张量模型并行¶

class VocabParallelEmbedding(Layer):
    def __init__(self,
                 num_embeddings,  # Embedding 参数的行数
                 embedding_dim,   # Embedding 参数的列数
                 weight_attr=None,
                 name=None):
        super(VocabParallelEmbedding, self).__init__()

class RowParallelLinear(Layer):
    def __init__(self,
                 in_features,
                 out_features,
                 weight_attr=None,
                 has_bias=True,
                 input_is_parallel=False, #输入是否是并行输入，为否的话需要按列切分输入参数
                 name=None):
        super(RowParallelLinear, self).__init__()

class ColumnParallelLinear(Layer):
    def __init__(self,
                 in_features,
                 out_features,
                 weight_attr=None,
                 has_bias=None,
                 gather_output=True, # 是否在该算子后汇聚所有卡的输出
                 name=None):


RNG_STATE_TRACKER = RNGStatesTracker() # RNGStatesTracker 是全局随机性控制器，用于保证张量模型并行的随机性
def get_rng_state_tracker():
   return RNG_STATE_TRACKER

class RNGStatesTracker(object):
   def __init__(self):
      self.states_ = {}     # 存储随机状态
      self.seeds_ = set()   # 可以设置多个随机种子

import paddle.distributed.fleet as fleet
word_embeddings = fleet.meta_parallel.VocabParallelEmbedding(
    vocab_size,
    hidden_size,
    weight_attr=paddle.ParamAttr(initializer=nn.initializer.Normal(
                  mean=0.0, std=initializer_range)))

dist_strategy = paddle.distributed.fleet.DistributedStrategy()
strategy.hybrid_configs = {
    "mp_degree": 2,
    "dp_degree": 1,
}
fleet.init(is_collective=True, strategy=strategy)
hcg = fleet.get_hybrid_communicate_group()
global_rank = hcg.get_global_rank() # 全局 rank
mp_rank = hcg.get_model_parallel_rank() # 模型并行组 rank

from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
local_seed = basic_seed + mp_rank * 11
global_seed = basic_seed
tracker.add('global_seed', global_seed)
tracker.add('local_seed', local_seed)

# For local dropout
import paddle.nn.functional as F
from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
with get_rng_state_tracker().rng_state('local_seed'):
    weights = F.dropout(
                      weights,
                      dropout_rate,
                      training=True,
                      mode='upscale_in_train')

# For global dropout
with get_rng_state_tracker().rng_state('global_seed'):
    weights = F.dropout(
                      weights,
                      dropout_rate,
                      training=True,
                      mode='upscale_in_train')

import paddle
import numpy as np
import random
import paddle.distributed as dist
import paddle.fluid as fluid
import paddle.distributed.fleet as fleet

vocab_size = 20
hidden_size = 10
inner_size = 8
output_size = 10
seq_length = 2
batch_size = 4

class SimpleMPNet(fluid.dygraph.Layer):
   def __init__(self, vocab_size, hidden_size, inner_size, output_size):
      super(SimpleMPNet, self).__init__()
      self.linear1 = fleet.meta_parallel.ColumnParallelLinear(
            hidden_size,
            inner_size,
            gather_output=False,
            has_bias=True)

      self.linear2 = fleet.meta_parallel.RowParallelLinear(
            inner_size,
            hidden_size,
            input_is_parallel=True,
            has_bias=True)

      self.linear3 = paddle.nn.Linear(hidden_size, output_size)

      self.embedding = fleet.meta_parallel.VocabParallelEmbedding(
                        vocab_size,
                        hidden_size)

   def forward(self, x):
      x = self.embedding(x)
      x = self.linear1(x)
      x = self.linear2(x)
      x = self.linear3(x)
      return x

def set_random_seed(seed, rank_id):
   random.seed(seed)
   np.random.seed(seed)
   paddle.seed(seed + rank_id)

strategy = fleet.DistributedStrategy()

# 设置两路张量模型并行
model_parallel_size = 2
data_parallel_size = 1
strategy.hybrid_configs = {
   "dp_degree": data_parallel_size,
   "mp_degree": model_parallel_size,
   "pp_degree": 1
}
# 注意 strategy 是这里传递的，动态图只能这里，静态图还可以在 distributed_optimizer 里传
fleet.init(is_collective=True, strategy=strategy)

hcg = fleet.get_hybrid_communicate_group()
mp_id = hcg.get_model_parallel_rank()
rank_id = dist.get_rank()
set_random_seed(1024, rank_id)

model = SimpleMPNet(vocab_size, hidden_size, inner_size, output_size)

optimizer = paddle.optimizer.SGD(learning_rate=0.001, parameters=model.parameters())
model = fleet.distributed_model(model)
optimizer = fleet.distributed_optimizer(optimizer)


for _ in range(5):
   np_data = np.random.randint(0, vocab_size, (batch_size, seq_length, ))

   output = model(paddle.to_tensor(np_data))
   loss = output.mean()
   loss.backward()
   optimizer.step()
   optimizer.clear_grad()
   print("loss", loss.numpy())