update megatron

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/models/automodel/train.py

@@ -547,14 +547,15 @@ def __call__(
                 logits, self.device_mesh, self.cp_mesh, sequence_dim
             )
 
-        # Wrap loss function for sequence packing if needed
-        wrapped_prepare_loss_input = partial(
+        # Wrap prepare_loss_input with sampling_params
+        prepare_loss_input_wrapped = partial(
             prepare_loss_input, sampling_params=self.sampling_params
         )
+        # Wrap loss function for sequence packing if needed
         if self.enable_seq_packing:
             loss_fn = SequencePackingLossWrapper(
                 loss_fn=self.loss_fn,
-                prepare_fn=wrapped_prepare_loss_input,
+                prepare_fn=prepare_loss_input_wrapped,
                 cu_seqlens_q=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
                 cu_seqlens_q_padded=processed_inputs.flash_attn_kwargs.cu_seqlens_q,
             )
@@ -565,7 +566,7 @@ def __call__(
                 global_valid_toks,
             )
         else:
-            loss_input, mb = wrapped_prepare_loss_input(logits, mb, self.loss_fn)
+            loss_input, mb = prepare_loss_input_wrapped(logits, mb, self.loss_fn)
             loss, loss_metrics = self.loss_fn(
                 data=mb,
                 global_valid_seqs=global_valid_seqs,

nemo_rl/models/megatron/config.py

@@ -21,6 +21,8 @@
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.transformer import MegatronModule
 
+from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
+
 
 class MegatronGenerationConfig(TypedDict):
     # Total GPU memory (in GB) allocated for KV cache buffers
@@ -55,6 +57,7 @@ class RuntimeConfig(NamedTuple):
     optimizer_cpu_offload: bool
     offload_optimizer_for_logprob: bool
     is_generation_colocated: Optional[bool]
+    sampling_params: Optional[TrainingSamplingParams]
     final_padded_vocab_size: int
 
 

nemo_rl/models/megatron/setup.py

@@ -66,6 +66,7 @@
 except ImportError:
     HAVE_FSDP2 = False
 
+from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.models.megatron.community_import import import_model_from_hf_name
 from nemo_rl.models.megatron.config import ModelAndOptimizerState, RuntimeConfig
@@ -194,7 +195,6 @@ def validate_and_set_config(
     hf_model_name,
     pretrained_path,
     weights_path,
-    tokenizer,
 ):
     # Handle generation colocation
     is_generation_colocated = None
@@ -218,6 +218,16 @@ def validate_and_set_config(
     optimizer_cpu_offload = config["megatron_cfg"]["optimizer"]["optimizer_cpu_offload"]
     offload_optimizer_for_logprob = config["offload_optimizer_for_logprob"]
 
+    # Sampling parameters configuration
+    sampling_params = None
+    if "generation" in config and config["generation"] is not None:
+        generation_cfg = config["generation"]
+        sampling_params = TrainingSamplingParams(
+            top_k=generation_cfg.get("top_k", None),
+            top_p=generation_cfg.get("top_p", 1.0),
+            temperature=generation_cfg.get("temperature", 1.0),
+        )
+
     # Reward models are not yet supported with Megatron.
     if "reward_model_cfg" in config and config["reward_model_cfg"]["enabled"]:
         raise NotImplementedError(
@@ -242,6 +252,7 @@ def validate_and_set_config(
         optimizer_cpu_offload,
         offload_optimizer_for_logprob,
         is_generation_colocated,
+        sampling_params,
         final_padded_vocab_size,
     )
 

nemo_rl/models/megatron/train.py

@@ -29,6 +29,7 @@
 from megatron.core.pipeline_parallel import get_forward_backward_func
 from megatron.core.utils import StragglerDetector
 
+from nemo_rl.algorithms.logits_sampling_utils import TrainingSamplingParams
 from nemo_rl.algorithms.loss import (
     SequencePackingLossWrapper,
     prepare_loss_input,
@@ -56,7 +57,6 @@
 def model_forward(
     model: GPTModel,
     data_dict: BatchedDataDict[Any],
-    cfg: PolicyConfig,
     input_ids_cp_sharded: torch.Tensor,
     position_ids: torch.Tensor,
     attention_mask: torch.Tensor,
@@ -106,27 +106,26 @@ def model_forward(
 
 
 def apply_temperature_scaling(
-    logits: torch.Tensor,
-    cfg: PolicyConfig,
+    logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams]
 ) -> torch.Tensor:
     """Apply temperature scaling to logits.
 
     Args:
         logits: Logits tensor to scale
-        cfg: Policy configuration containing generation settings
+        sampling_params: Sampling parameters
 
     Returns:
         torch.Tensor: Temperature-scaled logits
     """
-    if "generation" in cfg and cfg["generation"] is not None:
-        logits.div_(cfg["generation"]["temperature"])
+    if sampling_params is not None and sampling_params.temperature != 1.0:
+        logits.div_(sampling_params.temperature)
     return logits
 
 
 def forward_with_post_processing_fn(
     data_iterator: Iterator[ProcessedMicrobatch],
     model: GPTModel,
-    cfg: PolicyConfig,
+    sampling_params: TrainingSamplingParams,
     post_processing_fn: PostProcessingFunction,
     defer_fp32_logits: Optional[bool] = False,
     global_valid_seqs: Optional[torch.Tensor] = None,
@@ -142,7 +141,7 @@ def forward_with_post_processing_fn(
     Args:
         data_iterator: Iterator yielding ProcessedMicrobatch objects (already processed)
         model: The model to run forward pass on
-        cfg: Policy configuration dictionary
+        sampling_params: Sampling parameters
         post_processing_fn: Post-processing function to post-process the logits
         defer_fp32_logits: Whether to defer FP32 conversion of logits
         global_valid_seqs: Global valid sequence count for loss normalization
@@ -169,7 +168,6 @@ def forward_with_post_processing_fn(
     output_tensor = model_forward(
         model=model,
         data_dict=data_dict,
-        cfg=cfg,
         input_ids_cp_sharded=input_ids_cp_sharded,
         position_ids=position_ids,
         attention_mask=attention_mask,
@@ -187,7 +185,7 @@ def forward_with_post_processing_fn(
         # Temperature scaling is element-wise, directly applying it here.
         # Other sampling parameters like top-k and top-p need the logits from whole vocabulary,
         # so applying them when gathering logits from vocab parallel (called in LossPostProcessor and LogprobsPostProcessor).
-        apply_temperature_scaling(output_tensor, cfg)
+        apply_temperature_scaling(output_tensor, sampling_params)
 
     # Use type checking to dispatch to the correct post-processing method
     if isinstance(post_processing_fn, LossPostProcessor):
@@ -218,7 +216,7 @@ def forward_with_post_processing_fn(
 
 def megatron_forward_backward(
     model: GPTModel,
-    cfg: PolicyConfig,
+    sampling_params: TrainingSamplingParams,
     data_iterator: Iterator[ProcessedMicrobatch],
     num_microbatches: int,
     seq_length: int,
@@ -238,7 +236,7 @@ def megatron_forward_backward(
 
     Args:
         model: The model to train
-        cfg: Policy configuration dictionary
+        sampling_params: Sampling parameters
         data_iterator: Iterator yielding ProcessedMicrobatch objects (already processed)
         num_microbatches: Number of microbatches to process
         seq_length: Sequence length
@@ -255,7 +253,7 @@ def megatron_forward_backward(
     """
     forward_step = partial(
         forward_with_post_processing_fn,
-        cfg=cfg,
+        sampling_params=sampling_params,
         post_processing_fn=post_processing_fn,
         defer_fp32_logits=defer_fp32_logits,
         global_valid_seqs=global_valid_seqs,
@@ -282,11 +280,13 @@ def __init__(
         cfg: PolicyConfig,
         num_microbatches: int = 1,
         cp_normalize: bool = True,
+        sampling_params: Optional[TrainingSamplingParams] = None,
     ):
         self.loss_fn = loss_fn
         self.cfg = cfg
         self.num_microbatches = num_microbatches
         self.cp_normalize = cp_normalize
+        self.sampling_params = sampling_params
 
     def __call__(
         self,
@@ -310,12 +310,17 @@ def __call__(
         Returns:
             Callable: Function that takes output tensor and returns (loss, metrics) tuple
         """
+        # wrap prepare_loss_input with sampling_params
+        prepare_loss_input_wrapped = partial(
+            prepare_loss_input, sampling_params=self.sampling_params
+        )
+
         # wrap loss function with loss input preparation
         pack_sequences = self.cfg["sequence_packing"]["enabled"]
         if pack_sequences and packed_seq_params is not None:
             loss_fn_wrapped = SequencePackingLossWrapper(
                 loss_fn=self.loss_fn,
-                prepare_fn=prepare_loss_input,
+                prepare_fn=prepare_loss_input_wrapped,
                 cu_seqlens_q=packed_seq_params.cu_seqlens_q,
                 cu_seqlens_q_padded=packed_seq_params.cu_seqlens_q_padded,
                 vocab_parallel_rank=get_tensor_model_parallel_rank(),
@@ -326,7 +331,7 @@ def __call__(
             loss_fn_wrapped = partial(
                 wrap_loss_fn_with_input_preparation,
                 loss_fn=self.loss_fn,
-                prepare_fn=prepare_loss_input,
+                prepare_fn=prepare_loss_input_wrapped,
                 vocab_parallel_rank=get_tensor_model_parallel_rank(),
                 vocab_parallel_group=get_tensor_model_parallel_group(),
                 context_parallel_group=get_context_parallel_group(),
@@ -365,8 +370,9 @@ def _counteract_mcore_loss_averaging(*args, **kwargs):
 
 
 class LogprobsPostProcessor:
-    def __init__(self, cfg: PolicyConfig):
+    def __init__(self, cfg: PolicyConfig, sampling_params: TrainingSamplingParams):
         self.cfg = cfg
+        self.sampling_params = sampling_params
 
     def __call__(
         self,
@@ -406,6 +412,7 @@ def processor_fn_inner(output_tensor):
                     inference_only=True,
                     cp_group=get_context_parallel_group(),
                     chunk_size=logprob_chunk_size,
+                    sampling_params=self.sampling_params,
                 )
             else:
                 token_logprobs = from_parallel_logits_to_logprobs(
@@ -416,6 +423,7 @@ def processor_fn_inner(output_tensor):
                     tp_group=tp_grp,
                     inference_only=True,
                     chunk_size=logprob_chunk_size,
+                    sampling_params=self.sampling_params,
                 )
 
             # Prepend 0 logprob for first token to maintain same sequence length as input

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -781,8 +781,10 @@ def get_topk_logits(
     def use_reference_model(self) -> Generator[None, None, None]:
         """Context manager that temporarily swaps the reference model and active model.
 
-        On entry: Moves model to CPU, moves reference_model to CUDA. Swaps the references
-        On exit: Restores original references and re-flips cuda/cpu
+        On entry: Moves model to CPU, moves reference_model to CUDA. Swaps the references.
+                  Also disables top-k/top-p filtering since the reference policy's distribution
+                  is different from the current policy, making filtered logprobs incompatible.
+        On exit: Restores original references and re-flips cuda/cpu, restores sampling_params.
         """
         with torch.no_grad():
             try:
@@ -796,10 +798,11 @@ def use_reference_model(self) -> Generator[None, None, None]:
                     val = to_local_if_dtensor(v)
                     val.copy_(self.reference_model_state_dict[k])
 
-                # - self.model is the original reference_model, now on CUDA
-                # - curr_state_dict is the train model, now on CPU
-
-                # Save and adjust sampling_params for reference model
+                # Temporarily disable top-k/top-p filtering for reference policy logprobs.
+                # The reference policy has different weights, so its top-k/top-p set is
+                # inherently different from the current policy. Using filtered logprobs
+                # would cause -inf mismatches that cannot be resolved by masking.
+                # Note: We keep temperature scaling since it was applied to prev_logprobs.
                 saved_sampling_params = self.sampling_params
                 if saved_sampling_params is not None:
                     self.sampling_params = TrainingSamplingParams(
@@ -810,6 +813,8 @@ def use_reference_model(self) -> Generator[None, None, None]:
                 else:
                     self.sampling_params = None
 
+                # - self.model is the original reference_model, now on CUDA
+                # - curr_state_dict is the train model, now on CPU
                 yield
 
             finally: