Feature/amem nccl rl

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.

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