ViperEkura
76aa4edc9f
- _LinearFp8.backward computed g2.sum(0) unconditionally and dropped it when needs_input_grad[2] was false; now the column-sum only runs when the bias actually requires grad - saves ~327 reduce kernels per train step on bias-free LLMs (215M GQA: end-to-end 1.08x -> 1.13x vs bf16)
472 lines
17 KiB
Python
472 lines
17 KiB
Python
"""FP8 training: scaling recipes, per-tensor state, and aten::linear dispatch.
|
|
|
|
Layered (see ``ops/fp8.py`` for the CUDA interface adapter):
|
|
1. ``ops.fp8`` — the only module touching the pybind.
|
|
2. This module (strategy layer): scaling *recipes* (TE-style delayed scaling
|
|
or dynamic current-amax scaling), per-tensor scales + amax history, and the
|
|
``fp8_autocast`` context manager (like ``torch.autocast``).
|
|
3. aten::linear integration: registers the CUDA + AutogradCUDA impls.
|
|
|
|
Usage::
|
|
|
|
from astrai.extension.fp8 import fp8_autocast
|
|
with fp8_autocast(enabled=True, fp8_format="hybrid"):
|
|
logits = model(input_ids)
|
|
loss.backward() # fp8 backward runs anywhere; fwd captured state on the node
|
|
|
|
Format defaults follow the ecosystem consensus: E4M3 forward / E5M2 backward
|
|
("hybrid"); every operand's scale is a quantization step derived from its amax
|
|
history by the active recipe.
|
|
|
|
The context mirrors ``torch.autocast`` (``autocast_mode.py``): the active
|
|
``(enabled, recipe, fp8_format)`` triple is thread-local (a ``contextvars``
|
|
``ContextVar``, absent outside any region), and the manager is class-based and
|
|
reentrant with nested ``enabled=False`` disabling dispatch inside it. The module
|
|
targets *training*: every step quantizes x/w/g fresh (no weight-cast cache — the
|
|
optimizer bumps the weight version each step, so a torch-style cached_cast would
|
|
miss anyway), and the per-operand scales come from the delayed/dynamic recipe.
|
|
"""
|
|
|
|
import functools
|
|
from contextvars import ContextVar, Token
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Dict, List, Optional
|
|
|
|
import torch
|
|
from torch.library import Library
|
|
|
|
from astrai.extension.ops.fp8 import mm_fp8, quantize
|
|
|
|
# Max representable value per FP8 format (E4M3: 448, E5M2: 57344).
|
|
FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
|
|
|
|
|
|
class FP8Format(str, Enum):
|
|
"""Per-direction FP8 format. HYBRID = E4M3 forward / E5M2 backward."""
|
|
|
|
E4M3 = "e4m3"
|
|
E5M2 = "e5m2"
|
|
HYBRID = "hybrid"
|
|
|
|
def fwd(self) -> str:
|
|
return "e4m3" if self is FP8Format.HYBRID else self.value
|
|
|
|
def bwd(self) -> str:
|
|
return "e5m2" if self is FP8Format.HYBRID else self.value
|
|
|
|
|
|
class FP8Recipe:
|
|
"""Scale-from-amax policy: ``scale = (amax / FP8_MAX[fmt]) / 2^margin``.
|
|
|
|
``scale_from_history`` receives the operand's amax tensor (a ring window for
|
|
delayed scaling, the current amax for dynamic scaling) and returns the
|
|
quantization step. Subclasses set ``history_len`` / ``margin``.
|
|
"""
|
|
|
|
history_len: int = 16
|
|
margin: int = 0
|
|
|
|
def scale_from_history(self, amax: torch.Tensor, fmt: str) -> torch.Tensor:
|
|
peak = amax.max()
|
|
return ((peak / FP8_MAX[fmt]) / (2**self.margin)).clamp_min(1e-12)
|
|
|
|
|
|
@dataclass
|
|
class DelayedScaling(FP8Recipe):
|
|
"""TE-style delayed scaling: max over the amax history window (amax from
|
|
*previous* steps; the window trades responsiveness against stability)."""
|
|
|
|
history_len: int = 16
|
|
margin: int = 0
|
|
|
|
|
|
@dataclass
|
|
class DynamicScaling(FP8Recipe):
|
|
"""Current-amax scaling (torchao DYNAMIC): measure, then quantize. No
|
|
history — the scale is derived from the same-step amax, at an extra pass."""
|
|
|
|
history_len: int = 1
|
|
margin: int = 0
|
|
|
|
|
|
class _ScaleRing:
|
|
"""One operand's delayed-scaling state: a float32 buffer
|
|
``[hist[n] | scale | counter]`` (views). ``update`` folds the amax
|
|
returned by the quantize primitive into ``hist[idx]`` and publishes the
|
|
next scale from the window; ``idx`` advances host-side each step. The
|
|
trailing slot is a legacy counter kept for state-buffer compatibility.
|
|
"""
|
|
|
|
__slots__ = ("recipe", "state", "hist", "scale", "idx", "initialized")
|
|
|
|
def __init__(self, device: torch.device, recipe: FP8Recipe):
|
|
self.recipe = recipe
|
|
n = recipe.history_len
|
|
self.state = torch.zeros(n + 2, device=device, dtype=torch.float32)
|
|
self.hist = self.state[:n]
|
|
self.scale = self.state[n : n + 1]
|
|
self.idx = 0
|
|
self.initialized = False
|
|
|
|
def advance(self) -> None:
|
|
"""Rotate to the next history slot after metadata update."""
|
|
self.idx = (self.idx + 1) % self.hist.numel()
|
|
|
|
def seed(self, t: torch.Tensor, fmt: str) -> None:
|
|
amax = t.abs().amax().to(torch.float32).clamp_min(1e-12)
|
|
self.hist.fill_(amax)
|
|
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
|
|
self.initialized = True
|
|
|
|
def update(self, amax: torch.Tensor, fmt: str) -> None:
|
|
self.hist[self.idx].copy_(amax.reshape(()))
|
|
self.scale.copy_(self.recipe.scale_from_history(self.hist, fmt))
|
|
|
|
|
|
class FP8TensorMeta:
|
|
"""Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
|
|
|
|
DynamicScaling never allocates a meta; it measures the current amax inline.
|
|
"""
|
|
|
|
__slots__ = ("w", "x", "g")
|
|
|
|
def __init__(self, device: torch.device, recipe: FP8Recipe):
|
|
self.w = _ScaleRing(device, recipe)
|
|
self.x = _ScaleRing(device, recipe)
|
|
self.g = _ScaleRing(device, recipe)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ActiveConfig:
|
|
"""The immutable (enabled, recipe, format) triple of one open region."""
|
|
|
|
enabled: bool
|
|
recipe: FP8Recipe
|
|
fp8_format: FP8Format
|
|
|
|
|
|
# Thread-local active configuration (torch's autocast TLS analog): set by
|
|
# fp8_autocast on __enter__, absent outside any region. Autograd engine
|
|
# threads run backwards with their own empty context — fine, since backward
|
|
# only reads state captured on ctx at forward time.
|
|
_active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
|
|
"astrai_fp8_active_config", default=None
|
|
)
|
|
|
|
|
|
class FP8State:
|
|
"""Global fp8 training state: per-tensor metas + out-of-region defaults.
|
|
|
|
The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar`` set
|
|
by ``fp8_autocast``. The properties below read that active config when a
|
|
region is open and the global defaults otherwise; the setters (and
|
|
``fp8_linear_enable``) write the global defaults — the persistent switch
|
|
applying outside any region. The metas registry is shared across threads
|
|
(GIL-protected); fp8 backward runs on autograd engine threads and only
|
|
touches metas captured on ``ctx`` at forward time.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.default_enabled = False
|
|
self.default_recipe: FP8Recipe = DelayedScaling()
|
|
self.default_format: FP8Format = FP8Format.HYBRID
|
|
self._metas: Dict[tuple, FP8TensorMeta] = {}
|
|
|
|
# Active-config views (region config if open, else the defaults).
|
|
@property
|
|
def enabled(self) -> bool:
|
|
cfg = _active_config.get()
|
|
return cfg.enabled if cfg is not None else self.default_enabled
|
|
|
|
@property
|
|
def recipe(self) -> FP8Recipe:
|
|
cfg = _active_config.get()
|
|
return cfg.recipe if cfg is not None else self.default_recipe
|
|
|
|
@property
|
|
def fp8_format(self) -> FP8Format:
|
|
cfg = _active_config.get()
|
|
return cfg.fp8_format if cfg is not None else self.default_format
|
|
|
|
# Persistent (out-of-region) defaults.
|
|
@enabled.setter
|
|
def enabled(self, value: bool) -> None:
|
|
self.default_enabled = bool(value)
|
|
|
|
@recipe.setter
|
|
def recipe(self, value: FP8Recipe) -> None:
|
|
self.default_recipe = value
|
|
|
|
@fp8_format.setter
|
|
def fp8_format(self, value: FP8Format) -> None:
|
|
self.default_format = FP8Format(value)
|
|
|
|
def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
|
|
key = (w.data_ptr(), w.shape, w.dtype)
|
|
meta = self._metas.get(key)
|
|
if meta is None:
|
|
meta = FP8TensorMeta(w.device, self.recipe)
|
|
self._metas[key] = meta
|
|
return meta
|
|
|
|
def reset(self) -> None:
|
|
"""Restore construction defaults (switch, recipe, format) and drop all
|
|
per-weight metas — a full state reset for tests / reconfiguration."""
|
|
self.default_enabled = False
|
|
self.default_recipe = DelayedScaling()
|
|
self.default_format = FP8Format.HYBRID
|
|
self._metas.clear()
|
|
|
|
|
|
# Process-wide singleton; per-thread/per-region state lives in _active_config.
|
|
_state = FP8State()
|
|
|
|
|
|
def fp8_state() -> FP8State:
|
|
return _state
|
|
|
|
|
|
def _active() -> Optional[_ActiveConfig]:
|
|
"""The active config when fp8 dispatch is on, else ``None`` (fast guard).
|
|
|
|
A region config wins (honoring nested ``enabled=False`` regions); with no
|
|
region open this falls back to the persistent global switch
|
|
(``fp8_linear_enable``), so that flag still routes aten::linear to fp8.
|
|
"""
|
|
cfg = _active_config.get()
|
|
if cfg is not None:
|
|
return cfg if cfg.enabled else None
|
|
if _state.default_enabled:
|
|
return _ActiveConfig(True, _state.default_recipe, _state.default_format)
|
|
return None
|
|
|
|
|
|
def _current_config() -> _ActiveConfig:
|
|
"""Like ``_active()`` but always returns a config (disabled regions and
|
|
out-of-region direct calls resolve to the global defaults)."""
|
|
cfg = _active_config.get()
|
|
if cfg is not None:
|
|
return cfg
|
|
return _ActiveConfig(
|
|
_state.default_enabled, _state.default_recipe, _state.default_format
|
|
)
|
|
|
|
|
|
class fp8_autocast:
|
|
"""Autocast-style context: fp8 linear dispatch on this thread.
|
|
|
|
Mirrors ``torch.autocast`` — a class-based, reentrant, nestable context
|
|
over thread-local state::
|
|
|
|
with fp8_autocast(enabled=True, fp8_format="hybrid"):
|
|
logits = model(input_ids) # aten::linear -> fp8 path
|
|
loss.backward() # fp8 backward; state was captured at forward time
|
|
|
|
Nesting follows torch: each ``__enter__`` pushes the new active config, each
|
|
``__exit__`` restores the previous one, and a nested ``enabled=False`` region
|
|
simply disables dispatch inside it. The instance doubles as a decorator.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
enabled: bool = True,
|
|
update_interval: int = 16,
|
|
recipe: Optional[FP8Recipe] = None,
|
|
fp8_format: str = "hybrid",
|
|
margin: int = 0,
|
|
):
|
|
if recipe is None:
|
|
recipe = DelayedScaling(history_len=update_interval, margin=margin)
|
|
self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
|
|
self._tokens: List[Token] = []
|
|
|
|
def __enter__(self) -> "fp8_autocast":
|
|
self._tokens.append(_active_config.set(self._config))
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
|
|
token = self._tokens.pop()
|
|
_active_config.reset(token)
|
|
return False
|
|
|
|
def __call__(self, func):
|
|
@functools.wraps(func)
|
|
def decorate(*args, **kwargs):
|
|
with self:
|
|
return func(*args, **kwargs)
|
|
|
|
return decorate
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Strategy-level forward / backward (called from the aten::linear impl)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _dynamic_scale(t: torch.Tensor, recipe: FP8Recipe, fmt: str) -> torch.Tensor:
|
|
amax = t.abs().amax().to(torch.float32).clamp_min(1e-12)
|
|
return recipe.scale_from_history(amax, fmt)
|
|
|
|
|
|
def _is_fp8(dtype: torch.dtype) -> bool:
|
|
"""A pre-quantized weight takes the GEMM directly (no re-quantize)."""
|
|
return dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
|
|
|
|
|
|
def fp8_linear_forward(
|
|
x: torch.Tensor, w: torch.Tensor, bias=None, cfg: Optional[_ActiveConfig] = None
|
|
):
|
|
"""Scaled fp8 linear forward (called from the aten::linear impl).
|
|
|
|
Composed from the two stateless primitives: quantize x/w with the active
|
|
scales, run the pre-quantized GEMM with the bias fused into its epilogue.
|
|
Delayed scaling folds
|
|
the returned amax into the history ring and publishes the next scale;
|
|
dynamic scaling measures the current amax itself. Training quantizes the
|
|
weight every step (the optimizer bumps its version, so there is no cast
|
|
cache, matching ``cached_cast``-less behavior).
|
|
"""
|
|
state = fp8_state()
|
|
if cfg is None:
|
|
cfg = _current_config()
|
|
fmt = cfg.fp8_format.fwd()
|
|
if isinstance(cfg.recipe, DynamicScaling):
|
|
sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
|
|
sw = _dynamic_scale(w, cfg.recipe, fmt)
|
|
x8, _ = quantize(x, sx.reciprocal(), fmt)
|
|
w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
|
|
# Bias fuses into the GEMM epilogue (fp32 add before the single bf16
|
|
# rounding — one rounding fewer than the separate out + bias pass);
|
|
# None passes through to the kernel's no-bias path.
|
|
out = mm_fp8(
|
|
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
|
|
).reshape(*x.shape[:-1], w.size(0))
|
|
return out, sx, sw
|
|
|
|
meta = state.get_weight_meta(w)
|
|
if not meta.w.initialized:
|
|
meta.w.seed(w, fmt)
|
|
if not meta.x.initialized:
|
|
meta.x.seed(x, fmt)
|
|
sx, sw = meta.x.scale.clone(), meta.w.scale.clone()
|
|
x8, amax_x = quantize(x, sx.reciprocal(), fmt)
|
|
if _is_fp8(w.dtype):
|
|
w8, amax_w = w, None
|
|
else:
|
|
w8, amax_w = quantize(w, sw.reciprocal(), fmt)
|
|
out = mm_fp8(
|
|
x8.reshape(-1, x8.size(-1)), w8, sx * sw, trans_b=True, bias=bias
|
|
).reshape(*x.shape[:-1], w.size(0))
|
|
meta.x.update(amax_x, fmt)
|
|
if amax_w is not None:
|
|
meta.w.update(amax_w, fmt)
|
|
meta.x.advance()
|
|
if amax_w is not None:
|
|
meta.w.advance()
|
|
return out, sx, sw
|
|
|
|
|
|
class _LinearFp8(torch.autograd.Function):
|
|
"""The fp8 linear forward/backward pair (standard Function style).
|
|
|
|
The forward runs inside ``fp8_autocast`` and captures the active
|
|
fmt/recipe/meta on ``ctx``; the backward reads only that captured state, so
|
|
``loss.backward()`` may run after the context exits. The gradient is
|
|
quantized once (E5M2 in hybrid) and both dX/dW GEMMs share it; the output
|
|
masks come from ``needs_input_grad``.
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, w, bias):
|
|
cfg = _current_config()
|
|
out, sx, sw = fp8_linear_forward(x, w, bias, cfg)
|
|
ctx.save_for_backward(x, w, sx, sw)
|
|
ctx.fmt_bwd = cfg.fp8_format.bwd()
|
|
ctx.recipe = cfg.recipe
|
|
ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
|
|
ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w)
|
|
return out
|
|
|
|
@staticmethod
|
|
@torch.autograd.function.once_differentiable
|
|
def backward(ctx, g):
|
|
x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors
|
|
fmt = ctx.fmt_bwd
|
|
# Flatten leading dims (the forward GEMMs ran on [-1, N] / [-1, K]
|
|
# views; the kernels only accept 2D operands).
|
|
g2 = g.reshape(-1, g.size(-1))
|
|
if ctx.is_dynamic:
|
|
sg = _dynamic_scale(g2, ctx.recipe, fmt)
|
|
sw = _dynamic_scale(w, ctx.recipe, fmt)
|
|
sx = _dynamic_scale(x, ctx.recipe, fmt)
|
|
else:
|
|
meta = ctx.meta
|
|
if not meta.g.initialized:
|
|
meta.g.seed(g2, fmt)
|
|
sg = meta.g.scale.clone()
|
|
sw, sx = _sw_fwd, _sx_fwd
|
|
g8, amax_g = quantize(g2, sg.reciprocal(), fmt)
|
|
x8, _ = quantize(x.reshape(-1, x.size(-1)), sx.reciprocal(), fmt)
|
|
w8 = w if _is_fp8(w.dtype) else quantize(w, sw.reciprocal(), fmt)[0]
|
|
grad_x = mm_fp8(g8, w8, sg * sw).reshape(x.shape) # g8[m,n] @ w8[n,k]
|
|
grad_w = mm_fp8(g8, x8, sg * sx, trans_a=True) # g8.T @ x8
|
|
# bias-free linears must not pay the column-sum
|
|
# reduce: g2.sum(0) is another full read of the gradient.
|
|
grad_b = g2.sum(0).to(torch.bfloat16) if ctx.needs_input_grad[2] else None
|
|
if not ctx.is_dynamic:
|
|
meta.g.update(amax_g, fmt)
|
|
meta.g.advance()
|
|
return grad_x, grad_w, grad_b
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# aten::linear integration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def fp8_linear_enable(enabled: bool = True) -> None:
|
|
"""Toggle fp8 dispatch for aten::linear globally (the out-of-region default;
|
|
``fp8_autocast`` regions override it thread-locally)."""
|
|
fp8_state().default_enabled = enabled
|
|
|
|
|
|
def fp8_linear_enabled() -> bool:
|
|
"""Whether fp8 dispatch is active right now (region config or global)."""
|
|
return _active() is not None
|
|
|
|
|
|
def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
|
|
"""Shape guard for the fp8 path. Unlike a strict 16-alignment requirement,
|
|
the kernels handle unaligned M/N via boundary checks (slower but correct) —
|
|
so no whole-call bf16 fallback for small decode batches. Only the K-dimension
|
|
contraction must match and the weight must be 2D."""
|
|
return x.dim() >= 2 and w.dim() == 2 and x.size(-1) == w.size(1)
|
|
|
|
|
|
def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
|
|
if (
|
|
_active() is not None
|
|
and x.dtype is torch.bfloat16
|
|
and w.dtype is torch.bfloat16
|
|
and _fp8_supported(x, w)
|
|
):
|
|
return _LinearFp8.apply(x, w, bias)
|
|
return torch.ops.aten.linear.default.redispatch(
|
|
torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd),
|
|
x,
|
|
w,
|
|
bias,
|
|
)
|
|
|
|
|
|
_lib = Library("aten", "IMPL", "CUDA")
|
|
_lib.impl("linear", _linear_cuda_impl)
|
|
# Also replace torch's generated linear autograd formula (which would call
|
|
# aten::linear_backward after the fp8_autocast region exits). The fp8 backward
|
|
# is owned by _LinearFp8 with state captured at forward time, so loss.backward()
|
|
# works wherever it is called; the CUDA registration still covers inference_mode.
|
|
_lib_autograd = Library("aten", "IMPL", "AutogradCUDA")
|
|
_lib_autograd.impl("linear", _linear_cuda_impl)
|