From 432dfec3c2e35e16ec2504fffd6e979fc25d946f Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 29 Aug 2026 20:03:32 +0800 Subject: [PATCH] refactor: collapse fp8 recipe hierarchy and state property layers - 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 --- astrai/extension/fp8.py | 110 ++++++++++---------------------- tests/extension/test_fp8_mma.py | 49 ++++++++------ 2 files changed, 63 insertions(+), 96 deletions(-) diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index f2ca308..e7ac16d 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -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 diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index 4c3dbbb..bb56983 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -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