Fix GroupedLinear FP8 calibration loop

[fallintoplace]

Summary

• remove the unused outer calibration loop in GroupedLinear
• calibrate each input and weight quantizer once per GEMM
• avoid repeating calibration num_gemms times for every GEMM

Validation

• ran python3 -m py_compile transformer_engine/pytorch/module/grouped_linear.py
• ran git diff --check -- transformer_engine/pytorch/module/grouped_linear.py

[greptile-apps]

Greptile Summary

This PR fixes a bug in _GroupedLinear.forward where FP8 calibration was redundantly applied num_gemms × num_gemms times instead of once per GEMM. The original code had an outer for i in range(num_gemms) loop whose variable was immediately shadowed by two identical inner loops, making the outer loop a no-op wrapper that caused each calibration to repeat num_gemms times for every GEMM.

• Root cause removed: the outer for i loop (with its shadowed variable) and the two separate inner loops are collapsed into a single correct loop that calls input_quantizers[i].calibrate(inputmats[i]) and weight_quantizers[i].calibrate(weights[i]) exactly once each.
• Scope: change is confined to the non-grouped-tensor forward path under fp8_calibration=True; the grouped-tensor path is unaffected.

Confidence Score: 5/5

Safe to merge — the change is a minimal, well-scoped removal of a redundant nested loop in the FP8 calibration path with no behavioural side effects beyond fixing the over-calibration.

The diff is six lines touching a single code path that only runs during FP8 calibration. The original loop structure was unambiguously buggy (outer loop variable shadowed by two inner loops), and the replacement correctly calls each quantizer exactly once. The surrounding GEMM code, the grad bookkeeping, and every other code path are untouched.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Removes redundant nested calibration loops; single loop now calibrates each input and weight quantizer exactly once per GEMM — straightforward and correct fix.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[forward called with fp8_calibration=True] --> B[Build inputmats via torch.split]
    B --> C[Build weights_fp8]
    C --> D[general_grouped_gemm]
    D --> E{fp8_calibration?}
    E -- Yes --> F["for i in range(num_gemms)"]
    F --> G["input_quantizers[i].calibrate(inputmats[i])"]
    G --> H["weight_quantizers[i].calibrate(weights[i])"]
    H --> F
    F -- done --> I[cpu_offloading / grad bookkeeping]
    E -- No --> I

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into fix/grouped-lin..." | Re-trigger Greptile}

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

LGTM