[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

• 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 not 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"
• buffer_quantized_scaling_factors turns on the feature, and by default it is deactivated.
  • 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.)
• quantized_scaling_factor_buffering_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.
• We do not save blockwise scaling factors (or CustomRecipe), 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.

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-quantized-scaling-factors)

[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 (--quantized-scaling-factor-buffering-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]

Greptile Summary

The PR adds opt-in, nonpersistent PTQ calibration metadata buffers across Transformer Engine’s quantized PyTorch modules.

• Buffers activation and weight scaling metadata for Linear, LayerNormLinear, LayerNormMLP, and GroupedLinear.
• Supports decaying-max accumulation for activation metadata while retaining current weight metadata.
• Centralizes built-in quantization recipe identification and safely skips unsupported custom quantizers.
• Adds focused unit coverage for metadata selection, unsupported recipes, grouped GEMM naming, and activation decay.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the LayerNormMLP path now accumulates both activation inputs separately from weights, and unsupported custom quantizers are safely skipped by calibration buffering.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/_common.py	Adds shared metadata selection and update helpers; the revised role-specific update calls fix the previously reported LayerNormMLP activation accumulation issue.
transformer_engine/pytorch/module/layernorm_mlp.py	Buffers both FC activation inputs through the configured decay path and handles weights separately without duplicate checkpoint-recomputation updates.
transformer_engine/pytorch/tensor/utils.py	Adds stable built-in recipe detection and safely returns no recipe for arbitrary custom quantizers, resolving the previously reported buffering crash.
transformer_engine/pytorch/module/grouped_linear.py	Adds per-GEMM activation and weight calibration metadata buffering for both grouped execution paths.
transformer_engine/pytorch/module/linear.py	Adds opt-in input and weight scaling-factor buffering to Linear forwards.
transformer_engine/pytorch/module/layernorm_linear.py	Adds opt-in normalized-input and weight scaling-factor buffering.
tests/pytorch/test_ptq_calibration_metadata_buffering.py	Covers recipe metadata selection, unsupported blockwise formats, grouped buffer naming, and decaying activation maxima.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Quantized module forward] --> B[Identify quantization recipe]
  B --> C{Global FP32 metadata available?}
  C -->|No| D[Skip calibration buffering]
  C -->|Yes| E[Read activation and weight metadata]
  E --> F[Apply decaying maximum to activations]
  E --> G[Retain current weight metadata]
  F --> H[Register nonpersistent module buffers]
  G --> H

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[cspades]

/te-ci pytorch

[cspades]

/te-ci pytorch