Skip to content

Improve readability of dgrad overlap variable #3165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Prachi-kushwaha wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Improve readability of dgrad overlap variable #3165

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6f15e94
Improve readability of dgrad overlap variable
Prachi-kushwaha Jul 2, 2026
7a3a724
Improve readability of dgrad overlap variable
Prachi-kushwaha Jul 2, 2026
299e7cf
Merge branch 'main' into rename_dgrad_reduce_scatter_overlap
Prachi-kushwaha Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions transformer_engine/pytorch/module/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ def initialize_ub(
"fc2_wgrad",
]
layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"]
dgrad_reduce_scatter_overlap = ["qkv_dgrad", "fc1_dgrad"]
layers_dgrad_ag_and_rs_overlap = ["qkv_dgrad", "fc1_dgrad"]

# Default overlap methods for layers
methods = {
Expand Down Expand Up @@ -490,7 +490,7 @@ def add_ub(

for quantization_mode, user_ub_cfg in zip(quantization_modes, ub_cfgs):
if user_ub_cfg is not None:
for name in dgrad_reduce_scatter_overlap:
for name in layers_dgrad_ag_and_rs_overlap:
if (
name in user_ub_cfg
and "method" in user_ub_cfg[name]
Expand Down Expand Up @@ -523,7 +523,7 @@ def add_ub(
methods["external"] = []
external_gemm_to_overlap.clear()

for name in dgrad_reduce_scatter_overlap:
for name in layers_dgrad_ag_and_rs_overlap:
wgrad_name = name.replace("dgrad", "wgrad")
if name not in layers_reduce_scatter_overlap:
layers_reduce_scatter_overlap.append(name)
Expand Down