FP8 support on H100

navmarri: <p>with H100 supporting FP8, is there any plan to support FP8 training in pytorch?<br>
Looks like the only alternative is to use FP8 through <a href="https://github.com/NVIDIA/TransformerEngine" rel="noopener nofollow ugc">nvidia transformer engine</a>. Even transformer engine doesn’t support all nn modules like conv layers for FP8 support.</p>

marksaroufim: <p><a class="mention" href="/u/drisspg">@drisspg</a> has been working on this</p>

navmarri: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> is there a timeline for the feature to be available?</p>

drisspg: <p>Hey <a class="mention" href="/u/navmarri">@navmarri</a> we actually have landed some fp8 primitives already in PyTorch. Specifically we have landed fp8e4m3 and fp8e5m2 datatypes and a private function (with no BC/FC guarantees) for doing scaled matmul on h100 machines.</p>
<p>For a quick example script of how you can use it, check out:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/malfet/7874d96b99670c3da83cbb779ab770c6">
  <header class="source">

      <a href="https://gist.github.com/malfet/7874d96b99670c3da83cbb779ab770c6" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/malfet/7874d96b99670c3da83cbb779ab770c6" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/malfet/7874d96b99670c3da83cbb779ab770c6</a></h4>

  <h5>scale_mm_example.py</h5>
  <pre><code class="Python">import torch
import torch.nn.functional as F

def to_float8(x, dtype=torch.float8_e4m3fn):
    finfo = torch.finfo(dtype)
    # Calculate the scale as dtype max divided by absmax
    scale = finfo.max / x.abs().max().clamp(min=1e-12)
    # scale and clamp the tensor to bring it to
    # the representative range of float8 data type
    # (as default cast is unsaturated)</code></pre>
    This file has been truncated. <a href="https://gist.github.com/malfet/7874d96b99670c3da83cbb779ab770c6" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


navmarri: <p><a class="mention" href="/u/drisspg">@drisspg</a> awesome!<br>
does this work with FP8 precision training using FSDP on H100s?</p>

KohakuBlueleaf: <p>As I know (and also be mentioned in Nvidia/TransformerEngine)<br>
RTX 4000 series (compute capibility 8.9) should also support FP8 computing.<br>
Does pytorch support that too?<br>
or for now it is only for H100 and need further works for sm8.9 cards?</p>

bhack: <p>Is the mixed precision support planned for FP8?</p>
<p>There was an announcement in JAN in the dev forum:</p><aside class="onebox discoursetopic" data-onebox-src="https://dev-discuss.pytorch.org/t/float8-in-pytorch-1-x/1815">
  <header class="source">
      <img src="https://global.discourse-cdn.com/standard10/uploads/pytorch1/optimized/1X/6cd7da56682d360e2c6006ff3e31eb250c5a8675_2_32x32.png" class="site-icon" width="" height="">

      <a href="https://dev-discuss.pytorch.org/t/float8-in-pytorch-1-x/1815" target="_blank" rel="noopener nofollow ugc" title="10:36PM - 16 January 2024">PyTorch Dev Discussions – 16 Jan 24</a>
  </header>

  <article class="onebox-body">
    <img src="https://global.discourse-cdn.com/standard10/uploads/pytorch1/original/2X/0/0ce30c0f4225e400e2fc6e6423739c4e2c78d822.png" class="thumbnail" width="" height="">

<div class="title-wrapper">
  <h3><a href="https://dev-discuss.pytorch.org/t/float8-in-pytorch-1-x/1815" target="_blank" rel="noopener nofollow ugc">Float8 in PyTorch [1/x]</a></h3>
</div>

  <p>This note introduces the experimental support for float8 training in PyTorch. The float8 data formats (as laid out in https://arxiv.org/pdf/2209.05433.pdf) have been added in recent hardware such as NVIDIA Hopper. They have theoretical advantages...</p>

  <p>
    <span class="label1">Reading time: 4 mins 🕑</span>
      <span class="label2">Likes: 7 ❤</span>
  </p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


ptrblck: <p><a href="https://github.com/NVIDIA/TransformerEngine"><code>TransformerEngine</code></a> supports FP8 for a long time now. Alternatively you can also play around with <a href="https://github.com/pytorch-labs/float8_experimental"><code>pytorch-labs/float8_experimental</code></a>.</p>

bhack: <p>Thanks, my point was more related to pytorch AMP support.</p>