perf(megatron): skip redundant weight load

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
 

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,13 +81,17 @@ 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
 
         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.

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

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

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

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.
@@ -478,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()