fix: Add division by temperature in training model

nemo_rl/models/policy/__init__.py

@@ -37,7 +37,7 @@ class PolicyConfig(TypedDict):
     train_micro_batch_size: int
     learning_rate: float
     logprob_batch_size: int
-    generation: GenerationConfig
+    generation: Optional[GenerationConfig]
     precision: str
     dtensor_cfg: DTensorConfig
     make_sequence_length_divisible_by: int

nemo_rl/models/policy/dtensor_policy_worker.py

@@ -335,6 +335,10 @@ def train(
                     else:
                         logits = outputs.logits
 
+                    # Divide logits by temperature
+                    if "generation" in self.cfg and self.cfg["generation"] is not None:
+                        logits.div_(self.cfg["generation"]["temperature"])
+
                     loss, loss_metrics = loss_fn(logits, mb)
                     num_valid_samples = loss_metrics["num_valid_samples"]
                     loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"]

nemo_rl/models/policy/fsdp1_policy_worker.py

@@ -290,6 +290,10 @@ def train(
                         else:
                             logits = outputs.logits
 
+                    # Divide logits by temperature
+                    if "generation" in self.cfg and self.cfg["generation"] is not None:
+                        logits.div_(self.cfg["generation"]["temperature"])
+
                     loss, loss_metrics = loss_fn(logits, mb)
                     num_valid_samples = loss_metrics["num_valid_samples"]
                     loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"]

tests/unit/models/generation/test_vllm_generation.py

@@ -36,7 +36,7 @@
     },
     "dtype": "bfloat16",
     "max_new_tokens": 10,
-    "temperature": 1.0,
+    "temperature": 0.8,
     "top_p": 1.0,
     "top_k": None,
     "stop_token_ids": None,
@@ -85,6 +85,9 @@ def get_basic_hf_test_config(enable_dtensor: bool = False) -> PolicyConfig:
         },
         "max_grad_norm": 1.0,
         "make_sequence_length_divisible_by": 1,
+        "generation": {
+            "temperature": 0.8,
+        },
     }
 
 

tests/unit/utils/test_native_checkpoint.py

@@ -61,6 +61,9 @@
         "tensor_parallel_size": 1,
     },
     "max_grad_norm": 1.0,
+    "generation": {
+        "temperature": 1.0,
+    },
 }