NVIDIA Megatron-LM

Megatron-RL is a library designed to add native reinforcement learning (RL) based post-training to Megatron-LM. It provides a flexible framework for defining RL environments and agents, extending the Megatron-LM training loop with RL algorithm support. It is specifically optimized for research teams performing RL post-training of LLMs at scale on NVIDIA hardware.

Note: As of August 2025, Megatron-RL is actively under development and may not be fully released for external use. It is intended for research rather than as an enterprise framework. For out-of-the-box support for open models, NVIDIA recommends Nemo RL.

Megatron-FSDP is a high-performance implementation of Fully Sharded Data Parallelism (FSDP) written in native PyTorch. It is designed for training and inference of extremely large models at data-center scale on NVIDIA GPUs.

Key capabilities include:

• Scalability: Near-linear scaling from single nodes to entire data centers.
• Flexible Sharding: Supports sharding optimizer states, gradients, and model parameters. It includes hierarchical strategies like Hybrid-Sharded Data Parallelism (HSDP) and Hybrid-FSDP (HFSDP).
• Framework Compatibility: Works with PyTorch (DeviceMesh, DTensor, DCP), Megatron-Core, TransformerEngine, and can be integrated into HuggingFace Transformers or TorchTitan.
• Zero-Copy Communication: Uses NCCL User Buffer Registration and Symmetric Memory to achieve true zero-COPY data movement.

Megatron Core is a GPU-optimized library designed for training large language models (LLMs) at scale. It provides modular and composable building blocks (such as attention and MLP layers) that allow developers to create custom training frameworks.

While Megatron Core provides the foundation, Megatron-LM is the reference implementation that uses these components to train massive models (billions to trillions of parameters) across distributed GPU clusters.

Megatron Core MoE is a production-ready stack for training large-scale Mixture-of-Experts models. It supports a wide range of architectures and advanced parallelism strategies.

Supported Model Architectures

• DeepSeek: V2, V3 (including Multi-Token Prediction/MTP).
• Qwen: Qwen2-57B-A14B, Qwen3-30B-A3B, Qwen3-235B-A22B.
• Mixtral: 8x7B, 8x22B.

Core Capabilities

• Routing: Token dropless MoE (dMoE), Top-K Router with flexible K, and load balancing losses.
• Parallelism: Expert Parallel (EP) integrated with 3D parallelism (DP, TP, PP, SP), Context Parallel (CP) for long sequences, and Distributed Optimizer (ZeRO-1 equivalent).
• Optimizations: Fine-grained recomputations (mla, moe, mlp, moe_act, norm), GroupedGEMM (including FP8/MXFP8), and various communication/kernel fusions.

Megatron-FSDP is an NVIDIA-developed distributed parallelism library written in native PyTorch. It provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP) designed for training and inference of extremely large PyTorch models at data-center scale on NVIDIA GPUs. It is compatible with:

• PyTorch DeviceMesh, DTensor, and Distributed Checkpoint (DCP)
• Megatron Core
• TransformerEngine
• NVIDIA NeMo Framework Container

The models.gpt package provides an implementation of the GPT model architecture. It is designed for high-performance large-scale training and inference, supporting advanced distributed computing techniques and precision formats.

Key supported features include:

• Parallelization Strategies: Tensor Parallel (TP), Pipeline Parallel (PP), and Data Parallel (DP).
• Advanced Architectures: Mixture of Experts (MoE).
• Precision & Optimization: FP8 training/inference and Distributed Optimizer.

This package is part of the Megatron-LM ecosystem and is frequently updated with new features.

The pipeline_parallel package provides implementations for different pipeline parallelism schedules used in large-scale model training. It supports:

1. Non-interleaved pipeline parallelism: A standard pipeline schedule.
2. Interleaved pipeline parallelism: An optimized schedule based on the Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM paper.
3. No-pipelining schedule: A default schedule that does not use pipeline parallelism.

Additionally, the package includes the necessary point-to-point communication methods required to facilitate data transfer between different pipeline stages.

The repository is organized into the following main directories:

• megatron/core/: The core library containing kernels, parallelism implementations (Tensor, Pipeline, Distributed), transformer building blocks, optimizers, and inference engines.
• megatron/training/: Training scripts.
• megatron/post_training/: Tools for quantization, distillation, and pruning.
• megatron/rl/: Reinforcement learning components (including RLHF).
• examples/: Ready-to-use training examples.
• tools/: Utility tools.
• tests/: The test suite.