feat(training): MXFP8 混合精度训练落地(等 mlx-lm fp8 支持)

[dahai80]

需求

#402 落地了 QLoRA,但 MXFP8 混合精度训练仍未实现 —— service.py:463-468 对 cfg.mxfp8=True 直接 raise ValueError:

if cfg.mxfp8:
    raise ValueError(
        "mxfp8 mixed-precision training is not yet supported "
        "(mlx-lm 0.31.3 has no fp8 training path). Use qlora for "
        "memory savings; mxfp8 will be landed in a later phase."
    )

下游 fusion-trainer 的 use_mxfp8=True 因此只能告警 + 退回默认精度训练(fail-visible),preset 里的显存峰值估算(MXFP8 假设)偏高,大模型(32B)实际可能 OOM。

现状

• FineTuneConfig.mxfp8 字段已存在(service.py:66),但训练路径未实现
• 阻塞根因:外部依赖 —— mlx-lm 0.31.3 无 fp8 训练路径(已确认 mlx_lm 源码无 mxfp8/fp8 训练支持)

提议

跟踪 issue,待 mlx-lm 上游支持 fp8 训练后再落地。在此之前:

1. 保持当前 fail-visible raise(不静默降级)
2. 监控 mlx-lm 版本,fp8 训练支持落地后:
   • 在 _apply_training 接入 fp8 量化训练路径
   • 对 mxfp8=True 走 fp8 前向 + LoRA 梯度更新
   • 更新 preset 显存估算
3. 期间文档明确:用 QLoRA 替代以省显存

关联

• feat(training): FineTuneConfig 增加 mxfp8 精度开关 + qlora (quantize_base) fine_tune_type #402(QLoRA 已落地,MXFP8 部分未落地)
• 下游 fusion-trainer feat(video-gen): VideoGenEngine + POST /v1/videos/generate (#454) #38 复测确认 MXFP8 仍 blocked

[dahai80]

Re-verification (2026-08-08)

Re-checked the upstream blocker for this issue:

• mlx-lm latest = 0.31.3 (pip index versions mlx-lm — 0.31.3 is the newest published release; no newer version available)
• mlx_lm/tuner/ (trainer.py, lora.py, losses.py, utils.py) — no fp8 / mxfp8 / float8 references in any training module
• mxfp8 appears only in mlx_lm/convert.py (quantization export, --quantize mxfp8) and mlx_lm/utils.py (weight loading of pre-quantized mxfp8 checkpoints) — both are inference paths, not training
• Confirmed: MLX core has mx.float8 types and quantize supports mxfp8, but nn.value_and_grad / the LoRA trainer have no fp8 mixed-precision forward+backward path

Conclusion

The upstream dependency has not landed fp8 training support. The current fail-visible raise ValueError in service.py:463 remains the correct behavior (no silent downgrade — per the project's fail-visible rule).

Action

Keeping this issue open as a tracking issue until mlx-lm adds an fp8 training path. No code change possible in fusion-mlx until then. Workaround for memory savings: qlora=true (#402, already shipped).

If/when upstream lands fp8 training:

1. Replace the raise in _apply_training with an fp8 forward + LoRA gradient path
2. Update preset memory estimates (currently assume MXFP8, over-estimate for 32B)
3. Add real-model training test (per project rule: 大模型测试须真实加载模型)

🤖 Verified by Claude Code