Hugging Face Accelerate: Harmonizing DeepSpeed and FSDP Precision

Hugging Face Accelerate has introduced updates to align the precision handling of PyTorch Fully Sharded Data Parallel (FSDP) and Microsoft DeepSpeed. This change, integrated into the Accelerate 0.30.0 release, allows users to switch between these two ZeRO Redundancy Optimizer implementations without encountering the convergence discrepancies previously caused by differing internal precision defaults.

Precision Discrepancies and Convergence

Differences in how DeepSpeed and FSDP handle parameter precision can lead to divergent training results. In tests using a Mistral-7B base model loaded in bfloat16, DeepSpeed showed stable convergence while FSDP failed to decrease loss unless the learning rate was manually scaled by the number of GPUs or lowered to 1e-5.

This behavior is attributed to internal upcasting. DeepSpeed's DeepSpeedZeroOptimizer_Stage3 automatically upcasts trainable parameter groups to float32 via _create_fp32_partitions, maintaining master weights in full precision by design. This allows the optimizer to converge at learning rates that would be unstable in lower precision.

Comparison of Precision Workflows

DeepSpeed and FSDP differ fundamentally in how they manage "flattened" parameters and optimizer initialization:

Process FSDP DeepSpeed
Preparation Utilizes torch_dtype Disregards torch_dtype; creates in float32
Optimizer Initialization Creates parameters in torch_dtype Creates parameters in float32
Optimizer (Pre-step) Upcasting (if any) to torch_dtype Upcasts everything to float32
Optimizer (Actual step) Occurs in torch_dtype Occurs in float32

While DeepSpeed's native upcasting ensures convergence, it can increase memory consumption by 2x, which is significant when training on a small number of GPUs. Conversely, the torch-native FSDP implementation does not force upcasting, providing more flexibility for memory-constrained scenarios by allowing optimizers to operate in low precision.

New FSDP Modes in Accelerate 0.30.0

To harmonize these frameworks, Hugging Face Accelerate now supports two distinct FSDP modes to match the user's priority between convergence stability and memory efficiency:

  1. Mixed-Precision Mode: Aligns with DeepSpeed by upcasting parameters to fp32 during preparation and optimizer steps, while maintaining bf16 for training.
  2. Memory-Constrained Mode: Operates entirely in low precision (bf16) across preparation, training, and optimizer steps.

Framework Comparison Table

Framework Model Loading Mixed Precision Preparation Training Optimizer
FSDP (Memory-Constrained) bf16 None bf16 bf16 bf16
FSDP (Mixed Precision) bf16 bf16 fp32 bf16 fp32
DeepSpeed bf16 bf16 fp32 bf16 fp32

Throughput Performance

Benchmarking with the IBM Granite 7B model on four A100 GPUs demonstrates that FSDP (in aligned mode) and DeepSpeed (Zero3) provide nearly identical throughput performance:

Framework Tokens / sec / device Step time (s) Model Flops Utilization (MFU)
FSDP (Aligned) 3158.7 10.4 0.41
DeepSpeed 3094.5 10.6 0.40

Migration and Configuration

Hugging Face has released a concept guide to assist users in migrating between FSDP and DeepSpeed. Switching is primarily handled via the Accelerate config file or through the DeepSpeed and FSDP plugin classes. Key considerations for migration include:

  • Sharding Strategies: Achieving equivalent sharding across frameworks.
  • Model Loading: Implementing efficient loading patterns.
  • Weight Prefetching: Managing how weights are moved between GPUs.
  • Checkpointing: Handling differences in how each framework saves and loads model states.

Sources