refit latest update based on ahmads refactor

examples/configs/grpo_math_1B.yaml

@@ -146,6 +146,11 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "allgather"
     moe_shared_expert_overlap: false
+    moe_pad_experts_for_cuda_graph_inference: false
+    cuda_graph_impl: "local"
+    cuda_graph_scope: null
+    use_te_rng_tracker: true
+    inference_rng_tracker: true
 
     optimizer:
       optimizer: "adam"
@@ -252,13 +257,16 @@ policy:
     stop_strings: null
     mcore_generation_config:
       buffer_size_gb: 20  # Total GPU memory (in GB) allocated for KV cache buffers
-      buffer_guaranteed_fraction: 0.1  # Fraction of buffer reserved for guaranteed active requests
       num_cuda_graphs: 16  # Number of CUDA graphs to pre-compile for different batch sizes
       block_size_tokens: 256  # Size of each KV cache block in tokens (affects memory granularity)
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
-      enable_chunked_prefill: true  # Split long prefills into chunks for better memory management
-      unified_memory_level: 0  # Unified memory usage level (0=disabled, higher values enable more aggressive paging)
+      unified_memory_level: 0  # Unified memory usage level (0=disabled, 1+=enables unified memory )
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
+      enable_chunked_prefill: false
+      kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute"
+      static_kv_memory_pointers: false # Relevant only for offload and recompute modes
+      materialize_only_last_token_logits: false
+
     vllm_cfg:
       async_engine: false
       precision: ${policy.precision}

examples/configs/grpo_math_1B_megatron.yaml

@@ -100,6 +100,12 @@ policy:
     moe_shared_expert_overlap: false
     #gives ~20% training perf speedup with sequence packing 
     apply_rope_fusion: True
+    moe_pad_experts_for_cuda_graph_inference: false
+    cuda_graph_impl: "local"
+    cuda_graph_scope: "full_iteration_inference"
+    use_te_rng_tracker: true
+    inference_rng_tracker: true
+    batch_invariant_mode: false
 
     optimizer:
       optimizer: "adam"
@@ -125,6 +131,7 @@ policy:
       clip_grad: ${policy.max_grad_norm}
 
     scheduler:
+      override_opt_param_scheduler: true
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       weight_decay_incr_style: "constant"
@@ -151,9 +158,12 @@ policy:
       num_cuda_graphs: 16  # Number of CUDA graphs to pre-compile for different batch sizes
       block_size_tokens: 256  # Size of each KV cache block in tokens (affects memory granularity)
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
-      enable_chunked_prefill: false  # Split long prefills into chunks for better memory management
-      unified_memory_level: 0  # Unified memory usage level (0=disabled, higher values enable more aggressive paging)
+      unified_memory_level: 0  # Unified memory usage level (0=disabled, 1+=enables unified memory )
       max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
+      enable_chunked_prefill: false
+      kv_cache_management_mode: "persist" # Can be "persist", "offload", or "recompute"
+      static_kv_memory_pointers: false # Relevant only for offload and recompute modes
+      materialize_only_last_token_logits: false
 
     vllm_cfg:
       tensor_parallel_size: 1

examples/configs/grpo_math_8B_megatron.yaml

@@ -17,7 +17,7 @@ policy:
   train_global_batch_size: 512
   train_micro_batch_size: 1
   generation_batch_size: 32 # Only used when generating using HF backend
-  logprob_batch_size: 4
+  logprob_batch_size: ${policy.train_micro_batch_size}
   max_total_sequence_length: 4096
   precision: "bfloat16"
 
@@ -48,6 +48,7 @@ policy:
       params_dtype: "float32"
 
     scheduler:
+      override_opt_param_scheduler: true
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       weight_decay_incr_style: "constant"

nccl_timeout.md

@@ -0,0 +1,161 @@
+# NCCL Timeout During CUDA Graph Warmup in MoE RL Training
+
+## Symptom
+
+After several successful GRPO training steps (anywhere from step 5 to step 20+), the job crashes with NCCL collective operation timeouts during the generation phase. The errors look like:
+
+```
+(MegatronPolicyWorker[rank=1] pid=151407) [rank1]:[E228 00:10:27.202494521 ProcessGroupNCCL.cpp:2057] [PG ID 12 PG GUID 99(EXPERT_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1145527, OpType=ALLTOALL_BASE, NumelIn=206438400, NumelOut=206438400, Timeout(ms)=600000) ran for 600016 milliseconds before timing out.
+(MegatronPolicyWorker[rank=1] pid=151407)
+(MegatronPolicyWorker[rank=1] pid=151407) [2026-02-28 00:10:27,258 E 151407 153013] logging.cc:118: Unhandled exception: N3c1016DistBackendErrorE. what(): [PG ID 12 PG GUID 99(EXPERT_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=1145527, OpType=ALLTOALL_BASE, NumelIn=206438400, NumelOut=206438400, Timeout(ms)=600000) ran for 600016 milliseconds before timing out.
+(MegatronPolicyWorker[rank=1] pid=151407)
+(MegatronPolicyWorker[rank=1] pid=151407)
+(MegatronPolicyWorker[rank=7] pid=151429) [rank7]:[E228 00:10:27.168785312 ProcessGroupNCCL.cpp:2057] [PG ID 5 PG GUID 36(TENSOR_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=2036074, OpType=_REDUCE_SCATTER_BASE, NumelIn=3225600, NumelOut=1612800, Timeout(ms)=600000) ran for 600000 milliseconds before timing out.
+(MegatronPolicyWorker[rank=7] pid=151429)
+(MegatronPolicyWorker[rank=7] pid=151429) [2026-02-28 00:10:27,224 E 151429 152962] logging.cc:118: Unhandled exception: N3c1016DistBackendErrorE. what(): [PG ID 5 PG GUID 36(TENSOR_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=2036074, OpType=_REDUCE_SCATTER_BASE, NumelIn=3225600, NumelOut=1612800, Timeout(ms)=600000) ran for 600000 milliseconds before timing out.
+```
+
+Key signatures:
+- Different ranks report **different NCCL operation types** (ALLTOALL_BASE, REDUCE_SCATTER_BASE, ALLREDUCE, COALESCED) -- a collective mismatch
+- The crash always happens during `cuda graph warmup` at the start of generation
+- A new NCCL communicator is lazily initialized at the failing step (`NCCL version 2.27.5+cuda12.9` printed mid-warmup)
+- Steps 1 through N-1 complete normally; the crash is non-deterministic
+
+## Background: The Training-Inference Cycle
+
+In the RL training loop, each step does:
+
+```
+generate() {
+    _wake()          // resume inference engine (realloc KV cache, rebuild CUDA graphs)
+    <submit requests, run generation>
+    _sleep()         // suspend inference engine (dealloc KV cache, delete CUDA graphs)
+}
+```
+
+With `static_kv_memory_pointers=false` and `kv_cache_management_mode=recompute`, every suspend/resume cycle **destroys and recreates CUDA graphs**. Graph warmup runs forward passes through the model, which for MoE models includes NCCL alltoall collectives across expert-parallel (EP) ranks. All EP/TP ranks must execute the same sequence of NCCL operations in lockstep during this warmup.
+
+## Architecture: Two Communication Systems on One Event Loop
+
+The `DynamicInferenceEngine` runs an async engine loop on a dedicated event loop thread. This single event loop handles two different communication systems:
+
+| System | Purpose | Mechanism |
+|--------|---------|-----------|
+| **EP consensus** (`_ep_group_has_work`) | Coordinate EP ranks on work availability | Async ZMQ all-reduce |
+| **CUDA graph warmup** (inside `resume()`) | Capture model forward passes into graphs | Blocking NCCL collectives |
+
+Both run on the **same event loop thread**. This is the root of the problem.
+
+## Root Cause
+
+The engine loop has this structure (simplified from `run_engine_with_coordinator`):
+
+```python
+while True:
+    self.schedule_requests()                          # read ZMQ messages
+    ep_group_has_work = await self._ep_group_has_work(...)  # ZMQ all-reduce across EP ranks
+    if not ep_group_has_work:
+        if self.suspend_signal:
+            self.suspend()      # no-op when already suspended
+        else:
+            self.resume()       # CUDA graph warmup -- blocks with NCCL!
+        await asyncio.sleep(0.02)
+```
+
+When the coordinator sends `RESUME + UNPAUSE` to all engines, the signals arrive asynchronously. EP ranks process them at different times depending on ZMQ delivery and event loop scheduling. This leads to a **divergence**:
+
+```
+Rank A (received RESUME):  suspend_signal=False  -->  calls resume()  -->  NCCL alltoall BLOCKS event loop
+Rank B (not yet received): suspend_signal=True   -->  calls suspend() (no-op)  -->  sleeps 20ms
+```
+
+On the next iteration, Rank B calls `_ep_group_has_work()` which does an async ZMQ all-reduce. This requires Rank A to respond. But Rank A's event loop is **blocked inside NCCL** (graph warmup forward pass). Rank A can never respond to ZMQ while NCCL is blocking its event loop.
+
+**Deadlock: Rank A waits for Rank B in NCCL. Rank B waits for Rank A in ZMQ.**
+
+After 10 minutes, the NCCL watchdog times out and kills the process.
+
+### Why it's non-deterministic
+
+The deadlock only occurs when at least one EP rank enters `resume()` before all other EP ranks have received the `RESUME` signal. When all ranks happen to process the signals within the same ~20ms engine loop cycle, they all enter `resume()` together and the warmup succeeds. This timing depends on ZMQ delivery, event loop scheduling, and OS thread scheduling -- hence the non-determinism.
+
+
+### Implementation 1 (This causes delay of 25%)
+
+```python
+def _wake(self):
+    # Phase 1: Unpause the engine loop (async, event loop stays free for ZMQ)
+    asyncio.run_coroutine_threadsafe(self._unpause_engine(), self._inference_loop).result()
+
+    # Phase 2: Synchronized resume on the main thread
+    self._synchronized_resume()
+
+async def _unpause_engine(self):
+    # Send only UNPAUSE (not RESUME) -- keeps suspend_signal=True so the engine
+    # loop never calls resume() on its own
+    if torch.distributed.get_rank() == 0:
+        self.inference_client.unpause_engines()
+    await self.dynamic_inference_engine.running.wait()
+
+def _synchronized_resume(self):
+    engine = self.dynamic_inference_engine
+
+    # Guard: replace suspend() with a no-op while we resume
+    original_suspend = engine.suspend
+    engine.suspend = lambda: None
+
+    try:
+        torch.distributed.barrier()        # all ranks ready
+        engine.resume()                     # CUDA graph warmup (NCCL collectives)
+        engine.suspend_signal = False       # let engine loop transition to normal mode
+        torch.distributed.barrier()        # all ranks done
+    finally:
+        engine.suspend = original_suspend
+```
+
+### Why this works
+
+**No event-loop blocking.** The NCCL barriers and `resume()` run on the main thread. The event loop thread continues running the engine loop, freely handling ZMQ communication for EP consensus. No rank's ZMQ is ever starved.
+
+**No RESUME signal divergence.** We never send the `RESUME` header to the coordinator. Instead, we send only `UNPAUSE` (which restarts the engine loop) and keep `suspend_signal=True`. The engine loop sees `suspend_signal=True`, calls `suspend()` (no-op since already suspended), and idles. It never calls `resume()` on its own. We control exactly when `resume()` happens -- after the barrier on the main thread.
+
+**No thread-safety race.** When `resume()` runs on the main thread, it sets `is_suspended=False`. Without the guard, the engine loop's next `suspend()` call (on the event loop thread) would read `is_suspended=False`, enter the suspend body, and **deallocate buffers while the main thread is still creating CUDA graphs**. The `suspend()` guard (replacing it with `lambda: None`) prevents this. The guard is removed only after `suspend_signal=False` is set, so the engine loop transitions to calling `resume()` (which is a no-op since we already resumed) instead of `suspend()`.
+
+### Thread interaction timeline
+
+```
+Main Thread                              Event Loop Thread (engine loop)
+-----------                              --------------------------------
+_unpause_engine() ----sends UNPAUSE---->
+                                         schedule_requests(): reads UNPAUSE
+                                         _ep_group_has_work(): ZMQ all-reduce
+                                         suspend_signal=True -> suspend() [no-op]
+                                         asyncio.sleep(0.02)
+
+barrier() ............all ranks sync.... (ZMQ continues running freely)
+
+engine.suspend = no-op                   suspend() -> no-op [guarded]
+engine.resume()                          (ZMQ continues, no GPU conflict)
+  -> reinitialize buffers
+  -> create_cuda_graphs() [NCCL]
+engine.suspend_signal = False            (ZMQ continues)
+
+barrier() ............all ranks done....
+
+engine.suspend = original                suspend_signal=False -> resume() [no-op]
+                                         (engine is ready for requests)
+```
+
+## Affected Configuration
+
+This bug affects MoE models using the megatron generation backend with:
+- `moe_token_dispatcher_type=alltoall` (NCCL alltoall inside CUDA graph warmup)
+- `static_kv_memory_pointers=false` (CUDA graphs deleted/recreated each cycle)
+- `kv_cache_management_mode=recompute` (full dealloc on suspend)
+- `num_cuda_graphs > 0`
+- Expert parallelism (EP) > 1
+
+Dense models or configs with `static_kv_memory_pointers=true` are not affected because CUDA graphs are not recreated on resume.
+
+### Implementation 2 
+In dynamic_engine.py you set asyncio.sleep(0) instead of 0.02. This works 

nemo_rl/algorithms/grpo.py

@@ -60,6 +60,7 @@
 from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
 from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster
 from nemo_rl.environments.interfaces import EnvironmentInterface
+from nemo_rl.models.generation.megatron import MegatronGeneration
 from nemo_rl.experience.rollouts import (
     run_async_multi_turn_rollout,
     run_async_nemo_gym_rollout,
@@ -395,11 +396,6 @@ def setup(
         )
 
     else:
-        assert generation_config["backend"] != "megatron", (
-            "Non-colocated inference is not supported for Megatron generation backends. "
-            "Please use vLLM backend for generation."
-        )
-
         # train resources will be updated through overall and inference resources below
         train_gpus_per_node = cluster_config["gpus_per_node"]
         train_nodes = policy_nodes
@@ -539,6 +535,18 @@ def init_sglang():
         pg.finish_generation()
         return pg, time.perf_counter() - t0
 
+    def init_megatron_generation():
+        """Initialize Megatron generation workers for non-colocated inference."""
+        t0 = time.perf_counter()
+        mg = MegatronGeneration(
+            cluster=inference_cluster,
+            config=policy_config,
+            tokenizer=tokenizer,
+            processor=processor,
+            weights_path=weights_path,
+        )
+        return mg, time.perf_counter() - t0
+
     def initialize_generation_with_policy(
         init_generation_fn,
         generation_name: str,
@@ -603,14 +611,38 @@ def initialize_generation_with_policy(
     # Handle generation-specific setup
     if backend == "megatron":
         # Megatron generation: policy_generation is None, only initialize policy
-        policy_generation = None
-        print(
-            f"  ✓ Using {backend} backend for generation with {policy_config['model_name']}",
-            flush=True,
-        )
+        if colocated_inference:
+            policy_generation = None
+            print(
+                f"  ✓ Using {backend} backend for generation with {policy_config['model_name']}",
+                flush=True,
+            )
 
-        policy, policy_time = init_policy()
-        worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            policy, policy_time = init_policy()
+            worker_init_timing_metrics["policy_init_time_s"] = policy_time
+        else:
+            # Non-colocated Megatron backend: separate inference workers
+            print(
+                "  ⚡ Using parallel worker initialization (non-colocated Megatron mode)",
+                flush=True,
+            )
+
+            # Execute both initializations in parallel
+            parallel_start_time = time.perf_counter()
+            with ThreadPoolExecutor(max_workers=2) as executor:
+                megatron_gen_future = executor.submit(init_megatron_generation)
+                policy_future = executor.submit(init_policy)
+                policy_generation, megatron_gen_time = megatron_gen_future.result()
+                policy, policy_time = policy_future.result()
+            parallel_wall_time = time.perf_counter() - parallel_start_time
+
+            # Store timing metrics
+            worker_init_timing_metrics["megatron_generation_init_time_s"] = (
+                megatron_gen_time
+            )
+            worker_init_timing_metrics["policy_init_time_s"] = policy_time
+            worker_init_timing_metrics["parallel_wall_time_s"] = parallel_wall_time
+            worker_init_timing_metrics["parallel_init_enabled"] = True
 
     elif backend == "vllm":
         # vLLM generation: setup config, then initialize with policy

nemo_rl/models/generation/megatron/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2025, 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.
+
+from nemo_rl.models.generation.megatron.megatron_generation import (
+    MegatronGeneration,
+)
+
+__all__ = ["MegatronGeneration"]