Hugging Face Accelerate has made significant strides in harmonizing two popular implementations of the ZeRO Redundancy Optimizer (Zero) algorithm: DeepSpeed and PyTorch's FSDP. The company has upstreamed a precision-related change and created a concept guide to enable users to seamlessly switch between these backends, which are exposed through Accelerate for training and tuning models.

What Happened

A recent experiment by Hugging Face revealed that running a training pipeline with DeepSpeed and PyTorch FSDP resulted in differing outcomes. The specific model used was the Mistral-7B base, loaded in half-precision (bfloat16). While the DeepSpeed loss had converged well, the FSDP loss did not decrease, as seen in Figure 1.

The team hypothesized that the learning rate may need scaling by the number of GPUs and bumped up the learning rate by 4x since they were using 4 GPUs. This led to a desired behavior, where the FSDP loss began decreasing, as shown in Figure 2. However, when trying a different learning rate (1e-5) without scaling, similar loss and gradient norm characteristics were observed for both frameworks, as seen in Figure 3.

Background and Context

The ZeRO Redundancy Optimizer (Zero) algorithm is a popular implementation of the ZeRO optimizer, which is used for distributed training of large-scale models. DeepSpeed and PyTorch's FSDP are two different implementations of this algorithm, each with its own strengths and weaknesses. Hugging Face Accelerate exposes both these frameworks for users to train and tune their models.

Inside the DeepSpeed codebase, specifically in the implementation of DeepSpeedZeroOptimizer_Stage3, it was noticed that the trainable_param_groups pass through an internal _setup_for_real_optimizer function call, which calls another function called _create_fp32_partitions. This upcasting to full precision meant that the optimizer could converge at learning rates that it would not converge in lower precision.

Why It Matters

The differences between DeepSpeed and FSDP have significant implications for the industry. The ability to seamlessly switch between these backends is crucial for users who need to train and tune their models efficiently. The harmonization of these two frameworks will enable users to take advantage of the strengths of each implementation, leading to better performance and faster training times.

The precision-related change upstreamed by Hugging Face Accelerate allows FSDP to operate in two modes: a "mixed-precision" mode like DeepSpeed's counterpart and a low precision mode for memory-constrained scenarios. This will enable users to choose the most suitable mode based on their specific needs, leading to improved performance and efficiency.

What Comes Next

Hugging Face Accelerate plans to follow up with throughput comparisons at scale and techniques to better utilize GPUs for tuning and alignment jobs while maintaining model quality. The company will also continue to work on harmonizing other frameworks and improving the overall user experience.

Key Facts

  • Hugging Face Accelerate has upstreamed a precision-related change to enable seamless switching between DeepSpeed and FSDP.
  • The change allows FSDP to operate in two modes: mixed-precision and low precision.
  • DeepSpeed performs upcasting internally, keeping master weights in full precision by design.
  • FSDP uses different dtypes for "flattened" parameters, which has ramifications for PyTorch optimizers.
  • Hugging Face Accelerate plans to follow up with throughput comparisons at scale and techniques to improve GPU utilization.