Request
Add an fp8 (MXFP8) mixed-precision training path to mlx_lm.training — both full finetune and LoRA — so downstream tools can run memory-reduced training on large models on Apple Silicon.
Background
mlx.core already exposes fp8 primitives (mx.from_fp8 / mx.to_fp8), and mlx_lm supports MXFP8 for inference (weight-only quantization). But there is no fp8 path for training: mlx_lm/training/ has no fp8/mxfp8 forward+backward mode. The mlx_lm package (verified on 0.31.3) has no training/ module exposing an fp8 mixed-precision option.
For LoRA specifically, the attractive setup is: frozen base weights in fp8 (MXFP8) for the forward pass, with LoRA adapter gradients accumulated in a higher precision — analogous to how QLoRA uses 4-bit weights + bf16/fp16 adapter updates, but pushed to fp8 for a smaller memory footprint.
Why
On unified-memory Apple Silicon, the dominant cost of finetuning a large model (e.g. 32B) is the resident weight memory. QLoRA 4-bit helps but can still OOM on the largest models at the preset's estimated peak. An fp8 mixed-precision training path would:
- cut base-weight memory ~2x vs bf16 (fp8 = 1 byte/param),
- keep gradient/optimizer state in higher precision for stability,
- unlock finetuning of larger models than QLoRA currently permits on the same machine.
Current workaround
Downstream projects (e.g. fusion-mlx, issue #425) currently fail visibly with raise ValueError("mxfp8 mixed-precision training is not yet supported (mlx-lm 0.31.3 has no fp8 training path)") and fall back to QLoRA. We do not want to silently degrade; we would rather land the real path once mlx-lm supports it.
Ask
- Is an fp8 / MXFP8 mixed-precision training path on the roadmap for
mlx_lm?
- If not, would a contribution be welcome? If so, what shape would the maintainers prefer — a
quantization_mode="mxfp8" flag on the trainer, a separate trainer class, or something else — so a PR fits the project's design?
- Any known blocker in
mlx.core's fp8 ops that would prevent a training backward pass today (autograd support through from_fp8/to_fp8, gradient scaling)?
Happy to contribute a PR if there's appetite and a preferred shape. Thanks!
Request
Add an fp8 (MXFP8) mixed-precision training path to
mlx_lm.training— both full finetune and LoRA — so downstream tools can run memory-reduced training on large models on Apple Silicon.Background
mlx.corealready exposes fp8 primitives (mx.from_fp8/mx.to_fp8), andmlx_lmsupports MXFP8 for inference (weight-only quantization). But there is no fp8 path for training:mlx_lm/training/has no fp8/mxfp8 forward+backward mode. Themlx_lmpackage (verified on 0.31.3) has notraining/module exposing an fp8 mixed-precision option.For LoRA specifically, the attractive setup is: frozen base weights in fp8 (MXFP8) for the forward pass, with LoRA adapter gradients accumulated in a higher precision — analogous to how QLoRA uses 4-bit weights + bf16/fp16 adapter updates, but pushed to fp8 for a smaller memory footprint.
Why
On unified-memory Apple Silicon, the dominant cost of finetuning a large model (e.g. 32B) is the resident weight memory. QLoRA 4-bit helps but can still OOM on the largest models at the preset's estimated peak. An fp8 mixed-precision training path would:
Current workaround
Downstream projects (e.g. fusion-mlx, issue #425) currently fail visibly with
raise ValueError("mxfp8 mixed-precision training is not yet supported (mlx-lm 0.31.3 has no fp8 training path)")and fall back to QLoRA. We do not want to silently degrade; we would rather land the real path once mlx-lm supports it.Ask
mlx_lm?quantization_mode="mxfp8"flag on the trainer, a separate trainer class, or something else — so a PR fits the project's design?mlx.core's fp8 ops that would prevent a training backward pass today (autograd support throughfrom_fp8/to_fp8, gradient scaling)?Happy to contribute a PR if there's appetite and a preferred shape. Thanks!