Add opt-in rowwise-only quantized primary weights

[xiuhu17]

Description

quantized_model_init normally allocates both rowwise and columnwise storage for primary quantized weights when autograd is enabled. The columnwise representation is required by the default quantized dgrad, but it is not consumed when a fixed backward_override="high_precision" or "dequantized" mode computes dgrad from a high-precision or dequantized weight.

This matters for LoRA-style training with a large frozen quantized base: the base still participates in dgrad, while an unused columnwise copy remains attached to every primary weight.

This PR adds an explicit experimental opt-in:

with te.quantized_model_init(
    enabled=True,
    recipe=recipe,
    omit_columnwise_primary_weight_storage=True,
):
    model = ...

When enabled, primary quantized weights allocate only the rowwise representation required by forward. The option requires an initialization recipe whose backward_override is "high_precision" or "dequantized".

Compatibility and safety

• The default is False, so existing callers retain bidirectional primary-weight storage, including callers that switch from a backward override to default quantized backward at runtime.
• Opt-in models fail early with a descriptive error if a later autocast recipe requests default quantized backward.
• The option is intended for fixed-mode workloads such as frozen LoRA base weights. Optimizer or distributed master-weight writeback paths that require columnwise primary storage must leave it disabled.

For block-scaled formats, omitting one persistent direction saves approximately 1.03 bytes/parameter for MXFP8 and 0.56 bytes/parameter for 1D NVFP4 before shape padding. The NVFP4 test explicitly uses one-dimensional scaling, which cannot synthesize columnwise storage from rowwise storage.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Add omit_columnwise_primary_weight_storage to quantized_model_init and preserve it across nested context state.
• Apply the opt-in when allocating primary-weight storage.
• Validate the initialization recipe and reject an incompatible runtime switch before a quantized dgrad requests missing storage.
• Add MXFP8 and 1D NVFP4 coverage for default storage, rowwise-only storage, complete override forward/backward, preserved default runtime switching, and fail-fast behavior.

Validation

• Repository-configured Black 24.4.2 passed for all three changed Python files.
• Python syntax compilation passed for all three changed Python files.
• Pylint passed for both changed Transformer Engine source files (10.00/10).
• git diff --check passed.
• CUDA tests were not run locally because this host has no CUDA/Blackwell runtime; the added tests use the existing MXFP8/NVFP4 availability gates.
• Full recursive Python lint is blocked on this host by the repository-pinned Astroid version parsing Python 3.14 TemplateStr; the changed source files pass the pinned Pylint independently.

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 public API documentation
• My changes generate no new warnings
• I have added tests that prove the feature works and preserves default compatibility
• New and existing GPU unit tests pass locally (CUDA/Blackwell runtime unavailable locally)

[greptile-apps]

Greptile Summary

Adds an experimental quantized_model_init option that omits columnwise storage from primary quantized weights when dgrad uses a high-precision or dequantized override.

• Captures the storage policy when modules are initialized and applies it during primary-weight allocation.
• Validates initialization recipes and rejects later switches to quantized backward.
• Adds MXFP8 and one-dimensional NVFP4 coverage for storage layouts, successful overridden backward, default compatibility, and incompatible recipe errors.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete blocking or independently actionable non-blocking issue identified.

The opt-in is constrained to backward modes that do not consume columnwise primary-weight storage, default behavior remains unchanged, nested state is restored, and incompatible runtime recipe switches are rejected during forward preparation.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/quantization.py	Adds the public opt-in, validates its initialization recipe, and preserves the new global-state field across nested contexts.
transformer_engine/pytorch/module/base.py	Captures the opt-in per module, suppresses columnwise primary-weight allocation, and rejects incompatible runtime recipes before normal execution.
tests/pytorch/test_backward_override.py	Covers default and rowwise-only layouts for MXFP8 and one-dimensional NVFP4, overridden backward, runtime switching, and validation failures.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[quantized_model_init] --> B{Omit columnwise storage?}
    B -- No --> C[Allocate rowwise and columnwise weight]
    B -- Yes --> D{Backward override configured?}
    D -- No --> E[Reject initialization]
    D -- Yes --> F[Allocate rowwise-only weight]
    F --> G[Module forward preparation]
    G --> H{Current recipe uses override?}
    H -- No --> I[Reject runtime recipe switch]
    H -- Yes --> J[Forward and overridden backward]

Loading

_{Reviews (1): Last reviewed commit: "Merge branch 'NVIDIA:main' into fp8_para..." | Re-trigger Greptile}