fp8 is slower than fp16

[WangQiangItachi]

Hi experts, I tried to use Transformer Engine to do the inference test of llama2 on H800 and found that the speed of fp8 was much slower than fp16, bellow is a small reproduction that only contains Linear. The results show that the speed of fp8.Linear < te.linear < nn.linear, can anyone tell me if this test is reasonable? Thanks.

Docker image: nvcr.io/nvidia/pytorch:23.08-py3

Test code:

import os
import numpy as np
import torch
import transformer_engine.pytorch as te
from time import perf_counter
from torch import nn
from transformer_engine.common import recipe

os.environ['CUDA_VISIBLE_DEVICES'] = "1"

batch_sizes = [1024, 10240, 10240, 10240, 51200]
hidden_sizes = [1024, 1024, 5120, 10240, 10240]
output_sizes = [1024, 1024, 5120, 10240, 10240]

fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)
fp8_recipe.reduce_amax = False
use_te_fp8s = [[False, False], [True, False], [True, True]]

class Model(torch.nn.Module):
    def __init__(self, hidden_size, output_size, use_te):
        super().__init__()
        if not use_te:
            self.linear = nn.Linear(hidden_size, output_size)
        else:
            self.linear = te.Linear(hidden_size, output_size)

    def forward(self, x):
        for i in range(10):
            x = self.linear(x)
        return x

for i in range(len(batch_sizes)):
    for use_te_fp8 in use_te_fp8s:
        use_te = use_te_fp8[0]
        use_fp8 = use_te_fp8[1]
        batch_size = batch_sizes[i]
        hidden_size = hidden_sizes[i]
        output_size = output_sizes[i]
        model = Model(hidden_size, output_size, use_te)
        model.to('cuda')
        input = torch.randn(batch_size, hidden_size).to('cuda')

        times = []
        for j in range(20):
            if j >= 5:
                start = perf_counter()
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
            if j >= 5:
                end = perf_counter()
                times.append(end - start)
        print(f'use_te: {str(use_te):5}, use_fp8: {str(use_fp8):5}, bs: {batch_size:5}, hid_size: {hidden_size:5}, out_size: {output_size:5}, cost_time: {1000*np.mean(times):5f}ms')
    print('=============================================')

Test result:

use_te: False, use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 0.231356ms
use_te: True , use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 0.874557ms
use_te: True , use_fp8: True , bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 3.224556ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 0.238261ms
use_te: True , use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 0.903459ms
use_te: True , use_fp8: True , bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 3.308409ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 0.248760ms
use_te: True , use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 0.917531ms
use_te: True , use_fp8: True , bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 3.316881ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 0.256033ms
use_te: True , use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 0.922597ms
use_te: True , use_fp8: True , bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 11.672506ms
=============================================
use_te: False, use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 0.238444ms
use_te: True , use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 0.908310ms
use_te: True , use_fp8: True , bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 52.545521ms

[ptrendx]

Hi @WangQiangItachi, in order for this test to give you proper results, you need to synchronize with the GPU after computation (since GPU computation is asynchronous with respect to the CPU), by adding torch.cuda.synchronize() after the computation. You can see that this is a problem since the nn.Linear that you got does not increase at all with the increase of the problem size.
Adding synchronization introduces some overheads, so in order to get the accurate timing you should change the way your timing loop works a little bit, by measuring multiple invocations, synchronizing, and then averaging the times. Basically something like this:

        for j in range(5):
            # warmup
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        N = 20
        start = perf_counter()
        for j in range(N):
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        end = perf_counter()
        time = (end - start) / N
        print(f'use_te: {str(use_te):5}, use_fp8: {str(use_fp8):5}, bs: {batch_size:5}, hid_size: {hidden_size:5}, out_size: {output_size:5}, cost_time: {1000*time:5f}ms')

[WangQiangItachi]

Hi @ptrendx, thanks for your suggestion, I modified the code and re-tested it, the results are as follows. It can be seen that when the linear size reaches a certain level, fp8 will have a greater performance improvement. Does this mean that using fp8 for inference with LLM only yields speed benefits when the batch size is relatively large? Because I tested chatglm-6B/llama2-13b on a single H800 card with batch sizes of 8/16/32, and the inference speed with fp8 was still slower than fp16.

use_te: False, use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 0.271229ms
use_te: True , use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 1.009548ms
use_te: True , use_fp8: True , bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 3.510086ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 0.729535ms
use_te: True , use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 1.056175ms
use_te: True , use_fp8: True , bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 3.574828ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 15.609458ms
use_te: True , use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 15.755959ms
use_te: True , use_fp8: True , bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 5.614840ms
=============================================
use_te: False, use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 62.995147ms
use_te: True , use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 59.850148ms
use_te: True , use_fp8: True , bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 20.039033ms
=============================================
use_te: False, use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 320.128337ms
use_te: True , use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 324.473535ms
use_te: True , use_fp8: True , bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 97.020478ms

[ptrendx]

For inference (and especially for H800, since it has lower memory throughput) you should try the trick with caching FP8 weights. The forward method of TE modules has the option is_first_microbatch. When set with the default value of None, it will recompute the FP8 values of the weights every time. In order to avoid that recomputation, you can set it once to True (this will do the casting to FP8 and cache the results) and then do the actual inference with it set to False (this will use the cached values of the weights without casting again). That should limit the overheads and so boost the performance of the FP8 path.

Another tip for boosting the performance is to use the larger APIs vs just Linear layer - APIs like LayerNormLinear hide the cost of casting activations to FP8 in the LayerNorm, which is a double win - LN is faster and you do not have the overhead of the cast. Not sure if you used those in your full inference test.

Finally, the results you posted are still a little strange to me - I would have expected to see te.Linear and nn.Linear to be ~same speed for non-FP8 cases, whereas there seems to be some overhead in the small cases. Could you post the full benchmark code here? I will test it on my side and see where the problem is.

[WangQiangItachi]

Hi @ptrendx, I will try the trick with caching FP8 weights and use LayerNormLinear to replace Linear + LayerNorm, as you mentioned, my current tests simply involve replacing nn.Linear/nn.LayerNorm with te.Linear/te.LayerNorm, as shown below.

def replace_layers_recursively(model, child, name=''):
    if isinstance(child, nn.Linear):
        has_bias = child.bias is not None
        te_module = te.Linear(child.in_features, child.out_features, bias=has_bias)
        te_module.weight.data.copy_(child.weight.data)
        if has_bias:
            te_module.bias.data.copy_(child.bias.data)
        setattr(model, name, te_module)
        child = child.cpu()
    elif isinstance(child, nn.LayerNorm):
        te_module = te.LayerNorm(child.normalized_shape[0], eps=child.eps)
        te_module.weight.data.copy_(child.weight.data)
        te_module.bias.data.copy_(child.bias.data)
        setattr(model, name, te_module)
        child = child.cpu()

    for name1, module in child.named_children():
        replace_layers_recursively(child, module, name=name1)

And the full benchmark code is as follows.

import os
import numpy as np
import torch
import transformer_engine.pytorch as te
from time import perf_counter
from torch import nn
from transformer_engine.common import recipe

os.environ['CUDA_VISIBLE_DEVICES'] = "1"

batch_sizes = [1024, 10240, 10240, 10240, 51200]
hidden_sizes = [1024, 1024, 5120, 10240, 10240]
output_sizes = [1024, 1024, 5120, 10240, 10240]

fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)
fp8_recipe.reduce_amax = False
use_te_fp8s = [[False, False], [True, False], [True, True]]

class Model(torch.nn.Module):
    def __init__(self, hidden_size, output_size, use_te):
        super().__init__()
        if not use_te:
            self.linear = nn.Linear(hidden_size, output_size)
        else:
            self.linear = te.Linear(hidden_size, output_size)

    def forward(self, x):
        for i in range(10):
            x = self.linear(x)
        return x


for i in range(len(batch_sizes)):
    for use_te_fp8 in use_te_fp8s:
        use_te = use_te_fp8[0]
        use_fp8 = use_te_fp8[1]
        batch_size = batch_sizes[i]
        hidden_size = hidden_sizes[i]
        output_size = output_sizes[i]
        model = Model(hidden_size, output_size, use_te)
        model.to('cuda')
        input = torch.randn(batch_size, hidden_size).to('cuda')

        for j in range(5):
            # warmup
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        N = 20
        start = perf_counter()
        for j in range(N):
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        end = perf_counter()
        time = (end - start) / N
        print(f'use_te: {str(use_te):5}, use_fp8: {str(use_fp8):5}, bs: {batch_size:5}, hid_size: {hidden_size:5}, out_size: {output_size:5}, cost_time: {1000*time:5f}ms')
    print('=============================================')

[ptrendx]

Hi @WangQiangItachi , wanted to give you an update. I ran the test on the H100 and the results are similar to what you reported on H800:

use_te: False, use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 0.175725ms                                                                                                                                             
use_te: True , use_fp8: False, bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 0.573835ms                                                                                                                                             
use_te: True , use_fp8: True , bs:  1024, hid_size:  1024, out_size:  1024, cost_time: 2.072206ms                                                                                                                                             
=============================================                                                                                                                                                                                                 
use_te: False, use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 0.716221ms                                                                                                                                             
use_te: True , use_fp8: False, bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 0.666609ms                                                                                                                                             
use_te: True , use_fp8: True , bs: 10240, hid_size:  1024, out_size:  1024, cost_time: 2.116560ms                                                                                                                                             
=============================================                                                                                                                                                                                                 
use_te: False, use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 16.341905ms                                                                                                                                            
use_te: True , use_fp8: False, bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 16.344693ms                                                                                                                                            
use_te: True , use_fp8: True , bs: 10240, hid_size:  5120, out_size:  5120, cost_time: 5.692444ms                                                                                                                                             
=============================================                                                                                                                                                                                                 
use_te: False, use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 66.155760ms                                                                                                                                            
use_te: True , use_fp8: False, bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 62.284888ms                                                                                                                                            
use_te: True , use_fp8: True , bs: 10240, hid_size: 10240, out_size: 10240, cost_time: 20.845784ms                                                                                                                                            
=============================================                                                                                                                                                                                                 
use_te: False, use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 333.154325ms                                                                                                                                           
use_te: True , use_fp8: False, bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 335.041633ms                                                                                                                                           
use_te: True , use_fp8: True , bs: 51200, hid_size: 10240, out_size: 10240, cost_time: 99.745531ms                                                                                                                                            
=============================================                                                                          

I'm still looking into why the BF16 te.Linear is slower than nn.Linear for the smallest case. The initial investigation shows that the issue is on the CPU side - somehow the TE version for that smallest case has higher CPU overhead than the regular call, the GPU kernels launched have the same speed (and since in both TE and non-TE cases that case is CPU-limited, this corresponds to the slower reported time).

[WangQiangItachi]

Hi @ptrendx, thanks for your update. Is this considered a bug in your opinion? If it is, will there be attempts to resolve it in the future? If not, are there any alternative methods to work around this issue?

[ptrendx]

Hi @WangQiangItachi, I'm sorry for the delayed reply. Yes, in general we consider this a bug, although it is unclear how much we can really do to mitigate it. One of the approaches we are considering is making it easier to use CUDA graphs with TE, which should significantly lower the CPU overhead.

[WangQiangItachi]

Hi @ptrendx, thank you for helping to clarify my doubts.

[ftmtk]

I am observing the same issue. Would be good to know if there is an update.

[cslvjt]

Hi @ptrendx, I will try the trick with caching FP8 weights and use LayerNormLinear to replace Linear + LayerNorm, as you mentioned, my current tests simply involve replacing nn.Linear/nn.LayerNorm with te.Linear/te.LayerNorm, as shown below.

def replace_layers_recursively(model, child, name=''):
    if isinstance(child, nn.Linear):
        has_bias = child.bias is not None
        te_module = te.Linear(child.in_features, child.out_features, bias=has_bias)
        te_module.weight.data.copy_(child.weight.data)
        if has_bias:
            te_module.bias.data.copy_(child.bias.data)
        setattr(model, name, te_module)
        child = child.cpu()
    elif isinstance(child, nn.LayerNorm):
        te_module = te.LayerNorm(child.normalized_shape[0], eps=child.eps)
        te_module.weight.data.copy_(child.weight.data)
        te_module.bias.data.copy_(child.bias.data)
        setattr(model, name, te_module)
        child = child.cpu()

    for name1, module in child.named_children():
        replace_layers_recursively(child, module, name=name1)

And the full benchmark code is as follows.

import os
import numpy as np
import torch
import transformer_engine.pytorch as te
from time import perf_counter
from torch import nn
from transformer_engine.common import recipe

os.environ['CUDA_VISIBLE_DEVICES'] = "1"

batch_sizes = [1024, 10240, 10240, 10240, 51200]
hidden_sizes = [1024, 1024, 5120, 10240, 10240]
output_sizes = [1024, 1024, 5120, 10240, 10240]

fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)
fp8_recipe.reduce_amax = False
use_te_fp8s = [[False, False], [True, False], [True, True]]

class Model(torch.nn.Module):
    def __init__(self, hidden_size, output_size, use_te):
        super().__init__()
        if not use_te:
            self.linear = nn.Linear(hidden_size, output_size)
        else:
            self.linear = te.Linear(hidden_size, output_size)

    def forward(self, x):
        for i in range(10):
            x = self.linear(x)
        return x


for i in range(len(batch_sizes)):
    for use_te_fp8 in use_te_fp8s:
        use_te = use_te_fp8[0]
        use_fp8 = use_te_fp8[1]
        batch_size = batch_sizes[i]
        hidden_size = hidden_sizes[i]
        output_size = output_sizes[i]
        model = Model(hidden_size, output_size, use_te)
        model.to('cuda')
        input = torch.randn(batch_size, hidden_size).to('cuda')

        for j in range(5):
            # warmup
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        N = 20
        start = perf_counter()
        for j in range(N):
            with te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
                out = model(input)
        torch.cuda.synchronize()
        end = perf_counter()
        time = (end - start) / N
        print(f'use_te: {str(use_te):5}, use_fp8: {str(use_fp8):5}, bs: {batch_size:5}, hid_size: {hidden_size:5}, out_size: {output_size:5}, cost_time: {1000*time:5f}ms')
    print('=============================================')

I run this code on L20. And the result is as follows.

transformer engine version is 2.2.0

[sbhavani]

Closing this issue as the original question has been addressed:

• The initial benchmark was corrected (missing torch.cuda.synchronize())
• FP8 shows expected speedups for compute-bound workloads (large batch/hidden sizes)
• For small problem sizes, CPU overhead dominates - this is a known limitation, and using CUDA graphs can help

@cslvjt - Your results on L20 with TE 2.2.0 appear to show a different issue. If you're still experiencing unexpected performance, please open a new issue with full environment details (container, CUDA version, TE version). This will help us investigate your specific setup properly.