Moving between FSDP And DeepSpeed

🤗 Accelerate offers flexibilty of training frameworks, by integrating two extremely powerful tools for distributed training, namely Pytorch FSDP and Microsoft DeepSpeed. The aim of this tutorial is to draw parallels, as well as to outline potential differences, to empower the user to switch seamlessly between these two frameworks.

Configuring Functionalities

Model tensors are split into different GPUs in an attempt to scale up model sizes; this is termed sharding in FSDP, and partitioning in DeepSpeed. FSDP sharding and DeepSpeed ZeRO (partitioning) stages are configured by --fsdp_sharding_strategy, and --zero_stage, respectively. In particular, FSDP FULL_SHARD maps to DeepSpeed ZeRO stage 3; see this comprehensive mapping between FSDP sharding and DeepSpeed ZeRO settings. The below table summarizes and groups similar settings:

Group	Framework	Configuration	Example	Restrictions (if any)
sharding / partitioning	FSDP DeepSpeed	--fsdp_sharding_strategy --zero_stage	1 (FULL_SHARD) 3
offload	FSDP DeepSpeed	--fsdp_offload_params --offload_param_device --offload_optimizer_device	true cpu cpu	all or nothing
model loading	FSDP DeepSpeed	--fsdp_cpu_ram_efficient_loading --zero3_init_flag	true true	only ZeRO 3
efficient checkpointing	FSDP DeepSpeed	--fsdp_state_dict_type --zero3_save_16bit_model	SHARDED_STATE_DICT true	only ZeRO 3
weights prefetching	FSDP DeepSpeed	--fsdp_forward_prefetch --fsdp_backward_prefetch None	true BACKWARD_PRE
model	FSDP DeepSpeed	--fsdp_auto_wrap_policy --fsdp_transformer_layer_cls_to_wrap None	TRANSFORMER_BASED_WRAP <Layer Class>	Usually not needed Transparent to user.
parameters summoning	FSDP DeepSpeed	--fsdp_use_orig_params None	true	required for torch.compile Transparent to user
parameters syncing	FSDP DeepSpeed	--fsdp_sync_module_states None	true
training	FSDP DeepSpeed	None --gradient_accumulation_steps --gradient_clipping	auto auto	Transparent to user

For detailed descriptions of the above, refer to 🤗 Accelerate launch documentation.

Checkpointing

Do note that while FSDP can be configured via --fsdp_state_dict_type to save either full / sharded checkpoints.

Offloading

FSDP only allows all-or-nothing offload (i.e., either offload parameters, gradients, and optimizer, or keep them all in GPU), but DeepSpeed can offload parameters and optimizer differently. Furthermore, DeepSpeed also supports offloading to NVME.

Prefetching

FSDP allows two prefetching configurations --fsdp_forward_prefetch and --fsdp_backward_prefetch to improve overlap of comms / computation at a cost of extra memory, see FSDP documentation. For DeepSpeed, the prefetching will be turned on when needed, and it turns on depending on certain hyper-params like stage3_param_persistence_threshold, stage3_max_reuse_distance, etc, that can be configured for Zero3; 🤗 accelerate may set these hyper-params automatically if you don’t set those explicitly in the deepspeed config file.

Model Loading

While FSDP require an explicit --fsdp_cpu_ram_efficient_loading true to activate efficient model loading, 🤗 transformers will activate the similar feature whenever DeepSpeed Zero3 is used.

Model

FSDP requires an explicit --fsdp_auto_wrap_policy for the algorithm to decide how to schedule the all-gather and reduce-scatter operations. But for DeepSpeed this is transparent to the user.

Parameters Summoning

FSDP requires an explicit --fsdp_use_orig_params flag if using torch.compile, see the pytorch documenation. For DeepSpeed this is transparent to the user.

Training

Deepspeed requires explicit --gradient_accumulation_steps and --gradient_clipping flags. For FSDP this is transparent to the user.

On Differences in Data Precision Handling

To discuss the how data precision is handled in both FSDP and Deepspeed, it is instructive to first give an overview of how model parameters are handled in these frameworks. Before the model / optimizer parameters are distributed across GPUs, parameter preparation is involved to first “flatten” them to one-dimensional torch.Tensor. The implementation of FSDP / DeepSpeed varies in the respect of the dtype in which these “flattened” parameters are stored, and there are ramifications with regards to how torch.Optimizer allocate their dtypes. The table below outlines the processes for both frameworks; the “Local” column indicates the process occurring at a per-gpu level, therefore any memory overheads by upcasting should be understood to be amortized by the number of gpus used.

Process	Local	Framework	Details
Loading, i.e., AutoModel.from_pretrained(..., torch_dtype=torch_dtype)
Preparation, i.e., creation of “flat params”	✅	FSDP DeepSpeed	created in torch_dtype. disregards torch_dtype, created in float32.
Optimizer initialization	✅	FSDP DeepSpeed	creates parameters in torch_dtype creates parameters in float32
Training Step, i.e, forward, backward, reduction		FSDP DeepSpeed	follows MixedPrecision follows deepspeed_config_file mixed precision settings.
Optimizer (Pre-Step)	✅	FSDP DeepSpeed	upcasting (if any) to torch_dtype upcasted to float32
Optimizer (Actual Step)	✅	FSDP DeepSpeed	occurs in torch_dtype occurs in float32.

To clarify the above table consider the concrete examples below; the optimizer pre- and actual step combined for brevity. With FSDP it is possible to operate in the two modes shown below, but DeepSpeed can only operate in one.

Framework	Model Loading (torch_dtype)	Mixed Precision	Preparation (Local)	Training	Optimizer (Local)
FSDP	bf16	default (none)	bf16	bf16	bf16
FSDP	bf16	bf16	fp32	bf16	fp32
DeepSpeed	bf16	bf16	fp32	bf16	fp32

< > Update on GitHub