FP8 for norm inputs and residuals?

[cbcase]

Question for you guys: as best I can tell, there is no support at present for keeping activations in fp8 between the "output" matmul (of either an attention block or MLP block) and the next norm (layernorm or rmsnorm). The missing features to make that work are:

• FP8 residual input/output (along with the necessary scales) to fp8 gemm
• FP8 input (with scale) to {rms,layer}norm

Is this something you have considered, or is there a more fundamental limitation (kernel-wise or accuracy-wise) that means you always want to keep the residual path in a 16-bit format. (Fwiw, my own interest here is for inference.)

Thanks,
Carl

[timmoon10]

Matmuls are ideal for FP8 compute since they can take advantage of Tensor Cores and they're less sensitive to quantization error. While other operations might benefit (especially from reduced memory usage), we haven't prioritized them and performed the necessary convergence testing to trust them.

One other constraint is that TE modules usually treat FP8 activations as internal state, so it's cumbersome to make them handle complicated networks. For more flexibility, you may want to look into the experimental operation-based API (see transformer_engine/pytorch/ops):

fc1 = te.ops.Sequential(te.ops.Linear(...), te.ops.Quantize())
fc2 = te.ops.Sequential(te.ops.LayerNorm(...), te.ops.Linear(...), te.ops.Quantize())

with te.fp8_autocast():
    y = fc1(x)  # y is a Float8Tensor
    z = fc2(y) + y

This could implement your desired behavior once we have LayerNorm kernels that can accept FP8 input. However, right now it would cast FP8 to higher precision before the LayerNorm and the addition.