FSDP

Overview

FSDP, short for fully sharded data parallel, is PyTorch's native way to train a model that is too large to fit on one GPU while still using data parallelism. Instead of keeping a full copy of the model on every GPU, FSDP splits the parameters, the gradients and the optimizer states into shards. Each GPU owns one shard. When a layer needs to run, the GPUs briefly gather the full weights of that layer, do the work, and then drop the pieces they do not own. The result is that memory per GPU falls roughly in proportion to the number of GPUs.

For readers who work with tables, the practical message is simple. If you want to pre-train a transformer that learns from rows and columns, the model and its optimizer can quickly outgrow a single card. FSDP is one of the standard tools that lets you grow the model without leaving the normal PyTorch training loop. It is a close relative of DeepSpeed ZeRO stage 3, and it sits next to other scaling tools such as tensor parallelism, pipeline parallelism and Megatron-LM.

What it is

FSDP is a module wrapper and a runtime that lives inside PyTorch. You wrap a model, or parts of it, and FSDP takes over three jobs:

• It stores each wrapped module's parameters as a flat, sharded piece spread across all workers (ranks).
• It runs communication at the right moments, so that full weights exist only for the layer that is currently computing.
• It keeps gradients and optimizer states sharded too, so the optimizer step also runs on shards.

The key idea is that data parallelism is wasteful in memory. In classic data parallel training, every GPU holds the whole model, all gradients, and all optimizer states. With an Adam-style optimizer, the optimizer states can be larger than the weights themselves. FSDP removes this duplication.

Two generations

PyTorch ships two FSDP APIs. The first is the FullyShardedDataParallel class, which wraps modules and flattens parameters into large buffers. The second, often called FSDP2, uses a function called fully_shard and represents each parameter as a sharded tensor type. FSDP2 keeps the idea but changes the internals so that parameters stay as individual tensors. This makes it easier to mix with other features, to handle frozen parameters, and to save and load checkpoints. New code should look at the newer API first, while older tutorials still show the class-based one.

What it is not

FSDP is not a model parallel method in the strict sense. The math of each layer is still done whole, on each GPU, for its own slice of the batch. It does not split a single matrix multiply across GPUs. That is the job of tensor parallelism. FSDP also does not reduce the number of floating point operations. It trades extra communication for lower memory.

Where it came from

The story starts with data parallel training, where the model is copied to every worker and gradients are averaged. This works until the model no longer fits. Researchers then looked for ways to remove the redundancy in that copy.

The ZeRO work, published in 2019 and 2020, showed that optimizer states, gradients and parameters can be partitioned across data parallel workers in three stages, with the third stage sharding everything. That idea was first available in DeepSpeed ZeRO. A library called FairScale then offered a PyTorch style implementation of a fully sharded wrapper. The PyTorch team later brought a native version into the main library, and described the design in a paper about PyTorch FSDP published in 2023.

Since then the API has been refined. Features such as mixed precision policies, CPU offload, hybrid sharding, and better checkpoint formats were added. The second generation API followed, built around per-parameter sharding. Today FSDP is widely used in open training recipes and is supported by higher level tools such as PyTorch Lightning, Hugging Face training libraries and LLaMA-Factory.

How it works

Think of a team of cooks sharing one huge cookbook. In data parallel training, every cook buys a full copy of the book. It works, but the kitchen shelves overflow. With FSDP, the book is cut into chapters and each cook keeps only one chapter on the shelf. When it is time to cook a dish from chapter five, the cook who owns it reads it out loud, everyone copies the page for a moment, cooks, and then throws the copy away. Shelf space is saved, and the cost is some talking.

Step by step

1. Wrap the model. You choose which submodules become FSDP units. A typical choice is one unit per transformer block, plus one outer unit for embeddings and the final head.
2. Shard at construction. Each unit takes its parameters, flattens them, and splits them evenly across ranks. Each rank keeps only its piece.
3. All-gather before forward. Just before a unit runs its forward pass, all ranks do an all-gather. Each rank sends its piece to the others, so every rank has the full weights for that unit.
4. Compute forward. The unit runs on the rank's own micro-batch. After the forward pass, the rank frees the gathered weights and keeps only its own shard.
5. All-gather again for backward. During the backward pass, the same all-gather happens again, because the weights were freed to save memory.
6. Reduce-scatter gradients. After computing gradients for a unit, ranks run a reduce-scatter. This sums the gradients across ranks and leaves each rank with only the slice that matches its parameter shard.
7. Sharded optimizer step. Each rank updates its own shard of parameters using its own shard of optimizer state. No rank ever needs the full optimizer state.

Why overlap matters

All-gather and reduce-scatter move a lot of data. If they ran in sequence with compute, the GPUs would sit idle. FSDP uses separate communication streams so that the all-gather for the next unit can run while the current unit computes. This is called prefetching. With a fast interconnect, much of the communication hides behind compute. With a slow network, it does not, and training slows down. Communication is done through a library such as NCCL on NVIDIA GPUs.

Memory accounting

Suppose a model has N parameters and you train with mixed precision and an Adam-style optimizer. A rough count per parameter is: low precision weights, low precision gradients, and optimizer states that include a full precision master copy plus two moment estimates. Under plain data parallel training, each GPU pays this cost for all N parameters. Under FSDP with W workers, each GPU pays roughly one W-th of it, plus the temporary space for the gathered weights of one or two units, plus activations. Activations are not sharded by FSDP, so they often become the next bottleneck. That is why FSDP is often paired with gradient checkpointing.

Sharding strategies

Hybrid sharding is useful on clusters where links between GPUs in the same machine are much faster than links between machines. You shard inside the machine and replicate across machines.

Wrapping policy

The wrapping policy decides the size of each unit. If units are too large, the temporary full weights take too much memory. If units are too small, there are too many small collectives and the overhead grows. Wrapping by transformer block is a common sweet spot because it matches the natural repeating structure of the model.

Mixed precision and offload

FSDP lets you set separate dtypes for parameters, gradient reduction and buffers. A common setup computes in a 16-bit format and reduces gradients in the same or a higher precision. See mixed precision for the trade-offs. CPU offload moves the sharded parameters, and often the optimizer states, to host memory. This lets you train bigger models on fewer GPUs, at a large cost in speed because data must cross the PCIe link.

Checkpointing

Saving a sharded model has two flavors. A full state dict gathers everything to one rank and writes a single file, which is easy to share but needs a lot of memory on that rank. A sharded state dict has each rank write its own piece, which is fast and scalable, and distributed checkpoint tools can reload it even on a different number of GPUs. For large runs, sharded checkpoints are the norm.

Key concepts and terms

A worked example

Back-of-the-envelope numbers

Take a model with 1 billion parameters and an Adam-style optimizer, trained with a 16-bit compute format and a 32-bit master copy. Using the common rule of thumb of about 16 bytes per parameter for weights, gradients and optimizer state together, the training state is about 16 GB. This is a rough estimate, and real numbers vary with the setup.

• On one GPU, or with plain data parallel on 8 GPUs, each GPU holds about 16 GB of state.
• With full sharding on 8 GPUs, each GPU holds about 2 GB of state, plus temporary gathered weights and activations.

For a 10 billion parameter model the same rule gives about 160 GB of state. That cannot fit on one card, but spread over 16 GPUs it is about 10 GB each. This is the main reason FSDP exists.

A training script skeleton

The following sketch trains a small transformer that reads table rows, using the class-based API. Names are placeholders.

import os
import functools
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy


class RowEncoder(nn.Module):
    """Turns a row of numbers into tokens and runs a transformer."""

    def __init__(self, n_features=64, d_model=512, n_layers=12):
        super().__init__()
        self.embed = nn.Linear(1, d_model)
        layer = nn.TransformerEncoderLayer(d_model, nhead=8, batch_first=True)
        self.blocks = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.head = nn.Linear(d_model, 1)

    def forward(self, x):                # x: (batch, n_features)
        tokens = self.embed(x.unsqueeze(-1))
        h = self.blocks(tokens)
        return self.head(h.mean(dim=1))


def main():
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = RowEncoder().cuda()

    wrap = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={nn.TransformerEncoderLayer},
    )
    mp = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        buffer_dtype=torch.bfloat16,
    )
    model = FSDP(
        model,
        auto_wrap_policy=wrap,
        mixed_precision=mp,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        device_id=torch.cuda.current_device(),
    )

    # Create the optimizer AFTER wrapping.
    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for step in range(1000):
        x = torch.randn(32, 64, device="cuda")
        y = torch.randn(32, 1, device="cuda")
        loss = nn.functional.mse_loss(model(x), y)
        loss.backward()
        opt.step()
        opt.zero_grad(set_to_none=True)

    dist.destroy_process_group()


if __name__ == "__main__":
    main()

You launch it with the standard launcher, one process per GPU:

torchrun --nnodes 1 --nproc_per_node 8 train_rows.py

What to watch

• The optimizer must be built after the model is wrapped, so it sees the sharded parameters.
• The loss you print is local to each rank unless you reduce it.
• Memory use should drop clearly compared with a data parallel run. Check it with the peak memory counters.
• If throughput is poor, look at the gap between compute and communication in a profiler trace.

The newer API in short

With the per-parameter API, the pattern is to call fully_shard on each block and then on the whole model. Each call turns that module's parameters into sharded tensors and installs hooks for gathering and freeing. The rest of the training loop is unchanged. The ideas above all carry over.

A planning walk-through

Suppose you want to train a table transformer with 2 billion parameters on 8 GPUs with 40 GB each, using bf16 and an Adam-style optimizer, with about 16 bytes of model state per parameter.

The plain setup would leave too little room for activations. With sharding, there is plenty of space for a long context of table rows. The numbers are rough. Always confirm with a short run and read the peak memory report.

Checkpoint formats in more detail

A full state dict requires gathering all weights onto one rank, which can need a lot of temporary memory. For big models, save sharded during training and consolidate once at the end on a machine with enough memory, ideally using CPU memory.

Profiling a run

When throughput is lower than expected, look at a profiler trace and answer three questions.

1. Is the GPU idle waiting for communication? If yes, increase prefetching, use larger wrapped units, or use hybrid sharding.
2. Is the data loader the bottleneck? If yes, add workers, pre-generate data, or cache.
3. Are there many tiny kernels? If yes, try graph compilation and fused optimizers.

A simple health number is the share of step time spent in communication that is not overlapped with compute. Keep it small. Tokens per second per GPU should stay close to flat as you add GPUs on the same network; a steep drop signals a communication problem.

Multi-node notes

• Put the sharding group inside a node when the network between nodes is weaker, and replicate across nodes. This is the idea behind hybrid sharding.
• Keep the same software versions on all nodes.
• Test collective communication with a tiny script before a large run.
• Use placeholder names like node1.example.internal in shared config templates and keep real addresses in private settings.

torchrun --nnodes 2 --nproc_per_node 8 \
  --rdzv_backend c10d --rdzv_endpoint node1.example.internal:29500 \
  train.py --config configs/table_pretrain.yaml

Strengths and limitations

Strengths

• Big memory savings. Per-GPU memory for model state falls close to one over the number of GPUs.
• Native to PyTorch. No separate runtime or engine; it works with the usual training loop, optimizers, and profilers.
• Familiar programming model. You still write single-process-looking code. The model code does not need to be rewritten with parallel layers.
• Flexible. Works with mixed precision, offload, activation checkpointing, and hybrid sharding.
• Scales well on fast networks. With overlap and good wrapping, efficiency can stay high across many GPUs.
• Good ecosystem support. Many training frameworks expose it as a switch.

Limitations

• Communication cost. Every layer needs gathers, so slow interconnects hurt a lot.
• Activations are not sharded. Long sequences or large batches can still run out of memory.
• Single layer must fit. The largest wrapped unit, fully gathered, plus activations, has to fit on one GPU. Extremely wide layers may need tensor parallelism.
• Tuning is needed. Wrapping policy, prefetch settings and batch size all affect speed.
• Checkpoint handling is subtle. Full and sharded state dicts have different costs, and resuming on a different world size needs the right tools.
• Debugging is harder. Parameters are not full tensors when you inspect them, which can surprise people who print weights.
• Not a fit for small models. If the model fits easily, plain data parallel is simpler and often faster.

How it compares with alternatives

In practice these are often combined. A large run may use tensor parallelism inside a node, pipeline or hybrid sharding across nodes, and FSDP-style sharding for the data parallel dimension. FSDP also plays well with memory tools such as gradient checkpointing and fused attention kernels like FlashAttention.

Role in tabular foundation models

A tabular foundation model learns from raw tables. Its input is not text or pixels but a grid of rows and columns, where each column may be a number, a category, a date or a free text field. Many cells are missing, and column names may be absent or unhelpful. Such a model is pre-trained on a very large number of tables, often synthetic ones drawn from random generators, plus real ones when available. After pre-training it can do in-context learning: you show it a few labeled rows, and it predicts the rest without any gradient updates. The same backbone can be used for classification, regression, imputation of missing cells, and anomaly detection. Examples of this family include TabPFN and TabICL. For the research view see /model/research, and for the structure see /model/architecture.

FSDP fits into this picture as a training-time scaling tool. It does not change what the model learns. It changes how large the model and the batch of tables can be.

Why tabular pre-training needs sharding

• Attention across rows and columns. Table transformers often attend both over features inside a row and over rows inside a context. The model has to keep many cell tokens per table, so activations grow quickly. FSDP does not shard activations, but it frees memory from parameters and optimizer states so that more room is left for them.
• Large parameter counts. Moving from tens of millions to a billion or more parameters is a common path to better quality. Optimizer states for such models do not fit on a single card, and FSDP spreads them.
• Many small tables per step. Pre-training draws a stream of different tables, each with its own number of rows and columns. A step usually packs several tables into a batch. Sharded state lets each GPU process its own tables with plenty of memory left.
• Synthetic data on the fly. If the pre-training tables are generated during training, the data pipeline is CPU heavy. FSDP keeps the GPU memory budget predictable so the pipeline can be tuned separately.

Mixed types, missing values and no names

A model that handles mixed types usually has separate encoders for numbers and categories, a mask token or indicator for missing values, and column-position or learned column-identity embeddings instead of names. These pieces are small compared with the transformer blocks, so a block-level FSDP wrapping policy still applies cleanly. The encoders can sit in the outer unit, and each transformer block becomes its own unit. When column names are absent, the model relies on statistics and values, so the encoders matter, but they add few parameters.

Continual learning

Continual learning means adding new skills without erasing old ones. For a large table model, this can involve further training on a new domain with a small learning rate, or training small adapter weights. With FSDP, the sharded checkpoint format helps here: you can resume from a sharded checkpoint, change the world size, and continue. If you combine it with LoRA style adapters, only a small set of parameters carries gradients and optimizer states, which cuts memory further. Be aware that frozen and trainable parameters in the same unit were a pain point for the first FSDP generation, and the per-parameter design is better at it.

Imputation, anomaly detection and sector classification

These are downstream uses of the pre-trained backbone, not training-time concerns for FSDP. Still, the way the model is pre-trained shapes them. Imputation can be framed as masked cell prediction: hide some cells and ask the model to fill them. Pre-training with such a masked objective over many tables gives a ready imputer. Anomaly detection can use the model's likelihood or reconstruction error of a row. Semantic or sector classification, such as deciding what kind of business a table describes, needs the model to read hints from column values and headers when they exist. All of these benefit from a bigger backbone and a wider variety of pre-training tables, and FSDP is one way to make a bigger backbone trainable. See /model/benchmarks for how such abilities are measured.

Where FSDP is less needed

Many tabular tools, such as XGBoost and LightGBM, do not need FSDP at all. Even small table transformers that fit in memory are best trained with plain data parallel. FSDP is for the big end of the range. At inference time for in-context models, the cost is dominated by attention over the context rows, so inference-side tools such as vLLM-style serving ideas or fused kernels matter more than sharded training state.

Practical guidance and best practices

1. Start without FSDP. Measure single-GPU memory first. If the model and optimizer fit with room to spare, stay with data parallel.
2. Wrap by block. Use the repeating transformer block as the unit. Check that the policy finds the class you expect.
3. Use mixed precision. A 16-bit compute type with higher precision gradient reduction is a common safe default. See mixed precision.
4. Add activation checkpointing next. Since FSDP does not shard activations, gradient checkpointing is the natural partner.
5. Pick the strategy by network. Fast links across all GPUs favor full sharding. Slow links between nodes favor hybrid sharding.
6. Overlap communication. Keep prefetching on, make batches large enough to hide gathers, and profile before tuning.
7. Use sharded checkpoints for big runs. Save sharded state, and convert to a full file only when you need to share the weights.
8. Initialize large models smartly. Build the model on a meta device or use deferred initialization so the full model never needs to exist on one device.
9. Keep the seeds and data order stable. Each rank needs a distinct data shard, but the same random seed for any step that must match across ranks.
10. Log memory and throughput. Track peak memory, tokens or cells per second, and the share of time spent in communication.
11. Test on a few GPUs first. Bugs in wrapping and checkpointing show up at 2 GPUs just as well as at 200.
12. Pin library versions. The FSDP API has changed across releases, so record the PyTorch version with every checkpoint.

A short tuning loop

import torch

def report(tag):
    gb = torch.cuda.max_memory_allocated() / 1e9
    print(f"{tag}: peak {gb:.2f} GB")

torch.cuda.reset_peak_memory_stats()
# run a few steps here
report("after warmup")

Run the loop with different wrapping sizes and batch sizes. Pick the setting with the best throughput that leaves some memory headroom.

A rollout plan for a new sharded job

1. Run the model on a single GPU with a tiny batch to prove the code is correct.
2. Run two GPUs with plain data parallel and compare the loss curve to the single-GPU run.
3. Switch to FSDP and confirm the loss matches for the first hundred steps.
4. Scale to the full GPU count and record tokens per second.
5. Save a sharded checkpoint, stop the job, resume, and confirm the loss continues smoothly.
6. Start the long run with monitoring in place.

Debugging tips

Working with frozen and trainable parameters

Parameter-efficient methods such as LoRA freeze most weights. The older flat-parameter design had trouble mixing frozen and trainable parameters in one unit, because it flattened them together. The per-parameter design handles this more naturally. If you combine adapters with FSDP, check that your version supports it, wrap by block, and confirm that the optimizer only sees the trainable parameters.

Data loading for table workloads

Synthetic table pre-training creates data on the fly. Keep generation on CPU worker processes, return compact tensors, and pad by task size groups. Each rank should receive different tasks, so seed the generator with the rank. Check that no two ranks produce identical batches, which wastes compute and reduces diversity.

import torch
from torch.utils.data import IterableDataset, DataLoader

class SyntheticTaskStream(IterableDataset):
    def __init__(self, rank, world_size, seed=0):
        self.rank, self.world_size, self.seed = rank, world_size, seed

    def __iter__(self):
        gen = torch.Generator().manual_seed(self.seed + self.rank)
        while True:
            n_rows = int(torch.randint(32, 256, (1,), generator=gen))
            n_cols = int(torch.randint(3, 40, (1,), generator=gen))
            x = torch.randn(n_rows, n_cols, generator=gen)
            y = (x[:, 0] + 0.5 * x[:, 1] > 0).long()
            yield x, y

This is a toy sketch of a task stream. Real generators sample structure, noise, missing values and mixed types much more richly.

Cost thinking

Sharding lets you trade money for time or fit. Adding GPUs raises the hourly bill but can cut total hours if scaling is good. Measure scaling efficiency: take tokens per second per GPU at 8 GPUs and at 32 GPUs. If the 32-GPU figure is above roughly 80 percent of the 8-GPU figure, scaling is healthy for most teams. Below that, look for communication or data bottlenecks before buying more hardware.

Choosing a wrapping unit

The unit you wrap decides how much memory is gathered at once and how many communication calls happen. Use this as a guide.

For table transformers with row and column attention, treat one pair of row and column blocks as a natural unit if the blocks are small. Measure both ways, since the best choice depends on block size and network speed.

A short checklist before launch

• The global batch size and learning-rate schedule are written down.
• The wrapping policy matches the model's repeated block class.
• Mixed precision uses bf16 where available, with reductions in full precision if stability needs it.
• Activation checkpointing is on for long contexts.
• A sharded checkpoint save and resume has been tested.
• Peak memory per GPU is known from a short run.
• Tokens per second is recorded as a baseline for later comparisons.
• Library versions are pinned and identical on all nodes.

Common mistakes

• Building the optimizer before wrapping. The optimizer then points to unsharded parameters and training silently misbehaves or wastes memory.
• Wrapping the whole model as one unit. This gathers every weight at once and removes the memory gain.
• Wrapping every tiny layer. This creates thousands of small collectives and slows training.
• Ignoring activations. Teams see parameter memory fall and then hit an out-of-memory error from activations.
• Saving a full state dict on every rank. Memory spikes and the job crashes at checkpoint time.
• Comparing loss across ranks without reducing. Each rank sees different data, so local loss values differ.
• Mixing frozen and trainable parameters carelessly. In the first generation API this needed special flags.
• Running on a slow network without hybrid sharding. Communication then dominates and scaling stalls.
• Using FSDP for a model that fits. The extra gathers cost time for no benefit.
• Forgetting to set the device before init. Each process must pick its own GPU before creating the model.
• Not calling the sampler's epoch setter. Distributed data samplers need it to reshuffle each epoch.

Frequently asked questions

What does FSDP stand for and what does it do?

FSDP stands for fully sharded data parallel. It is a PyTorch feature that splits a model's parameters, gradients and optimizer states across all GPUs so that each GPU stores only a slice. When a layer needs to run, the full weights of that layer are gathered for a moment and then released. This lets you train models that are much larger than one GPU can hold.

How is FSDP different from DistributedDataParallel?

DistributedDataParallel keeps a full copy of the model on every GPU and only averages gradients between them. FSDP keeps just a shard on each GPU and gathers the rest when needed. The first is simpler and faster when the model fits. The second saves memory at the cost of more communication, and it is the choice when the model or optimizer state is too large for one card.

Is FSDP the same as DeepSpeed ZeRO stage 3?

They share the same core idea of sharding parameters, gradients and optimizer states across data parallel workers. The difference is mostly in packaging. FSDP is built into PyTorch and uses the usual training loop. ZeRO is part of a separate engine with its own configuration files and extra features. The two reach similar memory savings, and teams often pick the one that matches their existing stack.

Does FSDP reduce activation memory?

No. FSDP shards the model state but not the activations saved for the backward pass. For long sequences, many cell tokens, or large batches, activations can become the main memory user. The usual fix is activation checkpointing, which recomputes activations during the backward pass, and sometimes a smaller micro-batch with gradient accumulation.

When should I use hybrid sharding?

Use hybrid sharding when your GPUs inside one machine are connected by a much faster link than the links between machines. You shard the model within each machine, where gathers are cheap, and replicate across machines, where only gradient reduction is needed. This keeps the heavy traffic on the fast links and often improves scaling on multi-node clusters.

Can FSDP be combined with tensor and pipeline parallelism?

Yes. Large training systems often combine several forms of parallelism. Tensor parallelism splits wide layers, pipeline parallelism splits the model into stages, and FSDP-style sharding handles the data parallel dimension. The combination is more complex to configure, so most teams add one method at a time and measure the gain before adding the next.

How do I save and load a checkpoint with FSDP?

You can either gather a full state dict on one rank, which creates a single portable file but needs a lot of memory on that rank, or save a sharded state dict where every rank writes its own slice. Sharded saving is faster and scales better. Distributed checkpoint utilities can then load the shards, even onto a different number of GPUs.

Does FSDP help with training tabular foundation models?

It helps when the model becomes large enough that parameters and optimizer states do not fit on one GPU. A table transformer pre-trained on many synthetic and real tables may reach that size. FSDP does not change what the model learns, but it makes bigger backbones trainable, which can improve in-context learning, imputation and anomaly detection quality.

Is FSDP necessary for small models?

No. If the model, gradients, optimizer states and activations fit comfortably on one GPU, plain data parallel training is simpler and usually faster. FSDP adds communication for every layer, which only pays off when memory is the real limit. Measure first, and move to sharding only when you hit the memory wall.

What is the difference between FSDP and FSDP2?

FSDP2 is the newer API built around a fully_shard function and per-parameter sharding. It keeps the same core algorithm but represents each parameter as an individual sharded tensor rather than flattening many parameters into one buffer. This makes it easier to combine with other features, work with frozen parameters and handle checkpoints. The original class remains in use in many existing projects.

Glossary

• All-gather: A collective operation where every rank ends up with the concatenation of all ranks' pieces.
• All-reduce: A collective operation that sums a tensor across ranks and returns the sum to every rank.
• Activation: An intermediate output of a layer saved for the backward pass.
• Activation checkpointing: Recomputing activations during backward to save memory, also called gradient checkpointing.
• Data parallelism: Training with identical model copies on several devices, each on a different slice of the batch.
• FSDP unit: A wrapped module whose parameters are gathered and freed together.
• Hybrid sharding: Sharding within a group of GPUs and replicating across groups.
• Master weights: A higher precision copy of parameters kept by the optimizer in mixed precision training.
• Offload: Moving tensors to host memory or disk when not in use.
• Optimizer state: Extra per-parameter values kept by the optimizer, such as running moments.
• Rank: The integer id of one process in a distributed job.
• Reduce-scatter: A collective operation that sums tensors across ranks and gives each rank one slice of the result.
• Shard: A slice of a tensor owned by one rank.
• Sharded state dict: A checkpoint format where each rank stores its own shard.
• World size: The total number of ranks in a job.
• ZeRO: A family of memory optimizations that partition model state across data parallel workers.

Summary

FSDP is PyTorch's native way to shard parameters, gradients and optimizer states across GPUs, so each device stores only a slice and gathers full layer weights just in time. It saves large amounts of memory and keeps the normal training loop, at the price of extra communication and with no help for activation memory. It is closely related to DeepSpeed ZeRO and is often combined with tensor parallelism, pipeline parallelism and gradient checkpointing. For tabular foundation models, FSDP matters at the pre-training stage: it lets teams grow the backbone that learns from raw rows and columns, with mixed types and missing values, so that in-context learning, imputation, anomaly detection and classification can improve. Use it when the model does not fit, wrap by transformer block, pair it with mixed precision and activation checkpointing, prefer sharded checkpoints, and measure before you tune. Read more about the models themselves in /model/research, /model/architecture and /model/benchmarks.

Related

It belongs to Training tooling in my stack, under LLMs, retrieval, and evaluation. Others in the same group: