From e5d7be18f8d59e78d9b72f6832be8ccdd6b65a14 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 28 Jul 2026 21:06:00 -0500 Subject: [PATCH 1/3] perf(megatron): skip redundant weight load Closes #2615 Signed-off-by: Teodor-Dumitru Ene --- nemo_rl/algorithms/grpo.py | 1 + nemo_rl/models/generation/megatron/megatron_generation.py | 6 ++++++ nemo_rl/models/megatron/setup.py | 5 ++++- nemo_rl/models/policy/lm_policy.py | 3 +++ nemo_rl/models/policy/workers/megatron_policy_worker.py | 6 ++++++ 5 files changed, 20 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 4159e5d5b60..a646d7357b7 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1047,6 +1047,7 @@ def init_megatron_generation(policy=None): tokenizer=tokenizer, processor=processor, weights_path=weights_path, + skip_weight_load=True, ) return mg, time.perf_counter() - t0 diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 5565a6b3461..707f2230c6c 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -67,6 +67,7 @@ def __init__( name_prefix: str = "megatron_generation", processor: Optional[AutoProcessor] = None, weights_path: Optional[str] = None, + skip_weight_load: bool = False, ): """Initialize a MegatronGeneration instance. @@ -80,6 +81,7 @@ def __init__( 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. """ # Import here to avoid circular imports from nemo_rl.models.policy.lm_policy import Policy @@ -87,6 +89,9 @@ def __init__( assert (cluster is None) != (policy is None), ( "Provide exactly one of `cluster` or `policy`." ) + assert not (skip_weight_load and policy is not None), ( + "skip_weight_load only applies to the dedicated inference policy." + ) # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. # `self._policy_config` keeps a reference to the full PolicyConfig. @@ -120,6 +125,7 @@ def __init__( init_optimizer=False, init_reference_model=False, weights_path=weights_path, + skip_weight_load=skip_weight_load, ) # Start the persistent inference engine + HTTP server during construction. diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d5aaae2d6fc..88c6d2b362a 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1295,6 +1295,7 @@ def setup_model_and_optimizer( get_position_embedding_ranks=None, pre_load_checkpoint_hook: Optional[Callable] = None, additional_pre_wrap_hooks: Optional[list[Callable]] = None, + load_weights: bool = True, ): state = GlobalState() _patch_bridge_signal_handler_for_worker_threads() @@ -1478,7 +1479,9 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: print("Model, optimizer, and learning rate scheduler built") torch.distributed.barrier() - if megatron_cfg.peft is not None: + if not load_weights: + should_load_checkpoint = False + elif megatron_cfg.peft is not None: should_load_checkpoint = resume_checkpoint_exists if should_load_checkpoint: # The finetune toggle is explicitly set to True in order to avoid loading optimizer and RNG states diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 6f2a8aaccce..727254cec5f 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -97,6 +97,7 @@ def __init__( init_reference_model: bool = True, processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, + skip_weight_load: bool = False, ): if weights_path: weights_path = os.path.abspath(weights_path) @@ -258,6 +259,8 @@ def __init__( worker_sharding_annotations=self.sharding_annotations, pre_init_communication_queue=pre_init_queue, ) + if skip_weight_load: + worker_kwargs["skip_weight_load"] = True if use_v2: # DTensor v2 workers reconstruct tokenizer/processor locally to avoid diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 4052fe03774..6f0e940f4f5 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -319,6 +319,7 @@ def __init__( init_reference_model: bool = True, *, worker_sharding_annotations: NamedSharding, + skip_weight_load: bool = False, **kwargs: Any, ): """Initialize the MegatronPolicyWorker.""" @@ -440,11 +441,16 @@ def __init__( self.megatron_cfg.validate() # Step 4: Setup Megatron model and components + assert not (skip_weight_load and (init_optimizer or init_reference_model)), ( + "skip_weight_load is only valid for inference-only policies " + "(init_optimizer=False, init_reference_model=False)." + ) model_and_optimizer_state = setup_model_and_optimizer( config, self.megatron_cfg, init_optimizer, pre_load_checkpoint_hook=getattr(self, "_pre_load_checkpoint_hook", None), + load_weights=not skip_weight_load, ) self.mcore_state = model_and_optimizer_state.state From d13065a31e89eb9ee5823d7d2090fed7e155a3e3 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 4 Aug 2026 04:10:31 -0500 Subject: [PATCH 2/3] Add unit-test for skip_weight_load Signed-off-by: Teodor-Dumitru Ene --- .../generation/test_megatron_generation.py | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 787a051fa95..e655fb36428 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -430,10 +430,15 @@ def test_megatron_generation_colocated(cluster, test_input_data, tokenizer): @pytest.mark.mcore @pytest.mark.timeout(900) +@pytest.mark.parametrize("skip_weight_load", [False, True]) def test_megatron_generation_non_colocated_refit( - policy_cluster_separate, test_input_data, tokenizer + policy_cluster_separate, test_input_data, tokenizer, skip_weight_load ): - """Non-colocated Megatron generation.""" + """Non-colocated Megatron generation. + + With skip_weight_load the inference engine builds without loading the + checkpoint and must still generate correctly once refit delivers weights. + """ generation_cluster = RayVirtualCluster( bundle_ct_per_node_list=[1], use_gpus=True, @@ -455,8 +460,22 @@ def test_megatron_generation_non_colocated_refit( policy = Policy( cluster=policy_cluster_separate, config=config, tokenizer=tokenizer ) + + # construction guard: skip_weight_load requires a dedicated inference + # policy; wrapping an existing (colocated) policy must be rejected. + with pytest.raises(AssertionError): + MegatronGeneration( + config=config, + tokenizer=tokenizer, + policy=policy, + skip_weight_load=True, + ) + mg = MegatronGeneration( - config=config, tokenizer=tokenizer, cluster=generation_cluster + config=config, + tokenizer=tokenizer, + cluster=generation_cluster, + skip_weight_load=skip_weight_load, ) # init the refit collective on both sides. From 91ccce7f9a99a44e947c09e02733fe5f3ef89a07 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 4 Aug 2026 13:49:24 -0500 Subject: [PATCH 3/3] Assert logprob parity in unit-test Signed-off-by: Teodor-Dumitru Ene --- .../generation/test_megatron_generation.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index e655fb36428..873e4d419b3 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -497,12 +497,36 @@ def test_megatron_generation_non_colocated_refit( # refit the inference engine from the training weights, then generate refit_policy_generation(policy, mg, False) - outputs = mg.generate(test_input_data, greedy=True) + # Greedy needs to be false because processed logprobs doesn't handle it well. + outputs = mg.generate(test_input_data, greedy=False) _assert_valid_generation_output(outputs, test_input_data) generated_texts = tokenizer.batch_decode( outputs["output_ids"], skip_special_tokens=True ) assert all(len(t) > 0 for t in generated_texts), "Some texts are empty" + + # Training-policy logprobs must match generation-policy logprobs. + # A broken refit would fail this test. + fprop_data = BatchedDataDict( + { + "input_ids": outputs["output_ids"], + "input_lengths": outputs["unpadded_sequence_lengths"], + } + ) + policy.prepare_for_lp_inference() + train_logprobs = policy.get_logprobs(fprop_data)["logprobs"] + gen_mask = torch.zeros_like(outputs["logprobs"], dtype=torch.bool) + for i, (start, end) in enumerate( + zip(test_input_data["input_lengths"], outputs["unpadded_sequence_lengths"]) + ): + gen_mask[i, start:end] = True + abs_diff = (outputs["logprobs"] - train_logprobs).abs().masked_select(gen_mask) + avg_prob_mult_error = torch.exp(abs_diff).mean() + assert avg_prob_mult_error <= 1.05, ( + f"generation logprobs diverge from training-policy logprobs " + f"(avg prob mult error {avg_prob_mult_error:.4f}); inference weights " + f"do not match training weights after refit" + ) finally: if mg is not None: mg.shutdown()