SIGN IN SIGN UP

feat(model_init): add FP8 pre-flight check for the force_hf path (#2350)

* feat(model_init): add FP8 pre-flight check for the force_hf path

When force_hf is on and the quant config is FineGrainedFP8Config with dequantize=True, HF dequantizes the whole FP8 checkpoint to BF16 on every rank before anything gets sharded. For models bigger than one GPU's memory that just OOMs deep inside the HF loader with no useful error.

This adds a check that runs first in the force_hf branch. It estimates the BF16 footprint from the config, compares it against free CUDA memory, and stops early with a clear message if it won't fit, pointing the user at the streaming workaround instead of a silent OOM. Can be turned off with NEMO_AUTOMODEL_DISABLE_FP8_PREFLIGHT=1 if the estimate is ever wrong.

Refs #2114

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>

* fix(model_init): harden FP8 dequantize pre-flight gate math

Follow-up to the FP8 pre-flight check (issue #2114): make the memory gate
correct rather than just present.

- Budget against TOTAL device memory instead of FREE, so the verdict is
  deterministic across ranks and unaffected by transient allocations from
  co-located processes (a low FREE could false-trip a model that actually fits).
- Estimate parameters accurately: prefer an exact element count read from the
  local safetensors headers (dtype/GQA/MoE-agnostic), and make the fallback
  formula GQA-aware. The old flat 4*H*H attention term over-counted GQA models
  by ~13%, skewing the check too strict.
- Remove the dead num_parameters config lookup (it is a PreTrainedModel method,
  not a config attribute) and the speculative is_dequantize/dequant aliases
  that no transformers/Automodel code path sets.
- Run the same guard on the non-force_hf HF fallback branch, which reaches the
  identical full-materialize loader (e.g. PEFT-implicit dequantize=True).

Tests: 24 -> 32 (GQA vs MHA, MHA==legacy, total-not-free budgeting, exact
safetensors counting, non-force_hf wiring).

Refs #2114

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>

* fix: address review on the fp8 preflight guard

Gate the preflight on is_pretrained_init at both call sites, a from_config
build never materializes the checkpoint so there is nothing to dequantize
and it should not raise on fp8 metadata alone. The fallback path now passes
force_hf=False so the error stops recommending "drop force_hf" where that
was never on, and the size labels say GiB since the math divides by 1024^3.

Tests: config-only coverage for both branches, fallback wording, GiB label,
plus a guard that the force_hf wording keeps its hint.

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>

---------

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
Co-authored-by: Huiying <willwin.lee@gmail.com>
Co-authored-by: Yuhe Zhang <yuhez@nvidia.com>
S
stanley1208 committed
fbe1333e1d5ef5a66043466040795fca61e995f9
Parent: 82ece9f
Committed by GitHub <noreply@github.com> on 8/5/2026, 7:29:04 PM