[Distributed w/ Torchtitan] Enabling Float8 All-Gather in FSDP2

weifengpy: <p>with <a href="https://dev-discuss.pytorch.org/u/awgu" rel="noopener nofollow ugc">Andrew Gu</a>, <a href="https://dev-discuss.pytorch.org/u/wanchaol" rel="noopener nofollow ugc">Wanchao Liang</a>, <a href="https://dev-discuss.pytorch.org/u/drisspg" rel="noopener nofollow ugc">Driss Guessous</a>, <a href="https://dev-discuss.pytorch.org/u/vkuzo" rel="noopener nofollow ugc">Vasiliy Kuznetsov</a>, <a href="https://dev-discuss.pytorch.org/u/bdhirsh" rel="noopener nofollow ugc">Brian Hirsh</a></p>
<h2><a name="p-454012-tldr-1" class="anchor" href="#p-454012-tldr-1"></a>TL;DR</h2>
<ul>
<li>We focus on float8 because it speeds up large GEMMs on H100s and saves network bandwidth with reduced message size.</li>
<li>We enabled float8 all-gather in FSDP2. Readers can find training recipe for Llama3 in <a href="https://github.com/pytorch/torchtitan/blob/main/docs/float8.md" rel="noopener nofollow ugc">TorchTitan </a> and float8 dtype implementation in <a href="https://github.com/pytorch/ao/tree/main/torchao/float8" rel="noopener nofollow ugc">TorchAO/float8 </a>.</li>
<li>We observed 1.50x speedup with float8 compared with bfloat16 with on-par numerics. 20% of the speedup is contributed by float8 all-gather while the rest 80% come from float8 compute. The result is benchmarked by pretraining Llama3-70B using 128 H100s*.</li>
</ul>
<p>* Meta’s H100s are customized on <a href="https://engineering.fb.com/2024/03/12/data-center-engineering/building-metas-genai-infrastructure/" rel="noopener nofollow ugc">Grand Teton </a>. Specifications might be different to public ones.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/a/ea6302c70ab69344aa324f459d5d6c645d2e5a3a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/ea6302c70ab69344aa324f459d5d6c645d2e5a3a" title="Screenshot 2024-09-09 at 11.38.33 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea6302c70ab69344aa324f459d5d6c645d2e5a3a_2_690x214.jpeg" alt="Screenshot 2024-09-09 at 11.38.33 AM" data-base62-sha1="xrtRbAaN8jYJGs7cNdZaY2A0Wym" width="690" height="214" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea6302c70ab69344aa324f459d5d6c645d2e5a3a_2_690x214.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea6302c70ab69344aa324f459d5d6c645d2e5a3a_2_1035x321.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/ea6302c70ab69344aa324f459d5d6c645d2e5a3a_2_1380x428.jpeg 2x" data-dominant-color="DAE4F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.38.33 AM</span><span class="informations">2542×792 152 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<h2><a name="p-454012-float8-training-and-why-float8-all-gather-matters-2" class="anchor" href="#p-454012-float8-training-and-why-float8-all-gather-matters-2"></a>Float8 Training and Why Float8 All-Gather Matters</h2>
<p>Float8 data types are natively supported in NVIDIA H100. Float8 training enablement can be divided into 2 parts: float8 compute and float8 communication.</p>
<p><strong>Float8 compute</strong>: support float8 matrix multiplication with <a href="https://github.com/pytorch/pytorch/blob/618e2c9de42fbac3147e23e258122ebb76e2c041/aten/src/ATen/native/cuda/Blas.cpp#L935" rel="noopener nofollow ugc">torch._scaled_mm </a>. Different from bfloat16, float8 requires both raw tensor and scales to preserve numeric accuracy. Users need to maintain scales in the training loop. There are different scaling strategies including tensor-wise scaling, row/col-wise, group-wise, and block-wise. In this note, we focus on tensor-wise scaling and dynamic scaling <a href="https://arxiv.org/abs/2209.05433" rel="noopener nofollow ugc">ref</a>, where scales are computed from the current high-precision tensor.</p>
<p><strong>Float8 communication (Float8 all-gather)</strong>: with float8 compute, doing all-gathers in float8 is a almost a “free lunch” because we need to cast parameters before or after the all-gather. Casting before the all-gather saves 50% bandwidth (vs bfloat16) at the cost of one all-reduce for AMAX. Float8 can be applied to model weights, activations and gradients. We prioritized float8 weights since they are more stable numerically through the training loop and fit better with low-precision dtypes. We focus on Llama3 models in this note.</p>
<p>Readers can find training recipe for Llama3 in <a href="https://github.com/pytorch/torchtitan/blob/main/docs/float8.md" rel="noopener nofollow ugc">TorchTitan</a> and float8 dtype implementation in <a href="https://github.com/pytorch/ao/tree/main/torchao/float8" rel="noopener nofollow ugc">TorchAO/float8 </a>.</p>
<h2><a name="p-454012-applying-fsdp2-to-llama3-with-float8-weights-3" class="anchor" href="#p-454012-applying-fsdp2-to-llama3-with-float8-weights-3"></a>Applying FSDP2 to Llama3 with Float8 Weights</h2>
<p>Float8 Model (<a href="https://github.com/pytorch/torchtitan/blob/0f70507f1350679428ea64f90bc5a7db17b9c103/train.py#L239C5-L239C27" rel="noopener nofollow ugc">code</a>): PyTorch native float8 requires minimal changes to models. Taking Llama3-8B model as an example, we convert the bfloat16 model to a float8 model by swapping every <code>nn.Linear</code> with a <strong><code>Float8Linear</code></strong> , so that we can perform float8 compute.</p>
<pre><code class="lang-auto">TransformerBlock(
    (attention): Attention(
        (wq/wk/wv/wo): Float8Linear(in=4096, out=4096, bias=False) 
    )
    (feed_forward): FeedForward(
        (w1/w2/w3): Float8Linear(in=4096, out=14336, bias=False)
    )
    (attention_norm / ffn_norm): RMSNorm()
)
</code></pre>
<p><strong>Applying FSDP2</strong> (<a href="https://github.com/pytorch/torchtitan/blob/0f70507f1350679428ea64f90bc5a7db17b9c103/torchtitan/parallelisms/parallelize_llama.py#L501" rel="noopener nofollow ugc">code</a>): The UX of wrapping a float8 model is the same as wrapping a bfloat16 model. To keep track of scales efficiently, we call <code>precompute_float8_dynamic_scale_for_fsdp</code> after the optimizer step, so we can get replicated scales for float8 casting before float8 all-gather.</p>
<pre><code class="lang-auto"># wrapping each TransformerBlock, then root model
# the UX is the same across float8 model and bfloat16 model
for transformer_block in model.layers.values():
    fully_shard(transformer_block)
fully_shard(model)

# training loop
# ...
optimizer.step()
# all-reduce AMAX for Float8Linear.weight
precompute_float8_dynamic_scale_for_fsdp(model)
</code></pre>
<p><strong>FSDP2 extensions for float8 tensor subclass</strong>: We keep FSDP2 UX the same across bfloat16 models and float8 models because we implemented the float8 casting in FSDP2 extensions. The float8 linear module’s weight is a tensor subclass that knows how to cast to float8. We can customize the casting logic before and after all-gather, as shown by the following figure.</p>
<ul>
<li><strong>fsdp_pre_all_gather</strong> (<a href="https://github.com/pytorch-labs/float8_experimental/blob/0aca10aced1c4b3abdf00960d83316732cb08ed1/float8_experimental/fsdp_utils.py#L166" rel="noopener nofollow ugc">code</a>): casting the bfloat16 weight into a float8 weight according to the latest replicated AMAX/scale (requiring all-reduce). Note the bfloat16 weight here is sharded by 1/NGPU. Since we all-reduce to get the replicated AMAX and scale on all ranks, casting the sharded bfloat16 parameters to float8 before all-gather is equivalent to all-gathering bfloat16 parameters and then casting to float8.</li>
<li><strong>fsdp_post_all_gather</strong> (<a href="https://github.com/pytorch-labs/float8_experimental/blob/0aca10aced1c4b3abdf00960d83316732cb08ed1/float8_experimental/fsdp_utils.py#L196" rel="noopener nofollow ugc">code</a>): constructing Float8Tensor from all-gathered float8 data and replicated scale so they are ready for float8 compute in forward and backward.</li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/4/b441bb72d5b7bf3c0718c56ad928039ec7393d3a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b441bb72d5b7bf3c0718c56ad928039ec7393d3a" title="Screenshot 2024-09-09 at 11.41.37 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b441bb72d5b7bf3c0718c56ad928039ec7393d3a_2_690x161.jpeg" alt="Screenshot 2024-09-09 at 11.41.37 AM" data-base62-sha1="pICM4jdXoMYMxz9QDMIygnW4Jom" width="690" height="161" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b441bb72d5b7bf3c0718c56ad928039ec7393d3a_2_690x161.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b441bb72d5b7bf3c0718c56ad928039ec7393d3a_2_1035x241.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/4/b441bb72d5b7bf3c0718c56ad928039ec7393d3a_2_1380x322.jpeg 2x" data-dominant-color="E0E0DD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.41.37 AM</span><span class="informations">1920×449 69.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<h2><a name="p-454012-deep-dive-into-performance-4" class="anchor" href="#p-454012-deep-dive-into-performance-4"></a>Deep Dive Into Performance</h2>
<p>We discuss key optimizations in float8 to reach <strong>1.50x</strong> speed up vs bfloat16.</p>
<p><strong>Float8 Compute + Bfloat16 All-Gather</strong> (1.40x speed up, <a href="https://github.com/pytorch-labs/float8_experimental/blob/0aca10aced1c4b3abdf00960d83316732cb08ed1/float8_experimental/float8_linear.py#L439-L452" rel="noopener nofollow ugc">code</a>): When swapping <code>nn.Linear</code> with <code>Float8Linear</code>, it’s possible to keep the bfloat16 weight as is. We simply treat <code>Float8Linear</code> like a normal <code>nn.Linear</code> and perform bfloat16 all-gather in FSDP2 (stream 22). <code>Float8Linear.forward</code> is responsible for both bfloat16-to-float8 casting and float8 matmul (stream 7). This approach achieved 1.40x speed up and is a strong baseline to showcase the importance of float8 compute. However, it wasted 50% bandwidth to communicate bfloat16 parameters while those parameters will eventually get casted to float8 during forward.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94" title="Screenshot 2024-09-09 at 11.42.49 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94_2_690x113.jpeg" alt="Screenshot 2024-09-09 at 11.42.49 AM" data-base62-sha1="1BGuFjF8X09veDSr4AboXKsCU4I" width="690" height="113" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94_2_690x113.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94_2_1035x169.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b44b993d23fbf9f93f3a1edf832ed7bb9461c94_2_1380x226.jpeg 2x" data-dominant-color="D4D9DB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.42.49 AM</span><span class="informations">1920×316 36.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>Float8 All-Gather with Individual AMAX All-Reduce</strong> (+0.02x on top of 1.40x, <a href="https://github.com/pytorch/torchtitan/blob/0f70507f1350679428ea64f90bc5a7db17b9c103/torchtitan/float8_linear.py#L96" rel="noopener nofollow ugc">code </a>): We perform float8 casting before all-gather to save 50% bandwidth (stream 22). As a result, Float8Linear.forward uses float8 weights directly without the need for casting (stream 7). However, float8 casting requires a global AMAX (argmax of abs(max)) so we need to all-reduce partial AMAX (a scalar) across N ranks (stream 22 and 35). Each float8 parameter requires 1 all-reduce. Those small all-reduces degraded overall performance.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/6/d674e42fb30871a99257ff46695b82cd41d61b7a.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/d674e42fb30871a99257ff46695b82cd41d61b7a" title="Screenshot 2024-09-09 at 11.43.29 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d674e42fb30871a99257ff46695b82cd41d61b7a_2_690x107.jpeg" alt="Screenshot 2024-09-09 at 11.43.29 AM" data-base62-sha1="uBaCgGEDpvnajAJ6K1mk02kLPtw" width="690" height="107" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d674e42fb30871a99257ff46695b82cd41d61b7a_2_690x107.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d674e42fb30871a99257ff46695b82cd41d61b7a_2_1035x160.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d674e42fb30871a99257ff46695b82cd41d61b7a_2_1380x214.jpeg 2x" data-dominant-color="D6CEDD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.43.29 AM</span><span class="informations">1920×300 41.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>Combined AMAX AllReduce</strong> (+0.08x on top of 1.42x, <a href="https://github.com/pytorch/torchtitan/blob/0f70507f1350679428ea64f90bc5a7db17b9c103/torchtitan/float8_linear.py#L107" rel="noopener nofollow ugc">code </a>): We perform single all-reduce for all float8 parameters after the optimizer step. As a result, we avoided small all-reduces inside FSDP hooks (stream 47). We achieved horizontal fusion by calculating AMAX for all float8 parameters at once.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/4/843e07a710b73fe793cf0ae7fbb9544b27ae06ef.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/843e07a710b73fe793cf0ae7fbb9544b27ae06ef" title="Screenshot 2024-09-09 at 11.43.59 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/843e07a710b73fe793cf0ae7fbb9544b27ae06ef_2_690x130.jpeg" alt="Screenshot 2024-09-09 at 11.43.59 AM" data-base62-sha1="iRRVjTRUy5i1mOUXmDOvIduafhB" width="690" height="130" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/843e07a710b73fe793cf0ae7fbb9544b27ae06ef_2_690x130.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/843e07a710b73fe793cf0ae7fbb9544b27ae06ef_2_1035x195.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/4/843e07a710b73fe793cf0ae7fbb9544b27ae06ef_2_1380x260.jpeg 2x" data-dominant-color="DBDCDB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.43.59 AM</span><span class="informations">1920×362 39.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>SM contention between NCCL and Float8 compute</strong> : Depending on NCCL version and GPU total SMs, sometimes there are bubbles in float8 compute (stream 7). Both float8 compute (<code>sm90_xmm</code> ) and float8 all-gather (<code>ncclDevKernel</code> ) fight for SMs. The ideal case is to always prioritize float8 compute for layer k over float8 all-gather for layer k+1. In that case, if NCCL uses less SMs for slower communication or float8 compute uses less SMs. We find it useful to set <a href="https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-max-ctas" rel="noopener nofollow ugc">NCCL_MAX_CTAS </a>) to 16 or 8 during benchmarking to resolve contention.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/9/59105abf08c5952cd263b17cffda6a67183cce3e.png" data-download-href="https://discuss.pytorch.org/uploads/default/59105abf08c5952cd263b17cffda6a67183cce3e" title="Screenshot 2024-09-09 at 11.44.38 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/9/59105abf08c5952cd263b17cffda6a67183cce3e_2_690x73.png" alt="Screenshot 2024-09-09 at 11.44.38 AM" data-base62-sha1="cHTx7naA6g5KEEl5NBltdKzLcqa" width="690" height="73" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/9/59105abf08c5952cd263b17cffda6a67183cce3e_2_690x73.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/9/59105abf08c5952cd263b17cffda6a67183cce3e_2_1035x109.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/9/59105abf08c5952cd263b17cffda6a67183cce3e_2_1380x146.png 2x" data-dominant-color="C8CEDE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-09-09 at 11.44.38 AM</span><span class="informations">2532×268 251 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>Future Work</strong><br>
We are actively exploring the following directions (see more in <a href="https://dev-discuss.pytorch.org/t/meta-pytorch-team-2024-h2-roadmaps/2226" rel="noopener nofollow ugc">PyTorch roadmap </a>)</p>
<p><strong>Float8 in tensor parallel and pipeline parallel</strong>: for tensor parallel (including sequence parallel), we shard module input along sequence dim and would need float8 all-gather for inputs. For pipeline parallel, we are verifying if there are any performance gaps for float8.</p>
<p><strong>Delayed scaling</strong>: Comparing to dynamic scaling, delayed scaling gains perf by deriving AMAX from previous iterations. The cost is potential loss of numerical accuracy. In practice, float8 weights remain stable within adjacent iterations. We want to support delayed scaling to reach full performance.</p>
<p><strong>Row-wise scaling</strong>: Compared to tensor-wise scaling, row-wise scaling preserves better numerical accuracy by having fine-grained scales for each row. The cost is the complexity in the backward, because matrices are transposed from row-wise to column-wise. It requires special treatment for float8 all-gather in FSDP2. This is still a highly exploratory direction.</p>