From 6f15e947a65ae68b4a6fba711845dfa2104d1172 Mon Sep 17 00:00:00 2001 From: Prachi Kushwaha Date: Thu, 2 Jul 2026 07:57:19 +0530 Subject: [PATCH 1/2] Improve readability of dgrad overlap variable --- transformer_engine/pytorch/module/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5eab45b7ac..fdcfbc5c17 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -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 = ["qkv_dgrad", "fc1_dgrad"] # Default overlap methods for layers methods = { @@ -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: if ( name in user_ub_cfg and "method" in user_ub_cfg[name] @@ -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: wgrad_name = name.replace("dgrad", "wgrad") if name not in layers_reduce_scatter_overlap: layers_reduce_scatter_overlap.append(name) From 7a3a7244ac38a69f2b8c1e979aac24c324cd0a5f Mon Sep 17 00:00:00 2001 From: Prachi Kushwaha Date: Thu, 2 Jul 2026 08:05:55 +0530 Subject: [PATCH 2/2] Improve readability of dgrad overlap variable Signed-off-by: Prachi Kushwaha --- transformer_engine/pytorch/module/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index fdcfbc5c17..dcc2e776a8 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -319,7 +319,7 @@ def initialize_ub( "fc2_wgrad", ] layers_reduce_scatter_overlap = ["proj_fprop", "fc2_fprop", "qkv_wgrad", "fc1_wgrad"] - layers_dgrad_ag_and_rs = ["qkv_dgrad", "fc1_dgrad"] + layers_dgrad_ag_and_rs_overlap = ["qkv_dgrad", "fc1_dgrad"] # Default overlap methods for layers methods = { @@ -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 layers_dgrad_ag_and_rs: + for name in layers_dgrad_ag_and_rs_overlap: if ( name in user_ub_cfg and "method" in user_ub_cfg[name] @@ -523,7 +523,7 @@ def add_ub( methods["external"] = [] external_gemm_to_overlap.clear() - for name in layers_dgrad_ag_and_rs: + 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)