Megatron-LM

Overview

Megatron-LM is a library from NVIDIA for training very large transformer models across many GPUs and many machines. Its main idea is to split a model's math across devices in a careful way, so that thousands of GPUs can work on one network without waiting on each other too much. It combines tensor parallelism, pipeline parallelism, sequence parallelism and data parallelism, and it is tuned for NVIDIA hardware and fast interconnects.

If your model fits on one GPU, you do not need it. If your model has billions of parameters, and a normal data-parallel job runs out of memory, Megatron-LM is one of the standard answers. Its ideas also appear in many other tools, so learning how it works helps you understand DeepSpeed ZeRO, FSDP and Colossal-AI as well. This page explains the method in plain language, shows a worked example, and describes how the same scaling ideas connect to models that learn from tables.

What it is

Megatron-LM is an open source code base built on PyTorch. It started as a research project to train large language models and grew into a toolkit. Today it has two parts that people often talk about together.

• A reference training code base with ready-made scripts for GPT-style, BERT-style and T5-style models, and for newer decoder-only and mixture-of-experts designs.
• A modular library of building blocks, often called Megatron-Core, that other frameworks can import to get the parallel layers, schedules and optimizers without taking the whole code base.

The library gives you:

• Parallel versions of the main transformer layers: attention, the feed-forward block, the embedding and the output head.
• Schedules that decide the order in which micro-batches move through pipeline stages.
• A distributed optimizer that shards optimizer state across data-parallel workers.
• Fused GPU kernels for common operations, so that memory traffic is smaller.
• Support for mixed precision, including 16-bit and 8-bit floating point formats on supported hardware.
• Tools for checkpointing, data loading and converting checkpoints to other formats.

What it is not

It is not a model zoo and it is not a general training framework for small jobs. It does not hide the complexity of cluster training. You still need to choose a parallel layout, set up a job scheduler and manage data. It is a power tool for teams that train large models from scratch or continue pre-training them.

Who uses it

Research labs and companies that run long pre-training jobs use it, often through a wrapper framework. Cloud providers and hardware vendors also use it to test how well a cluster scales. Because the code is public, many other libraries copy its layer designs.

Ecosystem and wrappers

Many users meet Megatron-LM through another layer. Training frameworks for large models import its parallel layers, and conversion tools translate its checkpoints into formats that Transformers can read. Some projects combine it with the sharded optimizers of DeepSpeed. Others wrap it with launch and recipe tools so that you pick a model size and a cluster size and receive a tuned config. When you read a guide, check which layer it targets: the core library, the reference scripts or a wrapper.

Typical model sizes and layouts

The table below is a rough guide to how layouts grow with model size. It is a teaching aid, not a rule, because the right choice depends on sequence length, memory per GPU and the network.

Where it came from

Before 2019, most large language model training used data parallelism. Every GPU held a full copy of the model and processed a different slice of the data. That works until the model no longer fits on a single device.

The first Megatron-LM paper, published in 2019, showed a simple way to split the matrix multiplications inside a transformer layer across GPUs. It trained a language model with billions of parameters at a time when that was very large. The key point was that the split needed only a few small code changes and a small number of communication steps.

A later paper in 2021 described how to combine this tensor split with pipeline parallelism and data parallelism, and how to schedule the pipeline so that GPUs spend less time idle. It reported training runs on thousands of GPUs. This layered approach is often called 3D parallelism.

Another paper in 2022 introduced sequence parallelism and selective activation recomputation. These two ideas cut the memory used by activations, which are the intermediate values saved for the backward pass.

The library was also used in several very large public training efforts, including a joint effort with another large technology company to train a model with hundreds of billions of parameters. Over time the project added support for mixture-of-experts layers, context parallelism for long sequences, and lower precision number formats. The modular Megatron-Core library came from the need to reuse these parts outside the original scripts.

How it works

To understand Megatron-LM, start with the problem. A large model has too many weights for one GPU. Training it also needs space for gradients, optimizer state and activations. You must divide all of this across devices, and you must do it without making the devices spend most of their time talking.

Think of a big restaurant kitchen. One chef cannot cook a thousand meals an hour. You can add chefs who each cook the same dish for different customers. That is data parallelism. You can split one dish among chefs, so one chops, one fries and one plates. That is pipeline parallelism. You can also let several chefs work together on one big step, such as chopping a huge pile of vegetables. That is tensor parallelism. Megatron-LM mixes all three.

Step 1: data parallelism

The simplest layer of the design is data parallelism. Several groups of GPUs each hold a copy of the same model layout. Each group processes a different batch of data. After the backward pass, groups average their gradients with an all-reduce operation. Megatron-LM adds a distributed optimizer, which splits the optimizer state across these groups. This saves memory in the same spirit as ZeRO stage 1 in DeepSpeed ZeRO.

Step 2: tensor parallelism

Inside one transformer layer, the heavy work is matrix multiplication. Tensor parallelism splits a weight matrix across several GPUs.

Take the feed-forward block. It has two matrices: the first expands the hidden size, the second shrinks it back. The first matrix is split by columns. Each GPU computes its own slice of the output, and the activation function is applied locally with no talking. The second matrix is split by rows. Each GPU takes the slice it already holds, multiplies, and produces a partial result. A single all-reduce then adds the partial results.

The result is that each block needs only one all-reduce in the forward pass and one in the backward pass. Attention works the same way: each GPU takes a subset of the attention heads, since heads are independent, and the output projection is split by rows.

This is the same idea you can read about on the tensor parallelism page. It needs very fast links between the GPUs, so it is normally kept inside one server, where devices share a high-speed bus.

Step 3: pipeline parallelism

A deep model is cut into stages of consecutive layers. Stage one holds the first layers, stage two holds the next, and so on. A batch is split into micro-batches. The first micro-batch goes through stage one, then moves to stage two while stage one starts the second micro-batch. This overlap keeps all stages busy.

The cost is the pipeline bubble, the time at the start and end of a batch when some stages have nothing to do. Megatron-LM reduces the bubble with schedules. The one-forward-one-backward schedule interleaves forward and backward passes to limit memory. An interleaved schedule gives each device several small chunks of layers instead of one large chunk, which makes the bubble smaller at the cost of more messages. See pipeline parallelism for a general view.

Pipeline parallelism sends only activations between neighboring stages, which is much less data than tensor parallelism. So it is used across servers, where links are slower.

Step 4: sequence parallelism

Some operations inside a layer, such as layer normalization and dropout, are not split by tensor parallelism. Each GPU in a tensor group repeats them on the full activation. Sequence parallelism splits these operations along the sequence length. Each GPU handles a slice of the tokens for these parts. The all-reduce is replaced by a pair of operations, all-gather and reduce-scatter, which together cost about the same bandwidth. The gain is that activation memory in those parts falls in proportion to the tensor-parallel size.

Step 5: context parallelism

For very long inputs, even sequence parallelism is not enough, because attention cost grows with sequence length. Context parallelism splits the sequence across devices for the whole layer and exchanges keys and values in a ring or similar pattern. It lets a model train on long contexts that would not fit on one device. It pairs naturally with fast attention kernels such as FlashAttention.

Step 6: expert parallelism

For mixture-of-experts models, the expert networks are spread over devices. A router sends each token to a few experts. The tokens travel across the network with an all-to-all operation, are processed, and travel back. Megatron-LM supports this layout and can combine it with the other forms of parallelism. See mixture of experts for the model side.

Step 7: activation recomputation

Saving every activation for the backward pass costs a great deal of memory. Recomputation throws away some of them and recomputes them when needed. Full recomputation saves the most memory but costs about one extra forward pass. Selective recomputation targets only the parts that use a lot of memory but cost little to recompute, such as the attention score matrix. This is a refined form of gradient checkpointing.

Step 8: precision and kernels

Training uses mixed precision. Weights and updates are kept in higher precision while the heavy math runs in 16-bit or in newer 8-bit floating point formats on hardware that supports them. Fused kernels combine several small operations, such as bias, activation and dropout, into one GPU call. This cuts memory reads and writes. The kernels are written for CUDA, and the communication runs over NCCL.

Putting the layout together

The total number of GPUs is the product of the sizes of each form of parallelism.

world_size = tensor_parallel * pipeline_parallel * context_parallel * data_parallel

For example, 512 GPUs could be arranged as tensor size 8, pipeline size 8 and data size 8. The tensor groups live inside servers, the pipeline groups span a few servers, and the data groups span the rest of the cluster. Choosing these numbers is the main design task.

Key concepts and terms

A worked example

This example shows a launch of a small GPT-style run and a hand calculation of the layout. The numbers are placeholders. Check the options against the version you install, since flag names change between releases.

Hand calculation of a layout

Suppose you have 4 servers with 8 GPUs each, which is 32 GPUs. The model has 48 layers and a hidden size that makes one layer too large to train comfortably with data parallelism alone.

1. Tensor parallel size 4. Each group of 4 GPUs sits inside one server.
2. Pipeline parallel size 2. Each pipeline has 2 stages with 24 layers each.
3. Data parallel size = 32 / (4 * 2) = 4.

Check the global batch size. If each micro-batch has 2 sequences and a pipeline carries 8 micro-batches per step, each data-parallel replica sees 16 sequences. With 4 replicas, the global batch is 64 sequences per step.

Next estimate the pipeline bubble. With 2 stages and 8 micro-batches, the bubble fraction is about (stages - 1) / micro-batches, which is 1/8 or about 12 percent. Using more micro-batches shrinks it. This is why large pipelines need large batches.

A launch script

#!/usr/bin/env bash
# Launch on each node with the right NODE_RANK. Host name is a placeholder.
GPUS_PER_NODE=8
NNODES=4
MASTER_ADDR=node0.example.internal
MASTER_PORT=6000

torchrun \
  --nproc_per_node $GPUS_PER_NODE \
  --nnodes $NNODES \
  --node_rank $NODE_RANK \
  --master_addr $MASTER_ADDR \
  --master_port $MASTER_PORT \
  pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  --sequence-parallel \
  --use-distributed-optimizer \
  --num-layers 48 \
  --hidden-size 4096 \
  --num-attention-heads 32 \
  --seq-length 4096 \
  --max-position-embeddings 4096 \
  --micro-batch-size 2 \
  --global-batch-size 64 \
  --lr 3.0e-4 \
  --min-lr 3.0e-5 \
  --lr-decay-style cosine \
  --train-iters 100000 \
  --bf16 \
  --recompute-activations \
  --data-path /data/example.internal/corpus_text_document \
  --tokenizer-type GPT2BPETokenizer \
  --save /ckpt/example-run \
  --load /ckpt/example-run \
  --log-interval 10 \
  --save-interval 1000

What to look at in the log

• Throughput in tokens per second per GPU. This is the number to compare across layouts.
• Memory use per GPU. If it is close to the limit, add recomputation or raise the parallel sizes.
• Loss curve. It should fall smoothly. Spikes can mean a bad batch or a learning rate that is too high.
• Step time variance. Large swings point to a slow node or a network problem.

Converting a checkpoint

Training checkpoints are saved in a sharded layout that depends on the parallel sizes. To serve the model you often convert it to a standard format used by libraries such as Transformers. Converters exist for this. After conversion, run a small prompt on both versions and compare outputs to be sure the weights were mapped properly.

Strengths and limitations

Strengths

• Scales to very large clusters. The layered design lets you use thousands of GPUs.
• Efficient use of NVIDIA hardware. Fused kernels, fast collectives and support for new number formats keep utilization high.
• Memory control. Tensor, pipeline, sequence and context parallelism, plus recomputation and a sharded optimizer, give many ways to fit a model.
• Mature schedules. The interleaved pipeline schedule and overlap of communication with compute are well tested.
• Modular core. You can import parallel layers into your own training code.
• Active support for newer model types, such as mixture of experts and long context.

Limitations

• Complexity. You must understand the parallel layout, the cluster network and the job scheduler.
• Tight coupling to NVIDIA GPUs and a fast network. Other hardware may work through forks, but the main path is tuned for one vendor.
• Less friendly for small jobs. For a model that fits on one or two GPUs, simpler tools are faster to set up.
• Checkpoint formats depend on the layout, so changing the layout needs conversion or resharding tools.
• Flag names and defaults change. Scripts from older versions may not run unchanged.
• Debugging is hard. A hang in a collective on one rank can freeze the whole job with little information.
• It focuses on pre-training and continued pre-training. Fine-tuning and alignment recipes are often handled by wrapper frameworks.

How it compares with alternatives

A useful rule is to ask how big the model is compared with one GPU. If the model is a bit too large, sharding tools such as FSDP or ZeRO are often enough. If the model is far too large, or if you need the best throughput on a big cluster, adding tensor and pipeline parallelism through Megatron-LM becomes worthwhile.

Role in tabular foundation models

Tabular foundation models learn from raw tables. A table has rows and columns, mixed types such as numbers, categories, text and dates, missing values, and sometimes no useful column names. A model that reads such tables in context, a bit like TabPFN does, has to turn each cell into a vector, mix information across rows and across columns, and predict a missing value or a label. Megatron-LM does not define such a model. It is the engine that can train one when the model grows large.

Here is how the ideas connect, concretely.

Scaling the transformer over rows and columns

A table model often uses attention in two directions: across the cells of a row, and across the rows of a column. A large model of this kind has the same building blocks as a language model: attention blocks and feed-forward blocks. Tensor parallelism splits those blocks exactly as it does for text. Attention heads are spread over GPUs, and the feed-forward matrices are split by columns and rows. So a table model with billions of parameters can reuse the same layout that language models use.

Long contexts from many rows

In-context learning on tables puts many labeled rows into the input, then asks the model to predict a new row. More rows give better context but make the sequence longer. Row count times column count can reach tens of thousands of tokens. Context parallelism and sequence parallelism help here, because they split the long input across devices. Fast attention such as FlashAttention keeps the cost manageable.

Pre-training on synthetic and real tables

Many tabular foundation models are pre-trained on very large numbers of generated tables, drawn from random causal structures, and on collections of real tables. The data side of this is unusual: each training example is a whole table, not a sentence. Data loading must produce batches of tables with different shapes. A distributed trainer must handle padding and masks so that missing cells and variable column counts do not waste compute. The pipeline schedule works best when micro-batches have similar cost, so bucketing tables by size keeps stages balanced.

Missing values and masks

Missing values are normal in tables. A model treats them with a mask or a special token. Masking is also the training task for imputation: hide some cells, ask the model to fill them. This fits large-batch pre-training well, because masks can be generated on the fly. It also supports anomaly detection, where a cell that the model finds very unlikely given the rest of the row is flagged.

Continual learning

A large table model may need to learn new table families after the first pre-training run. Continued pre-training on new data is a standard use of Megatron-LM. The risk is forgetting. Teams mix old and new data, lower the learning rate and test on earlier tasks. Checkpoints saved with a known layout make it easy to resume.

Semantic and sector classification

Sometimes the column names are present and carry meaning. A model can read names and sample values to guess the type of a column or the industry sector a table belongs to. This links the table model to language models: text encoders may share the same distributed training stack. A single code base that handles both text and table blocks simplifies cluster operations.

Evaluating a table model during pre-training

Language model runs track perplexity. A table model needs different signals. Good ones include the masked-cell reconstruction error on held-out tables, accuracy and log loss on classification tasks, error on regression targets, and the rate of correct column-type guesses. Keep these tasks fixed across checkpoints so you can compare them. Run them on tables of several sizes, since a model may do well on small tables and poorly on wide ones.

Throughput tips specific to tables

• Group tables of similar width and length into the same micro-batch to reduce padding.
• Cap the number of rows in context per example, and sample rows at random each time, so the model sees many views of one table.
• Shuffle column order during training. This helps the model avoid relying on position, which is useful when column names are missing.
• Represent numbers with a scheme that handles wide ranges, such as per-column scaling, so that mixed units do not destabilize training.

When parallelism is not needed

Many tabular models are small. A model with tens of millions of parameters fits on a single GPU. For those, Megatron-LM is more than you need. It matters when you scale the model, widen the context, or add many experts. Even then, you should compare the cost with simpler sharding from FSDP. For research context, see /model/research, /model/architecture and /model/benchmarks.

Practical guidance and best practices

Pick the layout from the hardware up

• Keep tensor parallel groups inside one server. They talk the most.
• Use pipeline parallelism across servers when the model still does not fit.
• Fill the rest of the cluster with data parallelism.
• Prefer the smallest tensor and pipeline sizes that make the model fit, because more splitting means more communication.

Raise micro-batch count to shrink the bubble

The pipeline bubble shrinks as the number of micro-batches per step grows. If memory allows, use smaller micro-batches and more of them. If a larger global batch hurts convergence, use the interleaved schedule instead.

Use sequence parallelism with tensor parallelism

When tensor parallel size is above one, turn on sequence parallelism. It saves activation memory for almost no extra cost.

Prefer selective recomputation

Start with selective recomputation. Move to full recomputation only if memory is still short.

Overlap communication and compute

Enable options that overlap gradient reduction and parameter gathering with the backward pass. Check the profile to confirm that the overlap really happens.

Test small, then scale

1. Run a tiny model on one node to check data loading and loss.
2. Run the target layout on a small number of steps to check memory and speed.
3. Compare loss curves against a smaller layout. Same seed and same global batch should give similar loss early on.
4. Scale to the full cluster only after these checks pass.

Checkpoint often and test restore

Long jobs fail. Save checkpoints at a regular interval and test that a restore reproduces the same loss. Keep at least one older checkpoint in case the latest is corrupt. Use asynchronous saving if the version supports it, so that the job does not stop for long writes.

Monitor the cluster

Track GPU utilization, memory, temperature, network throughput and step time per rank. A single slow GPU slows every step. Replace or drain bad nodes quickly.

Data pipeline

• Tokenize and pack data ahead of time into a binary index format.
• Use a fixed random seed for data order and record it.
• Shuffle at the document level and verify that train and validation sets do not overlap.

Numerical stability

• Use bfloat16 where available. It has a wide range and avoids most scaling problems of float16.
• Clip gradients, and log the gradient norm.
• Watch for loss spikes. If they appear, lower the learning rate, check for bad data, or roll back to an earlier checkpoint and skip the batch.

Plan for evaluation

Run a small evaluation on a schedule during training. For table models, include held-out tables, imputation tasks and classification tasks, so that you notice regressions early.

A memory estimate you can do by hand

Before you launch, estimate memory per GPU. For mixed precision training with Adam, a common rule of thumb is about 16 bytes per parameter for the sum of weights, gradients and optimizer state: 2 bytes for the 16-bit weights, 2 for the 16-bit gradients, and 12 for the 32-bit master weights and the two Adam moments. Activations come on top.

Take a model with 20 billion parameters.

1. Total state is about 20 billion times 16 bytes, which is 320 billion bytes, or about 320 GB.
2. With tensor size 4 and pipeline size 2, the model is divided by 8 before data parallelism. That leaves about 40 GB per GPU.
3. A distributed optimizer then shards the 12-byte part across the data-parallel group. With a data-parallel size of 4, the optimizer part falls from 30 GB to about 7.5 GB per GPU, so the total drops to roughly 17.5 GB per GPU.
4. The rest of the memory budget is for activations, buffers and fragmentation.

This is an estimate, not a promise. It shows why combining several forms of parallelism works: each one cuts a different part of the cost.

Reading a profile

Capture a short profile of a few steps. Look for four things.

• Gaps between kernels. Gaps usually mean the CPU is slow to launch work or the data loader is late.
• Long collective calls. If an all-reduce is slow, the network or the group layout is the problem.
• Short kernels in large numbers. These are candidates for fusion.
• Imbalance across ranks. If one rank is always last, check its hardware and its share of the model.

Scaling efficiency

Measure weak scaling and strong scaling. In weak scaling, you grow the cluster and the global batch together and check that tokens per second per GPU stays flat. In strong scaling, you keep the global batch fixed and add GPUs, and you expect efficiency to fall as communication grows. A healthy large job keeps most of the per-GPU throughput as the cluster grows. When efficiency falls quickly, look at the communication share first.

Fault tolerance

At large scale, hardware failures are normal. A job on thousands of GPUs will see a failure in a long run. Good practice is:

• Use a scheduler that restarts the job from the latest checkpoint.
• Add health checks before the job starts, such as a quick all-reduce test and a memory test.
• Set timeouts on collectives so that a hung rank is detected rather than waited on forever.
• Keep a spare pool of nodes to swap in.

Mixing with other libraries

Some teams use the Megatron layers inside a larger framework that handles configuration, data and evaluation. Others combine its parallel layers with the sharded optimizers of DeepSpeed ZeRO. Another route is to use FSDP for the data-parallel dimension and tensor parallel layers for the model dimension. Whatever you choose, keep one source of truth for the layout, and test checkpoint conversion early.

A checklist before a long run

• The layer count divides by the pipeline size.
• The head count divides by the tensor size.
• The global batch size divides by micro-batch size times data-parallel size.
• Sequence length is the same in the data and in the config.
• The tokenizer and vocabulary size match the model's embedding size, with any padding for even splitting.
• Checkpoint restore has been tested.
• Logging reports loss, learning rate, gradient norm and throughput.
• Evaluation tasks are set up and run on a schedule.

Vocabulary padding

The embedding matrix is split across tensor-parallel ranks, so the vocabulary size must divide evenly. Libraries pad the vocabulary with unused entries to a convenient multiple. This is harmless, but remember it when you convert checkpoints, since the padded rows must be removed or ignored.

Seeds and determinism

Parallel training uses separate random streams for model-parallel regions and data-parallel regions. Dropout inside a tensor-parallel region must differ across ranks, while dropout on replicated activations must match. The library handles this with seeded state trackers. If you write custom layers, use the same mechanism, or you may create subtle bugs where ranks disagree.

Learning path

If you are new to large-scale training, learn in this order.

1. Train a small transformer on one GPU with PyTorch and mixed precision.
2. Add data parallelism across two GPUs and watch gradient averaging.
3. Try a sharding tool such as FSDP to see memory fall.
4. Read how tensor parallelism splits a feed-forward block, and reproduce the split by hand in a small script.
5. Add a pipeline of two stages and look at the bubble in a profile.
6. Only then move to Megatron-LM on a real cluster.

This order builds an intuition for each cost before you combine them. It also helps you decide whether the full library is needed at all.

Cost thinking

GPU hours are the main cost. Every percent of lost utilization is money wasted. Before a long run, spend a small share of the budget on layout search: try a few tensor, pipeline and micro-batch settings for a hundred steps each, and pick the best. The saving on the long run usually pays for the search many times over.

Common mistakes

• Putting tensor parallel groups across servers, which makes the all-reduce slow and wastes the cluster.
• Choosing large pipeline sizes with too few micro-batches, which creates a large bubble.
• Forgetting that the number of layers must divide evenly across pipeline stages, or that attention heads must divide evenly across tensor-parallel ranks.
• Changing the parallel layout between runs without converting the checkpoint.
• Comparing throughput between layouts at different global batch sizes.
• Ignoring the data loader. A slow loader leaves expensive GPUs idle.
• Using float16 without loss scaling and then blaming the model for not-a-number errors.
• Running with mismatched library, driver and communication library versions.
• Treating a hang as a model bug when it is a network or node failure.
• Skipping a restore test and learning that checkpoints were broken only after a crash.
• Over-splitting a small model. Extra communication can make a small job slower, not faster.
• Reading numbers from a paper and expecting the same utilization on different hardware.

Frequently asked questions

What is Megatron-LM used for?

Megatron-LM is used to pre-train very large transformer models across many GPUs. It splits the matrix math of each layer, the layers themselves and the data across devices, so that models too large for one GPU can be trained efficiently on a cluster.

What is the difference between tensor parallelism and pipeline parallelism?

Tensor parallelism splits the weight matrices inside one layer across several GPUs, and those GPUs must exchange results at every layer. Pipeline parallelism splits the stack of layers into stages and passes activations from one stage to the next. Tensor parallelism needs very fast links, so it stays inside a server, while pipeline parallelism works across servers.

Do I need Megatron-LM if I use FSDP or DeepSpeed ZeRO?

Not always. Sharding tools handle many models by spreading weights, gradients and optimizer state across data-parallel workers. Megatron-LM becomes attractive when the model is very large, when you want the highest throughput on a big cluster, or when you need tensor, pipeline and context parallelism in the same job.

What is the pipeline bubble?

The pipeline bubble is the idle time at the beginning and end of each batch, when some pipeline stages have no work yet or no work left. It shrinks when you use more micro-batches per step or an interleaved schedule, which gives each device several smaller chunks of layers.

What does sequence parallelism add?

Sequence parallelism splits operations such as layer normalization and dropout along the token axis, so each GPU in a tensor group stores only part of those activations. It cuts activation memory in proportion to the tensor-parallel size, and it costs about the same communication as plain tensor parallelism.

Can Megatron-LM train models other than language models?

Yes. The parallel layers are general transformer blocks, so they can be used for vision models, multimodal models and models that read tables. What matters is that the network is built from the same kinds of attention and feed-forward blocks, and that your data loader supplies the right inputs and masks.

Does it only run on NVIDIA GPUs?

The main code path is tuned for NVIDIA GPUs, their kernels and their communication library. Forks and ports for other accelerators exist in the wider community, but you should check their maturity and support before depending on them.

How do I choose the tensor and pipeline sizes?

Start from the hardware. Set the tensor size to fit inside one server, usually the number of GPUs per server or a divisor of it. Add pipeline stages only until the model fits in memory. Then fill the remaining GPUs with data parallelism and measure throughput. The best layout is the one with the highest tokens per second per GPU that still converges well.

How do I move a trained model to another framework?

Training checkpoints are sharded according to the parallel layout, so you convert them first. Converters merge the shards and rename the tensors to match the target format, such as the one used by Transformers. After conversion, compare outputs on a few prompts or inputs to check that the mapping is correct.

Is it useful for small tabular models?

Usually not. A small model fits on one GPU, and simple data parallelism or no parallelism at all is easier. Megatron-LM matters if you scale the model, extend the context to many rows, or use experts, and then it offers the same splitting methods that large language models use.

Glossary

• All-gather: a collective operation in which every device receives the slices held by all devices.
• All-reduce: a collective operation that sums values across devices and gives every device the sum.
• Bubble: idle time in a pipeline schedule.
• Context parallelism: splitting a long sequence across devices through the whole layer.
• Data parallelism: running copies of the model on different data and averaging gradients.
• Distributed optimizer: an optimizer whose state is sharded across data-parallel ranks.
• Expert parallelism: placing experts of a mixture-of-experts layer on different devices.
• Global batch size: the number of examples used for one optimizer update.
• Interleaved schedule: a pipeline schedule where each device holds several small chunks of layers.
• Micro-batch: a slice of the batch that flows through the pipeline as one unit.
• Model parallelism: a general name for splitting the model itself across devices.
• Rank: the index of one process in a distributed job.
• Recomputation: discarding activations and recomputing them in the backward pass.
• Reduce-scatter: a collective that sums data and gives each device a different slice of the sum.
• Sequence parallelism: splitting certain operations along the token axis in a tensor-parallel group.
• Stage: a group of consecutive layers assigned to one pipeline position.
• Tensor parallelism: splitting the matrices of a layer across devices.
• World size: the total number of processes in a distributed job.

Summary

Megatron-LM is NVIDIA's open source library for training very large transformer models on GPU clusters. It combines tensor, pipeline, sequence, context, expert and data parallelism, plus recomputation, a sharded optimizer and fused kernels, to fit huge models in memory and keep the hardware busy. Its strength is scale and efficiency on NVIDIA systems. Its cost is complexity, so it is best for large pre-training and continued pre-training, not for small jobs.

For tabular foundation models, it supplies the scaling layer: the same parallel methods that split language models can split table models that attend over rows and columns, handle long in-context row sets, train on masks for imputation and anomaly detection, and continue learning from new table families. Compare it with DeepSpeed ZeRO, FSDP and Colossal-AI, and pair it with FlashAttention, mixed precision and NCCL for best results. Further reading on the model side is at /model/research, /model/architecture and /model/benchmarks.

Related

It belongs to Training tooling in my stack, under LLMs, retrieval, and evaluation. Others in the same group: