pytorch-fsdp2
Category: ml-training · License: MIT · Version: 1.0.0
Tags: PyTorch, FSDP2, Fully Sharded Data Parallel, Distributed Training, DTensor, Device Mesh, Sharded Checkpointing, Mixed Precision, Offload, Torch Distributed
When to use it
Adds PyTorch FSDP2 (fully_shard) to training scripts with correct init, sharding, mixed precision/offload config, and distributed checkpointing. Use when models exceed single-GPU memory or when you need DTensor-based sharding with DeviceMesh.Example requests
Ask your agent in plain language, for example:- “convert my train.py to FSDP2”
- “I’m fine-tuning a 13B Llama in finetune.py on 8xA100 and get CUDA OOM during optimizer.step() with DDP. I’d like to shard params, grads and optimizer state with the new per-parameter DTensor approach (not the old FSDP1 wrapper), use bf16 with fp32 gradient reduction, and save checkpoints that I can reload on a different number of GPUs. Can you update the script?”
- “My model’s weights, gradients and Adam states don’t fit in one GPU’s memory, but I have a node with 4 GPUs. I want training to split the memory across them while still doing data parallelism, and later I might add tensor parallelism with a 2D device mesh. How should I set up my PyTorch training loop?”
Install
/plugin marketplace add KalarisLabs/research-agent-skills and install the whole ml-training category with /plugin install ml-training@research-agent-skills.