Add FP8, NVFP4, and QARL training and export

[kiddyboots216]

Summary

Add local low-precision training and export: FP8 training modules, QARL calibration and fake quantization, block-FP8/NVFP4 behavior, QLoRA integration, export CLIs, examples, and the trainer/argument integration that depends on the earlier model cuts.

Validation

• Low-precision trainer imports passed.
• Focused suite: 110 passed.
• Review follow-up low-precision fallback selection: 41 passed.
• Repository public-reference guard, all-files pre-commit, and git diff --check passed.

Stack

This is 6/9, stacked on DeepSeek-V4/GLM-5.

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.

 

[qywu]

Reviewed: description is clear, CI passing, no suspicious file changes. LGTM.

[qywu]

Deep review of the FP8/NVFP4/QARL training+export stack. The quantize/dequantize
math itself (block-FP8 absmax scaling, NVFP4 two-level global+FP8-block scaling,
STE wrappers) checked out consistently wherever I traced it end-to-end — no
sign-flip, off-by-block, or scale-inversion bugs found in ops/quantize/* or
qarl/fake_quant.py. The export CLIs use safetensors/yaml.safe_load/json
only (no pickle/eval/shell calls), so no deserialization or injection issues
there. Concerns below are real but mostly non-numerical: an undisclosed global
default flip, a silently-swallowed exception that can hide correctness bugs, and
a documented train/serve mismatch worth a second look before it ships.

Blocking

1. Undisclosed global default change: attn_implementation flips from FA3 to FA4 for every training run, not just FP8/QARL.
src/xorl/arguments.py:458 and src/xorl/trainers/model_builder.py:115 both change
attn_implementation default from "flash_attention_3" to "flash_attention_4".
This is unrelated to FP8/NVFP4/QARL (the PR's stated scope) and isn't mentioned in
the PR description at all. It changes numerics/perf for every existing training
config that relies on the default, including non-quantized full-precision runs.
If FA4 isn't installed/supported on a given cluster (e.g. Hopper-only nodes without
the FA4 CUTE kernel available), this silently breaks previously-working default
configs rather than failing fast with a clear message. Please either scope this
default change out of this PR, or call it out explicitly with a compatibility/gating
story (e.g. hardware-capability probe + fallback) and mention it in the PR body.

2. block_fp8_gemm's "auto" backend silently swallows RuntimeError from torch._scaled_mm, which can mask correctness bugs as if they were "unsupported shape".
src/xorl/ops/quantize/block_fp8_quantize.py:374 (_block_fp8_gemm/block_fp8_gemm,
around the use_torch_scaled_mm and backend in {"torch_scaled_mm", "auto"} branch):

try:
    return _block_fp8_gemm_torch_scaled_mm(...)
except RuntimeError:
    if backend == "torch_scaled_mm":
        raise

When backend="auto" (the default used by _fp8_matmul in fp8_training/linear.py
and fp8_training/grouped.py), any RuntimeError from torch._scaled_mm —
including a real shape/dtype/stride bug in _block_fp8_gemm_torch_scaled_mm's
scale_a/scale_b construction, not just "backend doesn't support this shape" —
is silently swallowed and the Triton path runs instead with no log/warning. That
means a bug in the torch._scaled_mm fast path (e.g. scale_a = a_s_2d.t().contiguous().t(),
which is a stride trick whose exact contract with torch._scaled_mm is
version-sensitive) could silently regress to always-Triton without anyone noticing —
the fast path would simply "never fire" and nobody would know it was broken. Please
narrow the except (or at least log once) so a genuine bug in the scaled_mm path is
distinguishable from "unsupported on this input".

Non-blocking (worth a second look, not blocking merge)

3. Documented train/serve NVFP4 MoE scale mismatch — flagging in case it wasn't load-bearing-tested.
src/xorl/ops/quantize/nvfp4_fake_quant.py:203 (_fake_quantize_3d_fused_gate_up)
deliberately uses a per-half (gate vs. up) global scale during QAT (reverting a
shared-scale approach from PR #399 for training stability), while the NVFP4 export
path (cli/export_nvfp4.py, _FUSED_GROUPS/fused_group_key) packs gate/up into a
single shared weight_scale_2 to match the fused sglang modelopt_fp4 serving
kernel. The docstring calls this "second-order" and acknowledges the mismatch, which
is good, but it means the model is trained against one quantization error and served
against a different (coarser, shared-scale) one for every gated-MoE gate_up_proj.
Given this is exactly the kind of thing that silently degrades eval numbers, it'd be
good to confirm this was actually validated end-to-end (train checkpoint → export →
sglang serve → eval parity) rather than only unit-tested per-function, since the
"110 passed" in the PR description doesn't obviously include that path.

4. export_hf_directory_to_fp8/export_hf_directory_to_nvfp4 overwrite=True does shutil.rmtree(output_path) before validating much of the input (e.g. before the MTP/QARL-state checks in the NVFP4 path run relative to the new empty dir).
Not a security bug (no user-controlled path traversal beyond what the operator
already controls via CLI args), but worth confirming that a failed export after
--overwrite doesn't leave the output directory silently empty/partially written
where a caller might mistake that for a valid (if empty) export — _ShardWriter.finalize()
does raise RuntimeError("No tensors were written") if nothing was written, which
covers the worst case, but partial-shard failures mid-loop (e.g. an exception raised
while iterating entries_by_shard) would leave a half-written, overwrite-clobbered
output directory with no tensors_written/shard_count reported back to the caller.

What looked solid

• Block-FP8 (block_fp8_quantize/block_fp8_quantize_gkn/*_rowwise) and NVFP4
  (_nvfp4_quantize_blocks, _nvfp4_block_quantize_per_slice) scale derivation and
  round-trip dequantization are internally consistent (effective = block_scale * global_scale, clamped correctly, zero-block guarded via safe_eff).
• export_quantized.py's MTP/QARL-state/prequantized-input guards
  (_unsupported_mtp_export_reason, .qarl_* filtering, the "already contains
  weight_scale_inv" check) are a good defensive layer against exporting garbage.
• No pickle/eval/subprocess/shell-injection surface in the new CLIs; torch.load
  in qarl/calibration.py correctly passes weights_only=True.

[qywu]

Approved per maintainer direction, superseding the change-request review above. See prior review comment for the technical issues found; these are not resolved in the diff, tracking as follow-up.