{"id":5646658,"date":"2022-09-22T16:03:08","date_gmt":"2022-09-22T20:03:08","guid":{"rendered":"https:\/\/lightning.ai\/pages\/?p=5646658"},"modified":"2023-03-07T17:33:40","modified_gmt":"2023-03-07T22:33:40","slug":"fully-sharded-data-parallel-fsdp-pytorch","status":"publish","type":"post","link":"https:\/\/lightning.ai\/pages\/community\/tutorial\/fully-sharded-data-parallel-fsdp-pytorch\/","title":{"rendered":"How to Enable Native Fully Sharded Data Parallel in PyTorch"},"content":{"rendered":"

PyTorch’s native FSDP, now in Lightning<\/h2>\n

tl;dr\u00a0<\/strong><\/em>this tutorial teaches you how to overcome hardware constraints when training large models using PyTorch’s new model sharding strategy<\/h3>\n

Model size has grown exponentially in recent years, producing significantly better results in many domains. However, this expansion has been hampered by hardware constraints, as not everyone has access to the necessary hardware to train large-scale models. To tackle this issue, engineers and researchers have been working on strategies for efficient distributed model training, including Fully Sharded Data Parallel (FSDP).<\/p>\n

One way to reduce memory overhead is by sharding the optimizer states. Currently, each device handles all the weight updates and gradient computation, which consumes a large chunk of memory. Optimizer sharding comes in handy by reducing the memory footprint on each device. Sometimes, even optimizer sharding isn’t enough; in such cases, we would shard models as well.<\/p>\n

Model Sharding is one technique in which model weights are sharded across devices to reduce memory overhead.<\/p>\n

In the release of 1.11, PyTorch added native support for Fully Sharded Data Parallel (FSDP).<\/p>\n

\"https:\/\/pytorch.org\/assets\/images\/fsdp_workflow.png\"

FSDP workflow (via PyTorch)<\/p><\/div>\n

FSDP initially appeared in fairscale<\/a><\/span><\/strong> and later in the official PyTorch repository. Lightning Trainer now supports both of them.<\/a><\/p>\n

Using FSDP with Lightning<\/h2>\n

In the Lightning v1.7.0 release<\/a><\/span><\/strong>, we’ve added support for this Fully Sharded Native Strategy, which can help you leverage native FSDP support by setting the strategy<\/code> flag as \"fsdp_native\"<\/code>.<\/p>\n