库 | TransformerEngine

xwhuan

应用介绍

Transformer Engine (TE) 是一个用于加速 NVIDIA GPU 上 Transformer 模型的库，支持在 Hopper、Ada 和 Blackwell GPU 上使用 8 位浮点数 (FP8) 精度，从而在训练和推理过程中提供更好的性能和更低的内存占用。TE 提供了一系列高度优化的构建块，适用于流行的 Transformer 架构，并提供类似自动混合精度的 API，可以与框架特定的代码无缝集成。TE 还包括一个与框架无关的 C++ API，可以与其他深度学习库集成，从而为 Transformer 模型启用 FP8 支持。

使用指南

官方推荐的使用方法是singularity镜像

镜像路径

/opt/app/sif/TransformerEngine.sif

使用方法

申请gpu资源
```
salloc -p <gpu-partition> --gres=gpu:1
```

登陆到申请的节点后，执行命令

/opt/app/singularity/bin/singularity  exec  /opt/app/sif/TransformerEngine.sif  bash

准备demo脚本 test.py 到个人目录

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048

# Initialize model and inputs.
model = te.Linear(in_features, out_features, bias=True)
inp = torch.randn(hidden_size, in_features, device="cuda")

# Create an FP8 recipe (not used if FP8 is disabled).
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

# Disable FP8 and use FP16/FP32
with te.fp8_autocast(enabled=False):
    out = model(inp)

loss = out.sum()
loss.backward()
# Display the results
print("Output shape: ", out.shape)
print("Output: ", out)  # Display the actual output tensor
print("Loss: ", loss)  # Display the loss value

运行测试脚本
```
which python
python $HOME/test.py 
```
可以看到测试结果

参考链接

Transformer Engine