Skip to content

Commit 925d022

Browse files
cspadescspades
committed
Lint.
Signed-off-by: Cory Ye <cye@nvidia.com>
1 parent 60b68f7 commit 925d022

9 files changed

Lines changed: 46 additions & 64 deletions

File tree

‎transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -575,14 +575,12 @@ def set_device_mesh(
575575
weight_mesh : Optional[DeviceMesh]
576576
Not used for DotProductAttention as there are no quantized weights.
577577
"""
578+
warnings.warn(f"weight_mesh not necessary for {self.__class__.__name__}: {weight_mesh}")
578579
if tp_mesh is not None:
579580
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
580-
assert (
581-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
582-
(
583-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
584-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
585-
),
581+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
582+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
583+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
586584
)
587585
# Set the tensor parallel group from the mesh.
588586
self.set_tensor_parallel_group(tp_mesh.get_group())

‎transformer_engine/pytorch/attention/multi_head_attention.py‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -201,10 +201,10 @@ class MultiheadAttention(torch.nn.Module):
201201
parameters and if the DTensor DeviceMesh includes dimensions that do not
202202
shard weights, such as in the case of HSDP (DP-Replicate x DP-Shard).
203203
For example:
204-
- device_mesh["dp"] for FSDP.
205-
- device_mesh["dp_cp"] if using CP ranks in FSDP.
206-
- device_mesh["tp"] if using TP.
207-
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
204+
- device_mesh["dp"] for FSDP.
205+
- device_mesh["dp_cp"] if using CP ranks in FSDP.
206+
- device_mesh["tp"] if using TP.
207+
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
208208
209209
Optimization parameters
210210
-----------------------
@@ -641,12 +641,9 @@ def set_device_mesh(
641641
"""
642642
if tp_mesh is not None:
643643
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
644-
assert (
645-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
646-
(
647-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
648-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
649-
),
644+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
645+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
646+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
650647
)
651648
# Set the tensor parallel group from the mesh.
652649
self.set_tensor_parallel_group(tp_mesh.get_group())
@@ -655,7 +652,7 @@ def set_device_mesh(
655652
# Iterate through child sub-modules without deep recursion.
656653
# Automatically detects TransformerEngine TP modules and
657654
# the capability to call this method at any level.
658-
for name, child in self.named_children():
655+
for child in self.children():
659656
if hasattr(child, "set_device_mesh"):
660657
child.set_device_mesh(tp_mesh, weight_mesh)
661658

‎transformer_engine/pytorch/module/grouped_linear.py‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -858,12 +858,9 @@ def set_device_mesh(
858858
"""
859859
if tp_mesh is not None:
860860
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
861-
assert (
862-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
863-
(
864-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
865-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
866-
),
861+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
862+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
863+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
867864
)
868865
# Set the tensor parallel group from the mesh.
869866
self.set_tensor_parallel_group(tp_mesh.get_group())

‎transformer_engine/pytorch/module/layernorm.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ def set_device_mesh(
168168
Quantized DTensor parameters are currently not supported for FusibleOperation(s),
169169
and this mesh is not used.
170170
"""
171+
warnings.warn(f"weight_mesh not necessary for {self.__class__.__name__}: {weight_mesh}")
171172
if tp_mesh is not None:
172173
# Construct TP-Replicate DTensors. Used to shim non-TP parameters for compatibility
173174
# with DTensor parameters in TP layers to support DTensor operations.

‎transformer_engine/pytorch/module/layernorm_linear.py‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1112,10 +1112,10 @@ class LayerNormLinear(TransformerEngineBaseModule):
11121112
parameters and if the DTensor DeviceMesh includes dimensions that do not
11131113
shard weights, such as in the case of HSDP (DP-Replicate x DP-Shard).
11141114
For example:
1115-
- device_mesh["dp"] for FSDP.
1116-
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1117-
- device_mesh["tp"] if using TP.
1118-
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
1115+
- device_mesh["dp"] for FSDP.
1116+
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1117+
- device_mesh["tp"] if using TP.
1118+
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
11191119
11201120
Optimization parameters
11211121
-----------------------
@@ -1492,12 +1492,9 @@ def set_device_mesh(
14921492
"""
14931493
if tp_mesh is not None:
14941494
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
1495-
assert (
1496-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
1497-
(
1498-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
1499-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
1500-
),
1495+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
1496+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
1497+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
15011498
)
15021499
# Set the tensor parallel group from the mesh.
15031500
self.set_tensor_parallel_group(tp_mesh.get_group())

‎transformer_engine/pytorch/module/layernorm_mlp.py‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1734,10 +1734,10 @@ class LayerNormMLP(TransformerEngineBaseModule):
17341734
parameters and if the DTensor DeviceMesh includes dimensions that do not
17351735
shard weights, such as in the case of HSDP (DP-Replicate x DP-Shard).
17361736
For example:
1737-
- device_mesh["dp"] for FSDP.
1738-
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1739-
- device_mesh["tp"] if using TP.
1740-
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
1737+
- device_mesh["dp"] for FSDP.
1738+
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1739+
- device_mesh["tp"] if using TP.
1740+
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
17411741
17421742
Optimization parameters
17431743
-----------------------
@@ -2058,12 +2058,9 @@ def set_device_mesh(
20582058
"""
20592059
if tp_mesh is not None:
20602060
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
2061-
assert (
2062-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
2063-
(
2064-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
2065-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
2066-
),
2061+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
2062+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
2063+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
20672064
)
20682065
# Set the tensor parallel group from the mesh.
20692066
self.set_tensor_parallel_group(tp_mesh.get_group())

‎transformer_engine/pytorch/module/linear.py‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1047,10 +1047,10 @@ class Linear(TransformerEngineBaseModule):
10471047
parameters and if the DTensor DeviceMesh includes dimensions that do not
10481048
shard weights, such as in the case of HSDP (DP-Replicate x DP-Shard).
10491049
For example:
1050-
- device_mesh["dp"] for FSDP.
1051-
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1052-
- device_mesh["tp"] if using TP.
1053-
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
1050+
- device_mesh["dp"] for FSDP.
1051+
- device_mesh["dp_cp"] if using CP ranks in FSDP.
1052+
- device_mesh["tp"] if using TP.
1053+
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
10541054
10551055
Optimization parameters
10561056
-----------------------
@@ -1383,12 +1383,9 @@ def set_device_mesh(
13831383
"""
13841384
if tp_mesh is not None:
13851385
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
1386-
assert (
1387-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
1388-
(
1389-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
1390-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
1391-
),
1386+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
1387+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
1388+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
13921389
)
13931390
# Set the tensor parallel group from the mesh.
13941391
self.set_tensor_parallel_group(tp_mesh.get_group())

‎transformer_engine/pytorch/module/rmsnorm.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ def set_device_mesh(
171171
Quantized DTensor parameters are currently not supported for FusibleOperation(s),
172172
and this mesh is not used.
173173
"""
174+
warnings.warn(f"weight_mesh not necessary for {self.__class__.__name__}: {weight_mesh}")
174175
if tp_mesh is not None:
175176
# Construct TP-Replicate DTensors. Used to shim non-TP parameters for compatibility
176177
# with DTensor parameters in TP layers to support DTensor operations.

‎transformer_engine/pytorch/transformer.py‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -258,10 +258,10 @@ class TransformerLayer(torch.nn.Module):
258258
parameters and if the DTensor DeviceMesh includes dimensions that do not
259259
shard weights, such as in the case of HSDP (DP-Replicate x DP-Shard).
260260
For example:
261-
- device_mesh["dp"] for FSDP.
262-
- device_mesh["dp_cp"] if using CP ranks in FSDP.
263-
- device_mesh["tp"] if using TP.
264-
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
261+
- device_mesh["dp"] for FSDP.
262+
- device_mesh["dp_cp"] if using CP ranks in FSDP.
263+
- device_mesh["tp"] if using TP.
264+
- device_mesh["dp_cp_tp"] if strided-sharding with FSDP-TP.
265265
266266
Optimization parameters
267267
-----------------------
@@ -629,12 +629,9 @@ def set_device_mesh(
629629
"""
630630
if tp_mesh is not None:
631631
# Validate TP DeviceMesh / Group. Must be consistent with tp_size.
632-
assert (
633-
tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(),
634-
(
635-
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
636-
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
637-
),
632+
assert tp_mesh.ndim == 1 and self.tp_size == tp_mesh.size(), (
633+
f"TransformerEngine {self.__class__.__name__} TP init size ({self.tp_size}) "
634+
f"does not match the size of the provided TP DeviceMesh ({tp_mesh.size()})."
638635
)
639636
# Set the tensor parallel group from the mesh.
640637
self.set_tensor_parallel_group(tp_mesh.get_group())
@@ -643,7 +640,7 @@ def set_device_mesh(
643640
# Iterate through child sub-modules without deep recursion.
644641
# Automatically detects TransformerEngine TP modules and
645642
# the capability to call this method at any level.
646-
for name, child in self.named_children():
643+
for child in self.children():
647644
if hasattr(child, "set_device_mesh"):
648645
child.set_device_mesh(tp_mesh, weight_mesh)
649646

0 commit comments

Comments
 (0)