Improve readability of dgrad overlap variable

[Prachi-kushwaha]

Description

Rename dgrad_reduce_scatter_overlap to better reflect its purpose.

Please include a brief summary of the changes, relevant motivation and context.

Fixes #3164

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

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

[greptile-apps]

Greptile Summary

This PR renames the local variable dgrad_reduce_scatter_overlap to layers_dgrad_ag_and_rs_overlap inside initialize_ub in base.py to better reflect that the listed layers require both an all-gather and a reduce-scatter overlap, not just reduce-scatter.

• The rename is applied consistently at all three sites (one definition on line 322, two for loop references on lines 493 and 526) with no missed occurrences elsewhere in the codebase.
• No logic, default values, or layer lists are changed — this is a pure identifier rename.

Confidence Score: 5/5

Safe to merge — this is a pure local variable rename with no logic changes.

All three occurrences of the old name are updated consistently and the renamed variable is local to initialize_ub, so there is no public API surface or external callsite affected.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Renames local variable dgrad_reduce_scatter_overlap to layers_dgrad_ag_and_rs_overlap in initialize_ub; all three occurrences (1 definition, 2 loop uses) are updated correctly.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["initialize_ub()"] --> B["Define layer lists"]
    B --> C["layers_all_gather_overlap"]
    B --> D["layers_reduce_scatter_overlap"]
    B --> E["layers_dgrad_ag_and_rs_overlap\n(renamed from dgrad_reduce_scatter_overlap)"]
    E --> F["Loop: validate user_ub_cfg methods\n(line 493)"]
    E --> G["Loop: append to layers_reduce_scatter_overlap\nor layers_all_gather_overlap\n(line 526)"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["initialize_ub()"] --> B["Define layer lists"]
    B --> C["layers_all_gather_overlap"]
    B --> D["layers_reduce_scatter_overlap"]
    B --> E["layers_dgrad_ag_and_rs_overlap\n(renamed from dgrad_reduce_scatter_overlap)"]
    E --> F["Loop: validate user_ub_cfg methods\n(line 493)"]
    E --> G["Loop: append to layers_reduce_scatter_overlap\nor layers_all_gather_overlap\n(line 526)"]

Loading

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

[denera]

layers_all_gather_overlap, layers_reduce_scatter_overlap and dgrad_reduce_scatter_overlap are all internal variables that are never exposed to the user. The first two layers_*_overlap lists encode the default comm overlap behavior while dgrad_reduce_scatter_overlap keeps track of which DGRAD GEMMs can optionally be overlapped with reduce-scatter instead of bulk all-gather.

I don't think the new layers_dgrad_ag_and_rs_overlap name in this PR is any more descriptive of how the variable is actually used in our internal workflow.

More importantly, focusing on the variable name is glossing over the actual deficiency — documentation.

The only reason why anyone goes digging in this source code and looks at these internal variables is because we are not telling the user what the default configuration is for each layer. The real solution to this problem is to update the initialize_ub() docstring with that missing information. That would enable users to supply partial config dictionaries that only modify what they want done differently than the default behavior.

So I would recommend that you either change this PR or close and file a new one to supply that docstring update. That would be a very valuable change to review and merge for us.