[PyTorch] Guard/document single parameter feature for grouped linear

[ksivaman]

Description

Guard/document single parameter feature for grouped linear.

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

• Guard/document single parameter feature for grouped linear

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 L0

[greptile-apps]

Greptile Summary

This PR guards the experimental single_grouped_weight / single_grouped_bias feature in GroupedLinear behind the NVTE_GROUPED_LINEAR_SINGLE_PARAM environment variable, emitting a UserWarning when the flags are requested without the env var or when the experimental path is enabled. CI scripts are updated to run three test suites with the feature enabled so the code path is exercised in automation. The implementation is straightforward and correct; the only nit is a slightly misleading warning message (see inline comment).

Confidence Score: 5/5

Safe to merge; the gating logic is correct, warnings are informative, and CI coverage of the new path is added.

Only P2 findings (minor warning-message wording). No logic bugs, no data-corruption risk, no security concerns. The function correctly short-circuits when neither flag is set, resolves at construction time, and the stacklevel is accurate for both direct-instantiation call sites.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/utils.py	Adds resolve_grouped_linear_single_param_flags gating logic; warning message says "is not set" even when the env var is explicitly set to "0", and the experimental-enabled warning fires unconditionally on every construction.
transformer_engine/pytorch/module/grouped_linear.py	Calls resolve_grouped_linear_single_param_flags early in init before assigning the instance attributes; docs updated consistently.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Same gating call added to ops-level GroupedLinear.init; docs updated; single_grouped_weight/single_grouped_bias are resolved before weight/bias registration loops.
qa/L0_pytorch_unittest/test.sh	Adds NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 to three test invocations to exercise the single-param code path in CI.
qa/L0_pytorch_debug_unittest/test.sh	Adds NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 to the debug sanity test run.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["GroupedLinear.__init__(single_grouped_weight, single_grouped_bias)"]
    A --> B["resolve_grouped_linear_single_param_flags()"]
    B --> C{Either flag True?}
    C -- No --> D["Return flags unchanged (no warning)"]
    C -- Yes --> E{NVTE_GROUPED_LINEAR_SINGLE_PARAM > 0?}
    E -- No --> F["Warn: env var not enabled\nForce both flags to False"]
    E -- Yes --> G["Warn: experimental feature active\nReturn flags unchanged"]
    F --> H["self.single_grouped_weight = False\nself.single_grouped_bias = False"]
    G --> I["self.single_grouped_weight = requested\nself.single_grouped_bias = requested"]
    D --> J["self.single_grouped_weight = False\nself.single_grouped_bias = False"]

Loading

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

[greptile-apps]

Non-integer env var value will raise ValueError

int(os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0")) will throw an uncaught ValueError if the variable is set to a non-numeric string (e.g. "true", "yes"). Wrapping in a try/except would give a cleaner error message. This is consistent with other similar env-var checks in the file, so this is a pre-existing pattern rather than a new regression — flagging for awareness.

Suggested change

	env_enabled = int(os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0")) > 0
	try:
	env_enabled = int(os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0")) > 0
	except ValueError:
	env_enabled = False

[timmoon10]

This logic breaks backward compatibility since we can no longer rely on the module kwargs to configure single grouped params. I guess we've already been treating single grouped params as an experimental feature. We should keep this instability in mind whenever we use this feature externally, e.g. in Mcore.

[timmoon10]

While we are breaking backward compatibility, we might consider consolidating these options together. Do we really want to take on the burden of supporting the case with a single grouped weight and discrete bias, or discrete weights and single grouped bias?

[timmoon10]

If we only respect the kwargs if an envvar is set, it doesn't really make sense to keep the envvars rather than just checking the envvar. I guess we're half-heartedly maintaining/preparing the stable API for this feature.

[ksivaman]

/te-ci pytorch

[ksivaman]

/te-ci pytorch