Merge remote-tracking branch 'origin' into yifu/vllm0112_bump

Signed-off-by: Yi-Fu Wu <[email protected]>

Author: yfw

docs/Makefile

@@ -46,7 +46,7 @@ ensure-docs-env:
 	@if [ ! -x "$(PYTHON)" ]; then \
 		echo "📦 Creating isolated docs environment..."; \
 		uv venv .venv; \
-		uv sync --no-config; \
+		uv sync --project ../pyproject.toml --group docs; \
 		echo "✅ Docs environment ready."; \
 		echo "📝 To activate it: $(ACTIVATE_CMD)"; \
 	fi

docs/fp8.md

@@ -66,9 +66,9 @@ FP8 generations are recommended to be configured with the following settings:
 
 ## Compatibility Note for Deepseek-Style FP8 Training
 
-When using FP8 training with Deepseek-style FP8 (sub-channel scaling), be aware of the following compatibility issue:
+The TransformerEngine implementation for this recipe requires **cuda version ≥ 12.9**. The latest nemo-rl depends on torch 2.8.0 + cuda 12.9 (since this [commit](https://github.com/NVIDIA-NeMo/RL/commit/3f36d14b53e906b27c01c06e36dbbd2b8eb300cd)). Users should check-out code to latest and build container from `docker/Dockerfile` ([instructions](docker.md)). 
 
-The TransformerEngine implementation for this recipe requires **cuBLAS version ≥ 12.9**. However, `nemo-rl` currently depends on **Torch 2.7.1**, which in turn requires **CUDA 12.8**. As a result, attempting to use the default setup will trigger the following error:
+If you are using nemo-rl before this [commit](https://github.com/NVIDIA-NeMo/RL/commit/3f36d14b53e906b27c01c06e36dbbd2b8eb300cd), you will see the following error when trying to use fp8 training
 
 ```
 File "/opt/ray_venvs/nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker/lib/python3.12/site-packages/transformer_engine/pytorch/fp8.py", line 646, in fp8_autocast
@@ -78,11 +78,6 @@ assert fp8_block_available, reason_for_no_fp8_block
            ^^^^^^^^^^^^^^^^^^^
 AssertionError: FP8 block scaled GEMM requires Hopper and CUDA >= 12.9.
 ```
-This issue will be resolved once the Torch version is upgraded to **≥ 2.8.0** (Please follow [#1122](https://github.com/NVIDIA-NeMo/RL/issues/1122) for more progress on the upgrade). In the meantime, you can enable Deepseek-style FP8 training using the following workaround:
-
-- **Build the NGC PyTorch container** from `docker/Dockerfile.ngc_pytorch`.  
-  This setup uses the system Python environment, which includes **CUDA version 12.9 or higher**, meeting the requirements for TransformerEngine’s FP8 implementation.
-
 
 
 ## Accuracy

docs/pyproject.toml

@@ -229,6 +229,8 @@ policy:
       use_deep_gemm: False
       num_last_layers_in_bf16: 0
       num_first_layers_in_bf16: 0
+      enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance
+      vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics
     vllm_kwargs: {}
     colocated:
       # true: generation shares training GPUs

docs/uv.lock

@@ -0,0 +1,59 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  val_period: -1
+loss_fn:
+  reference_policy_kl_penalty: 0.04
+  use_importance_sampling_correction: true
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo_megatron
+  save_period: 10000
+policy:
+  model_name: moonshotai/Moonlight-16B-A3B-Instruct
+  train_micro_batch_size: 1
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 8192
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    algorithm: modified_ffd
+  make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size}
+  optimizer: null
+  megatron_cfg:
+    enabled: true
+    moe_router_dtype: fp32
+    expert_model_parallel_size: 4
+    pipeline_model_parallel_size: 4
+    num_layers_in_first_pipeline_stage: 7
+    num_layers_in_last_pipeline_stage: 6
+    apply_rope_fusion: false
+    fp8_cfg:
+      enabled: true
+      fp8: e4m3
+      fp8_recipe: blockwise
+      fp8_param: false
+    optimizer:
+      lr: 1.0e-06
+      use_precision_aware_optimizer: false
+    scheduler:
+      lr_warmup_iters: 50
+    env_vars:
+      NVTE_FP8_BLOCK_SCALING_FP32_SCALES: '1'
+  generation:
+    vllm_cfg:
+      precision: fp8
+      use_deep_gemm: true
+      gpu_memory_utilization: 0.5
+      expert_parallel_size: 4
+      quantization_ignored_layer_kws: [
+        a_proj,
+        b_proj
+      ]
+logger:
+  monitor_gpus: false
+  wandb:
+    name: grpo-moonlight-16B-A3B-Instruct
+cluster:
+  gpus_per_node: 8
+  num_nodes: 4

examples/configs/grpo_math_1B.yaml

@@ -31,7 +31,7 @@ policy:
       PYTORCH_CUDA_ALLOC_CONF: expandable_segments:False
   generation:
     vllm_cfg:
-      tensor_parallel_size: 4
+      tensor_parallel_size: 2
 logger:
   log_dir: logs/grpo-qwen3-30ba3b-4n8g
   wandb_enabled: true

examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-fp8-e2e.yaml

@@ -1073,6 +1073,8 @@ def grpo_train(
 
                 dynamic_sampling_num_gen_batches += 1
                 with timer.time("generation"):
+                    # Clear vLLM logger metrics for each generation step
+                    policy_generation.clear_vllm_logger_metrics()
                     # Use penguin rollouts if enabled. We cascade penguin first since penguin requires async rollouts.
                     if _should_use_penguin(master_config):
                         generation_config = master_config["policy"]["generation"]
@@ -1122,6 +1124,9 @@ def grpo_train(
                             greedy=False,
                         )
                     policy_generation.finish_generation()
+                    # Collect vLLM logger metrics for performance reporting after each generation step
+                    # inflight batch sizes and num pending samples are collected from each vLLM worker
+                    vllm_logger_metrics = policy_generation.get_vllm_logger_metrics()
 
                 repeated_batch = scale_rewards(
                     repeated_batch, master_config["grpo"]["reward_scaling"]
@@ -1340,6 +1345,7 @@ def grpo_train(
                         metrics[k] = np.sum(v).item()
 
                 metrics.update(rollout_metrics)
+                metrics["vllm_logger_metrics"] = vllm_logger_metrics
                 total_valid_tokens += metrics["global_valid_toks"]
 
                 ## Checkpointing
@@ -1907,6 +1913,9 @@ def async_grpo_train(
 
     print("✅ All setup complete, starting buffer wait...")
 
+    # Clear vLLM logger metrics after at start of training
+    policy_generation.clear_vllm_logger_metrics()
+
     # Wait for initial buffer fill
     print(
         f"⏳ Waiting for replay buffer to have sufficient trajectories ({min_trajectories_needed} trajectories)..."
@@ -2145,12 +2154,17 @@ def async_grpo_train(
                     train_results = policy.train(train_data, loss_fn)
 
                 print("🔄 Synchronizing policy weights to trajectory collector…")
+                vllm_logger_metrics = None
                 if NEED_REFIT:
                     # Measure pending-generation wait as exposed_generation time
                     print("🔄 Coordinating with trajectory collector before refit...")
                     with timer.time("exposed_generation"):
                         ray.get(trajectory_collector.prepare_for_refit.remote())
 
+                    # Collect vLLM logger metrics for performance reporting
+                    # inflight batch sizes and num pending samples are collected from each vLLM worker
+                    vllm_logger_metrics = policy_generation.get_vllm_logger_metrics()
+
                     # Only the actual refit/weight transfer should be counted as weight_sync
                     print("🔄 Performing policy generation refit...")
                     with timer.time("weight_sync"):
@@ -2164,6 +2178,9 @@ def async_grpo_train(
                         trajectory_collector.set_weight_version.remote(weight_version)
                         trajectory_collector.resume_after_refit.remote()
 
+                # Clear vLLM logger metrics after each refit (weight sync), starting a new logging cycle
+                policy_generation.clear_vllm_logger_metrics()
+
                 # Validation
                 val_metrics, validation_timings = None, None
                 is_last_step = step + 1 == master_config["grpo"]["max_num_steps"]
@@ -2241,6 +2258,8 @@ def async_grpo_train(
                     else:
                         metrics[k] = np.sum(v).item()
                 metrics.update(rollout_metrics)
+                if vllm_logger_metrics is not None:
+                    metrics["vllm_logger_metrics"] = vllm_logger_metrics
                 total_valid_tokens += metrics["global_valid_toks"]
 
                 # Checkpointing (same as sync version)

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml

@@ -16,7 +16,7 @@
 import random
 import warnings
 from functools import partial, wraps
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 import torch
@@ -384,7 +384,7 @@ def maybe_pad_last_batch(batch: dict, dp_size: int, mbs: int) -> dict:
 
 def print_performance_metrics(
     train_results: dict[str, float],
-    metrics: dict[str, float],
+    metrics: dict[str, Any],
     timing_metrics: dict[str, float],
     master_config: dict,
 ) -> dict[str, float]:
@@ -400,13 +400,14 @@ def visualize_per_worker_load(per_worker_token_counts: dict[int, int]) -> float:
         per_worker_load_ratio = [
             v / max(per_worker_token_counts_list) for v in per_worker_token_counts_list
         ]
-        max_rows_to_print = 100
+        max_rows_to_print = 1000
+        bar_length = 20
         print("  • Visualizing Token Imbalance per Generation Worker:")
         for i in range(min(len(per_worker_token_counts_list), max_rows_to_print)):
             print(
                 f"    - Generated Tokens from Worker {i:3.0f}:"
-                f"{'■' * int(per_worker_load_ratio[i] * 10)}"
-                f"{'□' * (10 - int(per_worker_load_ratio[i] * 10))}"
+                f"{'■' * int(per_worker_load_ratio[i] * bar_length)}"
+                f"{'□' * (bar_length - int(per_worker_load_ratio[i] * bar_length))}"
                 f" Count: {per_worker_token_counts_list[i] / 1000:.1f}K"
             )
         estimated_idle_ratio = 1 - sum(per_worker_load_ratio) / len(
@@ -441,6 +442,125 @@ def visualize_per_worker_load(per_worker_token_counts: dict[int, int]) -> float:
             f"  • Mean Total Tokens per Sample: {metrics['mean_total_tokens_per_sample']:.2f}"
         )
 
+    # =====================================================
+    # vLLM Logger Metrics (inflight batch sizes, num pending samples, etc.)
+    # =====================================================
+    def resize_timeline(data, new_size):
+        old_size = len(data)
+        x_old = np.linspace(0, 1, old_size)
+        x_new = np.linspace(0, 1, new_size)
+        return np.interp(x_new, x_old, data)
+
+    def get_min_idle_time(
+        metric_dict: dict[int, list[int]], timeline_interval: float
+    ) -> float:
+        min_idle_time = float("inf")
+        for _, metric_values in metric_dict.items():
+            count_zeros = lambda x: sum(v == 0 for v in x)
+            idle_time = count_zeros(metric_values) * timeline_interval
+            min_idle_time = min(min_idle_time, idle_time)
+        return min_idle_time
+
+    def visualize_per_worker_timeline(
+        metric_dict: dict[int, list[int]],
+        metric_name: str,
+        timeline_interval: float | None,
+    ) -> None:
+        dp_ranks = list(metric_dict.keys())
+        max_rows_to_print = 1000
+        max_timeline_length = 50
+        marker = {0: "▃", 1: "▅", 2: "▆", 3: "▉"}
+        zero_marker = "▁"
+
+        max_value = max((max(v) if v else 0) for v in metric_dict.values())
+        bin_width = (max_value + 1) / len(marker)
+
+        print(f"  - {metric_name}:")
+        print(f"    - Max value: {max_value}")
+        if timeline_interval is not None:
+            print(
+                f"    - Min idle time: {get_min_idle_time(metric_dict, timeline_interval)} s"
+            )
+        print(
+            f"    - Timeline (0: {zero_marker}, {', '.join(f'{1.0 if k == 0 else k * (max_value / len(marker))}-{(k + 1) * (max_value / len(marker))}: {marker[k]}' for k in marker.keys())}):"
+        )
+        for dp_idx, metric_values in metric_dict.items():
+            if dp_idx > max_rows_to_print:
+                break
+            timeline = []
+            length = len(metric_values)
+            if timeline_interval is not None:
+                count_zeros = lambda x: sum(v == 0 for v in x)
+                idle = count_zeros(metric_values) * timeline_interval
+                active = length * timeline_interval - idle
+            if length > max_timeline_length:
+                resized_metric_values = resize_timeline(
+                    metric_values, max_timeline_length
+                )
+            else:
+                resized_metric_values = metric_values
+
+            for i, value in enumerate(resized_metric_values):
+                m = (
+                    zero_marker
+                    if value == 0
+                    else marker[min(int(value // bin_width), len(marker) - 1)]
+                )
+                timeline.append(m)
+            if timeline_interval is not None:
+                print(
+                    f"    - Generation Worker {dp_idx:3.0f}: {''.join(timeline)} (Active: {active:.2f} s, Idle: {idle:.2f} s)"
+                )
+            else:
+                print(f"    - Generation Worker {dp_idx:3.0f}: {''.join(timeline)}")
+
+    is_vllm_metrics_logger_enabled = master_config["policy"]["generation"].get(
+        "vllm_cfg", {}
+    ).get("enable_vllm_metrics_logger", False) and master_config["policy"][
+        "generation"
+    ].get("vllm_cfg", {}).get("async_engine", False)
+    if is_vllm_metrics_logger_enabled:
+        vllm_logger_metrics = metrics["vllm_logger_metrics"]
+        # vllm_logger_me    trics: dict[str (metric_name), dict[int (dp_idx), list[int] (metric_values)]]
+        # metric_name: "inflight_batch_sizes" or "num_pending_samples"
+
+        assert "inflight_batch_sizes" in vllm_logger_metrics, (
+            "inflight_batch_sizes not found in vllm_logger_metrics"
+        )
+        assert "num_pending_samples" in vllm_logger_metrics, (
+            "num_pending_samples not found in vllm_logger_metrics"
+        )
+        assert isinstance(vllm_logger_metrics["inflight_batch_sizes"], dict), (
+            "inflight_batch_sizes must be a dictionary"
+        )
+        assert isinstance(vllm_logger_metrics["num_pending_samples"], dict), (
+            "num_pending_samples must be a dictionary"
+        )
+
+        vllm_metrics_logger_interval = master_config["policy"]["generation"][
+            "vllm_cfg"
+        ]["vllm_metrics_logger_interval"]
+        print("  • vLLM Logger Metrics:")
+        # Visualize the inflight batch sizes timeline
+        if len(vllm_logger_metrics["inflight_batch_sizes"].values()) > 0:
+            visualize_per_worker_timeline(
+                vllm_logger_metrics["inflight_batch_sizes"],
+                "Inflight Batch Sizes",
+                vllm_metrics_logger_interval,
+            )
+        if len(vllm_logger_metrics["num_pending_samples"].values()) > 0:
+            max_num_pending_samples = max(
+                (max(v) if v else 0)
+                for v in vllm_logger_metrics["num_pending_samples"].values()
+            )
+            # If there is at least one pending sample, visualize the timeline
+            if max_num_pending_samples > 0:
+                visualize_per_worker_timeline(
+                    vllm_logger_metrics["num_pending_samples"],
+                    "Num Pending Samples",
+                    None,
+                )
+
     # =====================================================
     # Throughputs
     # =====================================================