Skip to content

Commit 537d470

Browse files
li126comli126com
committed
fix comment
1 parent 2fe2891 commit 537d470

3 files changed

Lines changed: 77 additions & 36 deletions

File tree

‎internlm/core/parallel/comm/tensor.py‎

Lines changed: 50 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,9 @@ def input_hook(
6666

6767
@abstractmethod
6868
def grad_output_hook(
69-
self, grad_output: torch.Tensor, async_op: bool = False, no_communication: bool = False
69+
self,
70+
grad_output: torch.Tensor,
71+
async_op: bool = False,
7072
) -> Tuple[torch.Tensor, AsyncCommHandle]:
7173
"""
7274
communication for grad_output when backward.
@@ -82,7 +84,9 @@ def grad_input_hook(self, grad_input: torch.Tensor, async_op: bool = False) -> T
8284

8385
@abstractmethod
8486
def output_hook(
85-
self, output: torch.Tensor, async_op: bool = False, no_communication: bool = False
87+
self,
88+
output: torch.Tensor,
89+
async_op: bool = False,
8690
) -> Tuple[torch.Tensor, AsyncCommHandle]:
8791
"""
8892
communication for output when forward.
@@ -95,13 +99,14 @@ class TensorParallelCommunicator(TPCommunicator):
9599
tensor parallel communicator for linear
96100
"""
97101

98-
def __init__(self, process_group: dist.ProcessGroup, role: LinearRole) -> None:
102+
def __init__(self, process_group: dist.ProcessGroup, role: LinearRole, last_block_layer=False) -> None:
99103
assert role in (LinearRole.COLUMN, LinearRole.ROW), f"Unknown linear role: {role}"
100104

101105
self._process_group = process_group
102106
self._role = role
103107

104108
self._save_total_input = False
109+
self.last_block_layer = last_block_layer
105110

106111
def save_total_input(self) -> bool:
107112
return self._save_total_input
@@ -120,8 +125,7 @@ def input_hook(
120125
def grad_output_hook(
121126
self,
122127
grad_output: torch.Tensor,
123-
async_op: bool = False,
124-
no_communication: bool = False, # pylint: disable=W0613
128+
async_op: bool = False, # pylint: disable=W0613
125129
) -> Tuple[torch.Tensor, AsyncCommHandle]:
126130
"""
127131
tensor parallel should do nothing for grad_output.
@@ -138,12 +142,18 @@ def grad_input_hook(self, grad_input: torch.Tensor, async_op: bool = False) -> T
138142
return all_reduce_raw(grad_input, process_group=self._process_group, async_op=async_op)
139143

140144
def output_hook(
141-
self, output: torch.Tensor, async_op: bool = False, no_communication: bool = False
145+
self,
146+
output: torch.Tensor,
147+
async_op: bool = False,
142148
) -> Tuple[torch.Tensor, AsyncCommHandle]:
143149
"""
144150
all reduce output only for row parallel linear when forward.
145151
"""
146-
if no_communication or dist.get_world_size(self._process_group) <= 1 or self._role == LinearRole.COLUMN:
152+
if (
153+
(self.last_block_layer and gpc.recompute_forward_no_comm)
154+
or dist.get_world_size(self._process_group) <= 1
155+
or self._role == LinearRole.COLUMN
156+
):
147157
return output, DUMMY_HANDLE_CONST
148158

149159
return all_reduce_raw(output, process_group=self._process_group, async_op=async_op)
@@ -155,14 +165,20 @@ class SequenceParallelCommunicator(TPCommunicator):
155165
"""
156166

157167
def __init__(
158-
self, process_group: dist.ProcessGroup, role: LinearRole, save_total_input_as_activation: bool = False
168+
self,
169+
process_group: dist.ProcessGroup,
170+
role: LinearRole,
171+
save_total_input_as_activation: bool = False,
172+
last_block_layer=False,
159173
) -> None:
160174
assert role in (LinearRole.COLUMN, LinearRole.ROW), f"Unknown linear role: {role}"
161175

162176
self._process_group = process_group
163177
self._role = role
164178

165179
self._save_total_input = save_total_input_as_activation
180+
self.last_block_layer = last_block_layer
181+
self.no_communication = False
166182

167183
def save_total_input(self) -> bool:
168184
return self._save_total_input
@@ -189,12 +205,19 @@ def input_hook(
189205
return all_gather_raw(_input, process_group=self._process_group, async_op=async_op, gather_dim=_GATHER_DIM)
190206

191207
def grad_output_hook(
192-
self, grad_output: torch.Tensor, async_op: bool = False, no_communication: bool = False
208+
self,
209+
grad_output: torch.Tensor,
210+
async_op: bool = False,
193211
) -> Tuple[torch.Tensor, AsyncCommHandle]:
194212
"""
195213
all gather grad_output only for row parallel linear when backward.
196214
"""
197-
if no_communication or dist.get_world_size(self._process_group) <= 1 or self._role == LinearRole.COLUMN:
215+
if (
216+
(self.last_block_layer and self.no_communication)
217+
or dist.get_world_size(self._process_group) <= 1
218+
or self._role == LinearRole.COLUMN
219+
):
220+
self.no_communication = False
198221
return grad_output, DUMMY_HANDLE_CONST
199222

200223
return all_gather_raw(grad_output, process_group=self._process_group, async_op=async_op, gather_dim=_GATHER_DIM)
@@ -211,12 +234,19 @@ def grad_input_hook(self, grad_input: torch.Tensor, async_op: bool = False) -> T
211234
)
212235

213236
def output_hook(
214-
self, output: torch.Tensor, async_op: bool = False, no_communication: bool = False
237+
self,
238+
output: torch.Tensor,
239+
async_op: bool = False,
215240
) -> Tuple[torch.Tensor, AsyncCommHandle]:
216241
"""
217242
reduce scatter output only for row parallel linear when forward.
218243
"""
219-
if no_communication or dist.get_world_size(self._process_group) <= 1 or self._role == LinearRole.COLUMN:
244+
self.no_communication = gpc.recompute_forward_no_comm
245+
if (
246+
(self.last_block_layer and self.no_communication)
247+
or dist.get_world_size(self._process_group) <= 1
248+
or self._role == LinearRole.COLUMN
249+
):
220250
return output, DUMMY_HANDLE_CONST
221251

222252
return reduce_scatter_raw(output, process_group=self._process_group, async_op=async_op, reduce_dim=_REDUCE_DIM)
@@ -236,8 +266,7 @@ def __init__(self, parallel_mode: ParallelMode, retain_out_sharded: bool = True)
236266
def grad_output_hook(
237267
self,
238268
grad_output: torch.Tensor,
239-
async_op: bool = False,
240-
no_communication: bool = False, # pylint: disable=W0613
269+
async_op: bool = False, # pylint: disable=W0613
241270
) -> Tuple[torch.Tensor, AsyncCommHandle]:
242271
"""
243272
split grad_output if retain_out_sharded is False.
@@ -248,7 +277,9 @@ def grad_output_hook(
248277
return _split(grad_output, parallel_mode=self._parallel_mode, dim=-1), DUMMY_HANDLE_CONST
249278

250279
def output_hook(
251-
self, output: torch.Tensor, async_op: bool = False, no_communication: bool = False # pylint: disable=W0613
280+
self,
281+
output: torch.Tensor,
282+
async_op: bool = False, # pylint: disable=W0613
252283
) -> Tuple[torch.Tensor, AsyncCommHandle]:
253284
"""
254285
all gather output for head layer if retain_out_sharded is False.
@@ -280,8 +311,7 @@ def __init__(
280311
def grad_output_hook(
281312
self,
282313
grad_output: torch.Tensor,
283-
async_op: bool = False,
284-
no_communication: bool = False, # pylint: disable=W0613
314+
async_op: bool = False, # pylint: disable=W0613
285315
) -> Tuple[torch.Tensor, AsyncCommHandle]:
286316
"""
287317
split grad_output if retain_out_sharded is False.
@@ -293,7 +323,9 @@ def grad_output_hook(
293323

294324
# rewrite ouput communication hook
295325
def output_hook(
296-
self, output: torch.Tensor, async_op: bool = False, no_communication: bool = False # pylint: disable=W0613
326+
self,
327+
output: torch.Tensor,
328+
async_op: bool = False, # pylint: disable=W0613
297329
) -> Tuple[torch.Tensor, AsyncCommHandle]:
298330
"""
299331
all gather output for head layer if retain_out_sharded is False.

‎internlm/model/modules/linear.py‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -45,12 +45,10 @@ def forward(
4545
bias: Optional[torch.Tensor],
4646
communicator: TPCommunicator,
4747
return_residual=False,
48-
no_communication=False,
4948
):
5049
ctx.compute_weight_gradient = weight.requires_grad
5150
ctx.return_residual = return_residual
5251
ctx.communicator = communicator
53-
ctx.no_communication = no_communication
5452

5553
if torch.is_autocast_enabled():
5654
x = x.to(dtype=torch.get_autocast_gpu_dtype())
@@ -79,7 +77,7 @@ def forward(
7977

8078
# parallel strategy-specific communication callback 2.
8179
# see more details in the communicator for different parallel strategies.
82-
output, _ = communicator.output_hook(output, async_op=False, no_communication=no_communication)
80+
output, _ = communicator.output_hook(output, async_op=False)
8381

8482
saved_x = None if ctx.compute_weight_gradient is False else total_x if communicator.save_total_input() else x
8583
ctx.save_for_backward(saved_x, weight)
@@ -93,9 +91,7 @@ def backward(ctx, grad_output, *args):
9391

9492
# parallel strategy-specific communication callback 3.
9593
# see more details in the communicator for different parallel strategies.
96-
grad_output, _ = communicator.grad_output_hook(
97-
grad_output, no_communication=ctx.no_communication, async_op=False
98-
)
94+
grad_output, _ = communicator.grad_output_hook(grad_output, async_op=False)
9995
grad_output = grad_output.contiguous()
10096

10197
if ctx.return_residual:
@@ -268,7 +264,6 @@ def fused_dense_func(
268264
module: Optional[nn.Module] = None,
269265
bias: Optional[torch.Tensor] = None,
270266
return_residual: bool = False,
271-
no_communication=False,
272267
):
273268
if communicator.communication_mode() == "wp":
274269
return WPFusedDenseFunc.apply(
@@ -286,7 +281,6 @@ def fused_dense_func(
286281
bias,
287282
communicator,
288283
return_residual,
289-
no_communication,
290284
)
291285

292286

@@ -349,19 +343,15 @@ def __init__(
349343
else:
350344
super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
351345

352-
self.last_block_layer = False
353-
354346
def forward(self, input: torch.Tensor) -> torch.Tensor: # pylint: disable=W0622
355347
_class_name = self.__class__.__name__
356348
assert self._communicator is not None, f"{_class_name} should register with a communicator first."
357-
no_communication = bool(gpc.recompute_forward_no_comm and self.last_block_layer)
358349
return fused_dense_func(
359350
input,
360351
self.weight,
361352
communicator=self._communicator,
362353
module=self,
363354
bias=self.bias,
364-
no_communication=no_communication,
365355
)
366356

367357

@@ -417,7 +407,6 @@ def __init__(
417407
multiple_of: int = 1,
418408
device: torch.device = None,
419409
dtype: torch.dtype = None,
420-
layer_name: str = "default",
421410
) -> None:
422411
if in_features % multiple_of:
423412
raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}")
@@ -434,9 +423,6 @@ def __init__(
434423
split_mode="row",
435424
)
436425

437-
if layer_name == "w2":
438-
self.last_block_layer = True
439-
440426

441427
class ScaleColumnParallelLinear(ParallelLinearWithCommExt):
442428
"""
@@ -602,15 +588,17 @@ def new_linear(
602588
dtype,
603589
)
604590
elif split_mode == "row":
605-
return RowParallelLinear(
591+
linear = RowParallelLinear(
606592
in_features,
607593
out_features,
608594
bias,
609595
multiple_of,
610596
device,
611597
dtype,
612-
layer_name=name,
613598
)
599+
if name == "w2":
600+
setattr(linear, "last_block_layer", True)
601+
return linear
614602
else:
615603
err_msg = (
616604
f"Parallel strategies for linear is unsupported, which is named as {name}.\n"

‎internlm/train/pipeline.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,15 @@ def initialize_parallel_communicator(model: Union[nn.Module, nn.ModuleList]):
277277
)
278278
_head_communicator = HeadTensorParallelCommunicator(ParallelMode.TENSOR, _retain_out_sharded)
279279
_embedding_communicator = EmbbedingTensorParallelCommunicator(ParallelMode.TENSOR)
280+
281+
# for tp recompute communication optimization, sign last block layer
282+
for row_parallel_linear in _submodule_filter(model, RowParallelLinear):
283+
if getattr(row_parallel_linear, "last_block_layer", False):
284+
row_parallel_linear.register_communicator(
285+
TensorParallelCommunicator(
286+
process_group=gpc.get_group(ParallelMode.TENSOR), role=LinearRole.ROW, last_block_layer=True
287+
)
288+
)
280289
# sequence parallel
281290
if gpc.config.parallel.tensor.mode in ("msp", "fsp"):
282291
save_total_input_as_activation = gpc.config.parallel.tensor.mode == "msp"
@@ -296,6 +305,18 @@ def initialize_parallel_communicator(model: Union[nn.Module, nn.ModuleList]):
296305
)
297306
)
298307

308+
# for tp recompute communication optimization, sign last block layer
309+
for row_parallel_linear in _submodule_filter(model, RowParallelLinear):
310+
if getattr(row_parallel_linear, "last_block_layer", False):
311+
row_parallel_linear.register_communicator(
312+
SequenceParallelCommunicator(
313+
gpc.get_group(ParallelMode.TENSOR),
314+
role=LinearRole.ROW,
315+
save_total_input_as_activation=save_total_input_as_activation,
316+
last_block_layer=True,
317+
)
318+
)
319+
299320
_head_communicator = HeadSequenceParallelCommunicator(
300321
ParallelMode.TENSOR, _retain_out_sharded, save_total_input_as_activation
301322
)

0 commit comments

Comments
 (0)