[PTQ] Store FP32 global scaling factors (absmax or scale_inv) for all quantized activations and weights.

[cspades]

Description

See issue for context: NVIDIA/Megatron-LM#5660
Related to: NVIDIA/Megatron-LM#6183

• Extended the Quantizer.calibrate() API to support decayed calibration and model buffering.
  • The public API autocast(calibrating: bool, calibration_config: QuantizationCalibrationConfig) (both arguments, kept the older one for legacy / backwards compatibility with decay defaulting to 0) turns on the feature, and by default it is deactivated.
  • calibrate will first look for calibration metadata according to the quantization recipe, but if that doesn't exist yet it will just take the non-quantized Tensor argument and compute its quantization metadata, since this is the original behavior of calibrate.
    • It is super un-performant to recompute the quantization metadata, usually it is fused into a TE kernel. So during training and inference, this path where we cannot find the quantization metadata should never be run.
  • calibrate now has implementations for all of the recipes that use FP32 global scaling factors, and the decayed calibration update is implemented in the base class Quantizer._update_calibration_value. It'll also buffer them automatically, which is a change in behavior but the extra buffering feature isn't a big deal for performance/memory (and frankly what's the point of storing calibration if we never export it), refer to my results below!
• To enable users to mine PTQ calibration data from their pre-training datasets, buffer scaling factors for quantized activations and weights for checkpointing in Megatron and inference in vLLM, Tensor-RT, ModelOpt, etc.
  • This is an exported common state artifact, it can be saved but (should) not (be) loaded, and will not appear in the model state dict. No risk of surprise overwriting of weight scaling factors, for instance.
  • Buffer names are linked to TE modules and are named pretty comprehensively: f"{tensor_name}_tensor_{metadata_name}_{recipe}_te_ptq_calibrated"
• QuantizationCalibrationConfig only has a single argument at the moment but in the future could hypothetically have way more, it's just that now we only buffer FP32 global scaling factors so it's simply just activation_scale_decay.
  • activation_scale_decay controls the decay of past max-accumulated activation scaling factors when your model or dataset is in-flight during training. Intends to capture steady-state maxima.
    • New AbsMax = Max(Old AbsMax * Decay, Observed AbsMax)
    • Setting it to 0 implies only taking the most recent scaling factor, while setting it to 1 implies taking an absmax over the entire history of your model and/or dataset.
    • You can do quite a lot of hacky customization on the MCore side, i.e. setting LR=0 means you're basically just calibrating on your calibration dataset with a frozen model.
  • Tiny performance loss to compute the accumulated max, but almost unnoticeable for both dense and MoE (GroupedLinear).
  • Tiny memory overhead to store the accumulated activation scaling factors. Since it's just 4 bytes per activation (FP32 global scale), it's not a lot at all.
    • Weights do not need to be accumulated, so the weight buffer points to the literal row-wise scaling factor in the quantized tensor storage and doesn't take any memory. (Set row-wise usage to True, on the user to do it.)
  • We do not save blockwise scaling factors (or CustomRecipe and HybridQuantizer, which will error out), those are usually computed during inference or using more advanced calibration techniques.
    • We do store the row-wise NVFP4 scaling factors, but in MCore I just max-reduce it to a scalar as well.
• TEAutocastState is another new data-class that may concern you since it changes the API of get_autocast_state and set_autocast_state which are public interfaces for the global quant state.
  • That being said, currently Megatron-LM and other NeMo repositories are not directly using this, and it gets really messy with checkpointing if I have to reset QuantizationCalibrationConfig so I think now's a good time to stop returning tuples...

Testing

• Really basic unit tests, tell me if you want to see more.
• Performance is only affected when decay is activated, and memory is only increased when buffering activations.

# Main / No scaling factor checkpointing.

[2026-07-27 13:24:57.339692] iteration       16/15258789 | consumed samples:         2048 | elapsed time per iteration (ms): 6282.9 | throughput per GPU (TFLOP/s/GPU): 2147.5 | learning rate: 7.864316E-08 | global batch size:   128 | lm loss: 1.207198E+01 | loss scale: 1.0 | grad norm: 20.134 | number of skipped iterations:   0 |number of nan iterations:   0 |

[Rank 0] (after 2 iterations) memory (MB) | allocated: 69607.20 | max allocated: 99633.89 | reserved: 72394.00| max reserved: 100310.00

# Scaling factor & absmax checkpointing. (--buffer-transformer-engine-calibration-metadata)

[2026-07-27 13:45:56.610449] iteration       16/15258789 | consumed samples:         2048 | elapsed time per iteration (ms): 6282.0 | throughput per GPU (TFLOP/s/GPU): 2147.8 | learning rate: 7.864316E-08 | global batch size:   128 | lm loss: 1.207239E+01 | loss scale: 1.0 | grad norm: 20.133 | number of skipped iterations:   0 | number of nan iterations:   0 |

[Rank 0] (after 2 iterations) memory (MB) | allocated: 69607.26 | max allocated: 99633.94 | reserved: 72394.00 | max reserved: 100310.00

# Decaying Activation Scaling Factors & Absmax (--transformer-engine-calibration-decay 0.99)

[2026-07-27 13:51:03.183943] iteration       16/15258789 | consumed samples:         2048 | elapsed time per iteration (ms): 6486.2 | throughput per GPU (TFLOP/s/GPU): 2080.2 | learning rate: 7.864316E-08 | global batch size:   128 | lm loss: 1.207229E+01 | loss scale: 1.0 | grad norm: 20.133 | number of skipped iterations:   0 |number of nan iterations:   0 |

[Rank 0] (after 2 iterations) memory (MB) | allocated: 69607.32 | max allocated: 99634.00 | reserved: 72396.00| max reserved: 100310.00

# Decaying Activation Scaling Factors + Blockwise Activation Scales (i.e. `_rowwise_scale_inv`)
# (Just FYI why we don't want to export blockwise scales, they significantly affect perf and memory, and aren't really necessary for calibration.)

[2026-07-27 13:57:09.949572] iteration       16/15258789 | consumed samples:         2048 | elapsed time per iteration (ms): 6857.8 | throughput per GPU (TFLOP/s/GPU): 1967.5 | learning rate: 7.864316E-08 | global batch size:   128 | lm loss: 1.207223E+01 | loss scale: 1.0 | grad norm: 20.134 | number of skipped iterations:   0 |number of nan iterations:   0 |

[Rank 0] (after 2 iterations) memory (MB) | allocated: 71158.36 | max allocated: 101184.57 | reserved: 74394.00 | max reserved: 101934.00

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

• Adds configurable decayed calibration and propagates it through autocast and CUDA-graph state.
• Buffers activation and weight metadata across Linear, LayerNormLinear, LayerNormMLP, and GroupedLinear.
• Adds recipe-specific calibration implementations, compatibility handling for custom quantizers, recomputation guards, and focused tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Autocast calibration config] --> B[TE module forward]
  B --> C{Activation recomputation?}
  C -- Yes --> D[Skip calibration update]
  C -- No --> E[Calibrate role-specific quantizers]
  E --> F[Apply decay to activation metadata]
  E --> G[Record current weight metadata]
  F --> H[Register nonpersistent PTQ buffers]
  G --> H

Loading

_{Reviews (8) · Last reviewed commit: "Merge branch 'main' into cye/scale_facto..."}

[cspades]

/te-ci pytorch

[cspades]

/te-ci pytorch

[ptrendx]

@cspades, there is already a calibrate API available in the quantizers (and invoked by modules), even though it only really works for the delayed scaling. Why couldn't that API be reused and extended in this PR instead of introducing calling new APIs in the modules?

[cspades]

already a calibrate API available in the quantizers (and invoked by modules)

@ptrendx It can! I didn't know that was there. Apologies for the late reply have been working on other tasks.

That being said, wanted to share that this max-reduce kind of works (vs. ModelOpt checkpoints, and I think Xiaokang did this with NVFP4 since that was the main topic of the issue):

but can definitely merge/fold this into the calibrate API. Will report back.

[cspades]

/te-ci pytorch

[cspades]

/te-ci L1 pytorch