Automatically omit unused columnwise primary weights for backward overrides

[xiuhu17]

Description

Quantized primary weights normally allocate rowwise and columnwise representations when autograd is enabled. With backward_override="high_precision" or "dequantized", the backward computation does not consume the columnwise primary weight. Keeping that representation wastes persistent memory, especially for a large frozen LoRA base.

Primary-weight initialization now derives its storage usage from the recipe automatically:

recipe = MXFP8BlockScaling(backward_override="dequantized")
with te.quantized_model_init(recipe=recipe):
    model = te.Linear(...)

No storage-policy argument is required. With an override recipe, primary weights are initialized rowwise-only. Without an override, autograd-enabled initialization retains both directions. The same allocation policy applies to the module API and the op fuser's BasicLinear and GroupedLinear.

The implementation records the inferred policy locally to reject an incompatible switch to default quantized backward before it requests missing columnwise storage. Weights initialized with both directions can still enter and leave override mode. There is no new global storage-policy state or public opt-in.

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 padding. The NVFP4 tests use one-dimensional scaling.

Compatibility

This changes storage automatically for override recipes, including trainable primary weights. External optimizers or distributed writeback paths that require both directions must initialize with a non-override recipe or add row-only support. In particular, current Megatron distributed MXFP8 master-weight writeback still assumes columnwise data and scales exist; this PR does not change that integration. Frozen LoRA bases excluded from optimizer parameter groups do not enter that writeback path.

The op fuser GroupedLinear change covers allocation and incompatible-recipe checking; it does not add grouped-op backward-override compute support.

Tests and validation

• MXFP8 and 1D NVFP4: recipe-derived layout and forward/backward coverage for Linear and BasicLinear, plus incompatible recipe switching.
• Bidirectional initialization followed by override and default quantized backward.
• GroupedLinear packed MXFP8 allocation with and without overrides, and incompatible recipe rejection.
• Repository-configured Black 24.4.2, Python syntax compilation, and git diff --check passed.
• Pylint passed for all four changed source files (10.00/10) using the local Python 3.14 lint environment.
• CUDA tests were not run locally; this host has no CUDA/Blackwell runtime.

[greptile-apps]

Greptile Summary

This PR changes quantized primary-weight allocation and adds guards and tests for rowwise-only weights.

• Derives rowwise-only storage automatically from backward_override.
• Rejects later quantized-backward use when columnwise storage is absent.
• Extends MXFP8 and one-dimensional NVFP4 coverage across module and basic-operation implementations.
• The automatic policy removes the stated opt-in and breaks callers that need to retain bidirectional storage while initializing with an override.

Confidence Score: 4/5

The PR is not safe to merge because it makes rowwise-only storage mandatory for override recipes, breaking supported runtime switching and storage-dependent writeback without an opt-out.

Override recipes now always omit columnwise primary-weight storage. Models that previously retained both directions by default can consequently fail either when switching to quantized backward or when a master-weight writeback path accesses the missing columnwise buffers.

Files Needing Attention: transformer_engine/pytorch/module/base.py, transformer_engine/pytorch/ops/basic/basic_linear.py, transformer_engine/pytorch/ops/basic/grouped_linear.py, transformer_engine/pytorch/quantization.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Automatically omits columnwise module-weight storage for override recipes and rejects later quantized backward, introducing a compatibility regression.
transformer_engine/pytorch/ops/basic/basic_linear.py	Applies the same automatic rowwise-only policy and runtime guard to BasicLinear.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Applies automatic rowwise-only allocation and guarding to grouped primary weights.
transformer_engine/pytorch/quantization.py	Removes the explicit storage-policy API and documents automatic omission based on the recipe.
tests/pytorch/test_backward_override.py	Covers the new automatic allocation policy and expected runtime rejection, but no longer preserves explicit opt-in compatibility.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[quantized_model_init with backward override] --> B[Initialize primary weight]
    B --> C[Allocate rowwise storage only]
    C --> D{Later consumer}
    D -->|Override backward| E[Forward and backward proceed]
    D -->|Quantized backward| F[Runtime guard raises]
    D -->|Master-weight writeback| G[Columnwise storage accessed]
    G --> H[Missing-storage failure]

Loading

_{Reviews (2): Last reviewed commit: "Infer primary weight storage from backwa..." | Re-trigger Greptile}

[timmoon10]

This is a nice memory optimization. I don't trust users to configure this correctly, so I'd prefer if we enabled it automatically.

[timmoon10]

Couldn't we deduce this case automatically?

Suggested change

	quantizer.set_usage(
	rowwise=True,
	columnwise=(
	torch.is_grad_enabled() and not self.omit_columnwise_primary_weight_storage
	),
	)
	columnwise_usage = torch.is_grad_enabled()
	if columnwise_usage:
	recipe = FP8GlobalStateManager.get_fp8_recipe()
	if recipe.backward_override in ("high_precision", "dequantized"):
	columnwise_usage = False
	quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)

[timmoon10]

This option is esoteric and low-level, enough that I don't think it's reasonable to expect users to set it correctly. Besides, we can deduce this automatically when constructing the primary weights.

Suggested change

omit_columnwise_primary_weight_storage: bool = False,

[timmoon10]

Similar changes should be made in the op fuser API:

TransformerEngine/transformer_engine/pytorch/ops/basic/basic_linear.py

Lines 331 to 334 in d1e9c33

	quantizer.set_usage(
	rowwise=True,
	columnwise=torch.is_grad_enabled(),
	)

TransformerEngine/transformer_engine/pytorch/ops/basic/grouped_linear.py

Line 466 in d1e9c33

with_columnwise_usage = torch.is_grad_enabled()

[greptile-apps]

Override Removes Required Storage

Initializing a quantized model with either backward override now always disables columnwise primary-weight storage. Previously, omitting this storage was an explicit opt-in, so callers could initialize under an override while retaining both directions. Existing callers that later switch to quantized backward now hit the runtime error at base.py:2072-2077. Distributed master-weight writeback can also access _columnwise_scale_inv and _columnwise_data unconditionally. Since the public storage option was removed, these callers cannot preserve the previous bidirectional layout. Please retain an explicit opt-in instead of deriving the storage layout solely from the recipe. The same automatic policy is applied in basic_linear.py:331-337 and grouped_linear.py:466-473.