[PyTorch] Avoid removing usages from quantized weight tensors

[timmoon10]

Description

We have experienced a correctness error in the grouped linear op when alternating between training and validation. During validation steps, we configure the weight quantizer without column-wise usage, which can cause quantized weight tensors to have stale column-wise data in the next training step. This PR makes sure to avoid disabling quantizer usages in quantized weight tensors, since a usage that was required in the past may be required in the future.

Related: #2222

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

Changes

• Avoid removing usages from quantized weight tensors in basic linear op
• Avoid removing usages from quantized weight tensors in grouped linear op
• Add tests for the linear op that alternate between training and inference steps

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

[timmoon10]

This comment was made incorrect in #1817.

[greptile-apps]

Greptile Summary

This PR fixes a correctness bug where quantized weight tensors could hold stale column-wise data after alternating between training and inference steps. The root cause was that pre_fuser_forward always set columnwise=False on the weight quantizer, and reset_recipe_state could overwrite an enabled column-wise usage with a disabled one; the fix changes the usage to columnwise=requires_grad and makes usage updates monotonically increasing (never disabling what was previously set).

Confidence Score: 5/5

Safe to merge — targeted, well-reasoned bug fix with no API-breaking changes and direct regression tests.

No P0 or P1 issues found. The monotonic-enable logic is self-consistent across both basic and grouped linear ops. The asymmetry between requires_grad (weight quantizer column-wise) and weight_requires_grad (input/grad-output quantizer column-wise) is intentional and correct: column-wise weight data is needed to compute input gradients regardless of whether the weight itself has requires_grad=True. The single_grouped_weight dead-code correction is a clean refactor with no side effects.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/basic_linear.py	Core fix: weight quantizer now uses columnwise=requires_grad in pre_fuser_forward (was always False), and reset_recipe_state monotonically enables — never disables — column-wise usage on the quantized weight tensor.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Same monotonic-usage fix applied; also corrects dead code path where the single_grouped_weight branch inside is_quantized_tensor(weight{group_idx}) was never reachable because weightN attributes don't exist when single_grouped_weight is True.
tests/pytorch/test_fusible_ops.py	Adds TestTrainingLoops with train/infer/train and infer/train/infer stage sequences; promotes to_cpu helper to module scope and extends it to handle QuantizedTensor dequantization.

Sequence Diagram

sequenceDiagram
    participant T as Training step
    participant I as Inference step
    participant PFF as pre_fuser_forward
    participant RRS as reset_recipe_state
    participant WQ as weight_quantizer
    participant WT as weight tensor

    Note over T,WT: Old behaviour (buggy)
    T->>PFF: requires_grad=True
    PFF->>WQ: set_usage(rowwise=True, columnwise=False)
    WQ->>WT: column-wise NOT cached
    I->>PFF: requires_grad=False
    PFF->>WQ: set_usage(rowwise=True, columnwise=False)
    WQ->>WT: column-wise NOT refreshed (stale)
    T->>PFF: requires_grad=True
    PFF->>WQ: set_usage(rowwise=True, columnwise=False)
    Note over WT: backward pass uses stale column-wise data — correctness error

    Note over T,WT: New behaviour (fixed)
    T->>PFF: requires_grad=True
    PFF->>WQ: set_usage(rowwise=True, columnwise=True)
    WQ->>WT: column-wise cached
    I->>PFF: requires_grad=False
    PFF->>WQ: set_usage(rowwise=True, columnwise=False)
    Note over RRS: reset_recipe_state preserves columnwise if previously set
    RRS->>WQ: columnwise stays True (monotonic)
    WQ->>WT: column-wise retained
    T->>PFF: requires_grad=True
    PFF->>WQ: set_usage(rowwise=True, columnwise=True)
    WQ->>WT: column-wise refreshed

Loading

_{Reviews (2): Last reviewed commit: "Restore pre-forward quantizer config in ..." | Re-trigger Greptile}

[timmoon10]

/te-ci pytorch L1

[vthumbe1503]

Shouldnt we update the quantized weight data itself as well to respect the quantizer usages here?

[vthumbe1503]

Do we really need this change of updating quantizer in case of quantized_model_init?

In general, for the quantized_model_init case, I am not sure if it even makes sense to modify the quantizer usages of the quantized_tensor ever during the lifecyle of a module being created.

For example lets say if we create module under torch.no_grad context manager. columnwise_usage will be set to False and if we enter a training loop, we ll be modifying the quantizer usage without modifying the parameters. Leading to quantizer and quantized tensor being in an inconsistent state.

Now in that case we can technically dequantize and quantize it back to have both usages, but we suffer dequantization errors in those case(which might be ok?).

I am wondering should we even touch the quantizer usages in case of quantized_model_init at all after the module is initialized?(for the quantized model parameters)