Any operator is supported on fp8 tensor?

cokespace2: <p>Since torch now supports <code>fp8_e5m2</code> and <code>fp8_e4m4fnz</code> data type, we could convert our fp16 tensor to fp8 like:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.randn(8, 16, dtype=torch.bfloat16, device='cuda')
a_f8 = a.to(torch.float8_e5m2)
</code></pre>
<p>but it seems like torch doesn’t provide any operator on float8 tensor, like we could not operate</p>
<pre data-code-wrap="python"><code class="lang-python">mask = a &lt; 0.5
</code></pre>
<p>so we have to utilize <code>torchao</code> for such dtype of tensors?</p>

ptrblck: <p>Yes, use <code>torchao</code> for native FP8 runs (or <code>TransformerEngine</code> as an alternative).<br>
CC <a class="mention" href="/u/marksaroufim">@marksaroufim</a> for visibility</p>

cokespace2: <p>Thanks for replying!<br>
BTW, does PyTorch has its plan on supporting such operations? Use <code>torchao</code> is kind a hacky way to do so.</p>

ptrblck: <p>I don’t know if the current derived dtype implementations will land in PyTorch directly. Could you describe what’s “hacky” about using <code>torchao</code>?</p>

marksaroufim: <p><a class="mention" href="/u/cokespace2">@cokespace2</a> mind sharing a bit more? We can certainly more stuff in core but just want to better understand what’s hacky? Is it the specific implementation in AO or is it any out of core implementation?</p>

cokespace2: <p>Well, for me I lean to use native PyTorch so I could have less library dependencies. The float8 dtype, as far as I understand, is more like float16 rather than qint8, while the latter contains attributes like <code>qscheme</code>, <code>qscales</code>.</p>
<p>For example, I believe that multiply 2 qint8 tensors needs a specific operator like <code>qadd</code>, since we need to compute the <code>qscales</code> of these input tensors and the output also come up with a new <code>qscale</code>.</p>
<p>However, the multiply of 2 float8 dtype tensors could be same as multiply 2 float16 tensors, so we could utilize <code>torch.matmul</code> directly.</p>
<p>Currently I have to <code>import torchao</code> and build a <code>Float8Linear</code> module to have a float8 matmul operator. Not mention that some other operations are not supported yet.</p>
<p>Maybe the word <strong>hacky</strong> is inaccurate, but I just thought that we can some basic operations supported in native PyTorch, so we could have more flexibility.</p>
<p>cc <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <aside class="quote no-group" data-username="cokespace2" data-post="6" data-topic="212371">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/cokespace2/48/51244_2.png" class="avatar"> cokespace2:</div>
<blockquote>
<p>However, the multiply of 2 float8 dtype tensors could be same as multiply 2 float16 tensors, so we could utilize <code>torch.matmul</code> directly.</p>
</blockquote>
</aside>
<p>No, that’s not the case as the scaling is needed for a proper training.<br>
Take a look at e.g. <a href="https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html">this TransformerEngine doc</a> explainig the workflow.</p>

cokespace2: <p>I know that during training with float8 we should scale out our float8 tensors, by doing so we can keep the stability of our model training.</p>
<p>But what I like to discuss is that the scaling and the computing can be separated. Just like we train model with float16, we use loss scaling, too. While we could still call many float16 functions in native PyTorch.</p>
<p>In a word, I think that float8 tensor could stand without any scaling factor and it can be used just like what we treat float16 tensors.</p>