FSDP › Wrapping Policies

A key decision in FSDP setup is how to group model submodules into FSDP units. Each FSDP unit is the granularity at which all-gather and reduce-scatter operations occur: finer granularity means smaller all-gather buffers and lower peak memory, but more frequent communication and higher overhead from kernel launches.

Getting the wrapping policy right matters more than many practitioners expect. A policy that is too fine-grained can reduce throughput by 2-3x compared to an optimal policy, because GPU kernels have significant launch overhead that becomes dominant when individual operations are small. A policy that is too coarse can cause out-of-memory errors because the all-gather buffer for a single large unit might not fit.

The transformer_auto_wrap_policy function in PyTorch's FSDP utilities wraps each transformer layer (self-attention block, FFN block, etc.) as a separate FSDP unit. This is the recommended starting point for transformer models.

The auto wrap policy accepts a set of module classes to wrap. For a standard transformer, you would typically pass the class corresponding to a single transformer block, for example GPTNeoXLayer for GPT-NeoX or LlamaDecoderLayer for Llama. The embedding and output head layers are wrapped at the top level (as part of the root FSDP module) unless you explicitly specify otherwise.

The reason this wrapping strategy works well for transformers is that each transformer layer is approximately equal in size (assuming the same hidden dimension and FFN expansion ratio across all layers), so all FSDP units have similar all-gather buffer sizes. Uniform unit sizes lead to uniform communication times, which maps well to overlapping communication with computation: while one layer runs its forward pass, FSDP can prefetch (all-gather) the next layer's parameters.

The size_based_auto_wrap_policy wraps modules based on their parameter count. Modules with fewer than a threshold number of parameters are merged with their parent, reducing the total number of FSDP units and thus the communication frequency.

This policy is useful for models that do not follow the uniform-transformer-layer pattern, such as encoder-decoder models with asymmetric components or models with large embedding tables followed by small classification heads.

For full control, you can manually wrap specific submodules by calling FSDP() on them before wrapping the parent model. This is useful when different parts of the model require different sharding strategies or mixed precision settings. For example, you might want the main transformer layers to use FULL_SHARD with bfloat16, but keep the embedding table in float32 with NO_SHARD because it is accessed sparsely and does not benefit from sharding.

A subtle but important performance feature of FSDP is prefetching: the ability to issue the all-gather for the next FSDP unit while the current unit is still computing. PyTorch FSDP implements this automatically for the forward pass by default, overlapping the communication for layer $k+1$ with the computation for layer $k$. This overlap is only possible when the wrapping granularity is at the layer level (one FSDP unit per layer), not at coarser granularity where there is no "next unit" to prefetch.

The backward_prefetch parameter controls whether the same overlap happens during the backward pass. By default, PyTorch prefetches during backward as well, which can significantly improve throughput on hardware where communication and compute can run concurrently.