[PyTorch] Make grouped weights opt-in

[ksivaman]

Description

Follow-up to #2654 that makes grouped weights opt-in via an envvar. The variable name is inaccurate for now but impending work will convert this opt-in option to use only a single-parameter.

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

• Make grouped weights opt-in.

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

[ksivaman]

/te-ci pytorch

[greptile-apps]

Greptile Overview

Greptile Summary

This PR makes the grouped weights feature opt-in by checking the environment variable NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS. When enabled, GroupedLinear will allocate weights in contiguous memory using GroupedTensor, otherwise it uses the default separate weight allocation.

Changes:

• Modified reset_parameters() in GroupedLinear to conditionally call make_grouped_weights() based on the envvar
• Added single_param test parameter to verify both grouped and non-grouped modes work correctly
• Test now sets/unsets the envvar to test the opt-in behavior

Note: The PR description mentions the variable name is temporary and will be replaced with a single-parameter option in future work.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk
• The changes are straightforward and well-tested. The implementation correctly gates the grouped weights feature behind an environment variable. The test coverage is comprehensive with the new single_param parameter testing both code paths. The only concern is the environment variable cleanup pattern in tests which lacks robustness (already noted in previous review thread).
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Added environment variable check to make grouped weights opt-in via NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS. The implementation is straightforward and correct.
tests/pytorch/test_sanity.py	Added single_param test parameter to test both grouped and non-grouped weight modes. Environment variable cleanup lacks try-finally protection for robustness.

Sequence Diagram

sequenceDiagram
    participant Test as test_sanity_grouped_linear
    participant Env as Environment
    participant GL as GroupedLinear
    participant GT as GroupedTensor
    
    alt single_param=True
        Test->>Env: Set NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS=1
    end
    
    Test->>GL: Create GroupedLinear instance
    GL->>GL: __init__() registers weight parameters
    GL->>GL: reset_parameters()
    
    alt Environment variable set
        GL->>Env: Check NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS
        Env-->>GL: "1" (enabled)
        GL->>GL: make_grouped_weights()
        GL->>GT: Create contiguous GroupedTensor storage
        GT-->>GL: Return grouped tensors
        GL->>GL: Re-register weights as GroupedTensor views
    else Environment variable not set
        GL->>Env: Check NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS
        Env-->>GL: "0" (disabled)
        Note over GL: Skip make_grouped_weights(), use separate weight tensors
    end
    
    Test->>GL: Forward and backward pass
    GL-->>Test: Results
    
    alt single_param=True
        Test->>Env: Delete NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS
    end

Loading

_{Last reviewed commit: 197f2e6}

[greptile-apps]

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[timmoon10]

We're not actually exposing a single param, but rather we're sneakily allocating the weights to be contiguous.

Suggested change

	if bool(int(os.getenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0"))):
	self.make_grouped_weights(defer_init=defer_init)
	if bool(int(os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0"))):
	self.make_grouped_weights(defer_init=defer_init)

Constructor kwargs are also much more Pythonic than these hidden envvars. It's not terrible for now since this is an internal implementation detail, but we should expose a proper kwarg when we actually add the option for a single param.

[zhongbozhu]

LGTM

[greptile-apps]

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

This parsing will raise ValueError if the envvar is set to non-numeric values like "true", "false", or empty string. Consider more robust parsing if users might set this directly:

Suggested change

	if bool(int(os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0"))):
	if os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0").lower() in ("1", "true", "yes"):