Skip to content

Commit a8bad4b

Browse files
yaoyu-33Mcore Bot
authored and
yanring
committed
ADLR/megatron-lm!4312 - [dev] Set tensor-parallel attributes irrespective of perform_initialization
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com> Co-authored-by: yaoyu-33 <yaoyu.094@gmail.com>
1 parent 12839ed commit a8bad4b

3 files changed

Lines changed: 128 additions & 0 deletions

File tree

‎megatron/core/tensor_parallel/layers.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,10 @@ def __init__(
248248
rank=get_pg_rank(self.tp_group),
249249
world_size=get_pg_size(self.tp_group),
250250
)
251+
else:
252+
set_tensor_model_parallel_attributes(
253+
tensor=self.weight, is_parallel=True, dim=0, stride=1
254+
)
251255
else:
252256
self.weight = Parameter(
253257
torch.empty(
@@ -259,6 +263,10 @@ def __init__(
259263
)
260264
if config.perform_initialization:
261265
_initialize_affine_weight_gpu(self.weight, init_method, partition_dim=0, stride=1)
266+
else:
267+
set_tensor_model_parallel_attributes(
268+
tensor=self.weight, is_parallel=True, dim=0, stride=1
269+
)
262270

263271
def forward(self, input_):
264272
"""Forward.
@@ -858,6 +866,10 @@ def __init__(
858866
rank=rank,
859867
world_size=world_size,
860868
)
869+
else:
870+
set_tensor_model_parallel_attributes(
871+
tensor=self.weight, is_parallel=True, dim=0, stride=stride
872+
)
861873
else:
862874
self.weight = Parameter(
863875
torch.empty(
@@ -875,6 +887,10 @@ def __init__(
875887
stride=stride,
876888
is_expert=self.is_expert,
877889
)
890+
else:
891+
set_tensor_model_parallel_attributes(
892+
tensor=self.weight, is_parallel=True, dim=0, stride=stride
893+
)
878894

879895
setattr(self.weight, "allreduce", not (self.is_expert and self.expert_parallel))
880896
else:
@@ -1170,6 +1186,10 @@ def __init__(
11701186
rank=rank,
11711187
world_size=world_size,
11721188
)
1189+
else:
1190+
set_tensor_model_parallel_attributes(
1191+
tensor=self.weight, is_parallel=True, dim=1, stride=stride
1192+
)
11731193
else:
11741194
self.weight = Parameter(
11751195
torch.empty(
@@ -1187,6 +1207,10 @@ def __init__(
11871207
stride=stride,
11881208
is_expert=self.is_expert,
11891209
)
1210+
else:
1211+
set_tensor_model_parallel_attributes(
1212+
tensor=self.weight, is_parallel=True, dim=1, stride=stride
1213+
)
11901214
setattr(self.weight, "allreduce", not (self.is_expert and self.expert_parallel))
11911215

11921216
if bias:

‎megatron/core/transformer/moe/experts.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from megatron.core.tensor_parallel.layers import (
3131
_initialize_affine_weight_cpu,
3232
_initialize_affine_weight_gpu,
33+
set_tensor_model_parallel_attributes,
3334
)
3435
from megatron.core.tensor_parallel.utils import divide
3536
from megatron.core.transformer.mlp import MLP, MLPSubmodules, apply_swiglu_sharded_factory
@@ -208,6 +209,14 @@ def activation_func_with_probs(x, probs):
208209
rank=tp_rank,
209210
world_size=tp_size,
210211
)
212+
else:
213+
# Ensure TP attrs are set even when not initializing
214+
set_tensor_model_parallel_attributes(
215+
tensor=self.weight1, is_parallel=True, dim=1, stride=1
216+
)
217+
set_tensor_model_parallel_attributes(
218+
tensor=self.weight2, is_parallel=True, dim=0, stride=1
219+
)
211220
else:
212221
self.weight1 = Parameter(
213222
torch.empty(
@@ -232,6 +241,14 @@ def activation_func_with_probs(x, probs):
232241
_initialize_affine_weight_gpu(
233242
self.weight2, config.output_layer_init_method, partition_dim=0, is_expert=True
234243
)
244+
else:
245+
# Ensure TP attrs are set even when not initializing
246+
set_tensor_model_parallel_attributes(
247+
tensor=self.weight1, is_parallel=True, dim=1, stride=1
248+
)
249+
set_tensor_model_parallel_attributes(
250+
tensor=self.weight2, is_parallel=True, dim=0, stride=1
251+
)
235252
setattr(self.weight1, 'allreduce', not self.expert_parallel)
236253
setattr(self.weight2, 'allreduce', not self.expert_parallel)
237254

‎tests/unit_tests/tensor_parallel/test_tp_attrs_without_init.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
import pytest
2+
import torch
3+
4+
from megatron.core.tensor_parallel.layers import (
5+
ColumnParallelLinear,
6+
RowParallelLinear,
7+
VocabParallelEmbedding,
8+
)
9+
from megatron.core.transformer.transformer_config import TransformerConfig
10+
from tests.unit_tests.test_utilities import Utils
11+
12+
13+
class TestTPAttributesWithoutInitialization:
14+
15+
def teardown_method(self, method):
16+
Utils.destroy_model_parallel()
17+
18+
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
19+
@pytest.mark.parametrize("use_cpu_init", [True, False])
20+
def test_vocab_parallel_embedding_tp_attrs_no_init(self, use_cpu_init):
21+
Utils.initialize_model_parallel(tensor_model_parallel_size=2)
22+
cfg = TransformerConfig(
23+
num_layers=1,
24+
hidden_size=8,
25+
num_attention_heads=4,
26+
use_cpu_initialization=use_cpu_init,
27+
perform_initialization=False,
28+
)
29+
30+
emb = VocabParallelEmbedding(
31+
num_embeddings=16, embedding_dim=8, init_method=cfg.init_method, config=cfg
32+
)
33+
w = emb.weight
34+
assert hasattr(w, "tensor_model_parallel") and w.tensor_model_parallel is True
35+
assert hasattr(w, "partition_dim") and w.partition_dim == 0
36+
assert hasattr(w, "partition_stride") and w.partition_stride == 1
37+
38+
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
39+
@pytest.mark.parametrize("use_cpu_init", [True, False])
40+
def test_column_parallel_linear_tp_attrs_no_init(self, use_cpu_init):
41+
Utils.initialize_model_parallel(tensor_model_parallel_size=2)
42+
cfg = TransformerConfig(
43+
num_layers=1,
44+
hidden_size=8,
45+
num_attention_heads=4,
46+
use_cpu_initialization=use_cpu_init,
47+
perform_initialization=False,
48+
)
49+
50+
layer = ColumnParallelLinear(
51+
input_size=8,
52+
output_size=8,
53+
init_method=cfg.init_method,
54+
bias=True,
55+
config=cfg,
56+
skip_bias_add=False,
57+
)
58+
w = layer.weight
59+
assert hasattr(w, "tensor_model_parallel") and w.tensor_model_parallel is True
60+
assert hasattr(w, "partition_dim") and w.partition_dim == 0
61+
assert hasattr(w, "partition_stride") and w.partition_stride == 1
62+
63+
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
64+
@pytest.mark.parametrize("use_cpu_init", [True, False])
65+
def test_row_parallel_linear_tp_attrs_no_init(self, use_cpu_init):
66+
Utils.initialize_model_parallel(tensor_model_parallel_size=2)
67+
cfg = TransformerConfig(
68+
num_layers=1,
69+
hidden_size=8,
70+
num_attention_heads=4,
71+
use_cpu_initialization=use_cpu_init,
72+
perform_initialization=False,
73+
)
74+
75+
layer = RowParallelLinear(
76+
input_size=8,
77+
output_size=8,
78+
init_method=cfg.init_method,
79+
bias=True,
80+
input_is_parallel=True,
81+
config=cfg,
82+
skip_bias_add=False,
83+
)
84+
w = layer.weight
85+
assert hasattr(w, "tensor_model_parallel") and w.tensor_model_parallel is True
86+
assert hasattr(w, "partition_dim") and w.partition_dim == 1
87+
assert hasattr(w, "partition_stride") and w.partition_stride == 1

0 commit comments

Comments
 (0)