test: [3/3] add L3 Nano 4B Gym training E2E

nemo_rl/algorithms/grpo.py

@@ -765,16 +765,22 @@ def init_train_dataloader(dataset, suffix: str = ""):
             "aggregation. Set policy.megatron_cfg.use_fused_linear_logprobs=false "
             "or policy.sequence_packing.enabled=false."
         )
+        sampling_params = TrainingSamplingParams(
+            top_k=generation_config["top_k"],
+            top_p=generation_config["top_p"],
+            temperature=generation_config["temperature"],
+        )
+        assert sampling_params.temperature == 1.0, (
+            "Linear CE fusion loss is not supported with non-unit training-time "
+            "temperature for GRPO. The fused path computes logprobs before "
+            "temperature scaling. Set policy.megatron_cfg.use_fused_linear_logprobs=false, "
+            "or set policy.generation.temperature to 1.0 (or 0.0 for greedy generation)."
+        )
         # The fused forward gathers the logprob of the realized token from the raw
         # (unfiltered) logits, so top-k/top-p training-time filtering cannot be
         # applied. This also keeps prev/reference logprobs (computed via the fused
         # get_logprobs path) consistent with the actor logprobs.
-        assert not need_top_k_or_top_p_filtering(
-            TrainingSamplingParams(
-                top_k=generation_config["top_k"],
-                top_p=generation_config["top_p"],
-            )
-        ), (
+        assert not need_top_k_or_top_p_filtering(sampling_params), (
             "Linear CE fusion loss is not supported with top-k/top-p training-time "
             "filtering for GRPO. The fused path computes logprobs from unfiltered "
             "logits. Set policy.megatron_cfg.use_fused_linear_logprobs=false, or "
@@ -2532,7 +2538,8 @@ def refit_policy_generation(
         kv_scales: Optional dictionary of KV cache scales for FP8 quantization.
 
     Returns:
-        Scalar metrics reported by the selected weight synchronizer.
+        Scalar metrics reported by the selected weight synchronizer, or the
+        number of generation worker groups that acknowledged a direct refit.
     """
     # Every SGLang deployment reaches its refit through this hook: `setup`
     # attaches an SGLang synchronizer that owns the whole lifecycle (phase
@@ -2562,7 +2569,7 @@ def refit_policy_generation(
     )
     with timer_context:
         # update weights
-        update_success = False
+        acknowledged_updates: list[bool] = []
         if colocated_inference:
             # get model param keys, which is grouped by size
             if _refit_buffer_size_gb is not None:
@@ -2587,7 +2594,7 @@ def refit_policy_generation(
             # wait for all futures to complete
             ray.get(futures_train)
             results = ray.get(futures_inference)
-            update_success = all(result for result in results if result is not None)
+            acknowledged_updates = [result is True for result in results]
         else:
             # update weights through nccl (vLLM)
             futures_train = policy.broadcast_weights_for_collective(
@@ -2597,10 +2604,10 @@ def refit_policy_generation(
             # wait for all futures to complete
             ray.get(futures_train)
             results = ray.get(futures_inference)
-            update_success = all(result for result in results if result is not None)
+            acknowledged_updates = [result is True for result in results]
 
         # check if update is successful
-        if not update_success:
+        if not acknowledged_updates or not all(acknowledged_updates):
             error_tag = "cuda-ipc" if colocated_inference else "nccl"
             error_message = (
                 "❌ Error: Updating weights for the generation policy failed during refit.\n"
@@ -2613,7 +2620,7 @@ def refit_policy_generation(
         policy.offload_after_refit()
         policy_generation.prepare_for_generation(tags=["kv_cache"])
 
-    return {}
+    return {"generation_workers_updated": float(len(acknowledged_updates))}
 
 
 def _initial_policy_generation_stale(

nemo_rl/algorithms/logits_sampling_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from dataclasses import dataclass
 from typing import Optional
 
@@ -43,6 +44,25 @@ class TrainingSamplingParams:
     top_p: float = 1.0
     temperature: float = 1.0
 
+    def __post_init__(self) -> None:
+        """Normalize generation parameters to vLLM's logprob semantics.
+
+        Generation backends use ``temperature=0`` to select tokens greedily,
+        while the selected tokens' policy logprobs still come from the unscaled
+        model distribution. Top-k and top-p do not constrain greedy selection,
+        so training-time filtering must be disabled as well.
+        """
+        if not math.isfinite(self.temperature):
+            raise ValueError(f"temperature must be finite, got {self.temperature}")
+        if self.temperature < 0.0:
+            raise ValueError(
+                f"temperature must be non-negative, got {self.temperature}"
+            )
+        if self.temperature == 0.0:
+            self.temperature = 1.0
+            self.top_k = None
+            self.top_p = 1.0
+
 
 def _need_top_k_filtering(top_k: int | None) -> bool:
     """Check if top-k filtering is needed."""
@@ -66,6 +86,15 @@ def need_top_k_or_top_p_filtering(
     return _need_top_k_filtering(top_k) or _need_top_p_filtering(top_p)
 
 
+def apply_temperature_scaling(
+    logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams]
+) -> torch.Tensor:
+    """Apply the effective training temperature to logits in place."""
+    if sampling_params is not None and sampling_params.temperature != 1.0:
+        logits.div_(sampling_params.temperature)
+    return logits
+
+
 @torch.no_grad()
 def _apply_top_k_only_fn(
     logits: torch.Tensor,

nemo_rl/models/automodel/train.py

@@ -44,6 +44,7 @@
 
 from nemo_rl.algorithms.logits_sampling_utils import (
     TrainingSamplingParams,
+    apply_temperature_scaling,
     apply_top_k_top_p,
     need_top_k_or_top_p_filtering,
 )
@@ -229,23 +230,6 @@ def extract_logits(
         return outputs.logits
 
 
-def apply_temperature_scaling(
-    logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams]
-) -> torch.Tensor:
-    """Apply temperature scaling to logits.
-
-    Args:
-        logits: Logits tensor to scale
-        sampling_params: Sampling parameters
-
-    Returns:
-        torch.Tensor: Temperature-scaled logits
-    """
-    if sampling_params is not None and sampling_params.temperature != 1.0:
-        logits.div_(sampling_params.temperature)
-    return logits
-
-
 def apply_top_k_top_p_filtering_for_local_logits(
     logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams]
 ) -> torch.Tensor:

nemo_rl/models/generation/__init__.py

@@ -22,6 +22,7 @@
 from nemo_rl.models.generation.vllm.config import VLLM_SPARSE_REFIT_TRANSPORTS
 
 TokenizerType = PreTrainedTokenizerBase
+_VLLM_MIN_NON_ZERO_TEMPERATURE = 1e-2
 
 
 def resolve_generation_class(
@@ -64,6 +65,16 @@ def configure_generation_config(
     trains_mtp: bool = False,
 ) -> GenerationConfig:
     """Apply specific configurations to generation config."""
+    # vLLM clamps tiny positive temperatures before sampling. Normalize the
+    # shared config at this backend boundary so policy logprob recomputation
+    # receives the same effective temperature without changing other backends.
+    temperature = config.get("temperature")
+    if (
+        config["backend"] in ("vllm", "dynamo")
+        and temperature is not None
+        and 0.0 < temperature < _VLLM_MIN_NON_ZERO_TEMPERATURE
+    ):
+        config["temperature"] = _VLLM_MIN_NON_ZERO_TEMPERATURE
     if (
         config["backend"] != "vllm"
         and config.get("worker_extension_cls_fqn") is not None

nemo_rl/models/generation/trtllm/trtllm_worker_async.py

@@ -362,10 +362,10 @@ async def update_weights_from_collective_async(
                 "update_weights_from_collective",
                 kwargs={"drain": drain, "recompute_kv": recompute_kv},
             )
-            worker_result = results[0] if results else True
-            if not worker_result:
+            if not results or not all(result is True for result in results):
                 print(
-                    f"Error: TRT-LLM worker failed to update weights. Result: {worker_result}"
+                    "Error: TRT-LLM worker failed to update weights. "
+                    f"Results: {results}"
                 )
                 return False
             return True
@@ -380,10 +380,10 @@ async def update_weights_via_ipc_zmq_async(self) -> bool:
         assert self.llm is not None
         try:
             results = await self.llm.collective_rpc("update_weights_via_ipc_zmq")
-            worker_result = results[0] if results else True
-            if not worker_result:
+            if not results or not all(result is True for result in results):
                 print(
-                    f"Error: TRT-LLM worker failed to update weights via IPC. Result: {worker_result}"
+                    "Error: TRT-LLM worker failed to update weights via IPC. "
+                    f"Results: {results}"
                 )
                 return False
             return True

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -606,7 +606,11 @@ def _load_model(self, bundle_indices, seed):
 
         # Override HF config for gpt-oss models to ensure compatibility with megatron
         # The megatron --> hf export is done in bf16, so we disable quantization
-        hf_config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True)
+        hf_config = AutoConfig.from_pretrained(
+            self.model_name,
+            trust_remote_code=True,
+            revision=vllm_kwargs.get("revision"),
+        )
         self.routed_experts_dtype = resolve_routed_experts_dtype(
             get_num_routed_experts(hf_config)
         )

nemo_rl/models/megatron/train.py

@@ -39,6 +39,7 @@
 
 from nemo_rl.algorithms.logits_sampling_utils import (
     TrainingSamplingParams,
+    apply_temperature_scaling,
     need_top_k_or_top_p_filtering,
 )
 from nemo_rl.algorithms.loss import (
@@ -233,23 +234,6 @@ def model_forward(
     return output_tensor
 
 
-def apply_temperature_scaling(
-    logits: torch.Tensor, sampling_params: Optional[TrainingSamplingParams]
-) -> torch.Tensor:
-    """Apply temperature scaling to logits.
-
-    Args:
-        logits: Logits tensor to scale
-        sampling_params: Sampling parameters
-
-    Returns:
-        torch.Tensor: Temperature-scaled logits
-    """
-    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,

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -52,6 +52,7 @@
 
 from nemo_rl.algorithms.logits_sampling_utils import (
     TrainingSamplingParams,
+    apply_temperature_scaling,
     apply_top_k_top_p,
     need_top_k_or_top_p_filtering,
 )
@@ -576,9 +577,7 @@ def create_context_parallel_ctx(
     # based on https://github.com/pytorch/torchtitan/blob/cddd7dc809f36fe0ed51cdaaea0671c084d75442/torchtitan/distributed/utils.py#L178
 
     def _apply_temperature_scaling(self, logits: torch.Tensor) -> torch.Tensor:
-        if self.sampling_params is not None and self.sampling_params.temperature != 1.0:
-            logits.div_(self.sampling_params.temperature)
-        return logits
+        return apply_temperature_scaling(logits, self.sampling_params)
 
     def _apply_top_k_top_p_filtering(self, logits: torch.Tensor) -> torch.Tensor:
         """Apply top-k and top-p filtering to the logits locally when TP is disabled."""

tests/functional/L1_Functional_Tests_GB200_Gym_Training.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..")
+
+cd "${PROJECT_ROOT}"
+
+if [[ "${FAST:-0}" == "1" ]]; then
+    echo "FAST: Skipping Nano 4B Gym training E2E on GB200"
+else
+    time uv run --no-sync bash ./tests/functional/grpo_nano4b_gym_training_e2e.sh
+fi

tests/functional/grpo_nano4b_gym_training_e2e.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..")
+EXP_NAME=$(basename "$0" .sh)
+EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}"
+LOG_DIR="${EXP_DIR}/logs"
+JSON_METRICS="${EXP_DIR}/metrics.json"
+RUN_LOG="${EXP_DIR}/run.log"
+CONFIG_PATH="${SCRIPT_DIR}/grpo_nano4b_gym_training_e2e.yaml"
+
+export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"
+
+rm -rf "${EXP_DIR}"
+mkdir -p "${LOG_DIR}"
+
+cd "${PROJECT_ROOT}"
+
+uv run coverage run -a --data-file="${PROJECT_ROOT}/tests/.coverage" --source="${PROJECT_ROOT}/nemo_rl" \
+    "${PROJECT_ROOT}/examples/nemo_gym/run_grpo_nemo_gym.py" \
+    --config "${CONFIG_PATH}" \
+    logger.log_dir="${LOG_DIR}" \
+    "$@" \
+    2>&1 | tee "${RUN_LOG}"
+
+grep -Fq "Running synchronous GRPO training" "${RUN_LOG}"
+
+uv run tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}"
+
+# The fixture intentionally contains one accepted and one rejected rollout. In
+# addition to testing both verifier outcomes, this gives Reinforce++ a non-zero
+# advantage so a finite nonzero grad norm and positive learning rate prove that
+# a trainable update signal reached the optimizer path.
+uv run tests/check_metrics.py "${JSON_METRICS}" \
+    'len(data["train/loss"]) == 1' \
+    'all_finite(data["train/loss"])' \
+    'all_finite(data["train/grad_norm"])' \
+    'min(data["train/grad_norm"]) > 0' \
+    'all_finite(data["train/lr"])' \
+    'min(data["train/lr"]) > 0' \
+    'all_finite(data["train/advantages/min"])' \
+    'all_finite(data["train/advantages/max"])' \
+    'min(data["train/advantages/min"]) < 0' \
+    'max(data["train/advantages/max"]) > 0' \
+    'data["train/total_reward/min"]["1"] == 0' \
+    'data["train/total_reward/max"]["1"] == 1' \
+    'data["train/total_reward/mean"]["1"] == 0.5' \
+    'all_finite(data["train/token_mult_prob_error"])' \
+    'max(data["train/token_mult_prob_error"]) < 1.05' \
+    'data["refit/generation_workers_updated"]["1"] > 0' \
+    'data["timing/train/generation"]["1"] > 0' \
+    'data["validation/accuracy"]["1"] == 0.5' \
+    'data["timing/validation/total_validation_time"]["1"] > 0'