megatron-fsdp

What it is and what it does

Megatron-FSDP is NVIDIA's PyTorch implementation of Fully Sharded Data Parallelism, a distributed training technique that splits model parameters, gradients, and optimizer states across multiple GPUs to reduce per-device memory usage. It lets you train models too large to fit on a single GPU by coordinating computation and communication across a cluster of NVIDIA GPUs. The library exposes a high-level API (fully_shard_model, fully_shard_optimizer) that wraps your model and optimizer, and supports multiple sharding strategies—from no sharding (similar to standard data parallelism) to aggressive parameter sharding (ZeRO-3 style)—so you can tune the trade-off between memory efficiency and communication overhead.

The package integrates with PyTorch's distributed checkpoint system, DeviceMesh, and DTensor abstractions, and is designed to work alongside Megatron-Core, TransformerEngine, and NVIDIA's NeMo framework. It handles the complexity of initializing extremely large models on meta-device to avoid out-of-memory errors during setup, and supports advanced parallelism patterns like tensor parallelism, context parallelism, and expert parallelism for mixture-of-experts models.

Use it for:

• Train large transformer models (billions of parameters) that cannot fit in a single GPU's memory.
• Reduce per-device memory footprint by distributing optimizer state and gradients across a multi-GPU cluster.
• Combine FSDP with tensor parallelism or expert parallelism for hybrid distributed training strategies.
• Save and restore distributed checkpoints of fully-sharded models and optimizers across training runs.
• Tune memory-communication trade-offs by switching between ZeRO-1, ZeRO-2, and ZeRO-3 sharding strategies.

Worth the install?

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

Megatron-FSDP implements Fully Sharded Data Parallelism (FSDP) in native PyTorch to distribute training and inference of large models across multiple GPUs with configurable memory-communication trade-offs.

Yes, if you are training large PyTorch models on multi-GPU NVIDIA clusters and need fine-grained control over distributed sharding strategies. The package is actively maintained, has no known vulnerabilities, installs with low friction, and integrates cleanly with PyTorch's distributed ecosystem. Not necessary for single-GPU training or if you are already satisfied with standard PyTorch FSDP or other parallelism frameworks.

Install

megatron-fsdp on PyPI

pip install megatron-fsdp

uv add megatron-fsdp

poetry add megatron-fsdp

Installing megatron-fsdp

Before you install

Low friction: pure Python wheel with only torch, einops, and packaging as runtime dependencies. Active maintenance with recent releases; repo shows 17428 stars and current development.

License in practice

Apache 2.0 is permissive; you can use this in commercial projects, modify it, and distribute it with minimal restrictions, though you must include a copy of the license and note any changes.

Quickstart

pip install megatron-fsdp

import torch
from megatron_fsdp import fully_shard_model, fully_shard_optimizer

torch.distributed.init_process_group()
model = torch.nn.Transformer()
fsdp_model = fully_shard_model(module=model, fsdp_unit_modules=[torch.nn.TransformerEncoder])
optimizer = fully_shard_optimizer(torch.optim.AdamW(fsdp_model.parameters()))

Requires torch.distributed setup and multiple GPUs; designed for data-center-scale training on NVIDIA GPUs.

Verify before relying

• Whether the package supports AMD or other non-NVIDIA GPU architectures despite NVIDIA authorship.
• Performance benchmarks or memory savings compared to standard PyTorch FSDP or other parallelism strategies.
• Compatibility with specific transformer frameworks beyond those listed (Megatron-Core, TransformerEngine, NeMo).

Package facts

Evidence: megatron_fsdp-0.5.1-py3-none-any.whl

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

Tags

More Libraries packages

Further reading