FP8 Training: H100 Native, RTX 4090'da Prematur — Transformer Engine Internals

# === Transformer Engine FP8 (H100 native) ===
import torch
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling
 
# FP8 recipe — TE auto-handles scaling factors
fp8_recipe = DelayedScaling(
    fp8_format=Format.HYBRID,           # e4m3 forward, e5m2 backward
    amax_history_len=16,                # scaling factor history
    amax_compute_algo="max",
)
 
# Model — TE layer'ları drop-in replace
class Llama_FP8(nn.Module):
    def __init__(self, hidden, ffn_hidden):
        super().__init__()
        self.attn_qkv = te.Linear(hidden, 3*hidden, bias=False)
        self.attn_o = te.Linear(hidden, hidden, bias=False)
        self.ffn_gate = te.Linear(hidden, ffn_hidden, bias=False)
        self.ffn_up = te.Linear(hidden, ffn_hidden, bias=False)
        self.ffn_down = te.Linear(ffn_hidden, hidden, bias=False)
 
# Training loop
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    output = model(input)
    loss = compute_loss(output)
loss.backward()
optimizer.step()
 
# H100'de bf16'ya göre +%30-40 throughput
# RTX 4090'da bf16'ya göre +%10-15 (kernel'lar Hopper-optimize)