megatron-core

What it is and what it does

Megatron Core is a composable library of GPU-optimized building blocks for training transformer models across distributed systems. It provides modular components for tensor parallelism, pipeline parallelism, data parallelism, expert parallelism, and context parallelism, along with support for mixed precision training (FP16, BF16, FP8, FP4) and model architectures including GPT, BERT, and Mamba-based models. The library is designed for framework developers and ML engineers building custom training pipelines, not as a standalone training script.

The package depends on torch, numpy, and packaging, and is actively maintained by NVIDIA with recent releases including dynamic context parallelism and multi-data center training support. Precompiled wheels are available for modern Python versions. Performance benchmarks show up to 47% Model FLOP Utilization on H100 clusters when training models from 2B to 462B parameters.

Use it for:

• Build custom distributed training frameworks that need composable parallelism strategies and transformer building blocks.
• Scale transformer model training across thousands of GPUs with optimized communication and computation overlap.
• Implement mixed-precision training (FP16, BF16, FP8, FP4) for large language models to reduce memory and compute costs.
• Train variable-length sequence models using dynamic context parallelism for adaptive efficiency gains.
• Integrate advanced model architectures (GPT, BERT, Mamba) into production training pipelines with fault tolerance.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Megatron Core provides GPU-optimized building blocks and parallelism strategies for training large transformer models at scale, including tensor parallelism, pipeline parallelism, and mixed precision support.

Yes, if you are building a custom distributed training framework or scaling transformer training across multiple GPUs. The library is production-stable, actively maintained, permissively licensed, and has no known vulnerabilities. Install friction is moderate due to GPU/CUDA requirements and torch dependency. Not recommended for simple single-GPU training or inference-only use cases.

Install

megatron-core on PyPI

pip install megatron-core

uv add megatron-core

poetry add megatron-core

Installing megatron-core

Before you install

Medium friction: precompiled wheels available for Python 3.11 and 3.12 on x86_64 and aarch64, but requires torch and numpy. Package now requires Python 3.12 or later.

License in practice

Apache 2.0 permissive license allows commercial use, modification, and redistribution with minimal restrictions, making it suitable for both research and production deployments.

Quickstart

pip install megatron-core

import torch
from megatron.core.models.gpt import GPTModel
from megatron.core.tensor_parallel import ColumnParallelLinear

Requires NVIDIA GPU(s) and torch compiled with CUDA support; building from source can consume significant memory—set MAX_JOBS environment variable if build fails.

Verify before relying

• Whether the package supports inference-only workflows or is primarily designed for training pipelines.
• Compatibility and integration requirements with specific NVIDIA GPU architectures beyond H100.
• Whether checkpoint conversion with Megatron Bridge is included or requires separate installation.

Package facts

Evidence: megatron_core-0.18.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; megatron_core-0.18.2-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; megatron_core-0.18.2-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; megatron_core-0.18.2-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl; megatron_core-0.18.2-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl; megatron_core-0.18.2-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

Keywords: NLP, NLU, deep, gpu, language, learning, machine, nvidia, pytorch, torch, transformer

Tags

More Libraries packages

Further reading