diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM index 23dd639cf3..b17248aa37 160000 --- a/3rdparty/Megatron-LM-workspace/Megatron-LM +++ b/3rdparty/Megatron-LM-workspace/Megatron-LM @@ -1 +1 @@ -Subproject commit 23dd639cf3de30f3b9d8d0fae71ee31180be9ddd +Subproject commit b17248aa37b454dcecbf8650b4267bebeb03997c diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 740f9ad24b..816b392dd0 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/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} diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index b240c6519c..3886ec0f8c 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/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 diff --git a/examples/configs/grpo_math_8B_megatron.yaml b/examples/configs/grpo_math_8B_megatron.yaml index 977ab394b5..94f14aaf2d 100644 --- a/examples/configs/grpo_math_8B_megatron.yaml +++ b/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" diff --git a/migration_notes_delete.md b/migration_notes_delete.md new file mode 100644 index 0000000000..161fcb3593 --- /dev/null +++ b/migration_notes_delete.md @@ -0,0 +1,675 @@ +# NOTES +### MegatronTokenizer Issue. +``` + from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer as MegatronTokenizer +ModuleNotFoundError: No module named 'megatron.core.datasets.megatron_tokenizer' +``` +Fix : +Create this file in megatron/core/datasets/megatron_tokenizer.py +``` +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Backward-compatibility shim. The legacy tokenizer base class was moved out of +# megatron.core.datasets in newer Megatron-LM versions. Megatron-Bridge still +# imports from this path, so we keep a thin re-export here. + +import json +import logging +from abc import ABC, abstractmethod +from collections import OrderedDict +from typing import Any + +import numpy + +logger = logging.getLogger(__name__) + + +class MegatronLegacyTokenizer(ABC): + """Abstract class for tokenizer + + Absent a config or class-specific tracking of which objects are uniquely identifying, we must + include all key word arguments as unique identifiers + + Args: + tokenizer_paths (Tuple[str]): All tokenizer source paths or prefixes + + tokenizer_options (Dict[str, Any]): All tokenizer options + """ + + def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any): + logger.warning( + "You're using the legacy tokenizer system, which is deprecated " + "and will be removed in a future release. Please migrate to the new tokenizer system " + "(`megatron.core.tokenizers.MegatronTokenizer`)." + ) + self.unique_identifiers = OrderedDict() + self.unique_identifiers["class"] = type(self).__name__ + self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths) + for option in tokenizer_options: + self.unique_identifiers[option] = str(tokenizer_options[option]) + + self.unique_description = json.dumps(self.unique_identifiers, indent=4) + + super().__init__() + + @abstractmethod + def tokenize(self, text: str) -> numpy.ndarray: + pass + + def detokenize(self, ids: numpy.ndarray) -> str: + raise NotImplementedError("{} has no method 'detokenize'".format(type(self).__name__)) + + def offsets(self, ids: list[int], text: str) -> list[int]: + raise NotImplementedError("{} has no method 'offsets'".format(type(self).__name__)) + + @property + @abstractmethod + def vocab(self): + pass + + @property + @abstractmethod + def inv_vocab(self): + pass + + @property + @abstractmethod + def vocab_size(self): + pass + + @property + def cls(self): + raise NotImplementedError("{} has no attribute 'cls'".format(type(self).__name__)) + + @property + def sep(self): + raise NotImplementedError("{} has no attribute 'sep'".format(type(self).__name__)) + + @property + def pad(self): + raise NotImplementedError("{} has no attribute 'pad'".format(type(self).__name__)) + + @property + def eod(self): + raise NotImplementedError("{} has no attribute 'eod'".format(type(self).__name__)) + + @property + def bos(self): + raise NotImplementedError("{} has no attribute 'bos'".format(type(self).__name__)) + + @property + def eos(self): + raise NotImplementedError("{} has no attribute 'eos'".format(type(self).__name__)) + + @property + def mask(self): + raise NotImplementedError("{} has no attribute 'mask'".format(type(self).__name__)) + + +# Older code imported this class under the name ``MegatronTokenizer``. +MegatronTokenizer = MegatronLegacyTokenizer + +``` + + +### DDP ISSUE +``` +@@ -369,12 +369,15 @@ class DistributedDataParallel(_BaseDataParallel): + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook ++ for module, handle in list(self.remove_forward_pre_hook_handles.items()): ++ handle.remove() ++ self.remove_forward_pre_hook_handles.clear() + +- for module in self.module.modules(): +- assert self.remove_forward_pre_hook_handles[module] is not None +- self.remove_forward_pre_hook_handles[module].remove() +- del self.remove_forward_pre_hook_handles[module] +- assert len(self.remove_forward_pre_hook_handles) == 0 + + # Force synchronize parameters. + if param_sync: +``` + +### EXPLANATION +## Why `disable_forward_pre_hook` is Called + +### The Context: `use_reference_model` + +The `use_reference_model` context manager (lines 1437-1486) temporarily **swaps the model weights** with the reference model weights: + +1. **On entry**: Copies the current model's state_dict to CPU, then loads the reference model's state_dict into the model +2. **On exit**: Restores the original model weights + +This allows running inference with the reference model's weights without having two full models in GPU memory. + +### What is the Forward Pre-Hook? + +Looking at the DDP code you attached, when **overlap_param_gather** is enabled with distributed optimizer: + +```376:386:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def enable_forward_pre_hook(self): + """ + Enable forward pre-hooks needed for param all-gather overlap with forward compute. + """ + assert self.use_forward_hook + assert len(self.remove_forward_pre_hook_handles) == 0 + # Register forward pre-hook for all sub-modules. + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +The forward pre-hook is used to **overlap parameter all-gather with forward compute**. Here's how it works: + +1. With **distributed optimizer**, model parameters are **sharded across data-parallel ranks** (each rank only holds a portion of the parameters) +2. Before forward pass, parameters need to be **all-gathered** to reconstruct full parameters +3. The forward pre-hook intercepts each module's forward call to **wait for the all-gather to complete** for that module's parameters before executing + +```411:437:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def hook(module, *unused): + // ... + # Make sure all parameters in this module have been all-gathered as necessary. + for param in module.parameters(recurse=False): + # Skip parameters without an associated buffer + if param not in self.param_to_bucket_group: + continue + // ... + self.param_to_bucket_group[param].finish_param_sync( + skip_next_bucket_dispatch=skip_next_bucket_dispatch + ) +``` + +### Why Disable It During Weight Swap? + +When swapping weights in `use_reference_model`: + +```1459:1459:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/nemo_rl/models/policy/workers/megatron_policy_worker.py + self.model.load_state_dict(self.reference_state_dict, strict=True) +``` + +**The forward pre-hook would interfere because:** + +1. The hook maintains state about **which parameters have been all-gathered** via `param_to_bucket_group` +2. When you swap weights using `load_state_dict`, the parameter tensors are replaced, but the **bucket mappings become stale/invalid** +3. If the hook tries to call `finish_param_sync` on parameters that weren't properly set up for the new weights, it could cause errors or incorrect behavior +4. Additionally, `disable_forward_pre_hook` calls `start_param_sync(force_sync=True)` which **synchronously gathers all parameters** first, ensuring a clean state before the weight swap + +After the weight swap is complete (in `finally` block), the hooks are re-enabled: + +```1484:1486:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/nemo_rl/models/policy/workers/megatron_policy_worker.py + ## re-enable overlap param gather after weight swap + if self.should_disable_forward_pre_hook: + self.enable_forward_pre_hook() +``` + +### Summary + +| Step | What Happens | +|------|-------------| +| 1. `disable_forward_pre_hook()` | Removes all forward hooks, forces a synchronous param all-gather to ensure all params are complete | +| 2. `load_state_dict()` | Safely swap weights without hooks interfering | +| 3. Forward passes | Run without overlap optimization (simpler, safer) | +| 4. `enable_forward_pre_hook()` | Re-register hooks after restoring original weights | + +This is essentially a safety measure: disable the async overlap optimization during the weight-swapping operation to avoid race conditions or stale state in the DDP infrastructure. + +```python + def enable_forward_pre_hook(self): + """ + Enable forward pre-hooks needed for param all-gather overlap with forward compute. + """ + assert self.use_forward_hook + assert len(self.remove_forward_pre_hook_handles) == 0 + # Register forward pre-hook for all sub-modules. + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +```python + def hook(module, *unused): + // ... + # Make sure all parameters in this module have been all-gathered as necessary. + for param in module.parameters(recurse=False): + # Skip parameters without an associated buffer + if param not in self.param_to_bucket_group: + continue + // ... + self.param_to_bucket_group[param].finish_param_sync( + skip_next_bucket_dispatch=skip_next_bucket_dispatch + ) +``` + +```python + self.model.load_state_dict(self.reference_state_dict, strict=True) +``` + +```python + ## re-enable overlap param gather after weight swap + if self.should_disable_forward_pre_hook: + self.enable_forward_pre_hook() +``` + +--- + +**User** + +When I run the application I get an error like this +``` + File "/opt/nemo-rl/examples/run_grpo_math.py", line 260, in + main() + File "/opt/nemo-rl/examples/run_grpo_math.py", line 243, in main + grpo_train( + File "/opt/nemo-rl/nemo_rl/algorithms/grpo.py", line 1339, in grpo_train + reference_logprobs = policy.get_reference_policy_logprobs( + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 390, in get_reference_policy_logprobs + self.worker_group.get_all_worker_results(futures) + File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 953, in get_all_worker_results + return future_bundle.get_results( + ^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/distributed/worker_groups.py", line 103, in get_results + all_results = ray.get(object_refs) + ^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper + return fn(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper + return func(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, in get + values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 968, in get_objects + raise value.as_instanceof_cause() +ray.exceptions.RayTaskError(AssertionError): ray::MegatronPolicyWorker.get_reference_policy_logprobs() (pid=710294, ip=10.65.29.29, actor_id=65a88a2a253f3b406b563ea701000000, repr=MegatronPolicyWorker[rank=0]) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/utils/nsys.py", line 88, in wrapper + ret = func(*args, **kwargs) + ^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/base_policy_worker.py", line 143, in get_reference_policy_logprobs + with self.use_reference_model(): + ^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/root/.local/share/uv/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/contextlib.py", line 137, in __enter__ + return next(self.gen) + ^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 1446, in use_reference_model + self.disable_forward_pre_hook() + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 944, in disable_forward_pre_hook + self.model.disable_forward_pre_hook(param_sync=param_sync) + File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py", line 399, in disable_forward_pre_hook + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +AssertionError: remove_forward_pre_hook_handles is not empty, it is: {CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): , CudaGraphManager(): } +``` +What does this error mean ? + +--- + +**Cursor** + +```plaintext +AssertionError: remove_forward_pre_hook_handles is not empty, it is: {CudaGraphManager(): ...} +``` + +```python +assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) + ... +``` + +This error is revealing a **mismatch between hook registration and deregistration** due to `CudaGraphManager` modules. Let me explain: + +## What the Error Means + +The assertion failure happens in `disable_forward_pre_hook`: + +```388:399:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +### The Problem + +1. **When `enable_forward_pre_hook` was called**: It iterated through `self.module.modules()` and registered hooks for **ALL** modules, including `CudaGraphManager` objects + +2. **When `disable_forward_pre_hook` is called**: It again iterates through `self.module.modules()` to remove hooks + +3. **The mismatch**: The `CudaGraphManager` modules that were present during registration are **no longer being returned** by `self.module.modules()` during deregistration. This leaves their handles in the dictionary, causing the assertion to fail. + +### Why This Happens + +Looking at the leftover keys in the error: +``` +{CudaGraphManager(): , + CudaGraphManager(): , + ...} +``` + +These `CudaGraphManager` objects were registered as submodules when `enable_forward_pre_hook` ran, but they're either: +- Dynamically removed from the module tree after registration +- Not yielded by `modules()` during the second iteration for some reason +- Created/destroyed dynamically as part of CUDA graph management + +### Root Cause + +This is likely a **bug or incompatibility** between: +1. **CUDA Graphs** being enabled in your configuration (`cuda_graph_impl = "local"` as I saw earlier in the code) +2. **Overlap param gather** being enabled (`overlap_param_gather = True`) + +When CUDA graphs are used, the model may dynamically add/remove `CudaGraphManager` wrapper modules, which breaks the assumption that `modules()` returns the same set both times. + +### Potential Fixes + +**Option 1**: Disable CUDA graphs when using reference model weight swapping (if that's acceptable for your use case) + +**Option 2**: Fix the `disable_forward_pre_hook` to be more robust by clearing ALL handles, not just those from current `modules()`: + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove all registered handles + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` + +**Option 3**: Skip calling `disable_forward_pre_hook` when reference model doesn't need it (if `init_reference_model=False` was used or reference model was never initialized) + +Would you like me to help you implement one of these fixes? + +```python + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all sub-modules. + for module in self.module.modules(): + assert self.remove_forward_pre_hook_handles[module] is not None + self.remove_forward_pre_hook_handles[module].remove() + del self.remove_forward_pre_hook_handles[module] + assert len(self.remove_forward_pre_hook_handles) == 0, "remove_forward_pre_hook_handles is not empty, it is: " + str(self.remove_forward_pre_hook_handles) +``` + +```plaintext +{CudaGraphManager(): , + CudaGraphManager(): , + ...} +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove all registered handles + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` + +```python +def toggle_cuda_graphs(model, set_to="none", reset_cuda_graphs=True): + ... + elif attribute == "cudagraph_manager": + for module in modules: + if set_to == "local": + if reset_cuda_graphs: + from megatron.core.transformer.cuda_graphs import CudaGraphManager + # If we are resetting cuda graphs we create a new cuda graph manager + setattr(module[0], attribute, CudaGraphManager(model.config)) + else: + # If we are not resetting cuda graphs we set it to its cached cuda graph + setattr(module[0], attribute, module[1]) + else: + for module in modules: + # If we are deleting the cuda graph, we delete its attribute + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") + ... +``` + +```python +setattr(module[0], attribute, CudaGraphManager(model.config)) +``` + +```python +delattr(module[0], "cudagraph_manager") +``` + +```python +toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) + + +``` + +```plaintext +Two potential solutions emerge: +1. Disable hooks before structural module changes +2. Modify hook removal logic to track references independently of module hierarchy +``` + +Let me check where `toggle_cuda_graphs` is being called in relation to the hook operations: + +**Yes, this is exactly the cause!** Let me explain the problem: + +## Root Cause Analysis + +### The Sequence of Events + +1. **During `generate()`** (line 1875): + ```python + toggle_cuda_graphs(lang_module, set_to="local", reset_cuda_graphs=True) + ``` + This **creates new `CudaGraphManager` objects** and adds them as submodules: + ```python + # From toggle_cuda_graphs when set_to="local": + setattr(module[0], attribute, CudaGraphManager(model.config)) + ``` + +2. **At some point**, `enable_forward_pre_hook()` is called, which registers hooks for **ALL** current modules, including the newly created `CudaGraphManager` objects: + + ```376:386:/lustre/fsw/portfolios/coreai/users/shanmugamr/RL/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py + def enable_forward_pre_hook(self): + # ... + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) + ``` + +3. **At end of `generate()`** (line 2045): + ```python + toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) + ``` + This **DELETES the `CudaGraphManager` attributes** from their parent modules: + ```python + # From toggle_cuda_graphs when set_to="none": + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") # ← REMOVES from module tree! + ``` + +4. **When `get_reference_policy_logprobs` → `use_reference_model` → `disable_forward_pre_hook()` is called**: + It iterates through `self.module.modules()`, but the `CudaGraphManager` objects are **no longer in the module tree** (they were deleted in step 3). + +5. **Result**: The hook handles for `CudaGraphManager` objects remain in `remove_forward_pre_hook_handles` but can't be removed → **Assertion fails!** + +## Visual Timeline + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ generate() starts │ +│ └─ toggle_cuda_graphs(set_to="local") │ +│ └─ Creates CudaGraphManager as submodule │ +├─────────────────────────────────────────────────────────────────────┤ +│ enable_forward_pre_hook() called somewhere │ +│ └─ Registers hooks for ALL modules (including CudaGraphManager) │ +│ └─ remove_forward_pre_hook_handles[CudaGraphManager] = handle │ +├─────────────────────────────────────────────────────────────────────┤ +│ generate() ends │ +│ └─ toggle_cuda_graphs(set_to="none") │ +│ └─ delattr(module, "cudagraph_manager") ← DELETES IT! │ +├─────────────────────────────────────────────────────────────────────┤ +│ get_reference_policy_logprobs() │ +│ └─ use_reference_model() │ +│ └─ disable_forward_pre_hook() │ +│ └─ for module in self.module.modules(): ← CudaGraphManager│ +│ # is NOT here anymore! │ +│ └─ ASSERTION FAILS: handles dict not empty! │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +## Solution Options + +**Option 1**: Fix `disable_forward_pre_hook` in DDP to be robust (best fix): + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove ALL registered handles (not just current modules) + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + if param_sync: + self.start_param_sync(force_sync=True) +``` + +**Option 2**: Call `disable_forward_pre_hook` BEFORE `toggle_cuda_graphs(set_to="none")` in your code. + +**Option 3**: Don't delete cuda graph managers, just disable them (modify `toggle_cuda_graphs`). + +Would you like me to implement Option 1 as a fix in the Megatron DDP code? + +```python + toggle_cuda_graphs(lang_module, set_to="local", reset_cuda_graphs=True) +``` + +```python + # From toggle_cuda_graphs when set_to="local": + setattr(module[0], attribute, CudaGraphManager(model.config)) +``` + +```python + def enable_forward_pre_hook(self): + # ... + for module in self.module.modules(): + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + self._make_forward_pre_hook() + ) +``` + +```python + toggle_cuda_graphs(lang_module, set_to="none", reset_cuda_graphs=True) +``` + +```python + # From toggle_cuda_graphs when set_to="none": + if hasattr(module[0], "cudagraph_manager"): + delattr(module[0], "cudagraph_manager") # ← REMOVES from module tree! +``` + +```plaintext +┌─────────────────────────────────────────────────────────────────────┐ +│ generate() starts │ +│ └─ toggle_cuda_graphs(set_to="local") │ +│ └─ Creates CudaGraphManager as submodule │ +├─────────────────────────────────────────────────────────────────────┤ +│ enable_forward_pre_hook() called somewhere │ +│ └─ Registers hooks for ALL modules (including CudaGraphManager) │ +│ └─ remove_forward_pre_hook_handles[CudaGraphManager] = handle │ +├─────────────────────────────────────────────────────────────────────┤ +│ generate() ends │ +│ └─ toggle_cuda_graphs(set_to="none") │ +│ └─ delattr(module, "cudagraph_manager") ← DELETES IT! │ +├─────────────────────────────────────────────────────────────────────┤ +│ get_reference_policy_logprobs() │ +│ └─ use_reference_model() │ +│ └─ disable_forward_pre_hook() │ +│ └─ for module in self.module.modules(): ← CudaGraphManager│ +│ # is NOT here anymore! │ +│ └─ ASSERTION FAILS: handles dict not empty! │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +```python +def disable_forward_pre_hook(self, param_sync: bool = True): + assert self.use_forward_hook + # Remove ALL registered handles (not just current modules) + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + if param_sync: + self.start_param_sync(force_sync=True) +``` + +AFTER FIX 1 +### What Changed + +| Before | After | +|--------|-------| +| Iterated over `self.module.modules()` | Iterates over `self.remove_forward_pre_hook_handles.items()` directly | +| Asserted each module exists in dict | No assertion needed - we're iterating over what we have | +| Deleted entries one by one | Uses `clear()` after removing all handles | +| Would fail if modules were removed from tree | Handles dynamically removed modules (like `CudaGraphManager`) | + +This fix ensures that all registered hook handles are properly removed regardless of whether the modules are still part of the module tree. You can now run your application again to test it. + +```python + def disable_forward_pre_hook(self, param_sync: bool = True): + """ + Disable forward pre-hooks needed for param all-gather overlap with forward compute. + Skip synchronous param all-gather if `param_sync` is False. + """ + assert self.use_forward_hook + # De-register forward pre-hook for all registered modules. + # We iterate over the dict keys directly (instead of self.module.modules()) to handle + # cases where modules may have been dynamically removed from the module tree + # (e.g., CudaGraphManager objects removed by toggle_cuda_graphs). + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + # Force synchronize parameters. + if param_sync: + self.start_param_sync(force_sync=True) +``` \ No newline at end of file diff --git a/nccl_timeout.md b/nccl_timeout.md new file mode 100644 index 0000000000..86f5250a80 --- /dev/null +++ b/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) + + _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 diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 6772739655..5398af6cdc 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/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 diff --git a/nemo_rl/models/generation/megatron/__init__.py b/nemo_rl/models/generation/megatron/__init__.py new file mode 100644 index 0000000000..cf9d2aa8e3 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py new file mode 100644 index 0000000000..24aab7de7d --- /dev/null +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -0,0 +1,205 @@ +# 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. + +"""MegatronGeneration: A GenerationInterface implementation for non-colocated +Megatron-based inference. + +This module wraps a Policy object (configured for inference only, without +optimizer or reference model) and exposes it through the GenerationInterface. +It enables non-colocated inference where training and generation run on +separate GPU clusters, with weights synchronized via NCCL collective +communication. + +The init_collective and update_weights_from_collective methods are currently +placeholders that will be implemented in a future PR. +""" + +from typing import Any, Optional + +import ray +from transformers import AutoProcessor +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.models.generation.interfaces import ( + GenerationDatumSpec, + GenerationInterface, + GenerationOutputSpec, +) +from nemo_rl.models.policy import PolicyConfig + + +class MegatronGeneration(GenerationInterface): + """Generation interface backed by Megatron for non-colocated inference. + + This class creates a Policy instance configured for inference only + (no optimizer, no reference model) on a dedicated inference cluster. + It implements the GenerationInterface so it can be used as a drop-in + replacement for VllmGeneration in the non-colocated inference flow. + + The weight synchronization methods (init_collective, update_weights_from_collective) + are placeholders that will be implemented in a future PR. + """ + + def __init__( + self, + cluster: RayVirtualCluster, + config: PolicyConfig, + tokenizer: PreTrainedTokenizerBase, + name_prefix: str = "megatron_generation", + processor: Optional[AutoProcessor] = None, + weights_path: Optional[str] = None, + ): + """Initialize a MegatronGeneration instance. + + Args: + cluster: The RayVirtualCluster to deploy inference workers on. + config: PolicyConfig for the Megatron model. + tokenizer: The tokenizer for the model. + name_prefix: Prefix for naming the worker group. + processor: Optional processor for VLMs. + weights_path: Optional path to model weights for initialization. + """ + # Import here to avoid circular imports + from nemo_rl.models.policy.lm_policy import Policy + + self.cfg = config + + # Create a Policy object configured for inference only: + # - No optimizer (not training on this cluster) + # - No reference model (not needed for generation) + self._policy = Policy( + cluster=cluster, + config=config, + tokenizer=tokenizer, + name_prefix=name_prefix, + processor=processor, + init_optimizer=False, + init_reference_model=False, + weights_path=weights_path, + ) + + def init_collective( + self, ip: str, port: int, world_size: int, *, train_world_size: int + ) -> list[ray.ObjectRef]: + """Initialize the collective communication for weight synchronization. + + This sets up NCCL communication between training workers and these + inference workers so that updated model weights can be broadcast + from the training cluster to the inference cluster. + + Uses init_collective_as_inference on the workers, which offsets each + worker's rank by train_world_size to avoid colliding with training + workers' ranks (rank = train_world_size + worker_rank). + + Args: + ip: IP address for the process group rendezvous. + port: Port for the process group rendezvous. + world_size: Total world size (train + inference workers). + train_world_size: Number of training workers (used to offset ranks). + + Returns: + List of Ray ObjectRefs for the collective init futures. + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "init_collective_as_inference", + ip=ip, + port=port, + world_size=world_size, + train_world_size=train_world_size, + ) + return futures + + def update_weights_from_collective(self) -> list[ray.ObjectRef]: + """Receive updated weights from the training cluster via collective communication. + + This method is called after the training side calls + policy.broadcast_weights_for_collective(). It receives the broadcast + weights and updates the local model parameters. + + TODO: This is a placeholder. The actual implementation will: + 1. Iterate over the model's state_dict info + 2. Use packed_broadcast_consumer to receive weights from the training side + 3. Update the local model parameters with the received weights + + Returns: + List of Ray ObjectRefs for the weight update futures. + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "update_weights_from_collective", + ) + return futures + + def generate( + self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False + ) -> BatchedDataDict[GenerationOutputSpec]: + """Generate a batch of data using the Megatron generation backend. + + Delegates to the internal Policy's generate method. + + Args: + data: BatchedDataDict containing input_ids and input_lengths. + greedy: Whether to use greedy decoding. + + Returns: + BatchedDataDict conforming to GenerationOutputSpec. + """ + return self._policy.generate(data, greedy=greedy) + + def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: + """Prepare the inference workers for generation. + + For Megatron generation, this is a no-op since the workers + are always ready for inference. + """ + return self._policy.prepare_for_generation(*args, **kwargs) + + def finish_generation(self, *args: Any, **kwargs: Any) -> bool: + """Clean up after generation. + + For Megatron generation, this is a no-op. + """ + return self._policy.finish_generation(*args, **kwargs) + + def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: + """Prepare state dict metadata for weight refitting. + + This stores the state dict info (tensor names, shapes, dtypes) on each + inference worker so that update_weights_from_collective knows what + tensors to expect during the weight broadcast. + + Note: This calls store_refit_info on workers (not prepare_refit_info), + because prepare_refit_info on MegatronPolicyWorker calculates and + returns metadata (training-side), while store_refit_info accepts and + stores metadata (inference-side). + + Args: + state_dict_info: Dictionary mapping tensor names to (shape, dtype) tuples, + as returned by the training-side prepare_refit_info(). + """ + futures = self._policy.worker_group.run_all_workers_single_data( + "store_refit_info", + state_dict_info=state_dict_info, + ) + ray.get(futures) + + def shutdown(self) -> bool: + """Shut down all inference workers and clean up resources.""" + return self._policy.shutdown() + + def __del__(self) -> None: + """Safety net to ensure workers are shut down.""" + if hasattr(self, "_policy"): + self._policy.shutdown() diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index e9fc2da9e1..ce03878307 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -318,6 +318,10 @@ def setup_model_config( # Apply performance settings _apply_performance_config(model_cfg, config) + # Apply generation settings + if config["generation"][f"backend"] == "megatron": + _apply_cuda_graph_and_rng_tracker_config(model_cfg, config) + # Validate optimizer configuration _validate_optimizer_config(config) @@ -367,6 +371,18 @@ def _apply_parallelism_config(model_cfg: Any, config: PolicyConfig) -> None: ) +def _apply_cuda_graph_and_rng_tracker_config(model_cfg: Any, config: PolicyConfig) -> None: + """Apply CUDA GRAPH and RNG TRACKER configuration.""" + model_cfg.cuda_graph_impl = config["megatron_cfg"]["cuda_graph_impl"] + model_cfg.cuda_graph_scope = config["megatron_cfg"]["cuda_graph_scope"] + model_cfg.use_te_rng_tracker = config["megatron_cfg"]["use_te_rng_tracker"] + model_cfg.inference_rng_tracker = config["megatron_cfg"]["inference_rng_tracker"] + model_cfg.batch_invariant_mode = config["megatron_cfg"]["batch_invariant_mode"] + if model_cfg.batch_invariant_mode: + from megatron.core.transformer.enums import AttnBackend + model_cfg.attention_backend = AttnBackend.flash + + def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: """Apply Mixture of Experts configuration.""" model_cfg.expert_tensor_parallel_size = config["megatron_cfg"][ @@ -399,6 +415,9 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.moe_token_dispatcher_type = config["megatron_cfg"][ "moe_token_dispatcher_type" ] + model_cfg.moe_pad_experts_for_cuda_graph_inference = config["megatron_cfg"][ + "moe_pad_experts_for_cuda_graph_inference" + ] model_cfg.moe_shared_expert_overlap = config["megatron_cfg"][ "moe_shared_expert_overlap" ] @@ -850,6 +869,10 @@ def setup_reference_model_state( ref_ckpt_context = init_checkpointing_context(ref_checkpoint_config) + megatron_cfg.model.cuda_graph_impl = "none" + megatron_cfg.model.use_te_rng_tracker = False + megatron_cfg.model.inference_rng_tracker = False + # Create a separate megatron config for the reference model ref_megatron_cfg = ConfigContainer( model=megatron_cfg.model, diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 29f034b065..c3787692c1 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -620,7 +620,13 @@ def train( def generate( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using the policy.""" + """Generate a batch of data using the policy. + + For coordinator-based inference (Megatron backend), all data is sent to DP rank 0 + only, which submits requests to the coordinator. The coordinator then distributes + work across all DP engines. Other DP ranks participate in the engine loop but + don't receive input data directly. + """ # Verify input data is right-padded assert isinstance(data, BatchedDataDict), ( f"data must be a BatchedDataDict, got type: {type(data)}" @@ -629,14 +635,38 @@ def generate( "Missing required input fields" ) - dp_size = self.sharding_annotations.get_axis_size("data_parallel") - sharded_data = data.shard_by_batch_size(dp_size, batch_size=None) + if self.cfg["generation"]['backend'] == "vllm": + dp_size = self.sharding_annotations.get_axis_size("data_parallel") + data = data.shard_by_batch_size(dp_size, batch_size=None) + in_sharded_axes = ["data_parallel"] + output_is_replicated = [ + "tensor_parallel", + "pipeline_parallel", + ] + elif self.cfg["generation"]['backend'] == "megatron": + # For coordinator-based inference: send ALL data to DP rank 0 only. + # Other DP ranks are called with data=None but still participate in the + # inference engine loop. The coordinator handles load balancing across DP ranks. + # + # With in_sharded_axes=[] and data_parallel not in replicate_on_axes, + # data_parallel becomes a "free axis". Only workers at DP coord 0 receive data, + # while workers at other DP coords get None (via make_dummy_calls_to_free_axes). + in_sharded_axes = [] + output_is_replicated = [ + "data_parallel", + "tensor_parallel", + "pipeline_parallel", + ] + else: + raise ValueError(f"Invalid generation backend: {self.cfg['generation']['backend']}, expected 'vllm' or 'megatron'") + futures = self.worker_group.run_all_workers_sharded_data( "generate", - data=sharded_data, - in_sharded_axes=["data_parallel"], + data=data, # Full data goes to DP=0 only (free axis behavior) + in_sharded_axes=in_sharded_axes, replicate_on_axes=["tensor_parallel", "pipeline_parallel"], - output_is_replicated=["tensor_parallel", "pipeline_parallel"], + output_is_replicated=output_is_replicated, + make_dummy_calls_to_free_axes=True, # Call all DP ranks, but only DP=0 gets data common_kwargs={"greedy": greedy}, ) assert self.cfg["generation"] is not None, "Generation config is not set" diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 5a6a683765..cc232aa9dd 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -11,9 +11,11 @@ # 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. +import asyncio import gc import os import re +import time import warnings from collections import defaultdict from contextlib import AbstractContextManager, contextmanager, nullcontext @@ -31,15 +33,13 @@ reduce_max_stat_across_model_parallel_group, ) from megatron.bridge.utils.common_utils import get_rank_safe +from megatron.core.transformer.utils import toggle_cuda_graphs from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( FullyShardedDataParallel as custom_FSDP, ) -from megatron.core.inference.config import InferenceConfig -from megatron.core.inference.text_generation_controllers.text_generation_controller import ( - TextGenerationController, -) +from megatron.core.inference.config import InferenceConfig, KVCacheManagementMode from megatron.core.optimizer import ChainedOptimizer from megatron.core.parallel_state import ( get_pipeline_model_parallel_group, @@ -230,6 +230,15 @@ def __init__( ## used for streaming update inference engine weights self._held_gather_buffer = None + + self.dynamic_inference_engine = None + self.inference_client = None + self.inference_context = None + self.inference_wrapped_model = None + self._inference_engine_initialized = False + self._inference_engine_alseep = True # Start paused since we begin with training + self._inference_loop = None # Event loop for inference operations + self._inference_thread = None # Thread running the event loop def enable_forward_pre_hook(self): assert isinstance(self.model, DistributedDataParallel) @@ -237,7 +246,14 @@ def enable_forward_pre_hook(self): def disable_forward_pre_hook(self, param_sync=True): assert isinstance(self.model, DistributedDataParallel) - self.model.disable_forward_pre_hook(param_sync=param_sync) + for module, handle in list(self.model.remove_forward_pre_hook_handles.items()): + handle.remove() + self.model.remove_forward_pre_hook_handles.clear() + if param_sync: + self.model.start_param_sync(force_sync=True) + + # TODO : Check why this doesnt work. + #self.model.disable_forward_pre_hook(param_sync=param_sync) @wrap_with_nvtx_name("megatron_policy_worker/train") def train( @@ -662,47 +678,23 @@ def get_topk_logits( [{"topk_logits": topk_logits.cpu(), "topk_indices": topk_indices.cpu()}] ) - @wrap_with_nvtx_name("megatron_policy_worker/generate") - def generate( - self, *, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False - ) -> BatchedDataDict[GenerationOutputSpec]: - """Generate a batch of data using huggingface framework generation. - - Args: - data: BatchedDataDict containing input_ids and input_lengths tensors - Returns: - BatchedDataDict conforming to GenerationOutputSpec: - - output_ids: input + generated token IDs - - logprobs: Log probabilities for each token - - generation_lengths: Lengths of each response - """ - # 512 bATCH SIZE (200 tokens) - no_grad = torch.no_grad() - no_grad.__enter__() - self.model.config.flash_decode = False - if self.should_disable_forward_pre_hook: - self.model = self.move_model( - self.model, "cuda", move_params=True, move_grads=False - ) - # Verify input is right padded - assert isinstance(data, BatchedDataDict), ( - f"data must be a BatchedDataDict, got type: {type(data)}" - ) - assert "input_ids" in data and "input_lengths" in data, ( - f"input_ids and input_lengths must be present in the BatchedDataDict, got keys: {data.keys()}" + def _get_lang_module(self): + """Get the underlying language module from the wrapped model.""" + return ( + self.model.module.module + if hasattr(self.model.module, "module") + else self.model.module ) - is_right_padded, error_msg = verify_right_padding( - data, pad_value=self.tokenizer.pad_token_id - ) - if not is_right_padded: - warnings.warn( - f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" - ) - model_cfg = self.megatron_cfg.model - mcore_generation_config = cast( - MegatronGenerationConfig, self.cfg["generation"]["mcore_generation_config"] - ) + def _initialize_inference_engine(self, mcore_generation_config: dict): + """Initialize the persistent inference engine and client. + + This method sets up the DynamicInferenceEngine, DynamicInferenceContext, + and InferenceClient for coordinator-based inference. The engine is created + once and reused across multiple generate() calls. + """ + if self._inference_engine_initialized: + return from megatron.core.inference.contexts.dynamic_context import ( DynamicInferenceContext, @@ -713,114 +705,342 @@ def generate( from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( GPTInferenceWrapper, ) - from megatron.core.inference.sampling_params import SamplingParams + from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, + ) - model_config = self.model.config - model_config.cuda_graph_impl = "local" + model_cfg = self.megatron_cfg.model + + + from megatron.core.utils import get_attr_wrapped_model + pg_collection = get_attr_wrapped_model(self.model, "pg_collection") + + buffer_size_gb = mcore_generation_config["buffer_size_gb"] + num_cuda_graphs = mcore_generation_config["num_cuda_graphs"] + block_size_tokens = mcore_generation_config["block_size_tokens"] + enable_chunked_prefill = mcore_generation_config["enable_chunked_prefill"] + use_cuda_graphs_for_non_decode_steps = mcore_generation_config[ + "use_cuda_graphs_for_non_decode_steps" + ] + max_tokens = mcore_generation_config["max_tokens"] + + # Level 0: No unified memory, CUDA graphs are deleted/recreated on pause/resume + # Level 1: Unified memory enabled, tensors maintain static addresses + unified_memory_level = mcore_generation_config["unified_memory_level"] + kv_cache_management_mode = mcore_generation_config["kv_cache_management_mode"] + static_kv_memory_pointers = mcore_generation_config["static_kv_memory_pointers"] + materialize_only_last_token_logits = mcore_generation_config["materialize_only_last_token_logits"] - local_rank = torch.cuda.current_device() - num_gpus_per_node = torch.cuda.device_count() - node_idx = self.rank // num_gpus_per_node if num_gpus_per_node > 0 else 0 - model_config.inference_sampling_seed = (node_idx * 1024) + local_rank + model_config = self.model.config inference_config = InferenceConfig( + block_size_tokens=block_size_tokens, + buffer_size_gb=buffer_size_gb, + num_cuda_graphs=num_cuda_graphs, + max_tokens=max_tokens, max_sequence_length=self.cfg["generation"]["max_new_tokens"], - buffer_size_gb=mcore_generation_config["buffer_size_gb"], - num_cuda_graphs=mcore_generation_config["num_cuda_graphs"], - block_size_tokens=mcore_generation_config["block_size_tokens"], - use_cuda_graphs_for_non_decode_steps=mcore_generation_config[ - "use_cuda_graphs_for_non_decode_steps" - ], - enable_chunked_prefill=mcore_generation_config["enable_chunked_prefill"], - unified_memory_level=mcore_generation_config["unified_memory_level"], - max_tokens=mcore_generation_config["max_tokens"], - materialize_only_last_token_logits=False, - use_flashinfer_fused_rope=False, - ) - - dynamic_context = DynamicInferenceContext(model_config, inference_config) - inference_wrapped_model = GPTInferenceWrapper(self.model, dynamic_context) - - inference_wrapped_model.prep_model_for_inference() - # Set pipeline parallel flag - inference_wrapped_model.model_is_pipeline_parallel = ( - self.cfg["megatron_cfg"]["pipeline_model_parallel_size"] > 1 + unified_memory_level=unified_memory_level, + kv_cache_management_mode=KVCacheManagementMode(kv_cache_management_mode), + static_kv_memory_pointers=static_kv_memory_pointers, + use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, + materialize_only_last_token_logits=materialize_only_last_token_logits, + enable_chunked_prefill=enable_chunked_prefill, + pg_collection=pg_collection, + ) + + # Create inference context + self.inference_context = DynamicInferenceContext(model_config, inference_config) + + # Create inference wrapper + self.inference_wrapped_model = GPTInferenceWrapper( + self.model, self.inference_context ) - + # Create text generation controller text_generation_controller = TextGenerationController( - inference_wrapped_model=inference_wrapped_model, + inference_wrapped_model=self.inference_wrapped_model, tokenizer=self.megatron_tokenizer, ) - dynamic_engine = DynamicInferenceEngine( + # Create the inference engine + self.dynamic_inference_engine = DynamicInferenceEngine( text_generation_controller, - dynamic_context, + self.inference_context ) - # Handle None values for top_k - convert to integer as required by Megatron - top_k_cfg = self.cfg["generation"]["top_k"] - top_k_val = 1 if greedy else (int(top_k_cfg) if top_k_cfg is not None else 0) + self._inference_engine_initialized = True + self._inference_engine_alseep = True # Engine starts in paused state + print(f"[Rank {self.rank}] Initialized persistent inference engine") - top_p_cfg = self.cfg["generation"]["top_p"] - top_p_val = ( - 0.0 if greedy else (float(top_p_cfg) if top_p_cfg is not None else 0.0) + async def _start_inference_coordinator(self, coordinator_port: int): + """Start the inference coordinator and engine loop. + + This is called once when the inference infrastructure is first needed. + The engine's start_listening_to_data_parallel_coordinator returns the + actual coordinator address (dp_addr) which is used to create the client. + """ + dp_addr = await self.dynamic_inference_engine.start_listening_to_data_parallel_coordinator( + inference_coordinator_port=coordinator_port, + launch_inference_coordinator=True, ) - # New API: SamplingParams now includes termination_id and uses num_tokens_total - sampling_params = SamplingParams( - temperature=self.cfg["generation"]["temperature"] if not greedy else 0, - top_k=top_k_val, - top_p=top_p_val, - skip_prompt_log_probs=False, - return_log_probs=True, - num_tokens_total=self.cfg["generation"]["max_new_tokens"], - num_tokens_to_generate=None, - termination_id=self.megatron_tokenizer.eod, + dist_rank = torch.distributed.get_rank() + if dist_rank == 0: + from megatron.core.inference.inference_client import InferenceClient + self.inference_client = InferenceClient(inference_coordinator_address=dp_addr) + await self.inference_client.start() + + self._inference_engine_alseep = False + + def _sleep(self): + """pause the inference engine to free GPU memory for training. + + This method should be called before training to: + 1. Deallocate KV cache and other inference-specific GPU memory + 2. Disable CUDA graphs for inference + 3. Toggle model configuration for training mode + + Uses the coordinator's pause mechanism to properly pause the engine loop + and then pause the engine (deallocate tensors, etc.). + + For coordinator-based inference: + - Only rank 0 sends pause signals via the coordinator + - The coordinator broadcasts to all DP engines + - Non-rank-0 workers wait for their engine to be paused via the event loop + """ + + future = asyncio.run_coroutine_threadsafe( + self._sleep_engine(), + self._inference_loop ) + future.result() + # Synchronize all ranks + torch.distributed.barrier() + + self._inference_engine_alseep = True + print(f"[Rank {self.rank}] paused inference engine") + + async def _sleep_engine(self): + """Send suspend signals via the coordinator and wait for acknowledgment. + + Mirrors MegatronLocal.suspend() from megatron/rl/inference/megatron.py: + 1. Rank 0 sends suspend (PAUSE + SUSPEND) to coordinator + 2. All ranks wait for engine to be paused + 3. All ranks call engine.suspend() to deallocate GPU state + """ + if torch.distributed.get_rank() == 0: + # Send PAUSE signals + self.inference_client.suspend_engines() + # Wait for the engine to acknowledge the pause + await self.dynamic_inference_engine.paused.wait() + self.dynamic_inference_engine.suspend() + + def _wake(self): + """Resume the inference engine after training. + + This method should be called before generation to: + 1. Reallocate KV cache and inference-specific GPU memory + 2. Enable CUDA graphs for inference + 3. Toggle model configuration for inference mode + + Uses the coordinator's resume mechanism to properly resume the engine loop. + + For coordinator-based inference: + - Only rank 0 sends resume signals via the coordinator + - The coordinator broadcasts to all DP engines + - Non-rank-0 workers wait for their engine to be running via the event loop + """ + + # Use the coordinator-based resume mechanism + # Only rank 0 sends the signal - coordinator broadcasts to all DP engines + future = asyncio.run_coroutine_threadsafe( + self._wake_engine(), + self._inference_loop + ) + future.result() + # Synchronize all ranks + torch.distributed.barrier() + + self._inference_engine_alseep = False + print(f"[Rank {self.rank}] Resumed inference engine") + + async def _wake_engine(self): + """Send resume signals via the coordinator and wait for acknowledgment. + + Mirrors MegatronLocal.resume() from megatron/rl/inference/megatron.py: + 1. Rank 0 sends resume (RESUME + UNPAUSE) to coordinator + 2. All ranks wait for engine to be running + 3. All ranks call engine.resume() to reallocate GPU state + """ + if torch.distributed.get_rank() == 0: + self.inference_client.resume_engines() + await self.dynamic_inference_engine.running.wait() + self.dynamic_inference_engine.resume() - input_ids = data["input_ids"] - prompt_tokens_tensor = input_ids.cuda() - prompt_lengths_tensor = data["input_lengths"] - request_id = 0 - # New API: add_request now takes sampling_params as a parameter - for p, prompt_len in zip( - prompt_tokens_tensor, prompt_lengths_tensor, strict=True - ): - dynamic_engine.add_request( - request_id, - p[:prompt_len], - sampling_params=sampling_params, + @wrap_with_nvtx_name("megatron_policy_worker/generate") + def generate( + self, *, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False + ) -> BatchedDataDict[GenerationOutputSpec]: + """Generate a batch of data using Megatron Core inference with coordinator. + + This method uses the coordinator-based inference pattern from Megatron Core, + which enables better parallelism across data-parallel ranks through a central + coordinator that routes requests to available engines. + + The inference engine is created once and reused across generate() calls. + The engine is paused between generate() calls to free GPU memory for training. + + For coordinator-based inference: + - Only DP rank 0 receives actual data and submits requests to the coordinator + - Other DP ranks receive data=None but still participate in the inference engine loop + - The coordinator distributes work across all DP engines + - Results are broadcast from rank 0 to all ranks + + Args: + data: BatchedDataDict containing input_ids and input_lengths tensors, + or None for non-DP-0 workers (they participate in engine loop only) + BatchedDataDict conforming to GenerationOutputSpec: + - output_ids: input + generated token IDs + - logprobs: Log probabilities for each token + - generation_lengths: Lengths of each response + """ + + from megatron.core.inference.sampling_params import SamplingParams + + def _log_gpu_memory(tag: str): + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + allocated = torch.cuda.memory_allocated() / (1024 ** 3) + reserved = torch.cuda.memory_reserved() / (1024 ** 3) + free, total = torch.cuda.mem_get_info() + free_gb, total_gb = free / (1024 ** 3), total / (1024 ** 3) + print( + f"[GPU Rank {rank}] {tag} | " + f"Allocated: {allocated:.2f} GB, Reserved: {reserved:.2f} GB, " + f"Free: {free_gb:.2f} GB, Total: {total_gb:.2f} GB" + ) + + _log_gpu_memory("generate START") + + self.model.config.flash_decode = False + if self.should_disable_forward_pre_hook: + self.model = self.move_model( + self.model, "cuda", move_params=True, move_grads=False + ) + + dist_rank = torch.distributed.get_rank() + is_request_submitter = (dist_rank == 0) + + # For non-rank-0 workers, data may be None (they participate in engine loop only) + if data is not None: + # Verify input is right padded + assert isinstance(data, BatchedDataDict), ( + f"data must be a BatchedDataDict, got type: {type(data)}" + ) + is_right_padded, error_msg = verify_right_padding( + data, pad_value=self.tokenizer.pad_token_id + ) + if not is_right_padded: + warnings.warn( + f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" + ) + + + mcore_generation_config = self.cfg["generation"]["mcore_generation_config"] + + lang_module = self._get_lang_module() + cuda_graph_impl = mcore_generation_config.get("cuda_graph_impl", "local") + was_training = lang_module.training + + lang_module.eval() + rotary_module = getattr(lang_module, "rotary_pos_emb", None) + has_lru_cache = rotary_module is not None and hasattr(rotary_module.forward, "cache_parameters") + if has_lru_cache: + rotary_module.forward.cache_clear() + + with torch.no_grad(): + + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) + + if not self._inference_engine_initialized: + self._initialize_inference_engine(mcore_generation_config) + coordinator_port = self.cfg["generation"].get( + "inference_coordinator_port", 5995 + ) + self._run_async_coordinator_start(coordinator_port) + + if self._inference_engine_alseep: + self._wake() + + top_k_cfg = self.cfg["generation"]["top_k"] + top_k_val = 1 if greedy else (int(top_k_cfg) if top_k_cfg is not None else 0) + + top_p_cfg = self.cfg["generation"]["top_p"] + top_p_val = ( + 0.0 if greedy else (float(top_p_cfg) if top_p_cfg is not None else 0.0) + ) + + sampling_params = SamplingParams( + temperature=self.cfg["generation"]["temperature"] if not greedy else 0, + top_k=top_k_val, + top_p=top_p_val, + skip_prompt_log_probs=False, + return_log_probs=True, + num_tokens_total=self.cfg["generation"]["max_new_tokens"], + num_tokens_to_generate=None, + termination_id=self.megatron_tokenizer.eod, ) - request_id += 1 - - result = [] - while dynamic_engine.has_unfinished_requests(): - result_step = dynamic_engine.step_modern() - result.extend(result_step["finished_request_records"]) - - # Sort results by request_id to maintain original batch order - result.sort(key=lambda x: x.request_id) - - out = { - "tokens": [ - x.requests[0].prompt_tokens.tolist() + x.requests[0].generated_tokens - for x in result - ], - "logprobs": [ - x.requests[0].prompt_log_probs + x.requests[0].generated_log_probs - for x in result - ], - } + + if is_request_submitter: + input_ids = data["input_ids"] + print(f"[Rank {dist_rank}] input_ids: {input_ids.shape}") + prompt_tokens_tensor = input_ids.cuda() + prompt_lengths_tensor = data["input_lengths"] + else: + print(f"[Rank {dist_rank}] Participating in engine loop (no data to submit)") + prompt_tokens_tensor = torch.empty(0, dtype=torch.long, device="cuda") + prompt_lengths_tensor = torch.empty(0, dtype=torch.long, device="cuda") + + result = self._run_async_generation_with_persistent_engine( + prompt_tokens_tensor, + prompt_lengths_tensor, + sampling_params, + ) + + if self._inference_engine_initialized and not self._inference_engine_alseep: + self._sleep() + + if cuda_graph_impl != "none": + toggle_cuda_graphs(lang_module, set_to="none") + + if has_lru_cache: + rotary_module.forward.cache_clear() + + if was_training: + lang_module.train() + + gc.collect() + torch.cuda.empty_cache() + + self.model.config.flash_decode = False + + # Only rank 0 needs to format and return results + # Other ranks return None (their results are ignored due to output_is_replicated) + if not is_request_submitter: + _log_gpu_memory("generate END (non-submitter)") + return BatchedDataDict({ + "output_ids": torch.empty(0, 0, dtype=torch.long), + "logprobs": torch.empty(0, 0, dtype=torch.float), + "generation_lengths": torch.empty(0, dtype=torch.long), + "unpadded_sequence_lengths": torch.empty(0, dtype=torch.long), + }).to("cpu") input_lengths = data["input_lengths"] - # pad the out "tokens" and "logprobs" and make them into tensors from lists batch_size = data["input_ids"].size(0) - max_gen_seq_len = max([len(x.requests[0].generated_tokens) for x in result]) + max_gen_seq_len = max([len(x.generated_tokens) for x in result]) padded_input_length = input_ids.size(1) max_seq_len = padded_input_length + max_gen_seq_len - # Create padded tensors for tokens and logprobs output_ids_padded = torch.full( (batch_size, max_seq_len), self.tokenizer.pad_token_id, @@ -834,7 +1054,6 @@ def generate( device=data["input_ids"].device, ) - # Fill in the padded tensors with actual values generation_lengths = torch.zeros( batch_size, dtype=torch.long, device=data["input_ids"].device ) @@ -842,15 +1061,17 @@ def generate( batch_size, dtype=torch.long, device=data["input_ids"].device ) for i in range(batch_size): - seq_len = len(out["tokens"][i]) + tokens = result[i].prompt_tokens.tolist() + result[i].generated_tokens + logprobs = result[i].prompt_log_probs + result[i].generated_log_probs + seq_len = len(tokens) output_ids_padded[i, :seq_len] = torch.tensor( - out["tokens"][i], dtype=torch.long, device=data["input_ids"].device + tokens, dtype=torch.long, device=data["input_ids"].device ) generation_lengths[i] = seq_len - input_lengths[i].item() unpadded_sequence_lengths[i] = seq_len - logprob_len = len(out["logprobs"][i]) + logprob_len = len(logprobs) logprobs_padded[i, 1 : logprob_len + 1] = torch.tensor( - out["logprobs"][i], + logprobs, dtype=torch.float, device=data["input_ids"].device, ) @@ -862,11 +1083,148 @@ def generate( "unpadded_sequence_lengths": unpadded_sequence_lengths, } - self.model.config.flash_decode = False - no_grad.__exit__(None, None, None) - + _log_gpu_memory("generate END") return BatchedDataDict.from_batches([out_dict]).to("cpu") + def _start_inference_loop_thread(self): + """Start a background thread with a persistent event loop for inference. + + This thread runs the event loop that hosts the engine loop task. + The loop runs forever until explicitly stopped. + """ + import threading + + def run_loop(): + asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) + self._inference_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._inference_loop) + # Run forever - the engine loop task will run in this loop + self._inference_loop.run_forever() + + self._inference_thread = threading.Thread(target=run_loop, daemon=True) + self._inference_thread.start() + + # Wait for the loop to be created + while self._inference_loop is None: + time.sleep(0.001) + + def _run_async_coordinator_start(self, coordinator_port: int): + """Start the coordinator and engine loop in the background thread. + + This is called once during the first generate() call to initialize + the persistent inference infrastructure. + """ + import concurrent.futures + + # Start the background thread with the event loop if not already running + if self._inference_loop is None: + self._start_inference_loop_thread() + + # Schedule the coordinator start in the inference loop + future = asyncio.run_coroutine_threadsafe( + self._start_inference_coordinator(coordinator_port), + self._inference_loop + ) + # Wait for completion + return future.result() + + def _run_async_generation_with_persistent_engine( + self, + prompt_tokens_tensor: torch.Tensor, + prompt_lengths_tensor: torch.Tensor, + sampling_params: "SamplingParams", + ) -> list: + """Run generation using the persistent inference engine. + + This method uses the pre-initialized engine and client to run generation. + Unlike the original method, it doesn't start/stop the coordinator each time. + The async operation runs in the persistent inference loop. + """ + if self._inference_loop is None: + raise RuntimeError("Inference loop not initialized. Call generate() first.") + + # Schedule the generation in the inference loop + future = asyncio.run_coroutine_threadsafe( + self._generate_with_persistent_engine( + prompt_tokens_tensor, + prompt_lengths_tensor, + sampling_params, + ), + self._inference_loop + ) + # Wait for completion and return the result + return future.result() + + async def _generate_with_persistent_engine( + self, + prompt_tokens_tensor: torch.Tensor, + prompt_lengths_tensor: torch.Tensor, + sampling_params: "SamplingParams", + ) -> list: + """Run generation using the persistent coordinator-based inference. + + This method uses the already-running engine and submits requests through + the persistent client. The engine loop continues running between calls. + + For coordinator-based inference with centralized request submission: + - Only rank 0 (the request submitter) submits requests and collects results + - Other ranks return early but their engine loops continue running in the + background, processing requests distributed by the coordinator + - No broadcast is needed since only rank 0's results are used by the caller + + Args: + prompt_tokens_tensor: Tensor of prompt token IDs [batch_size, seq_len] + prompt_lengths_tensor: Tensor of prompt lengths [batch_size] + sampling_params: Sampling parameters for generation + + Returns: + List of completed request records sorted by request_id (rank 0), + or empty list (other ranks) + """ + from megatron.core.inference.inference_request import DynamicInferenceRequestRecord + + dist_rank = torch.distributed.get_rank() + + if dist_rank == 0: + assert self.inference_client is not None, "Inference client not initialized" + + # Non-rank-0 workers: return immediately with empty results + # Their engine loops will continue processing requests from the coordinator + # in the background (the engine loop runs as a separate task in _inference_loop) + if dist_rank != 0: + print(f"[Rank {dist_rank}] Participating in engine loop only (not submitting requests)") + # Return empty results - the caller only uses rank 0's results + return [] + + # Rank 0: submit ALL requests and collect results + print(f"[Rank {dist_rank}] Submitting {prompt_tokens_tensor.size(0)} requests to coordinator") + + futures = [] + for request_id, (prompt_tokens, prompt_len) in enumerate( + zip(prompt_tokens_tensor, prompt_lengths_tensor, strict=True) + ): + # Extract the actual prompt tokens (without padding) and convert to list + prompt = prompt_tokens[: prompt_len.item()].tolist() + future = self.inference_client.add_request(prompt, sampling_params) + futures.append(future) + + # Wait for all requests to complete + # The coordinator distributes work to all DP engines, including this one + completed_records: list[DynamicInferenceRequestRecord] = await asyncio.gather( + *futures + ) + + # Extract the merged request from each record + results = [record.merge() for record in completed_records] + + # Sort by request_id to maintain original batch order + results.sort(key=lambda x: x.request_id) + + print(f"[Rank {dist_rank}] Completed {len(results)} requests") + + return results + + @torch.no_grad() @wrap_with_nvtx_name("megatron_policy_worker/prepare_refit_info") def prepare_refit_info(self) -> None: @@ -1023,6 +1381,70 @@ def broadcast_weights_for_collective( post_iter_func=lambda x: x[1], ) + @torch.no_grad() + def update_weights_from_collective(self) -> bool: + """Receive updated weights from collective communication (inference side). + + This method is the consumer counterpart of broadcast_weights_for_collective. + It receives weights broadcast by the training workers and updates the local + model parameters. + + TODO: Implement the actual weight update logic using packed_broadcast_consumer. + The implementation should: + 1. Iterate over the stored state_dict_info + 2. Use packed_broadcast_consumer to receive weights from the training side + 3. Update the local Megatron model parameters with the received weights + + Returns: + bool: True if weights were successfully updated. + """ + raise NotImplementedError( + "update_weights_from_collective for MegatronPolicyWorker is not yet implemented. " + "This placeholder will be replaced with actual NCCL collective weight reception logic." + ) + + def init_collective_as_inference( + self, ip: str, port: int, world_size: int, *, train_world_size: int + ) -> None: + """Initialize collective communication for inference-side workers. + + Unlike the base init_collective (used by training workers which use + self.rank directly), this method offsets the rank by train_world_size + so inference workers get globally unique ranks that don't collide + with training workers. + + Args: + ip: IP address for the process group rendezvous. + port: Port for the process group rendezvous. + world_size: Total world size (train + inference workers). + train_world_size: Number of training workers (used to offset ranks). + """ + from nemo_rl.distributed.stateless_process_group import StatelessProcessGroup + + # Offset rank by train_world_size so inference workers get unique global ranks + rank = train_world_size + self.rank + self.model_update_group = StatelessProcessGroup( + master_address=ip, port=port, rank=rank, world_size=world_size + ) + device = torch.cuda.current_device() + self.model_update_group.init_nccl_communicator(device=device) + + def store_refit_info(self, state_dict_info: dict[str, Any]) -> None: + """Store state dict metadata for weight refitting on the inference side. + + This is the inference-side counterpart of prepare_refit_info(). Instead of + calculating the metadata from the model, it accepts pre-computed metadata + from the training side. + + TODO: Implement proper storage and use of state_dict_info for + update_weights_from_collective. + + Args: + state_dict_info: Dictionary mapping tensor names to (shape, dtype) tuples, + as returned by the training-side prepare_refit_info(). + """ + self.state_dict_info = state_dict_info + def prepare_for_lp_inference(self): self.model = self.move_model(self.model, "cuda", move_grads=False) self.model.eval() diff --git a/pyproject.toml b/pyproject.toml index 7b1ced085c..c51bfacd30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,7 @@ mcore = [ # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 "flash-attn==2.8.1", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", + "torch-memory-saver", ] nemo_gym = ["nemo_gym"] @@ -350,7 +351,7 @@ exclude = ''' ''' [tool.pytest.ini_options] -addopts = "--durations=15 -s -rA -x" +addopts = "--durations=15 -s -rA" testpaths = ["tests"] python_files = "test_*.py" markers = [ diff --git a/tests/functional/grpo_non_colocated.sh b/tests/functional/grpo_non_colocated.sh index 8c65aedda2..4570e50279 100755 --- a/tests/functional/grpo_non_colocated.sh +++ b/tests/functional/grpo_non_colocated.sh @@ -20,6 +20,7 @@ mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo.py \ + --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ grpo.num_generations_per_prompt=4 \ @@ -27,7 +28,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE policy.train_micro_batch_size=1 \ policy.generation.colocated.enabled=false \ policy.generation.colocated.resources.gpus_per_node=1 \ - policy.generation.vllm_cfg.async_engine=true \ + policy.generation.backend=vllm \ + policy.generation.vllm_cfg.async_engine=false \ cluster.gpus_per_node=2 \ grpo.max_num_steps=2 \ logger.tensorboard_enabled=true \ diff --git a/uv.lock b/uv.lock index 8c95e36ee9..dbebe0eed2 100644 --- a/uv.lock +++ b/uv.lock @@ -4752,6 +4752,7 @@ mcore = [ { name = "flash-attn" }, { name = "megatron-bridge" }, { name = "megatron-core" }, + { name = "torch-memory-saver" }, { name = "transformer-engine", extra = ["pytorch"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" }, ] nemo-gym = [ @@ -4889,6 +4890,7 @@ requires-dist = [ { name = "tiktoken" }, { name = "torch", marker = "sys_platform != 'darwin'", specifier = "==2.9.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.9.0", index = "https://pypi.org/simple" }, + { name = "torch-memory-saver", marker = "extra == 'mcore'" }, { name = "torch-memory-saver", marker = "extra == 'sglang'" }, { name = "torchao", marker = "extra == 'sglang'" }, { name = "torchdata" },