fix: Use Float16Module even when defer_fp32_logits=True (#1537)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Author: yfw

nemo_rl/models/megatron/common.py

@@ -358,6 +358,7 @@ def forward_step_arbitrary_loss(
     pad_individual_seqs_to_multiple_of: int = 1,
     pad_packed_seq_to_multiple_of: int = 1,
     pad_full_seq_to: Optional[int] = None,
+    defer_fp32_logits: Optional[bool] = None,
     cp_normalize: bool = True,
     policy_cfg: Optional[dict] = None,
 ):
@@ -372,6 +373,9 @@ def forward_step_arbitrary_loss(
         loss_fn (LossFunction): Loss function to apply
         pack_sequences (bool): Whether to pack sequences for efficiency
         seq_length_key (Optional[str]): Key in data_dict containing actual sequence lengths
+        pad_individual_seqs_to_multiple_of (int): Pad individual sequences to a multiple of this value
+        pad_full_seq_to (Optional[int]): Pad packed sequences to this value
+        defer_fp32_logits (Optional[bool]): Whether to skip the conversion of logits to fp32
         cp_normalize (bool): Whether to normalize the loss by the cp_size
         policy_cfg (Optional[dict]): Policy configuration containing generation parameters
 
@@ -453,6 +457,9 @@ def forward_step_arbitrary_loss(
     if packed_seq_params is not None:
         additional_kwargs["packed_seq_params"] = packed_seq_params
 
+    if defer_fp32_logits:
+        additional_kwargs["fp32_output"] = False
+
     with straggler_timer:
         output_tensor = model(
             input_ids=input_ids_cp_sharded,

nemo_rl/models/policy/megatron_policy_worker.py

@@ -274,10 +274,6 @@ def freeze_moe_router(megatron_model):
         mixed_precision_wrapper = CustomFloat16Module
         pre_wrap_hook.extend([freeze_moe_router])
 
-    # If deferring fp32 logits, disable mixed-precision wrapper entirely
-    if policy_cfg["megatron_cfg"].get("defer_fp32_logits", None):
-        mixed_precision_wrapper = None
-
     # Model, optimizer, and learning rate.
     model = get_model(
         cfg.model,
@@ -663,6 +659,9 @@ def __init__(
             assert self.cfg["megatron_cfg"]["defer_fp32_logits"], (
                 "defer_fp32_logits must be True if logprob_chunk_size is set"
             )
+        self.defer_fp32_logits = self.cfg["megatron_cfg"].get(
+            "defer_fp32_logits", None
+        ) and (model_cfg.fp16 or model_cfg.bf16)
 
         checkpoint_config = CheckpointConfig(
             save_interval=100,
@@ -796,8 +795,6 @@ def __init__(
             ref_mixed_precision_wrapper = Float16Module
             if self.cfg["megatron_cfg"].get("freeze_moe_router", False):
                 ref_mixed_precision_wrapper = CustomFloat16Module
-            if self.cfg["megatron_cfg"].get("defer_fp32_logits", None):
-                ref_mixed_precision_wrapper = None
 
             reference_model = get_model(
                 self.megatron_cfg.model,
@@ -1068,6 +1065,7 @@ def train(
                             pad_individual_seqs_to_multiple_of=pad_factor,
                             pad_packed_seq_to_multiple_of=pad_packed_seq_to_multiple_of,
                             pad_full_seq_to=pad_full_seq_to,
+                            defer_fp32_logits=self.defer_fp32_logits,
                         ),
                         data_iterator=data_iterator,
                         model=self.model,
@@ -1284,6 +1282,9 @@ def forward_step_fn(
             if packed_seq_params is not None:
                 additional_kwargs["packed_seq_params"] = packed_seq_params
 
+            if self.defer_fp32_logits:
+                additional_kwargs["fp32_output"] = False
+
             output_tensor = model(
                 input_ids=input_ids_cp_sharded,
                 position_ids=position_ids,