NEW
A framework from NVIDIA for training large transformer models using tensor and pipeline parallelism efficiently across GPUs.
DESCRIPTION
Megatron-LM implements tensor-sliced model parallelism and pipeline parallelism to scale transformer training to billions of parameters. It integrates mixed precision, activation checkpointing, and optimized kernels for GPUs. Often combined with DeepSpeed for 3D parallelism, it underpins many large language model training runs. The project provides scripts and configs for pretraining GPT-like models. Effective use requires careful parallelism design and scheduling.
ALTERNATIVE DEFINITIONS
NVIDIA Megatron-LM
Megatron-LM is a framework for training transformer models at scale, leveraging model parallelism to efficiently train models with billions of parameters.
Highlights model-parallel training for very large transformers.
Shoeybi et al. (2019)
We present Megatron-LM, a model-parallel approach for training large transformer models by partitioning the weight matrices across multiple GPUs.
Academic description of tensor slicing across GPUs for scalability.
NVIDIA Developer Blog
Megatron-LM uses tensor and pipeline parallelism to train large transformer models efficiently across many GPUs.
Vendor blog articulating combined parallelism approach.
RELATED TERMS
Tensor Parallelism Pipeline Parallelism ZeRO Transformer Model GPT 3D Parallelism Model Parallelism
COMPETING TERMS
DeepSpeed FairScale Mesh TensorFlow
KNOWLEDGE AREAS
AI MLOps Infrastructure Platform & Tooling
TARGET AUDIENCE
AI Engineer MLOps Engineer Data Scientist Software Developer Hardware Engineer
EXTERNAL SOURCES
- Wikipedia - Transformer (machine learning model)
- NVIDIA - Official Megatron-LM repository
- Shoeybi et al. - Megatron-LM research paper
- NVIDIA - Overview of Megatron-LM parallelism
TOP QUESTIONS
When should I use Megatron-LM?
Use it to train very large transformers requiring tensor and pipeline parallelism for efficiency beyond data parallel alone.
Category: Requirements
How does tensor parallelism work in Megatron-LM?
It partitions weight matrices across GPUs, performing partial matmuls and all-reduces to assemble results.
Category: Mechanism
How does Megatron-LM compare to DeepSpeed?
Megatron focuses on model parallelism; DeepSpeed adds ZeRO and offloading. Together they enable 3D parallelism for maximal scale.
Category: Comparison
What are common pitfalls with Megatron-LM?
Mismatched micro-batch sizes, pipeline bubbles, and communication bottlenecks. Tune schedules and overlap compute with communication.
Category: Mistakes
Does Megatron-LM support inference?
Yes, it includes optimized inference paths, but many use Triton or DeepSpeed-Inference for serving.
Category: Capabilities
What hardware suits Megatron-LM best?
Multi-GPU servers with NVLink and cluster interconnects like InfiniBand. Sufficient GPU memory is critical.
Category: Infrastructure
Can I fine-tune models trained with Megatron-LM?
Yes, fine-tune with consistent parallelism settings or adapt with parameter-efficient methods to reduce cost.
Category: Application
How do I profile Megatron-LM performance?
Use GPU profilers, NCCL metrics, and pipeline timeline analysis to identify bubbles and hotspots.
Category: Observability
What datasets are commonly used with Megatron-LM?
Large text corpora such as The Pile or curated web datasets. Careful filtering improves downstream quality.
Category: Examples