diff --git a/DESIGN_AMEM_NCCL.md b/DESIGN_AMEM_NCCL.md new file mode 100644 index 00000000000..1dd6652e291 --- /dev/null +++ b/DESIGN_AMEM_NCCL.md @@ -0,0 +1,96 @@ +# Design: AMem NCCL Integration for RL Training + +## Overview + +This PR integrates optional AMem (Asynchronous Memory) NCCL-based offloading support for Megatron RL training workflows. The feature enables releasing NCCL-allocated communicator buffers during RL rollout phases to reduce GPU memory pressure. + +## When AMem Hooks Are Enabled + +AMem NCCL integration is activated when: +- `--rl-amem-offload-during-rollout` flag is set +- Training is in RL mode (rollout generation phase) +- NCCL backend is available for CPU-GPU memory transfers + +The integration is **completely optional** and disabled by default. No impact on existing training workflows. + +## Why RL Rollout Benefits + +RL training has two distinct phases: +1. **Rollout Phase**: Generate responses/trajectories (inference-heavy, memory-intensive) +2. **Training Phase**: Update model parameters (compute-heavy) + +During rollout: +- Large batches of sequences are generated +- Activation tensors accumulate rapidly +- GPU memory becomes a bottleneck before compute saturation +- NCCL buffer memory (communicator allocations) is idle during rollout generation + +AMem enables: +- Releasing GPU memory allocated by NCCL (communication buffers) +- Restoring NCCL memory before training collective operations resume +- Reducing peak GPU memory during rollout phases where NCCL buffers are otherwise idle + +## What Memory Is Offloaded + +AMem enables: +- Offloading GPU memory allocated by NCCL (communication buffers) +- Restoring NCCL memory before training collective operations resume +- Reducing peak GPU memory during rollout phases where NCCL buffers are otherwise idle + +AMem does NOT offload: +- Model parameters +- Optimizer states +- Gradient buffers +- Activations + +Parameter and optimizer offloading are handled separately via +`--rl-offload-optimizer-during-inference` and related flags. + +## Implementation Details + +### Entry Point +`megatron/training/initialize.py`: +- Checks `--rl-amem-offload-during-rollout` flag +- Initializes AMem NCCL backend if enabled +- Sets up environment variables (NCCL_ALGO=Ring for stability) + +### Core Logic +`megatron/core/amem_nccl.py`: +- Wraps NCCL operations for CPU-GPU memory transfers +- Manages memory pinning and buffer registration +- Provides async offload/prefetch primitives + +### RL Integration +`megatron/rl/rl_utils.py`: +- Hooks into rollout phase entry/exit +- Triggers offload before rollout starts +- Triggers prefetch before training phase begins + +### Configuration +`megatron/training/arguments.py`: +- Single flag: `--rl-amem-offload-during-rollout` +- No complex tuning parameters exposed initially +- Defaults designed for safety and stability + +## Non-Goals / Out of Scope + +- **No performance claims**: This PR establishes the integration. Performance tuning and benchmarks will be addressed separately. +- **No automatic memory management**: Offload is explicit, triggered by RL phase transitions. +- **No impact on non-RL training**: Code paths are isolated to RL workflows. + +## Testing Strategy + +- Unit tests validate NCCL offload/prefetch operations +- Functional tests ensure RL training completes successfully with flag enabled +- Backward compatibility: existing RL scripts work unchanged (flag defaults to off) + +## Future Work + +- Fine-tune offload granularity (which layers, when) +- Benchmark memory savings vs. transfer overhead +- Explore overlapping strategies for multi-stage pipelines +- Extend to other memory-constrained scenarios beyond RL + +--- + +**Note**: This design focuses on integration correctness and safety. Performance optimization will be data-driven based on real RL workloads. diff --git a/README.md b/README.md index 221dc0263f7..4219641f854 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,146 @@ Megatron-LM/ ā””ā”€ā”€ docs/ # Documentation ``` +### Megatron-LM: Reference Implementation + +**Reference implementation** that includes Megatron Core plus everything needed to train models. + +**Best for:** + +- **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware +- **Research teams** exploring new architectures and training techniques +- **Learning distributed training** concepts and best practices +- **Quick experimentation** with proven model configurations + +**What you get:** + +- Pre-configured training scripts for GPT, LLama, DeepSeek, Qwen, and more. +- End-to-end examples from data prep to evaluation +- Research-focused tools and utilities + +### Megatron Core: Composable Library + +**Composable library** with GPU-optimized building blocks for custom training frameworks. + +**Best for:** + +- **Framework developers** building on top of modular and optimized components +- **Research teams** needing custom training loops, optimizers, or data pipelines +- **ML engineers** requiring fault-tolerant training pipelines + +**What you get:** + +- Composable transformer building blocks (attention, MLP, etc.) +- Advanced parallelism strategies (TP, PP, DP, EP, CP) +- Pipeline schedules and distributed optimizers +- Mixed precision support (FP16, BF16, FP8) +- GPU-optimized kernels and memory management +- High-performance dataloaders and dataset utilities +- Model architectures (LLaMA, Qwen, GPT, Mixtral, Mamba, etc.) + +## Ecosystem Libraries + +**Libraries used by Megatron Core:** + +- **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** šŸ“£ **NEW!** - Multi-modal data loader (text, images, video, audio) with distributed loading and dataset blending +- **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - Optimized kernels and FP8 mixed precision support +- **[Resiliency Extension (NVRx)](https://github.com/NVIDIA/nvidia-resiliency-ext)** - Fault tolerant training with failure detection and recovery + +**Libraries using Megatron Core:** + +- **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ā†” Megatron checkpoint conversion, flexible training loops, and production-ready recipes +- **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods +- **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples +- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/). + +**Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed) + +# Installation + +## šŸ³ Docker (Recommended) + +We strongly recommend using the previous releases of [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) rather than the latest one for optimal compatibility with Megatron Core release and testing. Our releases are always based on the previous month's NGC container, so this ensures compatibility and stability. + +**Note:** The NGC PyTorch container constraints the python environment globally via `PIP_CONSTRAINT`. In the following examples we will unset the variable. + +This container comes with all dependencies pre-installed with compatible versions and optimized configurations for NVIDIA GPUs: + +- PyTorch (latest stable version) +- CUDA, cuDNN, NCCL (latest stable versions) +- Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs +- For best performance, use NVIDIA Turing GPU architecture generations and later + +```bash +# Run container with mounted directories +docker run --runtime --nvidia --gpus all -it --rm \ + -v /path/to/megatron:/workspace/megatron \ + -v /path/to/dataset:/workspace/dataset \ + -v /path/to/checkpoints:/workspace/checkpoints \ + -e PIP_CONSTRAINT= \ + nvcr.io/nvidia/pytorch:25.04-py3 +``` + +## Pip Installation + +Megatron Core offers support for two NGC PyTorch containers: + +- `dev`: Moving head that supports the most recent upstream dependencies +- `lts`: Long-term support of NGC PyTorch 24.01 + +Both containers can be combined with `mlm` which adds package dependencies for Megatron-LM on top of Megatron Core. + +```bash +# Install the latest release dependencies +pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" +pip install --no-build-isolation megatron-core[dev] +# For running an M-LM application: +pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" +pip install --no-build-isolation megatron-core[mlm,dev] +``` + +```bash +# Install packages for LTS support NGC PyTorch 24.01 +pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" +pip install --no-build-isolation megatron-core[lts] +# For running an M-LM application: +pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" +pip install --no-build-isolation megatron-core[mlm,lts] +``` + +For a version of Megatron Core with only torch, run: + +```bash +pip install megatron-core +``` + +### Optional MoE Dependencies + +For Mixture of Experts (MoE) training with Grouped GEMM support: + +```bash +pip install --no-build-isolation megatron-core[moe] +``` + +**Note:** The `nv-grouped-gemm` package requires: +- CUDA toolkit (nvcc) with CUTLASS headers +- On Ubuntu/Debian: `apt-get install libcutlass-dev` +- GPU with compute capability >= 8.0 + +If you encounter build errors, you can skip this optional dependency and use MoE without Grouped GEMM optimization. + +## System Requirements + +### Hardware Requirements + +- **FP8 Support**: NVIDIA Hopper, Ada, Blackwell GPUs +- **Recommended**: NVIDIA Turing architecture or later + +### Software Requirements + +- **CUDA/cuDNN/NCCL**: Latest stable versions +- **PyTorch**: Latest stable version +- **Transformer Engine**: Latest stable version +- **Python**: 3.12 recommended # Performance Benchmarking @@ -124,6 +264,22 @@ We also strong scaled the standard GPT-3 model (our version has slightly more th ![Strong scaling](images/strong_scaling.png) +### AMem NCCL Plugin for RL Training (Optional) + +For Reinforcement Learning scenarios, enable the AMem NCCL plugin to transparently offload NCCL-allocated GPU memory during inference/rollout phases. This can save up to 10GB+ memory per GPU card. + +```bash +--rl-amem-offload-during-rollout # Enable AMem NCCL plugin for memory offloading during rollout/inference (default: true) +--rl-amem-group-id GROUP_ID # Process group ID (if needed) +``` + +**Prerequisites:** +- Install AMem NCCL plugin from [asystem-amem](https://github.com/inclusionAI/asystem-amem) +- Set environment variables: `NCCL_CUMEM_ENABLE=1` and `AMEM_ENABLE=1` +- GPU compute capability ā‰„ sm80 (Ampere or newer) + +**ā†’ [Complete AMem Integration Guide](docs/amem_integration.md)** - Installation, configuration, and usage examples. + # Roadmaps - **[MoE Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements diff --git a/examples/rl/train_with_amem.sh b/examples/rl/train_with_amem.sh new file mode 100755 index 00000000000..cb44f7ccbb2 --- /dev/null +++ b/examples/rl/train_with_amem.sh @@ -0,0 +1,146 @@ +#!/bin/bash + +# Example script for running RL training with AMem NCCL plugin +# This demonstrates how to enable AMem for memory-efficient RL training + +set -e + +# ============================================================================ +# AMem NCCL Plugin Configuration +# ============================================================================ + +# Path to AMem installation (modify this to your installation path) +AMEM_PATH="${AMEM_PATH:-/path/to/asystem-amem}" + +# Required: Enable NCCL CUMEM +export NCCL_CUMEM_ENABLE=1 + +# Required: Enable AMem plugin +export AMEM_ENABLE=1 + +# Group ID for this training job (use different IDs for different process groups) +# For example: 100 for training, 200 for inference if running on shared GPUs +export AMEM_GROUPID=100 + +# Log level: 3=INFO, 4=DEBUG, 5=VERBOSE +export GMM_LOG=3 + +## TODO: expose AMEM_NCCL_OFFLOAD_FREE_TAG as a CLI option if needed + +# Path to AMem-enabled NCCL library +export AMEM_NCCL_LIB_PATH="${AMEM_PATH}/third_party/nccl/build/lib/libnccl.so.2" +export LD_LIBRARY_PATH="${AMEM_PATH}/third_party/nccl/build/lib:${LD_LIBRARY_PATH}" + +echo "=== AMem NCCL Plugin Configuration ===" +echo "NCCL_CUMEM_ENABLE: ${NCCL_CUMEM_ENABLE}" +echo "AMEM_ENABLE: ${AMEM_ENABLE}" +echo "AMEM_GROUPID: ${AMEM_GROUPID}" +echo "GMM_LOG: ${GMM_LOG}" +echo "AMEM_NCCL_LIB_PATH: ${AMEM_NCCL_LIB_PATH}" +echo "======================================" + +# ============================================================================ +# Training Configuration +# ============================================================================ + +# Model configuration +TENSOR_MODEL_PARALLEL_SIZE=2 +PIPELINE_MODEL_PARALLEL_SIZE=1 +CONTEXT_PARALLEL_SIZE=1 + +# RL configuration +GRPO_PROMPTS_PER_STEP=32 +GRPO_GROUP_SIZE=2 +GRPO_ITERATIONS=2 +GRPO_KL_BETA=0.001 + +# Training configuration +MICRO_BATCH_SIZE=1 +GLOBAL_BATCH_SIZE=8 +SEQ_LENGTH=2048 + +# Data paths (modify these to your data paths) +DATA_PATH="${DATA_PATH:-/path/to/your/data}" +TOKENIZER_PATH="${TOKENIZER_PATH:-/path/to/tokenizer}" +CHECKPOINT_PATH="${CHECKPOINT_PATH:-/path/to/checkpoints}" + +# RL environment config +RL_ENV_CONFIG="${RL_ENV_CONFIG:-/path/to/rl_env_config.yaml}" + +# ============================================================================ +# Distributed Training Setup +# ============================================================================ + +WORLD_SIZE=8 +NNODES=1 +NODE_RANK=0 +MASTER_ADDR=localhost +MASTER_PORT=6000 + +echo "=== Distributed Training Setup ===" +echo "WORLD_SIZE: ${WORLD_SIZE}" +echo "NNODES: ${NNODES}" +echo "==================================" + +# ============================================================================ +# Run Training with AMem Enabled +# ============================================================================ + +torchrun \ + --nproc_per_node=${WORLD_SIZE} \ + --nnodes=${NNODES} \ + --node_rank=${NODE_RANK} \ + --master_addr=${MASTER_ADDR} \ + --master_port=${MASTER_PORT} \ + pretrain_gpt.py \ + --perform-rl-step \ + --tensor-model-parallel-size ${TENSOR_MODEL_PARALLEL_SIZE} \ + --pipeline-model-parallel-size ${PIPELINE_MODEL_PARALLEL_SIZE} \ + --context-parallel-size ${CONTEXT_PARALLEL_SIZE} \ + --micro-batch-size ${MICRO_BATCH_SIZE} \ + --global-batch-size ${GLOBAL_BATCH_SIZE} \ + --seq-length ${SEQ_LENGTH} \ + --max-position-embeddings ${SEQ_LENGTH} \ + --num-layers 24 \ + --hidden-size 2048 \ + --num-attention-heads 32 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --lr 1.0e-5 \ + --min-lr 1.0e-6 \ + --lr-decay-style cosine \ + --train-iters 100000 \ + --lr-warmup-iters 1000 \ + --distributed-backend nccl \ + --data-path ${DATA_PATH} \ + --vocab-file ${TOKENIZER_PATH}/vocab.json \ + --merge-file ${TOKENIZER_PATH}/merges.txt \ + --save ${CHECKPOINT_PATH} \ + --load ${CHECKPOINT_PATH} \ + --save-interval 1000 \ + --eval-interval 100 \ + --eval-iters 10 \ + --log-interval 10 \ + --tensorboard-dir ${CHECKPOINT_PATH}/tensorboard \ + --fp16 \ + --use-flash-attn \ + --sequence-parallel \ + --rl-amem-offload-during-rollout \ + --rl-amem-group-id ${AMEM_GROUPID} \ + --rl-amem-offload-during-rollout \ + --grpo-prompts-per-step ${GRPO_PROMPTS_PER_STEP} \ + --grpo-group-size ${GRPO_GROUP_SIZE} \ + --grpo-iterations ${GRPO_ITERATIONS} \ + --grpo-kl-beta ${GRPO_KL_BETA} \ + --grpo-default-temperature 1.0 \ + --grpo-default-top-p 0.9 \ + --langrl-env-config ${RL_ENV_CONFIG} \ + --langrl-inference-server-type inplace_megatron \ + --rl-offload-optimizer-during-inference \ + --rl-offload-kv-cache-during-training \ + --rl-reset-cuda-graphs \ + --rl-use-sequence-packing \ + --rl-sequence-packing-bin-size 8192 \ + --rl-sequence-packing-algo round-robin + +echo "=== Training Complete ===" diff --git a/megatron/core/amem_nccl.py b/megatron/core/amem_nccl.py new file mode 100644 index 00000000000..20a867e0f93 --- /dev/null +++ b/megatron/core/amem_nccl.py @@ -0,0 +1,235 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +""" +AMem NCCL Plugin Integration + +This module provides integration with AMem NCCL plugin for transparent NCCL memory +offloading and restoration, particularly useful in RL scenarios where memory needs +to be freed between training and inference phases. + +AMem NCCL plugin enables saving up to 10GB+ GPU memory per card by offloading +NCCL-allocated memory during rollout/inference phases in RL training. + +For more information, see: https://github.com/inclusionAI/asystem-amem +""" + +import logging +import os +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +_AMEM_AVAILABLE = False +_AMEM_NCCL_LIB = None + + +def _try_load_amem(): + """Attempt to load the AMem NCCL library with the extended API.""" + global _AMEM_AVAILABLE, _AMEM_NCCL_LIB + + if _AMEM_NCCL_LIB is not None: + return _AMEM_AVAILABLE + + try: + import ctypes + + # Try to load the NCCL library with AMem extensions + # The AMem plugin provides libnccl.so.2 with extended APIs + nccl_lib_path = os.environ.get('AMEM_NCCL_LIB_PATH', 'libnccl.so.2') + + try: + _AMEM_NCCL_LIB = ctypes.CDLL(nccl_lib_path) + except OSError: + logger.debug(f"Could not load NCCL library from {nccl_lib_path}") + return False + + # Check if AMem APIs are available + required_funcs = ['ncclPause', 'ncclResume', 'ncclSetGroupID', 'ncclMemStats'] + for func_name in required_funcs: + if not hasattr(_AMEM_NCCL_LIB, func_name): + logger.debug(f"AMem API {func_name} not found in NCCL library") + _AMEM_NCCL_LIB = None + return False + + # Define function signatures + ncclResult_t = ctypes.c_int + ncclComm_t = ctypes.c_void_p + + _AMEM_NCCL_LIB.ncclPause.argtypes = [ctypes.POINTER(ncclComm_t)] + _AMEM_NCCL_LIB.ncclPause.restype = ncclResult_t + + _AMEM_NCCL_LIB.ncclResume.argtypes = [ctypes.POINTER(ncclComm_t)] + _AMEM_NCCL_LIB.ncclResume.restype = ncclResult_t + + _AMEM_NCCL_LIB.ncclSetGroupID.argtypes = [ctypes.c_int] + _AMEM_NCCL_LIB.ncclSetGroupID.restype = ncclResult_t + + _AMEM_NCCL_LIB.ncclMemStats.argtypes = [] + _AMEM_NCCL_LIB.ncclMemStats.restype = ncclResult_t + + _AMEM_AVAILABLE = True + logger.info("AMem NCCL plugin successfully loaded") + return True + + except Exception as e: + logger.debug(f"Failed to load AMem NCCL plugin: {e}") + _AMEM_NCCL_LIB = None + return False + + +def is_amem_available() -> bool: + """Check if AMem NCCL plugin is available. + + Returns: + bool: True if AMem is available and enabled, False otherwise. + """ + global _AMEM_AVAILABLE + + if _AMEM_NCCL_LIB is None: + _try_load_amem() + + # Check if AMem is enabled via environment variables + if not _AMEM_AVAILABLE: + return False + + # Check environment variable settings + nccl_cumem_enable = os.environ.get('NCCL_CUMEM_ENABLE', '0') + amem_enable = os.environ.get('AMEM_ENABLE', '0') + + return _AMEM_AVAILABLE and nccl_cumem_enable == '1' and amem_enable == '1' + + +def nccl_pause() -> bool: + """Offload NCCL-allocated GPU memory. + + This function releases all NCCL-allocated GPU memory in the current process, + moving data to CPU pinned buffers. Must be called before inference/rollout + phases in RL training to free up memory. + + Returns: + bool: True if successful, False otherwise. + """ + if not is_amem_available(): + return False + + try: + # ncclPause accepts NULL for the communicator parameter + result = _AMEM_NCCL_LIB.ncclPause(None) + if result == 0: # ncclSuccess + logger.debug("Successfully offloaded NCCL memory") + return True + else: + logger.warning(f"ncclPause returned error code: {result}") + return False + except Exception as e: + logger.warning(f"Failed to call ncclPause: {e}") + return False + + +def nccl_resume() -> bool: + """Restore NCCL-allocated GPU memory. + + This function restores all previously offloaded NCCL memory from CPU pinned + buffers back to GPU. Must be called before training phases in RL training. + + Returns: + bool: True if successful, False otherwise. + """ + if not is_amem_available(): + return False + + try: + # ncclResume accepts NULL for the communicator parameter + result = _AMEM_NCCL_LIB.ncclResume(None) + if result == 0: # ncclSuccess + logger.debug("Successfully restored NCCL memory") + return True + else: + logger.warning(f"ncclResume returned error code: {result}") + return False + except Exception as e: + logger.warning(f"Failed to call ncclResume: {e}") + return False + + +def nccl_set_group_id(group_id: int) -> bool: + """Set the process group ID for AMem. + + When multiple processes share GPUs (e.g., training and inference processes), + this function assigns a unique group ID to differentiate them. This must be + called before the first NCCL memory allocation. + + Args: + group_id: Unique identifier for the process group (e.g., 100 for training, 200 for inference). + + Returns: + bool: True if successful, False otherwise. + """ + if not is_amem_available(): + return False + + try: + result = _AMEM_NCCL_LIB.ncclSetGroupID(group_id) + if result == 0: # ncclSuccess + logger.info(f"Successfully set NCCL group ID to {group_id}") + return True + else: + logger.warning(f"ncclSetGroupID returned error code: {result}") + return False + except Exception as e: + logger.warning(f"Failed to call ncclSetGroupID: {e}") + return False + + +def nccl_mem_stats() -> bool: + """Print NCCL memory allocation statistics. + + This function reports the total NCCL memory usage and breakdown by allocation source. + Useful for debugging and monitoring memory usage. + + Returns: + bool: True if successful, False otherwise. + """ + if not is_amem_available(): + return False + + try: + result = _AMEM_NCCL_LIB.ncclMemStats() + return result == 0 # ncclSuccess + except Exception as e: + logger.warning(f"Failed to call ncclMemStats: {e}") + return False + + +def setup_amem_environment(enable: bool = True, group_id: Optional[int] = None) -> None: + """Setup environment variables for AMem NCCL plugin. + + This is a convenience function to set up the required environment variables + for AMem. Should be called early in the initialization process. + + Args: + enable: Whether to enable AMem (default: True). + group_id: Optional group ID to differentiate process groups. + """ + if enable: + # Enable NCCL CUMEM (required for AMem) + os.environ.setdefault('NCCL_CUMEM_ENABLE', '1') + + # Enable AMem plugin + os.environ.setdefault('AMEM_ENABLE', '1') + + # Set group ID if provided + if group_id is not None: + os.environ.setdefault('AMEM_GROUPID', str(group_id)) + + # Set log level (3 = INFO) + os.environ.setdefault('GMM_LOG', '3') + + # TODO: expose AMEM_NCCL_OFFLOAD_FREE_TAG as a CLI option if needed + + logger.info("AMem NCCL plugin environment configured") + else: + os.environ['AMEM_ENABLE'] = '0' + logger.info("AMem NCCL plugin disabled") diff --git a/megatron/core/transformer/moe/grouped_gemm_util.py b/megatron/core/transformer/moe/grouped_gemm_util.py new file mode 100644 index 00000000000..be096c817fd --- /dev/null +++ b/megatron/core/transformer/moe/grouped_gemm_util.py @@ -0,0 +1,35 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +try: + import grouped_gemm +except ImportError: + grouped_gemm = None + + +def grouped_gemm_is_available(): + """Check if grouped_gemm is available.""" + return grouped_gemm is not None + + +def assert_grouped_gemm_is_available(): + """Assert that grouped_gemm is available.""" + error_msg = ( + "Grouped GEMM is not available. To use MoE with grouped GEMM, you need to install " + "nv-grouped-gemm.\n\n" + "Installation options:\n" + "1. Install from PyPI (requires CUDA toolkit and CUTLASS headers):\n" + " pip install 'megatron-core[moe]'\n" + " or\n" + " pip install nv-grouped-gemm\n\n" + "2. Build from source:\n" + " pip install git+https://github.com/fanshiqing/grouped_gemm@v1.1.4\n\n" + "Note: Building from source requires:\n" + "- CUDA toolkit (nvcc)\n" + "- CUTLASS headers (can be installed via 'apt-get install libcutlass-dev' on Ubuntu)\n" + "- Compatible GPU with compute capability >= 8.0\n\n" + "If you don't need MoE functionality, you can continue without this package." + ) + assert grouped_gemm_is_available(), error_msg + + +ops = grouped_gemm.ops if grouped_gemm_is_available() else None diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 3d818b5b0cd..1659c540269 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1829,6 +1829,22 @@ def megatron_rl_inference_mode( loop = get_asyncio_loop() nvtx_range = get_nvtx_range() + amem_offload_during_rollout = getattr(args, 'rl_amem_offload_during_rollout', False) + + # Resolve use_amem here, not lazily + use_amem = False + amem_nccl = None + if amem_offload_during_rollout: + try: + from megatron.core import amem_nccl + use_amem = amem_nccl.is_amem_available() + if use_amem: + logger.info(f"[{dist.get_rank()}:DP] AMem NCCL plugin enabled for memory offloading") + else: + logger.warning(f"[{dist.get_rank()}:DP] AMem requested but not available") + except ImportError: + logger.warning(f"[{dist.get_rank()}:DP] AMem module not found, disabling") + logger.debug(f"[{dist.get_rank()}] Entering inference mode") # If we get a lower precision wrapper, we go one object deeper. @@ -1863,6 +1879,20 @@ def megatron_rl_inference_mode( ) optimizer.offload_to_cpu() + # Offload NCCL memory before inference if AMem is enabled + # NOTE: This is only safe if inference runs without collective comms. + # In Megatron's inplace inference mode, comms still occur ā€” use with caution. + if use_amem and amem_offload_during_rollout: + logger.warning( + "AMem NCCL pause before inference is experimental. " + "If inference uses collective communication, this may cause errors." + ) + with nvtx_range("amem-nccl-pause-before-inference"): + if amem_nccl.nccl_pause(): + logger.info(f"[{dist.get_rank()}:DP] Successfully offloaded NCCL memory") + else: + logger.warning(f"[{dist.get_rank()}:DP] Failed to offload NCCL memory") + if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: toggle_cuda_graphs(lang_module, cuda_graph_impl) @@ -1880,6 +1910,14 @@ def megatron_rl_inference_mode( if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: toggle_cuda_graphs(lang_module, 'none') + # Restore NCCL memory after inference if AMem is enabled + if use_amem and amem_offload_during_rollout: + with nvtx_range("amem-nccl-resume-after-inference"): + if amem_nccl.nccl_resume(): + logger.info(f"[{dist.get_rank()}:DP] Successfully restored NCCL memory") + else: + logger.warning(f"[{dist.get_rank()}:DP] Failed to restore NCCL memory") + # If this is a separate RL inference model, prefetch weights back to CPU so they # don't consume GPU memory during training. with nvtx_range("prefetch-inference-model-weights-to-cpu"): diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 91e26af99c6..b0f8cdf2f63 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2295,6 +2295,15 @@ def _add_rl_args(parser): help='Number of parallel generation tasks for RL inference.') group.add_argument('--rl-skip-bos-token', action=argparse.BooleanOptionalAction, type=bool, default=False, help='Skip BOS token at the beginning of the sequences. Default is False.') + + # AMem NCCL plugin arguments + # AMem NCCL plugin is now always gated by --rl-amem-offload-during-rollout + group.add_argument('--rl-amem-group-id', type=int, default=None, + help='Group ID for AMem NCCL plugin. Use different IDs for training and ' + 'inference processes when they share GPUs (e.g., 100 for training, ' + '200 for inference). Must be set before first NCCL memory allocation.') + group.add_argument('--rl-amem-offload-during-rollout', action=argparse.BooleanOptionalAction, default=True, + help='Enable AMem NCCL plugin for transparent NCCL memory offloading during rollout/inference. Default is True. When enabled, NCCL-allocated GPU memory will be offloaded during rollout/inference phases and restored before training. This can save up to 10GB+ GPU memory per card. Requires AMem NCCL plugin to be installed. See https://github.com/inclusionAI/asystem-amem') return parser def _add_training_args(parser): diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index c150ac3d5ca..7e190269318 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -128,6 +128,29 @@ def finish_mpu_init(): # Pytorch distributed. _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, store) + # Initialize AMem NCCL plugin if requested for RL scenarios + if getattr(args, 'rl_amem_offload_during_rollout', True): + try: + from megatron.core import amem_nccl + amem_nccl.setup_amem_environment( + enable=True, + group_id=getattr(args, 'rl_amem_group_id', None), + ) + group_id = getattr(args, 'rl_amem_group_id', None) + # Try to set group ID if specified + if group_id is not None: + if amem_nccl.nccl_set_group_id(group_id): + if args.rank == 0: + print(f"> AMem NCCL plugin initialized with group ID {group_id}") + else: + logger.warning(f"Failed to set AMem group ID {group_id}") + elif args.rank == 0: + print("> AMem NCCL plugin enabled (no group ID set)") + except ImportError as e: + if args.rank == 0: + print(f"> Warning: AMem NCCL plugin requested but not available: {e}") + logger.warning(f"AMem NCCL module import failed: {e}") + # Random seeds for reproducibility. print_rank_0("> setting random seeds to {} ...".format(args.seed)) _set_random_seed( diff --git a/pyproject.toml b/pyproject.toml index d4de4e40ca1..0e245a40d60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,7 +95,6 @@ dev = [ "mamba-ssm~=2.2", "causal-conv1d~=1.5", "flash-linear-attention~=0.4.0", - "nv-grouped-gemm~=1.1", "megatron-energon[av_decode]~=6.0", "av", "flashinfer-python~=0.5.0", @@ -109,6 +108,14 @@ dev = [ "openai", ] +# Optional MoE dependencies that may require building from source with CUTLASS headers +moe = [ + "nv-grouped-gemm~=1.1", +] + +# Optional AMem NCCL plugin for memory offloading in RL scenarios +amem = [] + lts = [ "tqdm", "einops~=0.8", @@ -118,7 +125,6 @@ lts = [ "opentelemetry-api~=1.33.1", "mamba-ssm~=2.2", "causal-conv1d~=1.5", - "nv-grouped-gemm~=1.1", "megatron-energon[av_decode]~=6.0", "av", "flashinfer-python~=0.5.0", diff --git a/scripts/check_api_backwards_compatibility.py b/scripts/check_api_backwards_compatibility.py index 3c66f00b619..2d05d88147b 100644 --- a/scripts/check_api_backwards_compatibility.py +++ b/scripts/check_api_backwards_compatibility.py @@ -367,9 +367,9 @@ def main(): print(f"\n{i}. {change.kind.value}\n Package: {package_name}{path_info}\n ā†’ {change.explain()}\n{'-'*80}") print(f"\n{'='*80}\nSUMMARY\n{'='*80}\nTotal breaking changes: {len(breaking_changes)}\n{'='*80}\n") - + return 1 - + except Exception as e: logger.error(f"\nāŒ Error: {e}") if args.verbose: