diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index ed97496ce5d..2b9ab8d5ee1 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1268,12 +1268,39 @@ def init_megatron_generation( cluster=None if colocated_inference else inference_cluster, policy=policy if colocated_inference else None, processor=processor, - weights_path=weights_path, skip_weight_load=not colocated_inference, reserved_http_server_port=reserved_http_server_port, ) return mg, time.perf_counter() - t0 + def init_megatron_weight_synchronizer( + policy: ColocatablePolicyInterface, + policy_generation: MegatronGeneration, + ) -> None: + """Initialize Megatron weight synchronizer. + + For non-colocated inference, also performs the initial weight sync. + """ + t0 = time.perf_counter() + weight_synchronizer = create_weight_synchronizer( + policy=policy, + generation=policy_generation, + generation_backend="megatron", + colocated=colocated_inference, + train_cluster=train_cluster, + inference_cluster=None if colocated_inference else inference_cluster, + ) + policy_generation.weight_synchronizer = weight_synchronizer + weight_synchronizer.init_communicator() + setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 + if not colocated_inference: + # The skip-load inference engine gets its final weight buffers here. + # Its first prepare_for_generation also starts the HTTP server only after the refit, + # so CUDA graphs can capture those persistent buffers. + t0 = time.perf_counter() + weight_synchronizer.sync_weights() + setup_timing_metrics.weight_sync_time_s = time.perf_counter() - t0 + def initialize_generation_with_policy( init_generation_fn, colocated_inference: bool, @@ -1349,41 +1376,55 @@ def initialize_generation_with_policy( setup_timing_metrics.generation_init_reserve_time_s = reserve_time print(f" ✓ Reserved Megatron server URL: {reserved_url}", flush=True) - def init_megatron_stack(): - """Init policy then generation; rank 0 holds the reserved port.""" - p, policy_t = init_policy( - reserved_http_server_port=reserved_http_server_port - if colocated_inference - else None - ) - pg, gen_t = init_megatron_generation( - p, - reserved_http_server_port=None - if colocated_inference - else reserved_http_server_port, - ) - return p, policy_t, pg, gen_t - def init_nemo_gym(): """Spin up NeMo Gym servers against the reserved URL.""" return _spinup_nemo_gym([reserved_url], generation_config["model_name"]) - init_tasks = { - "megatron": init_megatron_stack, - "nemo_gym": init_nemo_gym, - } - print(f" ⚡ Init tasks: {', '.join(init_tasks.keys())}", flush=True) + # Exactly one task adopts the reserved port: the policy when colocated + # (generation wraps it), else the dedicated generation policy. + policy_port, generation_port = ( + (reserved_http_server_port, None) + if colocated_inference + else (None, reserved_http_server_port) + ) + + def init_megatron_generation_task(policy_future): + """Colocated generation waits; non-colocated inits in parallel.""" + if colocated_inference: + p, _ = policy_future.result() + return init_megatron_generation(p) + return init_megatron_generation( + reserved_http_server_port=generation_port + ) + + print(" ⚡ Init tasks: policy, megatron_generation, nemo_gym", flush=True) + init_tasks_t0 = time.perf_counter() try: - with ThreadPoolExecutor(max_workers=len(init_tasks)) as executor: - submitted = {k: executor.submit(fn) for k, fn in init_tasks.items()} - results = {k: f.result() for k, f in submitted.items()} + with ThreadPoolExecutor(max_workers=3) as executor: + policy_future = executor.submit( + init_policy, reserved_http_server_port=policy_port + ) + generation_future = executor.submit( + init_megatron_generation_task, policy_future + ) + nemo_gym_future = executor.submit(init_nemo_gym) + policy, policy_time = policy_future.result() + policy_generation, megatron_gen_time = generation_future.result() + if not colocated_inference: + setup_timing_metrics.parallel_wall_time_s = ( + time.perf_counter() - init_tasks_t0 + ) + setup_timing_metrics.parallel_init_enabled = 1.0 + # NeMo Gym probes the pre-published endpoint before its future completes. + # A skip-load Megatron endpoint starts only during this initial refit, + # so it must happen while Gym is waiting rather than after it resolves. + init_megatron_weight_synchronizer(policy, policy_generation) + nemo_gym_actor, nemo_gym_time = nemo_gym_future.result() finally: ray.kill(port_holder) - policy, policy_time, policy_generation, megatron_gen_time = results[ - "megatron" - ] - nemo_gym_actor, nemo_gym_time = results["nemo_gym"] + if colocated_inference: + setup_timing_metrics.parallel_init_enabled = 0.0 setup_timing_metrics.policy_init_time_s = policy_time setup_timing_metrics.generation_init_time_s = ( reserve_time + megatron_gen_time @@ -1392,13 +1433,20 @@ def init_nemo_gym(): setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time else: - # Initialize training first so checkpoint conversion completes before inference starts. - policy, policy_time = init_policy() - setup_timing_metrics.policy_init_time_s = policy_time + if not colocated_inference: + policy_generation, policy = initialize_generation_with_policy( + init_megatron_generation, + colocated_inference, + setup_timing_metrics, + ) + else: + # Colocated generation wraps the training policy. + policy, policy_time = init_policy() + setup_timing_metrics.policy_init_time_s = policy_time - # Colocated wraps the training policy; non-colocated builds a dedicated inference policy. - policy_generation, megatron_gen_time = init_megatron_generation(policy) - setup_timing_metrics.generation_init_time_s = megatron_gen_time + policy_generation, megatron_gen_time = init_megatron_generation(policy) + setup_timing_metrics.generation_init_time_s = megatron_gen_time + setup_timing_metrics.parallel_init_enabled = 0.0 print( f" ✓ Using {backend} backend for generation with {policy_config['model_name']}", @@ -1643,22 +1691,8 @@ def init_dynamo(): ) if backend == "megatron": - t0 = time.perf_counter() - policy_generation.weight_synchronizer = create_weight_synchronizer( - policy=policy, - generation=policy_generation, - generation_backend=backend, - colocated=colocated_inference, - train_cluster=train_cluster, - inference_cluster=None if colocated_inference else inference_cluster, - ) - policy_generation.weight_synchronizer.init_communicator() - setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 - if not colocated_inference: - # Load the model weights now. - t0 = time.perf_counter() - policy_generation.weight_synchronizer.sync_weights() - setup_timing_metrics.weight_sync_time_s = time.perf_counter() - t0 + if policy_generation.weight_synchronizer is None: + init_megatron_weight_synchronizer(policy, policy_generation) if enable_nemo_gym: served_urls = policy_generation.dp_openai_server_base_urls if served_urls != [reserved_url]: diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index e7c62a74389..09753cc4f89 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -70,10 +70,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # parity checks should select raw_logprobs explicitly. logprobs_mode: Literal["processed_logprobs", "raw_logprobs"] - # FP8/MXFP8 for the dedicated (non-colocated) inference model; merged into its - # `megatron_cfg` by `merged_inference_megatron_cfg`. When enabled, the first - # refit re-quantizes the weight buffers, so engine construction (CUDA-graph - # capture) is deferred until after that refit (#3731). + # FP8/MXFP8 for the dedicated (non-colocated) inference model; + # merged into its `megatron_cfg` by `merged_inference_megatron_cfg`. fp8_cfg: NotRequired[Fp8Config] diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 2a4c878168c..604e136a8f6 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -147,7 +147,6 @@ def __init__( policy: Optional["Policy"] = None, name_prefix: str = "megatron_generation", processor: Optional[AutoProcessor] = None, - weights_path: Optional[str] = None, skip_weight_load: bool = False, reserved_http_server_port: Optional[int] = None, ): @@ -162,7 +161,6 @@ def __init__( policy: Existing training Policy to reuse for generation. name_prefix: Prefix for naming the worker group (non-colocated only). processor: Optional processor for VLMs (non-colocated only). - weights_path: Optional path to model weights (non-colocated only). skip_weight_load: Do not load the weights from the checkpoint; refit will do it. reserved_http_server_port: Driver-reserved OpenAI server port for non-colocated. """ @@ -214,16 +212,14 @@ def __init__( processor=processor, init_optimizer=False, init_reference_model=False, - weights_path=weights_path, skip_weight_load=skip_weight_load, reserved_http_server_port=reserved_http_server_port, ) - # MXFP8 inference re-quantizes weights at the first refit, so CUDA graphs must - # capture the post-refit buffers (#3731). Everything else starts the engine + - # HTTP server now: the NeMo-Gym overlap blocks on the server URL in setup (#3569). - gen_fp8_cfg = self.cfg["mcore_generation_config"].get("fp8_cfg") - if not (skip_weight_load and gen_fp8_cfg and gen_fp8_cfg["enabled"]): + # Skip-load models do not have their final refit weight buffers yet. + # Defer engine initialization so CUDA graphs capture the persistent buffers. + # The engine + HTTP server then first come up at the initial refit. + if not skip_weight_load: self.prepare_for_generation() def init_collective( diff --git a/nemo_rl/models/megatron/community_import.py b/nemo_rl/models/megatron/community_import.py index b162f2acb2b..6555b812b56 100644 --- a/nemo_rl/models/megatron/community_import.py +++ b/nemo_rl/models/megatron/community_import.py @@ -13,6 +13,9 @@ # limitations under the License. import os +import shutil +import threading +import uuid from collections.abc import Iterator from contextlib import contextmanager from typing import Any, Callable, Optional @@ -62,6 +65,76 @@ def to_torch_dtype(dtype: str | torch.dtype) -> torch.dtype: raise ValueError(f"Unknown dtype: {dtype}") +def megatron_conversion_is_complete(pretrained_path: str) -> bool: + """Whether a completed HF->Megatron conversion exists at `pretrained_path`.""" + return os.path.exists( + os.path.join(pretrained_path, "iter_0000000", "run_config.yaml") + ) + + +def publish_megatron_conversion( + staging_path: str, pretrained_path: str, *, overwrite: bool = False +) -> None: + """Atomically publish a staged HF->Megatron conversion. + + Args: + staging_path: Directory the conversion was saved into. + pretrained_path: Final conversion-cache path. + overwrite: Replace any existing complete conversion. + """ + displaced: Optional[str] = None + for attempt in range(2): + if not overwrite and megatron_conversion_is_complete(pretrained_path): + # A concurrent producer sharing the cache won; use its artifact. + print( + f"Completed conversion already published at {pretrained_path}; " + "discarding the staged copy.", + flush=True, + ) + threading.Thread( + target=shutil.rmtree, + args=(staging_path,), + kwargs={"ignore_errors": True}, + daemon=True, + ).start() + break + try: + os.rename(staging_path, pretrained_path) + break + except OSError: + if attempt == 1: + if not overwrite and megatron_conversion_is_complete(pretrained_path): + # A peer published into the slot we freed; use its artifact. + threading.Thread( + target=shutil.rmtree, + args=(staging_path,), + kwargs={"ignore_errors": True}, + daemon=True, + ).start() + break + raise + if not overwrite and megatron_conversion_is_complete(pretrained_path): + # A complete artifact raced in between the probe and the rename; + # concede to it on the next pass. + continue + # The final path is occupied (stale partial conversion, or a complete one under overwrite); + # displace it, retry the rename, then delete the displaced copy. + displaced = f"{pretrained_path}.displaced-{uuid.uuid4().hex[:8]}" + try: + os.rename(pretrained_path, displaced) + except FileNotFoundError: + # A concurrent publisher displaced it first; retry from the top. + displaced = None + if displaced is not None: + # Doomed uniquely-named copy nothing reads; delete without blocking startup. + threading.Thread( + target=shutil.rmtree, + args=(displaced,), + kwargs={"ignore_errors": True}, + daemon=True, + ).start() + + @contextmanager def _prefer_nvrx_for_dist_ckpt_save(): """Prefer NVRx async strategy for torch_dist save in HF->Megatron import. @@ -106,6 +179,8 @@ def import_model_from_hf_name( model_post_wrap_hook: Optional[Callable] = None, transformer_layer_spec: Optional[ModuleSpec | Callable] = None, mamba_stack_spec: Optional[ModuleSpec | Callable] = None, + *, + overwrite: bool = False, **config_overrides: Any, ): """Import a Hugging Face model into Megatron checkpoint format and save the Megatron checkpoint to the output path. @@ -123,6 +198,7 @@ def import_model_from_hf_name( mamba_stack_spec: Optional Megatron ``ModuleSpec`` (or callable returning one) overriding the default Mamba stack spec selected by Mamba model providers. + overwrite: Publish over an existing complete conversion. **config_overrides: Extra keyword arguments forwarded to ``AutoBridge.from_hf_pretrained``. """ @@ -225,8 +301,35 @@ def import_model_from_hf_name( config.num_layers_in_last_pipeline_stage = orig_num_layers_in_last_pipeline_stage config.pipeline_dtype = orig_pipeline_dtype + # Stage the save next to the final path, then atomically rename into place + # so concurrent readers of the shared cache never see a partial checkpoint. + output_path = os.path.normpath(output_path) + output_parent = os.path.dirname(output_path) + if output_parent: + os.makedirs(output_parent, exist_ok=True) + dist_active = ( + torch.distributed.is_available() and torch.distributed.is_initialized() + ) + staging_token = uuid.uuid4().hex[:8] + if dist_active: + token_box = [staging_token] + torch.distributed.broadcast_object_list(token_box, src=0) + staging_token = token_box[0] + staging_path = os.path.join( + output_parent, f".{os.path.basename(output_path)}.staging-{staging_token}" + ) + with _prefer_nvrx_for_dist_ckpt_save(): - bridge.save_megatron_model(megatron_model, output_path) + bridge.save_megatron_model(megatron_model, staging_path) + + # Every rank must finish writing before rank 0 publishes the staging dir, + # and no rank may read output_path before the rename lands. + if dist_active: + torch.distributed.barrier() + if not dist_active or torch.distributed.get_rank() == 0: + publish_megatron_conversion(staging_path, output_path, overwrite=overwrite) + if dist_active: + torch.distributed.barrier() # resetting mcore state import megatron.core.rerun_state_machine diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 399633cea52..05a92a7c25a 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -232,6 +232,7 @@ def _sync_distrib_opt(distrib_opt): from nemo_rl.models.megatron.community_import import ( import_model_from_hf_name, iter_vlm_config_overrides, + megatron_conversion_is_complete, ) from nemo_rl.models.megatron.config import ( ColocatedReshardPlan, @@ -342,6 +343,8 @@ def validate_and_set_config( pretrained_path, weights_path, optimizer_path, + *, + skip_weight_load: bool = False, ): # inference_optimized layers hard-require SP with TP>1; fail here with the config key. # This guards the training cfg; the inference cfg is guarded in @@ -418,6 +421,7 @@ def validate_and_set_config( pretrained_path, weights_path, optimizer_path, + skip_weight_load=skip_weight_load, ) final_padded_vocab_size = calculate_padded_vocab_size( @@ -584,9 +588,7 @@ def validate_model_paths(config: PolicyConfig) -> tuple[str, str, bool]: overrides_hash = _get_hf_config_overrides_hash(hf_config_overrides) hf_model_subdir = f"{hf_model_subdir}__hfovr_{overrides_hash}" pretrained_path = os.path.join(get_megatron_checkpoint_dir(), hf_model_subdir) - pt_checkpoint_exists = os.path.exists(pretrained_path) and os.path.exists( - os.path.join(pretrained_path, "iter_0000000") - ) + pt_checkpoint_exists = megatron_conversion_is_complete(pretrained_path) return hf_model_name, pretrained_path, pt_checkpoint_exists @@ -598,16 +600,30 @@ def setup_model_config( pretrained_path: str, weights_path: Optional[str] = None, optimizer_path: Optional[str] = None, + *, + skip_weight_load: bool = False, ) -> tuple[ConfigContainer, Any]: - """Handle all the model configuration logic.""" + """Handle all the model configuration logic. + + Args: + config: Policy config. + rank: Global rank (used in error messages). + dtype: Training dtype. + hf_model_name: HF model id (or local path). + pretrained_path: Path to the pretrained Megatron checkpoint. + weights_path: Path to save/load training weights. + optimizer_path: Path to the optimizer state (None if not resuming). + skip_weight_load: This policy never loads the checkpoint (weights arrive via refit). + """ pretrained_ckpt = config.get("pretrained_checkpoint") fmt = pretrained_ckpt["format"] if pretrained_ckpt is not None else None validate_router_replay_config(config) - if fmt == "megatron_lm": - # For megatron_lm format: build the model config from the HF architecture. - # pretrained_path has already been resolved to a specific iter dir by - # validate_model_paths, so no conversion step is needed. + derive_provider_from_hf = fmt == "megatron_lm" or ( + skip_weight_load and fmt != "megatron_bridge" + ) + + if derive_provider_from_hf: from transformers import AutoConfig hf_config_overrides = config.get("hf_config_overrides", {}) or {} @@ -717,8 +733,8 @@ def setup_model_config( # Reconstructed providers must be finalized so derived fields reflect the # merged config. Without overrides, preserve the existing checkpoint-load - # behavior: only megatron_lm providers need finalization here. - if fmt == "megatron_lm" or model_overrides: + # behavior: only HF-derived providers need finalization here. + if derive_provider_from_hf or model_overrides: model_cfg.finalize() model_cfg.__post_init__() @@ -730,23 +746,26 @@ def setup_model_config( fp8_cfg and fp8_cfg.get("enabled", False) and fp8_cfg.get("fp8_param", False) ) + # Refit-fed policies never read the pretrained checkpoint (weights arrive via refit). + ckpt_pretrained_path: Optional[str] = None if skip_weight_load else pretrained_path + # When fp8_param starts from a pretrained checkpoint, model params may already # be quantized before optimizer main params are initialized. Load main params # from the checkpoint state dict to preserve the original checkpoint precision. load_main_params_from_ckpt = ( fp8_param_enabled - and pretrained_path is not None + and ckpt_pretrained_path is not None and weights_path is None and optimizer_path is None ) - # Create checkpoint configs + # Create checkpoint configs. A refit-fed policy neither saves nor loads checkpoints. checkpoint_config = _create_checkpoint_config( - pretrained_path, + ckpt_pretrained_path, weights_path, optimizer_path, load_main_params_from_ckpt, - ckpt_cfg=config["megatron_cfg"].get("checkpoint"), + ckpt_cfg=None if skip_weight_load else config["megatron_cfg"].get("checkpoint"), ) # Validate training configuration @@ -1257,7 +1276,7 @@ def _validate_chunking_config(config: PolicyConfig) -> None: def _create_checkpoint_config( - pretrained_path: str, + pretrained_path: Optional[str], weights_path: Optional[str], optimizer_path: Optional[str], load_main_params_from_ckpt: bool = False, @@ -2052,6 +2071,7 @@ def handle_model_import( model_post_wrap_hook=model_post_wrap_hook, transformer_layer_spec=transformer_layer_spec, mamba_stack_spec=mamba_stack_spec, + overwrite=force_reconvert, **hf_config_overrides, ) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index cfcbca407d2..89d3eb0a069 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -497,16 +497,19 @@ def __init__( # worker) may set ``_model_import_post_wrap_hook`` and # layer-spec hooks on ``self`` before calling # super().__init__() to inject quantization hooks into HF->Megatron - # import. - handle_model_import( - config, - hf_model_name, - pretrained_path, - pt_checkpoint_exists, - model_post_wrap_hook=getattr(self, "_model_import_post_wrap_hook", None), - transformer_layer_spec=getattr(self, "_transformer_layer_spec", None), - mamba_stack_spec=getattr(self, "_mamba_stack_spec", None), - ) + # import. Refit-fed inference-only policies (skip_weight_load) skip the import entirely. + if not skip_weight_load: + handle_model_import( + config, + hf_model_name, + pretrained_path, + pt_checkpoint_exists, + model_post_wrap_hook=getattr( + self, "_model_import_post_wrap_hook", None + ), + transformer_layer_spec=getattr(self, "_transformer_layer_spec", None), + mamba_stack_spec=getattr(self, "_mamba_stack_spec", None), + ) log_gpu_memory_diagnostics( label="after_hf_import", worker_type="MegatronPolicyWorker" ) @@ -524,6 +527,7 @@ def __init__( pretrained_path, weights_path, optimizer_path, + skip_weight_load=skip_weight_load, ) self.megatron_cfg = runtime_config.megatron_cfg diff --git a/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh b/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh index 525dd628d30..16c9f7391b4 100644 --- a/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh +++ b/tests/functional/grpo_megatron_mxfp8_refit_gb200.sh @@ -97,18 +97,20 @@ uv run coverage run -a --data-file="$PROJECT_ROOT/tests/.coverage" --source="$PR uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" -# The setup timing proves that the initial load ran, while the per-step transfer -# timing proves that the post-update refit ran. Raw rollout log-probs are compared -# against BF16 policy recomputation; the bounds allow the established MXFP8 -# quantization delta while rejecting a bad refit. +# The setup weight-sync timing proves that the initial refit ran, while the +# per-step transfer timing proves that the post-update refit ran. Raw rollout +# log-probs are compared against BF16 policy recomputation; the bounds allow the +# established MXFP8 quantization delta while rejecting a bad refit. uv run tests/check_metrics.py "$JSON_METRICS" \ 'len(data["train/loss"]) == 2' \ - 'len(data["timing/setup/generation_init_load_time_s"]) == 1' \ - 'min(data["timing/setup/generation_init_load_time_s"]) > 0' \ + 'len(data["timing/setup/weight_sync_time_s"]) == 1' \ + 'min(data["timing/setup/weight_sync_time_s"]) > 0' \ 'len(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) == 1' \ 'min(data["timing/train/prepare_for_generation/transfer_and_update_weights"]) > 0' \ + 'len(data["train/gen_kl_error"]) == 2' \ 'max(data["train/gen_kl_error"]) < 0.15' \ - 'max(data["train/token_mult_prob_error"]) < 2.0' + 'len(data["train/token_mult_prob_error"]) == 2' \ + 'max(data["train/token_mult_prob_error"]) < 1.5' assert_grep 'cuda graph warmup' "$RUN_LOG" diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 21773c2ad15..bc5d4999577 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -15,6 +15,8 @@ import os from contextlib import ExitStack, contextmanager from pathlib import Path +from threading import Event +from types import SimpleNamespace from typing import Any from unittest.mock import MagicMock, patch @@ -3144,6 +3146,129 @@ def prepare_refit_info(self, _state): ) +def test_setup_refits_noncolocated_megatron_while_nemo_gym_waits( + monkeypatch, mock_grpo_components +): + """The initial refit must start a skip-load endpoint before Gym can finish.""" + from nemo_rl.algorithms import grpo as grpo_mod + + events = [] + gym_started = Event() + engine_ready = Event() + checkpointer = MagicMock() + checkpointer.get_latest_checkpoint_path.return_value = None + checkpointer.load_training_info.return_value = None + checkpointer.get_resume_paths.return_value = (None, None) + + reserved_url = "http://megatron.example/v1" + port_holder = object() + generation = SimpleNamespace( + weight_synchronizer=None, + dp_openai_server_base_urls=[], + ) + generation_cls = MagicMock(return_value=generation) + generation_cls.reserve_http_server_address.return_value = ( + reserved_url, + 1234, + port_holder, + ) + + synchronizer = MagicMock() + synchronizer.init_communicator.side_effect = lambda: events.append("init") + + def sync_weights(): + assert gym_started.wait(timeout=5), "NeMo Gym did not start concurrently" + events.append("sync") + generation.dp_openai_server_base_urls = [reserved_url] + engine_ready.set() + + synchronizer.sync_weights.side_effect = sync_weights + nemo_gym_actor = object() + + def spinup_nemo_gym_actor(**kwargs): + assert kwargs["base_urls"] == [reserved_url] + events.append("gym_started") + gym_started.set() + assert engine_ready.wait(timeout=5), ( + "NeMo Gym waited for an endpoint that the initial refit never started" + ) + events.append("gym_ready") + return nemo_gym_actor + + logger = MagicMock() + policy_cls = MagicMock(return_value=MagicMock()) + ray_kill = MagicMock() + monkeypatch.setattr(grpo_mod, "Logger", lambda *_args, **_kwargs: logger) + monkeypatch.setattr( + grpo_mod, "CheckpointManager", lambda *_args, **_kwargs: checkpointer + ) + monkeypatch.setattr( + grpo_mod, "ClippedPGLossFn", lambda *_args, **_kwargs: MagicMock() + ) + monkeypatch.setattr( + grpo_mod, "StatefulDataLoader", lambda *_args, **_kwargs: [None] + ) + monkeypatch.setattr(grpo_mod, "RayVirtualCluster", MagicMock) + monkeypatch.setattr(grpo_mod, "Policy", policy_cls) + monkeypatch.setattr(grpo_mod, "MegatronGeneration", generation_cls) + monkeypatch.setattr( + grpo_mod, "create_weight_synchronizer", lambda **_kwargs: synchronizer + ) + monkeypatch.setattr(grpo_mod, "spinup_nemo_gym_actor", spinup_nemo_gym_actor) + monkeypatch.setattr(grpo_mod.ray, "kill", ray_kill) + + master_config = mock_grpo_components["master_config"] + master_config.policy["model_name"] = "test-model" + master_config.policy["tokenizer"] = {"use_fastokens": False} + master_config.policy["dtensor_cfg"] = {"enabled": False} + master_config.policy["megatron_cfg"] = { + "enabled": False, + "pipeline_model_parallel_size": 1, + } + master_config.policy["generation"] = { + "backend": "megatron", + "temperature": 1.0, + "top_p": 1.0, + "top_k": None, + "val_temperature": 1.0, + "val_top_p": 1.0, + "val_top_k": None, + "colocated": { + "enabled": False, + "resources": {"gpus_per_node": 1, "num_nodes": 1}, + }, + "mcore_generation_config": {"expose_http_server": True}, + } + master_config.env = {"should_use_nemo_gym": True} + master_config.loss_fn = ClippedPGLossConfig(reference_policy_kl_penalty=0.0) + master_config.grpo.val_period = 0 + master_config.grpo.batch_multiplier = 1 + master_config.cluster["gpus_per_node"] = 2 + master_config.data["shuffle"] = False + master_config.data["num_workers"] = 0 + + dataset = MagicMock() + dataset.__len__ = MagicMock(return_value=1) + result = grpo_mod.setup(master_config, MagicMock(), dataset, None) + + assert generation_cls.call_args.kwargs["skip_weight_load"] is True + assert generation_cls.call_args.kwargs["reserved_http_server_port"] == 1234 + assert "reserved_http_server_port" not in policy_cls.call_args.kwargs + assert events.index("init") < events.index("sync") + assert events.index("gym_started") < events.index("sync") + assert events.index("sync") < events.index("gym_ready") + synchronizer.init_communicator.assert_called_once_with() + synchronizer.sync_weights.assert_called_once_with() + ray_kill.assert_called_once_with(port_holder) + setup_metrics = next( + call.args[0] + for call in logger.log_metrics.call_args_list + if call.kwargs.get("prefix") == "timing/setup" + ) + assert setup_metrics["weight_sync_time_s"] > 0 + assert result[2] is nemo_gym_actor + + def test_grpo_train_collects_generation_logger_and_seq_metrics( monkeypatch, mock_grpo_components ): diff --git a/tests/unit/models/megatron/test_community_import.py b/tests/unit/models/megatron/test_community_import.py index 21d18b03202..b2bdbc6fbab 100644 --- a/tests/unit/models/megatron/test_community_import.py +++ b/tests/unit/models/megatron/test_community_import.py @@ -15,9 +15,13 @@ """Unit tests for community_import checkpoint-save strategy shim.""" import importlib +import os import sys +import time from types import ModuleType, SimpleNamespace +import pytest + def _ensure_package(monkeypatch, name: str) -> ModuleType: """Create a minimal package module in sys.modules for import stubbing.""" @@ -43,6 +47,9 @@ def _load_community_import_module(monkeypatch): fake_torch.float32 = object() fake_torch.bfloat16 = object() fake_torch.float16 = object() + # import_model_from_hf_name guards its collectives on this; single-process + # tests take the non-distributed path. + fake_torch.distributed = SimpleNamespace(is_available=lambda: False) monkeypatch.setitem(sys.modules, "torch", fake_torch) # Stub megatron imports required by the module. @@ -188,7 +195,14 @@ def async_save(self, sharded_state_dict, checkpoint_dir, async_strategy): assert strategy.original_save_calls == [({"y": 2}, "/tmp/ckpt")] -def test_import_model_from_hf_name_calls_bridge_save(monkeypatch): +def _stage_conversion(path) -> None: + """Materialize a complete conversion layout (iter_0000000/run_config.yaml).""" + os.makedirs(os.path.join(str(path), "iter_0000000"), exist_ok=True) + with open(os.path.join(str(path), "iter_0000000", "run_config.yaml"), "w") as f: + f.write("{}\n") + + +def test_import_model_from_hf_name_calls_bridge_save(monkeypatch, tmp_path): module = _load_community_import_module(monkeypatch) _install_runtime_stubs_for_hf_import(monkeypatch) # Force this import path to stay unavailable even if real megatron modules @@ -231,6 +245,9 @@ def to_megatron_provider(self, load_weights): def save_megatron_model(self, megatron_model, output_path): self.saved_model = megatron_model self.saved_path = output_path + # The real save materializes the checkpoint; publish needs the + # staged directory (and its completion marker) to exist. + _stage_conversion(output_path) fake_bridge = FakeBridge() @@ -243,7 +260,48 @@ def from_hf_pretrained(hf_model_name, *args, **kwargs): monkeypatch.setattr(module, "AutoBridge", FakeAutoBridge) - module.import_model_from_hf_name("fake/hf-model", "/tmp/out") + output_path = tmp_path / "out" + module.import_model_from_hf_name("fake/hf-model", str(output_path)) assert fake_bridge.saved_model is not None - assert fake_bridge.saved_path == "/tmp/out" + # The save lands in a hidden staging sibling, then is renamed into place. + assert os.path.basename(fake_bridge.saved_path).startswith(".out.staging-") + assert module.megatron_conversion_is_complete(str(output_path)) + assert sorted(p.name for p in tmp_path.iterdir()) == ["out"] + + +@pytest.mark.parametrize( + ("occupant", "overwrite", "staged_wins"), + [ + # The empty-target happy path is covered end to end by + # test_import_model_from_hf_name_calls_bridge_save. + # A concurrent producer's complete artifact wins; the staged copy is discarded. + ("complete", False, False), + # A stale partial occupant (bare iter_0000000/, interrupted run) is displaced. + ("partial", False, True), + # force_reconvert_from_hf replaces even a complete artifact. + ("complete", True, True), + ], +) +def test_publish_conversion(monkeypatch, tmp_path, occupant, overwrite, staged_wins): + """Publish atomically renames the staging dir, resolving occupants of the final path.""" + module = _load_community_import_module(monkeypatch) + staging, final = tmp_path / ".ckpt.staging-abc", tmp_path / "ckpt" + _stage_conversion(staging) + (staging / "staged_marker").touch() + if occupant == "complete": + _stage_conversion(final) + elif occupant == "partial": + (final / "iter_0000000").mkdir(parents=True) + + module.publish_megatron_conversion(str(staging), str(final), overwrite=overwrite) + + assert module.megatron_conversion_is_complete(str(final)) + assert (final / "staged_marker").exists() is staged_wins + # Doomed copies are deleted on a background thread; wait for it. + deadline = time.monotonic() + 30.0 + while time.monotonic() < deadline and sorted( + p.name for p in tmp_path.iterdir() + ) != ["ckpt"]: + time.sleep(0.01) + assert sorted(p.name for p in tmp_path.iterdir()) == ["ckpt"] diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 8622fac910f..47bf43c8249 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -100,14 +100,17 @@ def test_model_name_is_local_path(self, tmp_path): assert "model_" in pretrained_path assert pt_checkpoint_exists is False - def test_checkpoint_exists(self, tmp_path): - """Test when a Megatron checkpoint already exists.""" + @pytest.mark.parametrize("complete", [True, False]) + def test_checkpoint_exists(self, tmp_path, complete): + """Only a completed conversion counts; a bare iter_0000000/ (interrupted) must not.""" from nemo_rl.models.megatron.setup import validate_model_paths # Create the checkpoint directory structure checkpoint_dir = tmp_path / "checkpoints" / "test-model" iter_dir = checkpoint_dir / "iter_0000000" iter_dir.mkdir(parents=True) + if complete: + (iter_dir / "run_config.yaml").write_text("{}\n") config = {"model_name": "test-model"} @@ -120,7 +123,7 @@ def test_checkpoint_exists(self, tmp_path): ) assert hf_model_name == "test-model" - assert pt_checkpoint_exists is True + assert pt_checkpoint_exists is complete def test_hf_config_overrides_change_hashed_pretrained_path(self, tmp_path): """Test that different hf_config_overrides map to different hashed paths.""" @@ -2508,8 +2511,18 @@ def _make_model_cfg_mock() -> MagicMock: model_cfg.__post_init__ = MagicMock() return model_cfg - def test_megatron_lm_passes_hf_config_overrides_to_autoconfig(self, request): - """hf_config_overrides must be forwarded to AutoConfig.from_pretrained for megatron_lm.""" + @pytest.mark.parametrize( + ("pretrained_checkpoint", "skip_weight_load"), + [ + ({"format": "megatron_lm", "path": "/ckpt"}, False), + # Refit-fed policies derive from the HF config even without a cache. + (None, True), + ], + ) + def test_hf_derived_provider_passes_hf_config_overrides_to_autoconfig( + self, request, pretrained_checkpoint, skip_weight_load + ): + """Both from_hf_config routes forward hf_config_overrides to AutoConfig.""" from nemo_rl.models.megatron.setup import setup_model_config self._apply_patches(request) @@ -2520,7 +2533,7 @@ def test_megatron_lm_passes_hf_config_overrides_to_autoconfig(self, request): overrides = {"rope_scaling": {"rope_type": "yarn", "factor": 4.0}} config = { - "pretrained_checkpoint": {"format": "megatron_lm", "path": "/ckpt"}, + "pretrained_checkpoint": pretrained_checkpoint, "hf_config_overrides": overrides, "megatron_cfg": {}, } @@ -2536,6 +2549,7 @@ def test_megatron_lm_passes_hf_config_overrides_to_autoconfig(self, request): dtype=torch.bfloat16, hf_model_name="test-model", pretrained_path="/ckpt/iter_0005000", + skip_weight_load=skip_weight_load, ) mock_ac.assert_called_once_with( @@ -2544,6 +2558,66 @@ def test_megatron_lm_passes_hf_config_overrides_to_autoconfig(self, request): rope_scaling={"rope_type": "yarn", "factor": 4.0}, ) + @pytest.mark.parametrize("fmt", [None, "megatron_bridge"]) + def test_skip_weight_load_has_no_pretrained_checkpoint_dependency( + self, tmp_path, request, fmt + ): + """Refit-fed policies carry no pretrained path or checkpoint knobs. + + hf format derives the provider from the HF config (the conversion cache + may not exist yet during parallel init); megatron_bridge keeps its + user-supplied serialized provider. + """ + from nemo_rl.models.megatron.setup import setup_model_config + + mocks = self._apply_patches(request) + + mock_model_cfg = self._make_model_cfg_mock() + mock_provider = MagicMock() + mock_provider.to_megatron_provider.return_value = mock_model_cfg + + if fmt == "megatron_bridge": + (tmp_path / "run_config.yaml").touch() + pretrained_ckpt = {"format": "megatron_bridge", "path": str(tmp_path)} + pretrained_path = str(tmp_path) + else: + pretrained_ckpt = None + pretrained_path = str(tmp_path / "never-published") + + config = { + "pretrained_checkpoint": pretrained_ckpt, + "megatron_cfg": {"checkpoint": {"async_save": True}}, + } + + with ( + patch("transformers.AutoConfig.from_pretrained"), + patch("nemo_rl.models.megatron.setup.AutoBridge") as mock_ab, + patch( + "nemo_rl.models.megatron.setup.load_model_config", + return_value=(self._make_model_cfg_mock(), None), + ) as mock_load, + ): + mock_ab.from_hf_config.return_value = mock_provider + setup_model_config( + config, + rank=0, + dtype=torch.bfloat16, + hf_model_name="test-model", + pretrained_path=pretrained_path, + skip_weight_load=True, + ) + + if fmt == "megatron_bridge": + mock_load.assert_called_once() + mock_ab.from_hf_config.assert_not_called() + else: + mock_load.assert_not_called() + # Freshly derived providers must be finalized before __post_init__. + mock_model_cfg.finalize.assert_called_once() + ckpt_call = mocks["_create_checkpoint_config"].call_args + assert ckpt_call.args[0] is None + assert ckpt_call.kwargs["ckpt_cfg"] is None + def test_model_overrides_are_finalized_and_serialized(self, tmp_path, request): """The reconstructed provider is the finalized, serializable config.""" from megatron.bridge.training.config import ConfigContainer @@ -2781,6 +2855,7 @@ def test_import_when_checkpoint_missing(self, mock_ps, mock_import, tmp_path): model_post_wrap_hook=None, transformer_layer_spec=None, mamba_stack_spec=None, + overwrite=False, ) @patch("nemo_rl.models.megatron.setup.import_model_from_hf_name") @@ -2845,6 +2920,8 @@ def test_force_reconvert_from_hf_when_checkpoint_exists( model_post_wrap_hook=None, transformer_layer_spec=None, mamba_stack_spec=None, + # The forced re-conversion must atomically replace the published cache. + overwrite=True, rope_scaling={ "rope_type": "yarn", "factor": 4.0,