Request for additional features from TE Debug modules for low precision training.

[jomitchellnv]

Is your feature request related to a problem? Please describe.

1. For LogFp4Tensor stats and LogFP8TensorStats we want to obtain the standard deviation of the scale factors.
2. LogFp8 and LogFP4 tensor stats can't currently be used at the same time, which is useful for situations when you have some layers as FP8 and some as FP4.
3. [most important] When using MXFP8 or NVFP4 micro-scaling methods we are only obtaining global statistics over the entire tensor instead of per each block.

Describe the solution you'd like

1. We want to add the standard deviation metric for the scale factors because simply having min and max doesn't tell us enough information. We could have serious clipping going on but won't know if we only have a max value.
2. Currently, I've had to hack in support by modifying yaml files online in order to perform FP8 and FP4 logging for networks that have mixed FP4/FP8 layers https://github.com/NVIDIA/bionemo-framework/blob/main/bionemo-recipes/recipes/esm2_native_te/quantization.py#L94 I am hoping this can be supported natively.
3. For LogFP8TensorStats and LogFP4TensorStats specifically for micro-block formats such as MXFP8 / NVFP4 we need an optional mechanism to extract block level metadata in addition to per-tensor level metadata including scale factor statistics per block of (scale_value, block_indices, tensor_chunk) etc saved to files. By default this should be able to work on the first block (0:32 for MXFP8, 0:16 for NVFP4) unless otherwise specified etc. It would be super useful in order to identify issues on a more fine grained scale.

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.

[JiwaniZakir]

The core limitation for points 1 and 3 is that the current LogFP8TensorStats/LogFP4TensorStats implementations appear to aggregate scale factor metadata at the tensor level only — adding std alongside min/max is straightforward, but per-block extraction for MXFP8 (32-element blocks) and NVFP4 (16-element blocks) would require reshaping the scale factor tensor by block indices before computing statistics, which is a non-trivial change to the logging pipeline. For point 2, the mutual exclusivity of FP8 and FP4 logging likely stems from the debug config YAML schema treating the quantization format as a single global field rather than a per-layer property, as seen in the workaround referenced at the bionemo-framework recipe where the format is hardcoded at the recipe level — native support would require the debug module to introspect per-layer quantization config at hook time. The relevant logic is likely in transformer_engine/debug/features/ under the tensor stats feature implementations and the corresponding YAML-driven config parsing.

[nvMelissa]

Hi @pggPL - I was checking if PR #2296 (merged) would resolve any of the request here. Leaving this comment here to help me track comparison and what is left:

PR #2296 — "[PyTorch Debug] NVFP4 debug stats support":

• Adds NVFP4 statistics including underflows and MSE (log_nvfp4_tensor_stats.py)
• Renames fp8-specific variables to quantization-generic names
• Adds stats_computation.py utilities and new disable_quantization_* features

Comparison: Partial overlap, but mostly no.

• std dev of scale factors: Not directly addressed — PR adds MSE and underflows for NVFP4, but the issue specifically asks for std dev of scale factors.
• simultaneous FP8 + FP4 logging: Not addressed in this PR.
• per-block statistics: Not addressed — the PR adds tensor-level NVFP4 stats, not per-block granularity.

The PR is complementary (expanding NVFP4 debug capabilities), and the variable renaming from fp8_* → quantization_* might be a prerequisite step toward enabling mixed FP8/FP4 logging (simultaneous FP8 + FP4), but it doesn't solve it yet.