megatron-fsdp

Name	megatron-fsdp JSON
Version	0.1.0rc0 JSON
	download
home_page	None
Summary	**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP)
upload_time	2025-08-22 17:13:31
maintainer	None
docs_url	None
author	None
requires_python	>=3.10
license	Apache 2.0
keywords	nlp nlu deep gpu language learning learning machine nvidia pytorch torch transformer
VCS
bugtrack_url
requirements	No requirements were recorded.
Travis-CI	No Travis.
coveralls test coverage	No coveralls.

            <div align="center">

# 🚀 Megatron-FSDP

</div>

<div align="center">

[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/release/python-3100/)

</div>

## ✨ What is Megatron-FSDP?

**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP). It offers seamless cross-compatibility with major deep learning frameworks and parallelism libraries, making it easy to scale your PyTorch models across multiple GPUs and nodes.

Megatron-FSDP can provide up to 25% speed up and 23% memory savings compared to FSDP2.

### Compatibility

- **[PyTorch DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**
- **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)**
- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)**

## ✨ Features

- **Easy Integration**: Simple `fully_shard` function for quick model parallelization
- **High Performance**: Optimized for NVIDIA GPUs with efficient memory management
- **Cross-Framework**: Works seamlessly with PyTorch, Huggingface Transformers, Megatron-LM, Megatron Bridge and TransformerEngine
- **Scalable**: Supports both single-node multi-GPU and multi-node distributed training
- **Flexible Configuration**: Configurable sharding strategies and process groups

## ⚡ Optimizations

- **Advanced Bucketing**: Data-type aware bucketing system to minimize the overhead of collective operations
- **Buffer Management**: Zero copy communication is achieved by reorganizing the storage of parameters and main grad with `ParamAndGradBuffer` class
- **Communication Overlapping**: Improved communication overlap of paramter all-gather and gradient reduce-scatter
- **User-Buffer-Registration NCCL communication**: Offload NCCL collective communication to NVL/IB Sharp to reduce GPU SM usage for communication
- **FP8 Mixed Precision with Transformer Engine**: Compatibility with Transformer Engine enables efficient FP8 mixed precision training
- **Gradient accumulate fusion support with Transformer Engine**: Remove the explicit gradient copy to the communication buffer in backwards pass

<!-- ## 📊 Performance  -->

<!-- ## 📦 Installation -->

## 🚀 Quick Start

### Basic Usage

Transform your PyTorch model to use Fully Sharded Data Parallelism with just a few lines:

```python
import torch
from megatron_fsdp import fully_shard

# Your existing model and optimizer
model = YourModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Enable FSDP with Megatron-FSDP
model, optimizer = fully_shard(
    model,
    optimizer,
    fsdp_unit_modules=[YourTransformerBlock], # Modules to shard
)

# Your model is now ready for distributed training!
```

### Comparison with FSDP-2

We provide a similar approach for sharding the model with `fully_shard` function:

- No need to call `fully_shard` on all the submodules.
- One liner for the sharding change

Here is an FSDP2 usage example for better comparison

```python
import torch
from torch.distributed.fsdp import fully_shard

# Your existing model and optimizer
model = YourModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Enable FSDP with FSDP2
for module in model.modules():
    if isinstance(module, YourTransformerBlock): # Sub-Modules to shard
        fully_shard(module)
fully_shard(model)

# Your model is now ready for distributed training!
```

            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "megatron-fsdp",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.10",
    "maintainer_email": "NVIDIA <nemo-toolkit@nvidia.com>",
    "keywords": "NLP, NLU, deep, gpu, language, learning, learning, machine, nvidia, pytorch, torch, transformer",
    "author": null,
    "author_email": "NVIDIA <nemo-toolkit@nvidia.com>",
    "download_url": "https://files.pythonhosted.org/packages/70/b1/d2bdbf15ecd2101aefa15afe758cc02e5d2bdddc9cce6c9f4fffa88b0efb/megatron_fsdp-0.1.0rc0.tar.gz",
    "platform": null,
    "description": "<div align=\"center\">\n\n# \ud83d\ude80 Megatron-FSDP\n\n</div>\n\n<div align=\"center\">\n\n[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/release/python-3100/)\n\n</div>\n\n## \u2728 What is Megatron-FSDP?\n\n**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP). It offers seamless cross-compatibility with major deep learning frameworks and parallelism libraries, making it easy to scale your PyTorch models across multiple GPUs and nodes.\n\nMegatron-FSDP can provide up to 25% speed up and 23% memory savings compared to FSDP2.\n\n### Compatibility\n\n- **[PyTorch DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**\n- **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)**\n- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)**\n\n## \u2728 Features\n\n- **Easy Integration**: Simple `fully_shard` function for quick model parallelization\n- **High Performance**: Optimized for NVIDIA GPUs with efficient memory management\n- **Cross-Framework**: Works seamlessly with PyTorch, Huggingface Transformers, Megatron-LM, Megatron Bridge and TransformerEngine\n- **Scalable**: Supports both single-node multi-GPU and multi-node distributed training\n- **Flexible Configuration**: Configurable sharding strategies and process groups\n\n## \u26a1 Optimizations\n\n- **Advanced Bucketing**: Data-type aware bucketing system to minimize the overhead of collective operations\n- **Buffer Management**: Zero copy communication is achieved by reorganizing the storage of parameters and main grad with `ParamAndGradBuffer` class\n- **Communication Overlapping**: Improved communication overlap of paramter all-gather and gradient reduce-scatter\n- **User-Buffer-Registration NCCL communication**: Offload NCCL collective communication to NVL/IB Sharp to reduce GPU SM usage for communication\n- **FP8 Mixed Precision with Transformer Engine**: Compatibility with Transformer Engine enables efficient FP8 mixed precision training\n- **Gradient accumulate fusion support with Transformer Engine**: Remove the explicit gradient copy to the communication buffer in backwards pass\n\n<!-- ## \ud83d\udcca Performance  -->\n\n<!-- ## \ud83d\udce6 Installation -->\n\n## \ud83d\ude80 Quick Start\n\n### Basic Usage\n\nTransform your PyTorch model to use Fully Sharded Data Parallelism with just a few lines:\n\n```python\nimport torch\nfrom megatron_fsdp import fully_shard\n\n# Your existing model and optimizer\nmodel = YourModel()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n# Enable FSDP with Megatron-FSDP\nmodel, optimizer = fully_shard(\n    model,\n    optimizer,\n    fsdp_unit_modules=[YourTransformerBlock], # Modules to shard\n)\n\n# Your model is now ready for distributed training!\n```\n\n### Comparison with FSDP-2\n\nWe provide a similar approach for sharding the model with `fully_shard` function:\n\n- No need to call `fully_shard` on all the submodules.\n- One liner for the sharding change\n\nHere is an FSDP2 usage example for better comparison\n\n```python\nimport torch\nfrom torch.distributed.fsdp import fully_shard\n\n# Your existing model and optimizer\nmodel = YourModel()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n# Enable FSDP with FSDP2\nfor module in model.modules():\n    if isinstance(module, YourTransformerBlock): # Sub-Modules to shard\n        fully_shard(module)\nfully_shard(model)\n\n# Your model is now ready for distributed training!\n```\n",
    "bugtrack_url": null,
    "license": "Apache 2.0",
    "summary": "**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP)",
    "version": "0.1.0rc0",
    "project_urls": {
        "Download": "https://github.com/NVIDIA/Megatron-LM/releases",
        "Homepage": "https://github.com/NVIDIA/Megatron-LM/megatron/core"
    },
    "split_keywords": [
        "nlp",
        " nlu",
        " deep",
        " gpu",
        " language",
        " learning",
        " learning",
        " machine",
        " nvidia",
        " pytorch",
        " torch",
        " transformer"
    ],
    "urls": [
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "70b1d2bdbf15ecd2101aefa15afe758cc02e5d2bdddc9cce6c9f4fffa88b0efb",
                "md5": "0ad54d1368cb67cd0f69c31fe71d478e",
                "sha256": "c1f6750659b849a2c1362a07200d60f1999e0d587deab1847ad2b1d76441431f"
            },
            "downloads": -1,
            "filename": "megatron_fsdp-0.1.0rc0.tar.gz",
            "has_sig": false,
            "md5_digest": "0ad54d1368cb67cd0f69c31fe71d478e",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.10",
            "size": 71622,
            "upload_time": "2025-08-22T17:13:31",
            "upload_time_iso_8601": "2025-08-22T17:13:31.066378Z",
            "url": "https://files.pythonhosted.org/packages/70/b1/d2bdbf15ecd2101aefa15afe758cc02e5d2bdddc9cce6c9f4fffa88b0efb/megatron_fsdp-0.1.0rc0.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2025-08-22 17:13:31",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "NVIDIA",
    "github_project": "Megatron-LM",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": true,
    "lcname": "megatron-fsdp"
}