From c13a4e0543026a7beedef70fe160efcbbc5b2964 Mon Sep 17 00:00:00 2001 From: Chinmaya Sahu <42250012+csking101@users.noreply.github.com> Date: Tue, 28 Apr 2026 20:46:43 +0530 Subject: [PATCH] feat(fp8): add rowwise and rowwise_with_gw_hp scaling recipes for FP8 training (#62) Add fp8_recipe config field to TrainingConfig with three torchao-backed scaling recipes: tensorwise (default, v0.28.0 behavior), rowwise (more accurate via CUTLASS), and rowwise_with_gw_hp (most accurate, grad_weight in high precision). Dispatches via Float8LinearConfig.from_recipe_name(). - schema.py: add fp8_recipe Literal field with validator requiring quantization_aware='fp8' for non-default recipes - fp8.py: update apply_fp8_training() to accept recipe parameter - sft.py: pass tcfg.fp8_recipe to apply_fp8_training() - README.md: document recipe options with comparison table - tests: 24 tests covering schema, dispatch, validation, backward compat --- README.md | 18 ++ soup_cli/config/schema.py | 18 ++ soup_cli/trainer/sft.py | 7 +- soup_cli/utils/fp8.py | 23 ++- tests/test_fp8_recipe.py | 386 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 445 insertions(+), 7 deletions(-) create mode 100644 tests/test_fp8_recipe.py diff --git a/README.md b/README.md index 1486643..ed2e6ef 100644 --- a/README.md +++ b/README.md @@ -426,6 +426,24 @@ training: quantization: none # FP8 converts linears directly; no bnb 4bit needed ``` +### FP8 Scaling Recipes (v0.28.1) + +Choose a scaling recipe to trade off speed vs accuracy: + +```yaml +training: + quantization_aware: fp8 + fp8_recipe: rowwise # tensorwise | rowwise | rowwise_with_gw_hp +``` + +| Recipe | Kernel | Scaling | Trade-off | +|---|---|---|---| +| `tensorwise` (default) | cuBLAS | Single scale per tensor | Fastest, good accuracy | +| `rowwise` | CUTLASS | Per-row scale, e4m3, power-of-2 scales | Slower, more accurate | +| `rowwise_with_gw_hp` | CUTLASS | Rowwise + grad_weight in high precision | Slowest, most accurate | + +Omitting `fp8_recipe` defaults to `tensorwise` (identical to v0.28.0 behavior). + Bool `true` stays on the int8 QAT path for backward compatibility. FP8 requires CUDA + Hopper+ (compute capability ≥ 9.0) and is rejected on unsloth/mlx backends. Wired across every transformer-backend trainer (SFT, DPO, GRPO, KTO, ORPO, SimPO, IPO, PPO, Reward-Model, Embedding, Pretrain). ## Cut Cross-Entropy (Large-Vocab Models) diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 6c051e8..13d5337 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -146,6 +146,14 @@ class TrainingConfig(BaseModel): "'fp8'=FP8 training on H100/B100 (v0.28.0)." ), ) + fp8_recipe: Literal["tensorwise", "rowwise", "rowwise_with_gw_hp"] = Field( + default="tensorwise", + description=( + "FP8 scaling recipe (only used when quantization_aware='fp8'). " + "'tensorwise' (fastest, default), 'rowwise' (more accurate, CUTLASS rowwise), " + "'rowwise_with_gw_hp' (most accurate, grad_weight in high precision). (v0.28.1)." + ), + ) optimizer: str = Field(default="adamw_torch", description="Optimizer name") scheduler: str = Field(default="cosine", description="LR scheduler type") save_steps: int = Field(default=100, description="Save checkpoint every N steps") @@ -545,6 +553,16 @@ class TrainingConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_fp8_recipe_requires_fp8(self) -> "TrainingConfig": + """fp8_recipe is only meaningful when quantization_aware='fp8'.""" + if self.fp8_recipe != "tensorwise" and self.quantization_aware != "fp8": + raise ValueError( + f"fp8_recipe='{self.fp8_recipe}' requires quantization_aware='fp8'. " + "Either set quantization_aware: 'fp8' or remove the fp8_recipe field." + ) + return self + class EvalConfig(BaseModel): """Evaluation configuration for auto-eval after training.""" diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index 5afba3c..929edac 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -455,10 +455,11 @@ class SFTTrainerWrapper: if tcfg.quantization_aware == "fp8": from soup_cli.utils.fp8 import apply_fp8_training - if apply_fp8_training(self.model): + fp8_recipe = getattr(tcfg, "fp8_recipe", "tensorwise") + if apply_fp8_training(self.model, recipe=fp8_recipe): console.print( - "[green]FP8 training enabled:[/] " - "converted linears to Float8Linear" + f"[green]FP8 training enabled:[/] " + f"converted linears to Float8Linear (recipe={fp8_recipe})" ) else: console.print( diff --git a/soup_cli/utils/fp8.py b/soup_cli/utils/fp8.py index c7c48b0..af866a0 100644 --- a/soup_cli/utils/fp8.py +++ b/soup_cli/utils/fp8.py @@ -57,14 +57,27 @@ def is_fp8_gpu_supported() -> bool: return False -def apply_fp8_training(model) -> bool: +def apply_fp8_training( + model, + recipe: str = "tensorwise", +) -> bool: """Convert eligible linear layers to FP8 for training. - Uses torchao's ``convert_to_float8_training`` with a tensorwise scaling - recipe (default / most widely supported). + Uses torchao's ``convert_to_float8_training`` with a scaling recipe + selected via :pydata:`Float8LinearConfig.from_recipe_name`. + + Supported recipes (from ``torchao.float8.config.Float8LinearRecipeName``): + + - ``"tensorwise"`` — single scale per tensor, cuBLAS kernel (fastest, + default, v0.28.0 behavior). + - ``"rowwise"`` — per-row scale, CUTLASS kernel, e4m3 everywhere, + power-of-2 scales (more accurate). + - ``"rowwise_with_gw_hp"`` — rowwise but grad_weight stays in high + precision (most accurate). Args: model: PyTorch model to convert (typically after LoRA has been applied). + recipe: Scaling recipe name. Default ``"tensorwise"``. Returns: True on success, False if FP8 is unavailable or conversion failed. @@ -74,8 +87,10 @@ def apply_fp8_training(model) -> bool: try: from torchao.float8 import convert_to_float8_training + from torchao.float8.config import Float8LinearConfig - convert_to_float8_training(model) + config = Float8LinearConfig.from_recipe_name(recipe) + convert_to_float8_training(model, config=config) return True except (ImportError, RuntimeError, ValueError): return False diff --git a/tests/test_fp8_recipe.py b/tests/test_fp8_recipe.py new file mode 100644 index 0000000..19ab1b9 --- /dev/null +++ b/tests/test_fp8_recipe.py @@ -0,0 +1,386 @@ +"""Tests for FP8 recipe support (v0.28.1). + +Covers: +- Schema: fp8_recipe field accepts valid literals, rejects invalid strings +- Schema: fp8_recipe requires quantization_aware='fp8' when non-default +- Schema: default recipe is 'tensorwise' (backward compat with v0.28.0) +- Dispatch: apply_fp8_training passes correct recipe to Float8LinearConfig +- Integration: fp8_recipe with non-SFT tasks rejected (via quantization_aware gate) +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import SoupConfig, TrainingConfig + +# ─── Schema: fp8_recipe field ────────────────────────────────────────────── + + +class TestFP8RecipeSchema: + """fp8_recipe accepts 'tensorwise', 'rowwise', 'rowwise_with_gw_hp'.""" + + def test_fp8_recipe_default_tensorwise(self): + cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"}) + assert cfg.training.fp8_recipe == "tensorwise" + + def test_fp8_recipe_tensorwise_explicit(self): + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": "fp8", "fp8_recipe": "tensorwise"}, + ) + assert cfg.training.fp8_recipe == "tensorwise" + + def test_fp8_recipe_rowwise(self): + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": "fp8", "fp8_recipe": "rowwise"}, + ) + assert cfg.training.fp8_recipe == "rowwise" + + def test_fp8_recipe_rowwise_with_gw_hp(self): + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "rowwise_with_gw_hp", + }, + ) + assert cfg.training.fp8_recipe == "rowwise_with_gw_hp" + + def test_fp8_recipe_invalid_string_rejected(self): + """Only the three literal values are accepted.""" + with pytest.raises(ValidationError) as exc: + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "delayed", + }, + ) + assert "fp8_recipe" in str(exc.value) + + def test_fp8_recipe_invalid_empty_string_rejected(self): + with pytest.raises(ValidationError): + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": "fp8", "fp8_recipe": ""}, + ) + + def test_fp8_recipe_invalid_int_rejected(self): + with pytest.raises(ValidationError): + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": "fp8", "fp8_recipe": 42}, + ) + + +# ─── Schema: fp8_recipe requires quantization_aware='fp8' ───────────────── + + +class TestFP8RecipeRequiresFP8: + """Non-default fp8_recipe without quantization_aware='fp8' is rejected.""" + + def test_default_recipe_allowed_without_fp8(self): + """tensorwise (default) is fine even without fp8 — it's the default.""" + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"fp8_recipe": "tensorwise"}, + ) + assert cfg.training.fp8_recipe == "tensorwise" + assert cfg.training.quantization_aware is False + + def test_rowwise_without_fp8_rejected(self): + with pytest.raises(ValidationError) as exc: + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"fp8_recipe": "rowwise"}, + ) + assert "quantization_aware" in str(exc.value) + + def test_rowwise_with_gw_hp_without_fp8_rejected(self): + with pytest.raises(ValidationError) as exc: + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"fp8_recipe": "rowwise_with_gw_hp"}, + ) + assert "quantization_aware" in str(exc.value) + + def test_rowwise_with_bool_true_qat_rejected(self): + """Bool True = int8 QAT, not FP8 — recipe should be rejected.""" + with pytest.raises(ValidationError) as exc: + SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={ + "quantization_aware": True, + "fp8_recipe": "rowwise", + }, + ) + assert "quantization_aware" in str(exc.value) + + +# ─── Dispatch: apply_fp8_training recipe parameter ──────────────────────── + + +class TestFP8RecipeDispatch: + """apply_fp8_training passes recipe to Float8LinearConfig.from_recipe_name.""" + + def test_apply_fp8_dispatches_tensorwise(self): + """Default recipe passes 'tensorwise' to from_recipe_name.""" + mock_config = MagicMock() + mock_from_recipe = MagicMock(return_value=mock_config) + mock_convert = MagicMock() + + fake_float8 = MagicMock() + fake_float8.convert_to_float8_training = mock_convert + fake_config_mod = MagicMock() + fake_config_mod.Float8LinearConfig.from_recipe_name = mock_from_recipe + + with patch.dict( + "sys.modules", + { + "torchao": MagicMock(), + "torchao.float8": fake_float8, + "torchao.float8.config": fake_config_mod, + }, + ): + # Need to reimport to pick up the mocked modules + import importlib + + import soup_cli.utils.fp8 as fp8_mod + + importlib.reload(fp8_mod) + + # Mock is_fp8_available to return True + with patch.object(fp8_mod, "is_fp8_available", return_value=True): + model = MagicMock() + result = fp8_mod.apply_fp8_training(model, recipe="tensorwise") + + mock_from_recipe.assert_called_once_with("tensorwise") + mock_convert.assert_called_once_with(model, config=mock_config) + assert result is True + + def test_apply_fp8_dispatches_rowwise(self): + """Rowwise recipe passes 'rowwise' to from_recipe_name.""" + mock_config = MagicMock() + mock_from_recipe = MagicMock(return_value=mock_config) + mock_convert = MagicMock() + + fake_float8 = MagicMock() + fake_float8.convert_to_float8_training = mock_convert + fake_config_mod = MagicMock() + fake_config_mod.Float8LinearConfig.from_recipe_name = mock_from_recipe + + with patch.dict( + "sys.modules", + { + "torchao": MagicMock(), + "torchao.float8": fake_float8, + "torchao.float8.config": fake_config_mod, + }, + ): + import importlib + + import soup_cli.utils.fp8 as fp8_mod + + importlib.reload(fp8_mod) + + with patch.object(fp8_mod, "is_fp8_available", return_value=True): + model = MagicMock() + result = fp8_mod.apply_fp8_training(model, recipe="rowwise") + + mock_from_recipe.assert_called_once_with("rowwise") + assert result is True + + def test_apply_fp8_dispatches_rowwise_with_gw_hp(self): + """rowwise_with_gw_hp recipe passes through correctly.""" + mock_config = MagicMock() + mock_from_recipe = MagicMock(return_value=mock_config) + mock_convert = MagicMock() + + fake_float8 = MagicMock() + fake_float8.convert_to_float8_training = mock_convert + fake_config_mod = MagicMock() + fake_config_mod.Float8LinearConfig.from_recipe_name = mock_from_recipe + + with patch.dict( + "sys.modules", + { + "torchao": MagicMock(), + "torchao.float8": fake_float8, + "torchao.float8.config": fake_config_mod, + }, + ): + import importlib + + import soup_cli.utils.fp8 as fp8_mod + + importlib.reload(fp8_mod) + + with patch.object(fp8_mod, "is_fp8_available", return_value=True): + model = MagicMock() + result = fp8_mod.apply_fp8_training( + model, recipe="rowwise_with_gw_hp" + ) + + mock_from_recipe.assert_called_once_with("rowwise_with_gw_hp") + assert result is True + + def test_apply_fp8_returns_false_when_unavailable(self): + """When FP8 deps are missing, apply_fp8_training returns False.""" + from soup_cli.utils.fp8 import apply_fp8_training + + with patch("soup_cli.utils.fp8.is_fp8_available", return_value=False): + model = MagicMock() + assert apply_fp8_training(model, recipe="rowwise") is False + + def test_apply_fp8_default_recipe_is_tensorwise(self): + """Calling without recipe= uses 'tensorwise' (v0.28.0 compat).""" + mock_config = MagicMock() + mock_from_recipe = MagicMock(return_value=mock_config) + mock_convert = MagicMock() + + fake_float8 = MagicMock() + fake_float8.convert_to_float8_training = mock_convert + fake_config_mod = MagicMock() + fake_config_mod.Float8LinearConfig.from_recipe_name = mock_from_recipe + + with patch.dict( + "sys.modules", + { + "torchao": MagicMock(), + "torchao.float8": fake_float8, + "torchao.float8.config": fake_config_mod, + }, + ): + import importlib + + import soup_cli.utils.fp8 as fp8_mod + + importlib.reload(fp8_mod) + + with patch.object(fp8_mod, "is_fp8_available", return_value=True): + model = MagicMock() + fp8_mod.apply_fp8_training(model) + + # Default should be tensorwise + mock_from_recipe.assert_called_once_with("tensorwise") + + +# ─── Integration: fp8_recipe + non-SFT tasks ────────────────────────────── + + +class TestFP8RecipeNonSFT: + """FP8 recipe on non-SFT tasks: accepted on transformer backends (v0.35.0).""" + + def test_fp8_recipe_on_dpo_accepted(self): + """DPO + fp8 + recipe is valid (all transformer trainers wired in v0.35.0).""" + cfg = SoupConfig( + base="m", + task="dpo", + data={"train": "./d.jsonl", "format": "dpo"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "rowwise", + }, + ) + assert cfg.training.fp8_recipe == "rowwise" + + def test_fp8_recipe_on_grpo_accepted(self): + cfg = SoupConfig( + base="m", + task="grpo", + data={"train": "./d.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "rowwise_with_gw_hp", + }, + ) + assert cfg.training.fp8_recipe == "rowwise_with_gw_hp" + + def test_fp8_recipe_on_mlx_rejected(self): + """MLX backend does not support FP8 — rejected at config level.""" + with pytest.raises(ValidationError): + SoupConfig( + base="m", + task="sft", + backend="mlx", + data={"train": "./d.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "rowwise", + }, + ) + + def test_fp8_recipe_on_sft_accepted(self): + """SFT + fp8 + recipe is valid.""" + cfg = SoupConfig( + base="m", + task="sft", + data={"train": "./d.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": "rowwise", + }, + ) + assert cfg.training.fp8_recipe == "rowwise" + assert cfg.training.quantization_aware == "fp8" + + def test_all_recipes_accepted_on_sft(self): + """All three recipes are valid on SFT with fp8.""" + for recipe in ("tensorwise", "rowwise", "rowwise_with_gw_hp"): + cfg = SoupConfig( + base="m", + task="sft", + data={"train": "./d.jsonl"}, + training={ + "quantization_aware": "fp8", + "fp8_recipe": recipe, + }, + ) + assert cfg.training.fp8_recipe == recipe + + +# ─── Backward compatibility ─────────────────────────────────────────────── + + +class TestFP8RecipeBackwardCompat: + """v0.28.0 configs without fp8_recipe still work (defaults to tensorwise).""" + + def test_v028_config_no_recipe_field(self): + """Config with quantization_aware='fp8' but no fp8_recipe is valid.""" + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": "fp8"}, + ) + assert cfg.training.quantization_aware == "fp8" + assert cfg.training.fp8_recipe == "tensorwise" + + def test_v028_bool_true_unaffected(self): + """Bool True (int8 QAT) is unaffected by fp8_recipe field.""" + cfg = SoupConfig( + base="test/model", + data={"train": "./data.jsonl"}, + training={"quantization_aware": True}, + ) + assert cfg.training.quantization_aware is True + assert cfg.training.fp8_recipe == "tensorwise" # default, unused + + def test_training_config_fp8_recipe_default(self): + """TrainingConfig alone defaults fp8_recipe to tensorwise.""" + tcfg = TrainingConfig() + assert tcfg.fp8_recipe == "tensorwise"