fix: Add division by temperature in training model

nemo_rl/models/policy/dtensor_policy_worker.py

@@ -335,6 +335,9 @@ def train(
                     else:
                         logits = outputs.logits
 
+                    # Divide logits by temperature
+                    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,9 @@ def train(
                         else:
                             logits = outputs.logits
 
+                    # Divide logits by temperature
+                    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,
+        },
     }