FSDP2 and DeepSpeed

The last several pages described an idea. Two pieces of software implement it, and the choice between them is a practical one.

DeepSpeed is the library Microsoft released alongside the ZeRO paper. It is configured with a single JSON file that names the stage and the options, and it has the most complete support for offload to system memory and to storage.

FSDP2 stands for Fully Sharded Data Parallel, version 2. It is PyTorch's own implementation, shipped as part of the framework, and it is the recommended path from PyTorch 2.6 onward. It corresponds to ZeRO stage 3.

How each one is applied

FSDP2 is applied by calling fully_shard() on parts of the model, which replaces the older approach of wrapping the whole model in a class.

from torch.distributed.fsdp import fully_shard

for block in model.layers:
    fully_shard(block)
fully_shard(model)

Sharding the blocks individually is what lets a layer be gathered and released on its own, rather than the whole model at once.

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param":     { "device": "cpu" }
  }
}

The stage number is the same one from the earlier pages, and the two offload lines are the subject of the last one.

The mental model

Before the configuration, the thing being configured. Four GPUs, and what each one permanently holds under each stage.

ZeRO-1
GPU 0: Params(all) + Grads(all) + Optimizer shard 0
GPU 1: Params(all) + Grads(all) + Optimizer shard 1
GPU 2: Params(all) + Grads(all) + Optimizer shard 2
GPU 3: Params(all) + Grads(all) + Optimizer shard 3

ZeRO-2
GPU 0: Params(all) + Grad shard 0 + Optimizer shard 0
GPU 1: Params(all) + Grad shard 1 + Optimizer shard 1
GPU 2: Params(all) + Grad shard 2 + Optimizer shard 2
GPU 3: Params(all) + Grad shard 3 + Optimizer shard 3

ZeRO-3
GPU 0: Param shard 0 + Grad shard 0 + Optimizer shard 0
GPU 1: Param shard 1 + Grad shard 1 + Optimizer shard 1
GPU 2: Param shard 2 + Grad shard 2 + Optimizer shard 2
GPU 3: Param shard 3 + Grad shard 3 + Optimizer shard 3

Read down the column that changes. Params(all) survives stages 1 and 2 and disappears at stage 3, and that is the whole of the difference. At ZeRO-3 no card permanently holds the complete parameter set. Parameters are gathered for the layer being executed and repartitioned immediately afterwards, which is the cycle from the stages page.

It is also where the extra communication comes from. Stages 1 and 2 never need to fetch a weight, because every card already has all of them. Stage 3 does, before every layer, twice per step.

A configuration in full

The earlier snippet showed only the ZeRO block. Here is a complete DeepSpeed file for stage 1, with the batch numbers from the data parallelism page.

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 2,

  "fp16": {
    "enabled": true
  },

  "zero_optimization": {
    "stage": 1
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 0.0001,
      "betas": [0.9, 0.95],
      "eps": 1e-8,
      "weight_decay": 0.1
    }
  }
}

The first three lines are the global batch formula from the data parallelism page, written out: 4 × 8 × 2 = 64, being sequences per GPU, GPUs, and accumulation steps. DeepSpeed checks that the three agree, so setting two of them fixes the third.

The fp16 block is the 16-bit arithmetic from the first page, and the AdamW block is what creates the two running averages that stage 1 is sharding. Every number in this file has appeared earlier in the session.

ZeRO-2, with the options that matter

Stage 2 is the same file with the stage number changed, and it is the first stage where the ZeRO block takes options worth setting deliberately.

{
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 2,

  "fp16": {
    "enabled": true
  },

  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_scatter": true
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 0.0001,
      "betas": [0.9, 0.95],
      "eps": 1e-8,
      "weight_decay": 0.1
    }
  }
}

Those three flags are pages 06 and 07 of this session, written as configuration.

The first two are the reason stage 2 costs 2P rather than more. Turning reduce_scatter off would have each card complete a full all-reduce and then throw away the parts it does not own, which is the waste the collectives page spent its length identifying. Turning overlap_comm off would leave the 2.40 seconds of InfiniBand traffic sitting outside the computation instead of inside it.

Stage 3 is one more integer

And two optional lines if the state should go further out than the GPU.

"zero_optimization": {
  "stage": 3,
  "overlap_comm": true,
  "contiguous_gradients": true
}

"zero_optimization": {
  "stage": 3,
  "offload_optimizer": { "device": "cpu" },
  "offload_param":     { "device": "cpu" }
}

Nothing else in the file moves. That is worth noticing: the difference between holding 447 GiB per card and holding 55.9 is one integer, and the difference between keeping the optimizer state on the GPU and keeping it in system memory is two lines. The reasoning behind the choice took eleven pages; expressing it takes almost nothing.

Why the representation matters

Each parameter is split along its first dimension and represented as a DTensor, which is a tensor that knows which part of itself lives on which GPU. That representation is what allows the compiler to work on a sharded model, which the earlier version did not support.

Choosing between them

The table decides most cases on its own. A run that needs stage 1 or stage 2, or that needs to offload as far as NVMe, wants DeepSpeed. A run that wants stage 3 and wants the rest of the PyTorch toolchain to keep working — the compiler in particular — wants FSDP2.