Topic 5: What is FSDP and YaFSDP?

This article we open for everybody but if you want to support us →

• It might not always match local training results exactly, especially with optimizers that rely on unsharded values or specific tensor structures.

• Handling shared parameters can be tricky, and if not done right, it can cause errors and use up more memory than necessary.

• Even with optimizations, there can still be a lot of communication overhead, particularly with larger models or more GPUs, which requires careful strategy adjustments.

• Setting up FSDP can be complicated, especially for large models, sometimes needing other methods that come with their own challenges.

• Mixing FSDP with other parallelism methods like pipeline or tensor parallelism can be tough and needs careful setup to avoid too much overhead.

• Additionally, because it integrates deeply with PyTorch, FSDP can be harder to debug and troubleshoot compared to simpler data parallelism methods.

As the authors of YaFSDP approach wrote: “Despite all these advantages, there are also issues that we faced:

1. FSDP dynamically allocates memory for layers and sometimes requires much more memory than is actually necessary.

2. During backward passes, we came across a phenomenon that we called the “give-way effect”.

The first line here is the computation stream, and the other lines represent communication streams.

So what’s happening in the profile? Before the reduce_scatter operation (blue), there are many preparatory computations (small operations under the communications). The small computations run in parallel with the main computation stream, severely slowing down communications. This results in large gaps between communications, and consequently, the same gaps occur in the computation stream.”

Here comes YaFSDP

Yet Another Fully Sharded Data Parallel (YaFSDP) YaFSDP was developed and open-sourced by Yandex in May 2024. It improves upon FSDP by offering more efficient memory management, reducing redundant computations, and optimizing communication and synchronization during the training of large language models.

Here’s how YaFSDP works

1. Layer Sharding: Instead of sharding individual parameters, YaFSDP shards entire layers. This approach maintains efficient communication and reduces redundancy. Each GPU handles a different shard of the model, minimizing memory usage.

1. Buffer Pre-allocation: YaFSDP pre-allocates buffers for all necessary data, ensuring that memory management by the Torch allocator does not introduce inefficiencies. This method uses two buffers for intermediate weights and gradients, alternating between odd and even layers.

Memory Consumption Optimization

YaFSDP significantly reduces memory consumption by:

1. Efficient Buffer Use: Buffers store intermediate values and consume a constant amount of memory.

2. Activation Checkpointing: stores only essential activations during the forward pass and recomputes them in the backward pass, significantly reducing memory usage. For instance, training a Llama 2 70B model with a batch size of 8192 tokens can reduce activation storage from over 110 GB to 5 GB. However, this technique adds computational overhead, which YaFSDP can mitigate by optimizing memory usage and avoiding activation checkpointing for some layers.

3. Sharding Weights, Gradients, and Optimizer States: These components' memory consumption tends to approach zero as the number of processes increases, minimizing duplication.

Communication Optimization

YaFSDP enhances communication efficiency by:

1. Overlapping Communication with Computation: Using CUDA streams, YaFSDP manages concurrent computations and communications effectively. Two streams are used: one for computation and one for communication, synchronized by events to ensure correct operation order.

2. Reducing Communication Overhead: By ensuring that data transfers occur only when necessary and using techniques to minimize redundant operations, YaFSDP improves overall efficiency.

What YaFSDP is especially good at? Performance Gains

YaFSDP demonstrates significant performance improvements. For a model with 70 billion parameters, it can save the resources of approximately 150 GPUs, translating to monthly cost savings of $0.5 to $1.5 million. Training time is reduced by up to 26% compared to existing methods like FSDP.

When compared to FSDP, the final speedup shown by YaFSDP on Llama 2 and Llama 3 demonstrates significant improvements in training efficiency.

YaFSDP can be used in conjunction with huggingface workflows and is up to 25% faster compared to FSDP.

Original resources

• Paper: PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

• GitHub: YaFSDP

• Blog post: Yandex develops and open-sources YaFSDP — a tool for faster LLM training and optimized GPU consumption

Thank you for reading! Share this article with three friends and get a 1-month subscription free! 🤍