ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: collapse fp8 recipe hierarchy and state property layers

Browse Source 
- Merge DelayedScaling/DynamicScaling and the abstract FP8Recipe base into one FP8Recipe dataclass with a dynamic flag; dispatch now reads cfg.recipe.dynamic instead of isinstance checks
- Drop the _ActiveOrDefault descriptor and the FP8State property views; the persistent defaults are plain default_* attributes and get_weight_meta takes the active recipe explicitly
- Convert FP8TensorMeta to a NamedTuple of the three per-operand rings
- Update tests to the new API; the autocast context test now asserts _active_config push/restore directly
This commit is contained in:
ViperEkura
2026-08-31 14:24:51 +08:00
parent e3c3e28a11
commit 432dfec3c2
2 changed files with 63 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+34 -76
View File
@@ -31,7 +31,7 @@ import functools
from contextvars import ContextVar, Token
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional
from typing import Dict, List, NamedTuple, Optional
import torch
from torch.library import Library
@@ -56,40 +56,27 @@ class FP8Format(str, Enum):
return "e5m2" if self is FP8Format.HYBRID else self.value
@dataclass
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``.
``dynamic=False`` (default) is TE-style delayed scaling: max over the
amax history window (amax from *previous* steps; the window trades
responsiveness against stability). ``dynamic=True`` is current-amax
scaling (torchao DYNAMIC): measure, then quantize — no history, at an
extra pass. ``scale_from_history`` receives the operand's amax tensor
(a ring window / the current amax) and returns the quantization step.
"""
history_len: int = 16
margin: int = 0
dynamic: bool = False
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 | legacy | amax | done]`` (views). The quantize
@@ -131,18 +118,15 @@ class _ScaleRing:
}
class FP8TensorMeta:
"""Per-weight delayed-scaling state for ``w``, ``x`` and ``g``.
class FP8TensorMeta(NamedTuple):
"""Per-weight delayed-scaling rings for ``w``, ``x`` and ``g``.
DynamicScaling never allocates a meta; it measures the current amax inline.
Dynamic scaling 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)
w: _ScaleRing
x: _ScaleRing
g: _ScaleRing
@dataclass(frozen=True)
@@ -166,55 +150,29 @@ _active_config: ContextVar[Optional[_ActiveConfig]] = ContextVar(
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.
The active ``(enabled, recipe, fp8_format)`` triple is a ``ContextVar``
set by ``fp8_autocast`` (see ``_active``/``_current_config``); these plain
attributes are the persistent defaults applied outside any region —
``fp8_linear_enable`` writes ``default_enabled``. 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_recipe: FP8Recipe = FP8Recipe()
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:
def get_weight_meta(self, w: torch.Tensor, recipe: FP8Recipe) -> FP8TensorMeta:
key = (w.data_ptr(), w.shape, w.dtype)
meta = self._metas.get(key)
if meta is None:
meta = FP8TensorMeta(w.device, self.recipe)
meta = FP8TensorMeta(
_ScaleRing(w.device, recipe),
_ScaleRing(w.device, recipe),
_ScaleRing(w.device, recipe),
)
self._metas[key] = meta
return meta
@@ -222,7 +180,7 @@ class FP8State:
"""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_recipe = FP8Recipe()
self.default_format = FP8Format.HYBRID
self._metas.clear()
@@ -285,7 +243,7 @@ class fp8_autocast:
margin: int = 0,
):
if recipe is None:
recipe = DelayedScaling(history_len=update_interval, margin=margin)
recipe = FP8Recipe(history_len=update_interval, margin=margin)
self._config = _ActiveConfig(bool(enabled), recipe, FP8Format(fp8_format))
self._tokens: List[Token] = []
@@ -339,7 +297,7 @@ def fp8_linear_forward(
if cfg is None:
cfg = _current_config()
fmt = cfg.fp8_format.fwd()
if isinstance(cfg.recipe, DynamicScaling):
if cfg.recipe.dynamic:
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)
@@ -352,7 +310,7 @@ def fp8_linear_forward(
).reshape(*x.shape[:-1], w.size(0))
return out, sx, sw
meta = state.get_weight_meta(w)
meta = state.get_weight_meta(w, cfg.recipe)
if not meta.w.initialized:
meta.w.seed(w, fmt)
if not meta.x.initialized:
@@ -392,8 +350,8 @@ class _LinearFp8(torch.autograd.Function):
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)
ctx.is_dynamic = cfg.recipe.dynamic
ctx.meta = None if ctx.is_dynamic else _state.get_weight_meta(w, cfg.recipe)
return out
@staticmethod
tests/extension/test_fp8_mma.py
+29 -20
View File
@@ -15,9 +15,8 @@ import torch.nn.functional as F
import astrai.extension.fp8 as f8mod
from astrai.extension.fp8 import (
DelayedScaling,
DynamicScaling,
FP8Format,
FP8Recipe,
FP8TensorMeta,
_ScaleRing,
fp8_autocast,
@@ -234,7 +233,7 @@ def test_delayed_scaling_forward_uses_snapshot_scale():
dev = torch.device("cuda")
state = f8mod.fp8_state()
state.reset()
state.default_recipe = DelayedScaling(history_len=1, margin=0)
state.default_recipe = FP8Recipe(history_len=1, margin=0)
state.default_format = FP8Format.E4M3
try:
m, n, k = 32, 16, 64
@@ -273,7 +272,7 @@ def test_fp8_linear_forward_and_backward():
state = f8mod.fp8_state()
state.reset()
state.default_recipe = DynamicScaling()
state.default_recipe = FP8Recipe(dynamic=True)
try:
out, _, _ = f8mod.fp8_linear_forward(x, weight, bias)
@@ -400,13 +399,13 @@ def test_mm_fp8_matches_scaled_mm():
def test_recipe_scale_from_history():
"""Delayed: max over the window + margin; dynamic: current amax."""
hist = torch.tensor([1.0, 2.0, 0.5])
d = DelayedScaling(history_len=3, margin=0)
d = FP8Recipe(history_len=3, margin=0)
assert torch.allclose(d.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0))
d_m = DelayedScaling(history_len=3, margin=2)
d_m = FP8Recipe(history_len=3, margin=2)
assert torch.allclose(
d_m.scale_from_history(hist, "e4m3"), torch.tensor(2.0 / 448.0 / 4.0)
)
dyn = DynamicScaling()
dyn = FP8Recipe(dynamic=True)
amax = torch.tensor([0.25])
assert torch.allclose(
dyn.scale_from_history(amax, "e4m3"), torch.tensor(0.25 / 448.0)
@@ -424,27 +423,37 @@ def test_fp8_format_enum():
def test_fp8_autocast_context():
"""fp8_autocast sets and restores recipe + format on the global state."""
"""fp8_autocast pushes and restores the thread-local active config."""
state = fp8_state()
prev = (state.enabled, state.recipe, state.fp8_format)
state.reset()
try:
with fp8_autocast(enabled=True, fp8_format="hybrid", update_interval=8):
assert state.enabled
assert isinstance(state.recipe, DelayedScaling)
assert state.recipe.history_len == 8
assert state.fp8_format is FP8Format.HYBRID
with fp8_autocast(enabled=True, recipe=DynamicScaling(), fp8_format="e4m3"):
assert isinstance(state.recipe, DynamicScaling)
assert state.fp8_format is FP8Format.E4M3
assert state.fp8_format is FP8Format.HYBRID # restored on exit
assert not state.enabled
cfg = f8mod._active_config.get()
assert cfg is not None and cfg.enabled
assert not cfg.recipe.dynamic
assert cfg.recipe.history_len == 8
assert cfg.fp8_format is FP8Format.HYBRID
with fp8_autocast(
enabled=True, recipe=FP8Recipe(dynamic=True), fp8_format="e4m3"
):
inner = f8mod._active_config.get()
assert inner.recipe.dynamic
assert inner.fp8_format is FP8Format.E4M3
assert f8mod._active_config.get() is cfg # restored on exit
assert f8mod._active_config.get() is None
assert not fp8_linear_enabled()
finally:
state.enabled, state.recipe, state.fp8_format = prev
state.reset()
def test_fp8_tensor_meta_delayed_update():
"""Meta seeds from data; hist/scale are packed views of one state buffer."""
meta = FP8TensorMeta(torch.device("cpu"), DelayedScaling(history_len=4, margin=0))
recipe = FP8Recipe(history_len=4, margin=0)
meta = FP8TensorMeta(
_ScaleRing(torch.device("cpu"), recipe),
_ScaleRing(torch.device("cpu"), recipe),
_ScaleRing(torch.device("cpu"), recipe),
)
w = torch.randn(8, 8)
meta.w.seed(w, "e4m3")
assert meta.w.initialized