Skip to content

GitLab

Unverified Commit 379c1ee3 authored by Quentin Anthony's avatar Quentin Anthony Committed by GitHub
Browse files

Update example to use new TE_DType path (#660) 


Signed-off-by: default avatarQuentin Anthony <qganthony@yahoo.com>
parent 91d52ac7
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
docs/examples/quickstart_utils.py
View file @ 379c1ee3
...@@ -205,8 +205,9 @@ def share_parameters_with_transformerlayer_te_model(te_model, basic_model): ...@@ -205,8 +205,9 @@ def share_parameters_with_transformerlayer_te_model(te_model, basic_model):
def cast_to_representable(inp, scale = 1., fp8_format='e4m3'): def cast_to_representable(inp, scale = 1., fp8_format='e4m3'):
import transformer_engine.pytorch.cpp_extensions as texcpp import transformer_engine.pytorch.cpp_extensions as texcpp
import transformer_engine_extensions as tex import transformer_engine_extensions as tex
from transformer_engine.pytorch.constants import TE_DType
fp8_type = tex.DType.kFloat8E4M3 if fp8_format == 'e4m3' else tex.DType.kFloat8E5M2 fp8_type = tex.DType.kFloat8E4M3 if fp8_format == 'e4m3' else tex.DType.kFloat8E5M2
input_type = texcpp.TE_DType[inp.dtype] input_type = TE_DType[inp.dtype]
meta = tex.FP8TensorMeta() meta = tex.FP8TensorMeta()
meta.scale = torch.ones(1,dtype=torch.float32, device="cuda") * scale meta.scale = torch.ones(1,dtype=torch.float32, device="cuda") * scale
meta.scale_inv = torch.ones(1, dtype=torch.float32, device="cuda") / scale meta.scale_inv = torch.ones(1, dtype=torch.float32, device="cuda") / scale
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment