fix: consolidate Megatron Inference tests

docs/design-docs/generation.md

@@ -247,9 +247,8 @@ policy:
       max_tokens: 16384
 ```
 
-For a complete example, see:
-- **Configuration**: `examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml`
-- **Test Script**: `tests/functional/grpo_megatron_generation.sh`
+For a complete example, see
+`examples/configs/recipes/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.yaml`.
 
 ## Extend with New Backends
 

docs/guides/async-grpo.md

@@ -24,7 +24,7 @@ loss_fn:
   use_importance_sampling_correction: true
 ```
 
-3. **Disable colocated inference** (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see `examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml`):
+3. **Disable colocated inference** (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see `examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml`):
 ```yaml
 policy:
   generation:

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick.yaml

@@ -0,0 +1,9 @@
+defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml
+# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only
+# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B
+# megatron-inference stack without paying for the full nightly leg.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym-quick

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml

@@ -0,0 +1,36 @@
+defaults: ../../../nemo_gym/grpo_nanov3.yaml
+# Async colocated GRPO + gym on Megatron generation with a resharded inference
+# layout: training runs TP2xCP2xEP8 across all 16 GPUs; every wake reshards the
+# weights into TP4xEP4 inference_optimized models on the same GPUs. One nightly
+# covers the async loop with in-flight updates, the gym HTTP ingress, colocated
+# reshard, and (via the suite's save_period=val_period=4) the save-bound
+# deferred-wake path.
+grpo:
+  async_grpo:
+    enabled: true
+    in_flight_weight_updates: true
+  overlong_filtering: false
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym
+policy:
+  model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
+  tokenizer:
+    name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  train_global_batch_size: 16
+  max_total_sequence_length: 2048
+  megatron_cfg:
+    pipeline_model_parallel_size: 1
+    context_parallel_size: 2
+  generation:
+    backend: megatron
+    mcore_generation_config:
+      tensor_model_parallel_size: 4
+      expert_model_parallel_size: 4
+logger:
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym
+cluster:
+  num_nodes: 2

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick.yaml

@@ -0,0 +1,9 @@
+defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml
+# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only
+# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B
+# megatron-inference stack without paying for the full nightly leg.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym-quick

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym.yaml → examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml

@@ -5,7 +5,7 @@ grpo:
     in_flight_weight_updates: true
   overlong_filtering: false
 checkpointing:
-  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym
+  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym
 policy:
   model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
   tokenizer:
@@ -17,6 +17,8 @@ policy:
     context_parallel_size: 2
   generation:
     backend: megatron
+    mcore_generation_config:
+      expert_model_parallel_size: 4
     colocated:
       enabled: false
       resources:
@@ -27,6 +29,6 @@ logger:
   tensorboard_enabled: true
   wandb:
     project: nemo-rl
-    name: grpo-nanov3-30BA3B-2n8g-megatron_generation-async-gym
+    name: grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym
 cluster:
   num_nodes: 2

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick.yaml

@@ -0,0 +1,9 @@
+defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml
+# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only
+# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B
+# megatron-inference stack without paying for the full nightly leg.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym-quick

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym.yaml

@@ -0,0 +1,14 @@
+defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-colocated-reshard-async-gym.yaml
+# GB200 sizing: the same 16-GPU layouts on 4-GPU nodes, so the training EP8
+# groups span two nodes (cross-node expert parallelism, the production
+# Blackwell shape) and placement runs through the NVLink-domain topology path;
+# inference TP4 occupies exactly one node.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-4n4g-megatron_generation-colocated-reshard-async-gym
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick.yaml

@@ -0,0 +1,9 @@
+defaults: ./grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml
+# Reduced mcore-lane smoke leg: 3 steps, no checkpointing, gen_kl_error-only
+# gate (see the suite). Exists so Megatron-Core/Bridge bumps exercise the 30B
+# megatron-inference stack without paying for the full nightly leg.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym-quick

examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym.yaml

@@ -0,0 +1,19 @@
+defaults: ./grpo-nanov3-30BA3B-2n8g-megatron_generation-noncolocated-async-gym.yaml
+# GB200 sizing: the same 8+8-GPU non-colocated split on 4-GPU nodes, so the
+# training group (CP2xEP8) and the dedicated inference group each span two
+# nodes and placement runs through the NVLink-domain topology path.
+checkpointing:
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym
+policy:
+  generation:
+    colocated:
+      resources:
+        gpus_per_node: 4
+        num_nodes: 2
+logger:
+  wandb:
+    name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-async-gym
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2

examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_generation.yaml → examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts.yaml

@@ -1,16 +1,19 @@
 defaults: ../../grpo_math_1B.yaml
+# BF16 training with MXFP8-quantized Megatron inference on a dedicated 8-GPU group:
+# every weight sync quantizes BF16 -> MXFP8 over the NCCL refit,
+# with EP8-sharded expert transfers across NVLink domains.
 grpo:
   num_prompts_per_step: 2
   num_generations_per_prompt: 8
 checkpointing:
-  checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_generation
+  enabled: false
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts
 policy:
   model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
   tokenizer:
     name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
   train_global_batch_size: 16
   train_micro_batch_size: 1
-  logprob_batch_size: 1
   max_total_sequence_length: 2048
   dtensor_cfg:
     enabled: false
@@ -27,24 +30,29 @@ policy:
     backend: megatron
     mcore_generation_config:
       transformer_impl: inference_optimized
-      inference_grouped_gemm_backend: "vllm"
-      moe_router_num_groups: null
-      moe_router_group_topk: null
-      activation_checkpointing: false
+      inference_grouped_gemm_backend: "torch"
+      inference_moe_token_dispatcher_type: nvls
       tensor_model_parallel_size: 2
       expert_model_parallel_size: 8
       sequence_parallel: true
+      logprobs_mode: raw_logprobs
+      fp8_cfg:
+        enabled: true
+        fp8: e4m3
+        fp8_recipe: mxfp8
+        fp8_param: true
     colocated:
       enabled: false
       resources:
-        gpus_per_node: 8
-        num_nodes: 1
+        gpus_per_node: 4
+        num_nodes: 2
 logger:
   wandb_enabled: true
   tensorboard_enabled: true
   wandb:
     project: nemo-rl
-    name: grpo-nanov3-30BA3B-2n8g-megatron_generation
+    name: grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-mxfp8-rollouts
 cluster:
-  gpus_per_node: 8
-  num_nodes: 2
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2