mindspore.sync_pipeline_shared_parameters

mindspore.sync_pipeline_shared_parameters(net)

在流水线并行场景下，部分参数可能会被不同的stage之间共享。例如 embedding table 被 VocabEmbedding 和 LMHead 两层共享，这两层通常会被切分到不同的stage上。 在流水线并行推理时，有必要 embedding table 变更后在stage之间进行权重同步。

说明

网络需要先编译，再执行stage之间权重同步。

参数：

• net (nn.Cell) - 推理网络。

样例：

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import nn, ops, Parameter, Tensor
>>> class VocabEmbedding(nn.Cell):
...     def __init__(self, vocab_size, embedding_size):
...         super().__init__()
...         self.embedding_table = Parameter(Tensor(np.ones([vocab_size, embedding_size]), ms.float32),
...                                          name='embedding')
...         self.gather = ops.Gather()
...
...     def construct(self, x):
...         output = self.gather(self.embedding_table, x, 0)
...         output = output.squeeze(1)
...         return output, self.embedding_table.value()
...
>>> class LMHead(nn.Cell):
...     def __init__(self):
...         super().__init__()
...         self.matmul = ops.MatMul(transpose_b=True)
...
...     def construct(self, state, embed):
...         return self.matmul(state, embed)
...
>>> class Network(nn.Cell):
...     @lazy_inline
...     def __init__(self):
...         super().__init__()
...         self.word_embedding = VocabEmbedding(vocab_size=4, embedding_size=4)
...         self.head = LMHead()
...
...     def construct(self, x):
...         x, embed = self.word_embedding(x)
...         x = self.head(x, embed)
...         return x
>>>
>>> net = Network()
>>> net.word_embedding.pipeline_stage = 0
>>> net.head.pipeline_stage = 1
>>> x = Tensor(np.ones((8, 4))
>>> net.compile()
>>> ms.sync_pipeline_shared_parameters(net)
>>> print(net.word_embedding.embedding_table.asnumpy())
>>> [[1. 1. 1. 1.]
     [1. 1. 1. 1.]
     [1. 1. 1. 1.]
     [1. 1. 1. 1.]]