fix

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

nemo_rl/models/policy/dtensor_init.py

@@ -134,7 +134,7 @@ def validate_and_set_config(
 
     # Get other configuration values
     cpu_offload = config["dtensor_cfg"]["cpu_offload"]
-    offload_optimizer_for_logprob = config["offload_optimizer_for_logprob"]
+    offload_optimizer_for_logprob = config.get("offload_optimizer_for_logprob", False)
     max_grad_norm = config["max_grad_norm"]
     enable_seq_packing = config["sequence_packing"]["enabled"]
     model_name = config["model_name"]

nemo_rl/models/policy/dtensor_policy_worker_v2.py

@@ -349,21 +349,23 @@ def train(
 
                     # Forward and backward pass
                     loss, loss_metrics = forward_backward(
-                        self.model,
-                        mb,
-                        loss_fn,
-                        global_valid_seqs,
-                        global_valid_toks,
-                        processed_inputs,
-                        self.dtype,
-                        self.cp_size,
-                        self.cp_mesh,
-                        self.device_mesh,
-                        self.enable_seq_packing,
-                        self._is_reward_model,
-                        self.allow_flash_attn_args,
-                        eval_mode,
-                        self._apply_temperature_scaling,
+                        model=self.model,
+                        mb=mb,
+                        loss_fn=loss_fn,
+                        global_valid_seqs=global_valid_seqs,
+                        global_valid_toks=global_valid_toks,
+                        processed_inputs=processed_inputs,
+                        dtype=self.dtype,
+                        cp_size=self.cp_size,
+                        cp_mesh=self.cp_mesh,
+                        device_mesh=self.device_mesh,
+                        enable_seq_packing=self.enable_seq_packing,
+                        is_reward_model=self._is_reward_model,
+                        allow_flash_attn_args=self.allow_flash_attn_args,
+                        is_hf_model=self.is_hf_model,
+                        is_moe_model=self.is_moe_model,
+                        eval_mode=eval_mode,
+                        apply_temperature_fn=self._apply_temperature_scaling,
                     )
 
                     # skip the update for dummy batches