NVIDIA / NVIDIA/Megatron-LM

[MCore][TE][NVFP4] Persist activation global input_scale/input_amax for offline FP4 inference

Open
#5,660 26 comments 0 reactions 1 assignee Assigned to @cspades View on GitHub
community-request enhancement module: checkpoint module: transformer engine waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

I am training/fine-tuning models with Megatron Core + Transformer Engine using NVFP4, for example SFT with:

- `--use-mcore-models`
- `--transformer-impl transformer_engine`
- `--bf16`
- `--fp4-recipe nvfp4`
- `--fp4-format e2m1`
- `--ckpt-format torch_dist`

After training finishes, I would like the checkpoint/export artifact to include the NVFP4 quantization metadata needed by inference frameworks such as sglang or TensorRT-LLM, especially the per-tensor global amax or the derived global scale used by NVFP4.

**Describe the solution you'd like**
I would like Megatron Core + Transformer Engine NVFP4 training/fine-tuning to save the **global NVFP4 input scale**, or the corresponding input amax, into the final checkpoint/export artifact.

The important metadata is the per-layer / per-GEMM **global input scale** used for activation quantization, for example the scale consumed by downstream inference runtimes as:

- `input_scale`
- `input_global_scale`
- `input_scale_inv`
- or equivalent `input_amax`

I am **not** requesting that per-block / block-size activation scales be saved. Those block scales depend on the actual runtime activation tensor and can still be computed online by the inference kernel. The missing piece is the offline/global input scale or amax, which can be collected during NVFP4 SFT/training and reused during inference.

Downstream runtimes such as SGLang and vLLM already support using offline/static input scales or amax metadata. This avoids computing global input scales repeatedly during model loading, conversion, or inference, which can slow down deployment and serving startup.

For SFT, the training job has already seen the target data distribution. Therefore, MCore/TE should be able to persist the observed/final NVFP4 global input scale directly, instead of requiring a separate PTQ calibration dataset and calibration pass after training.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.