SequenceParallelEnd¶
标识序列并行结束的策略,该策略应该标识在序列并行结束之后的最后一个 Layer 上。
注解
不要将应该进行序列并行的 Layer 用此策略标识。
参数¶
need_transpose (bool,可选) - 是否在策略中进行
transpose
操作。 如果为 True,本策略会讲形状为[s/mp, b, h]
的输入张量变为形状为[b, s, h]
的输出张量。 反之,本策略会将形状为[s/mp, b, h]
的输入张量变为形状为[s, b, h]
的输出张量。默认为 True。
代码示例
>>> import paddle
>>> import paddle.distributed as dist
>>> class MLP(paddle.nn.Layer):
... def __init__(self):
... super().__init__()
... self.fc1 = paddle.nn.Linear(8, 8)
... self.fc2 = paddle.nn.Linear(8, 8)
...
... def forward(self, input):
... return self.fc2(self.fc1(input))
>>> layer = MLP()
>>> mp_config = {
... 'fc1': dist.SequenceParallelEnd()
... }