Add blockwise FP8 training recipe for linears and MoE grouped experts

scripts/fp8_blockwise/compare_deepseek_v3_16b.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Compare BF16 and blockwise FP8 DeepSeek V3 16B loss convergence on 8 GPUs.
+#
+# Usage:
+#   ./scripts/fp8_blockwise/compare_deepseek_v3_16b.sh
+#   STEPS=1000 NGPU=8 ./scripts/fp8_blockwise/compare_deepseek_v3_16b.sh
+#
+# Requires a torchao build with torchao.prototype.blockwise_fp8_training and
+# torchao.prototype.moe_training.blockwise_fp8 (set TORCHAO_REPO to a checkout
+# to use it via PYTHONPATH).
+
+set -euo pipefail
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "${REPO_ROOT}"
+
+TORCHAO_REPO="${TORCHAO_REPO:-}"
+NGPU="${NGPU:-8}"
+STEPS="${STEPS:-1000}"
+SEED="${SEED:-42}"
+DUMP_ROOT="${DUMP_ROOT:-./outputs/fp8_blockwise_16b_compare}"
+LOG_DIR="${LOG_DIR:-${DUMP_ROOT}/logs}"
+
+if [[ -n "${TORCHAO_REPO}" ]]; then
+    export PYTHONPATH="${TORCHAO_REPO}:${PYTHONPATH:-}"
+fi
+
+mkdir -p "${LOG_DIR}"
+
+echo "Running BF16 DeepSeek V3 16B on ${NGPU} GPU(s)"
+NGPU="${NGPU}" MODULE=deepseek_v3 CONFIG=deepseek_v3_16b ./run_train.sh \
+    --training.steps "${STEPS}" \
+    --debug.seed "${SEED}" \
+    --dump_folder "${DUMP_ROOT}/bf16" "$@" \
+    2>&1 | tee "${LOG_DIR}/bf16.log"
+
+echo "Running blockwise FP8 DeepSeek V3 16B on ${NGPU} GPU(s)"
+NGPU="${NGPU}" MODULE=deepseek_v3 CONFIG=deepseek_v3_16b_fp8_blockwise ./run_train.sh \
+    --training.steps "${STEPS}" \
+    --debug.seed "${SEED}" \
+    --dump_folder "${DUMP_ROOT}/fp8_blockwise" "$@" \
+    2>&1 | tee "${LOG_DIR}/fp8_blockwise.log"
+
+echo "Done. Compare TensorBoard loss curves under ${DUMP_ROOT}/{bf16,fp8_blockwise}."

tests/unit_tests/test_quantization.py

@@ -64,3 +64,51 @@ def test_quantized_grouped_experts():
     assert issubclass(float8_cls, GptOssGroupedExperts)
     assert hasattr(mxfp8_cls.Config, "swiglu_limit")
     assert hasattr(float8_cls.Config, "swiglu_limit")
+
+
+def test_fp8_blockwise_applied_by_model_registry():
+    """Blockwise FP8 converters swap both linears and grouped experts."""
+    pytest.importorskip("torchao.prototype.blockwise_fp8_training.linear")
+    from torchtitan.components.quantization import Float8BlockwiseLinear
+    from torchtitan.tools.utils import has_cuda_capability
+
+    if Float8BlockwiseLinear is None:
+        pytest.skip("torchao blockwise FP8 training linear is unavailable")
+    if not has_cuda_capability(9, 0):
+        pytest.skip("blockwise FP8 training requires SM90 or later")
+
+    config_manager = ConfigManager()
+    config = config_manager.parse_args(
+        [
+            "--module",
+            "deepseek_v3",
+            "--config",
+            "deepseek_v3_debugmodel_fp8_blockwise",
+        ]
+    )
+    model_config = config.model_spec.model
+    assert has_quantization(model_config)
+
+    converted = [
+        (fqn, lc)
+        for fqn, lc, _parent, _attr in model_config.traverse(Linear.Config)
+        if isinstance(lc, Float8BlockwiseLinear.Config)
+    ]
+    assert len(converted) > 0
+    for _fqn, lc in converted:
+        assert lc.in_features % 128 == 0
+        assert lc.out_features % 128 == 0
+        assert not lc.bias
+    for fqn, lc, _parent, _attr in model_config.traverse(Linear.Config):
+        if fqn == "lm_head" or "router.gate" in fqn:
+            assert not isinstance(lc, Float8BlockwiseLinear.Config)
+
+    experts_configs = [
+        config
+        for _fqn, config, _parent, _attr in model_config.traverse(GroupedExperts.Config)
+    ]
+    assert len(experts_configs) > 0
+    assert all(
+        getattr(config, "recipe_name", None) == "fp8_blockwise"
+        for config in experts_configs
+    )

torchtitan/components/quantization/__init__.py

@@ -33,6 +33,7 @@ class Config(ModelConfigConverter.Config):
 
 # Re-export all public symbols so callers can import from the package directly.
 from .float8 import (  # noqa: F401, E402
+    Float8BlockwiseLinear,
     Float8GroupedExpertsConverter,
     Float8Linear,
     Float8LinearConverter,
@@ -44,6 +45,7 @@ class Config(ModelConfigConverter.Config):
 )
 
 __all__ = [
+    "Float8BlockwiseLinear",
     "Float8GroupedExpertsConverter",
     "Float8Linear",
     "Float8LinearConverter",