---
id: megatron-fsdp
version: "0.5.1"
license: Apache 2.0
license_treatment: permissive
maintenance: active
---

# megatron-fsdp — **Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP)

License: permissive · Maintenance: active · Downloads: 76.4K/mo

## 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 above — 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

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

- License: Apache 2.0 (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 76.4K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

distributed model training pytorch, fully sharded data parallelism, large language model parallelism, multi-gpu model sharding, fsdp pytorch implementation, zero redundancy optimizer, gpu memory efficient training, distributed-training, model-parallelism, gpu-optimization

[View on SkillFed](https://skillfed.io/packages/megatron-fsdp) · [View on PyPI](https://pypi.org/project/megatron-fsdp/)