Graph
Megatron-LM

🧱Megatron-LM

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.

Source

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.

Source

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.

Source

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