Graph Bert 模型搭建

[L1aoXingyu]

在本 issue 中主要记录 Graph Bert 模型搭建的流程以及中间的结果，搭建过程主要分为4步：

• 模型结构
• 模型训练流程
• 下游任务的 finetune 测试指标
• README 以及相关文档和工具代码

模型结构的正确性验证

模型结构的正确性验证中主要包括下面两个步骤：

• Graph 载入 Lazy 模型，在相同输入时，对齐输出结果和 loss (廖星宇）
• Eager, Graph 和 Lazy 模型使用相同的初始化，关掉模型中的随机性(Dropout) 和数据读取的随机性(random shuffle)，关掉 lr scheduler，使用 SGD 进行 1000 轮 loss 对齐 （廖星宇）

使用 dataset/bert_regression_test/0/part-0，训练配置如下

batch size: 32
iter_num: 1000
lr: 1e-3
optimizer: SGD, momentum=0.9
scheduler: None
oneflow commit: ba36de99c

1000 轮 loss 曲线

从上面的结果来看，eager, graph 和 lazy 在相同的数据集上训练 1000 轮，loss 曲线近似重合，可以认为三种方式的结果已经对齐。

模型结构的正确性验证 double check（程鹏）

目前看来lazy和graph版本几乎可以对齐，但是lazy和eager版本还是有差距

实验配置

• 采用相同的模型初始化（load from lazy），关掉dataloader中的随机元素，关掉dropout

batch size: 32
iter_num: 1000
lr: 5e-4
scheduler: None
optimizer：SGD, momentum=0.9

模型训练流程的正确性验证

模型的训练流程主要包括下三个流程

• 完成模型训练所需依赖，在单卡上进行 1000 轮 loss 对齐 Lazy 结果
  • 完成 AdamW with weight_decay_excludes (廖星宇)
    
  • 完成 polynominal lr scheduler (程鹏)
  • 完成 warmup lr scheduler (程鹏)
    • polynomial_scheduler+warmuplr, 在sgd上可以完美对齐，但是在adam和adamw上有细微偏差。
      
      
      
  • 完成 AMP(程鹏)
    • polynomial_scheduler+warmuplr+amp
      
  • 完成 clip gradient (程鹏)
  • 完成 gradient accumulation
  • 完成 Graph consistent 多卡
  • 完成 eager DDP 多卡
  • LAMB optimizer
• 使用 8 卡进行模型训练
• 进行多卡训练性能测试

下游任务 finetune 测试指标验证

下游任务的精度验证主要分为下面三个过程：

• 载入 train 好的 Lazy 模型权重，进行下游任务 finetune，在 SQuAD 上进行精度验证 (廖星宇)

• 根据 TensorFlow Bert 官方实现加入所有下游任务的评测 (廖星宇)
  Lazy Bert 使用不同的 pretrain model 结果如下

  • 使用 tf pretrain model 得到的结果 {"exact_match": 81.7123935666982, "f1": 89.07152353706256}
  • 使用 of pretrain model 得到的结果 {"exact_match": 73.50993377483444, "f1": 82.1944788581551}

  Graph Bert 使用不同的 pretrain model 结果如下

  • 使用 tf pretrain model 得到的结果 {"exact_match": 82.57332071901608, "f1": 89.63355726606137}
  • 使用 of pretrain model 得到的结果 {"exact_match": 73.30179754020814, "f1": 82.10048936661447}

  在 tf 官方实现中需要对齐的精度为 {"exact_match": 80.71901608325449, "f1": 88.4073447260652}，在使用和 tf 相同的 pretrain model 进行训练可以认为已经对齐了 SQuAD 的精度，of pretrain model 的问题等待后续 check

• 载入 train 好的 graph 模型权重，进行下游任务 finetune，在 SQuAD 上进行精度验证

[L1aoXingyu]

记录一下在扩展 bert 多卡 consistent 中遇到的一些坑，这些坑在单卡上并不会显现出来，甚至有的时候在多卡里面也并不会报错，不过最终会影响到 loss 的收敛和精度：

1. 在 nn.Module 的 forward 里面初始化一个 zero tensor 进行 broadcast add 来扩展维度，这个时候初始化的 zero tensor 的 sbp 需要是 sbp.broadcast，placement 和网络中的传递的 tensor 一致，否则 zero 的 sbp 会是 s(0)，在计算的时候会进行维度切分导致出错，参考代码如下

if output.is_consistent:
	zeros = flow.zeros(
        (from_seq_length, to_seq_length),
        dtype=flow.float32,
        placement=output.placement,
        sbp=flow.sbp.broadcast
    )
else:
    zeros = flow.zeros(
        (from_seq_length, to_seq_length),
        dtype=flow.float32,
        device=output.device,
    )

output = output + zeros

2. 在网络中如果要使用 slice 操作时，不能使用 [start:stop] 的操作，这样会造成 s->b 的转化，需要使用 flow.slice，参考代码如下

# position_ids = self.position_ids[:, : self.seq_length]
position_ids = flow.slice(self.position_ids, [[None, None, None], [0, self.seq_length, 1]])

3. 使用 expand 或者 repeat 操作都会造成 s->b 的转化，需要特别小心
   
   附上 loss 计算的代码

def get_masked_lm_loss(
        logit_blob,
        masked_lm_positions,
        masked_lm_labels,
        label_weights,
        max_predictions_per_seq,
    ):
        # NOTE: `repeat` and `expand` will convert `logit_blob` sbp from S(0) to B
        # logit_blob = flow.gather(
        #     logit_blob,
        #     index=masked_lm_positions.unsqueeze(2).repeat(1, 1, args.vocab_size),
        #     dim=1,
        # )
        if logit_blob.is_consistent:
            zeros = flow.zeros(
                (1, 1, args.vocab_size),
                dtype=masked_lm_positions.dtype,
                placement=masked_lm_positions.placement,
                sbp=flow.sbp.broadcast,
            )
        masked_lm_positions = masked_lm_positions.unsqueeze(2) + zeros

        # gather valid position indices
        logit_blob = flow.gather(logit_blob, index=masked_lm_positions, dim=1,)

        logit_blob = flow.reshape(logit_blob, [-1, args.vocab_size])
        label_id_blob = flow.reshape(masked_lm_labels, [-1])

        # The `positions` tensor might be zero-padded (if the sequence is too
        # short to have the maximum number of predictions). The `label_weights`
        # tensor has a value of 1.0 for every real prediction and 0.0 for the
        # padding predictions.
        pre_example_loss = mlm_criterion(logit_blob, label_id_blob)
        pre_example_loss = flow.reshape(pre_example_loss, [-1, max_predictions_per_seq])
        sum_label_weight = flow.sum(label_weights, dim=-1)
        sum_label_weight = sum_label_weight / label_weights.shape[0]
        numerator = flow.sum(pre_example_loss * label_weights)
        denominator = flow.sum(label_weights) + 1e-5
        loss = numerator / denominator
        return loss

[daquexian]

这样会造成 s->b 的转化

s->b 会带来什么问题呢 0。0

[L1aoXingyu]

这样会造成 s->b 的转化

s->b 会带来什么问题呢 0。0

会产生一些 unexpected 的行为，比如 loss 会变成 b，本来 loss 应该是 P 的 :<

[MARD1NO]

这样会造成 s->b 的转化

s->b 会带来什么问题呢 0。0

会产生一些 unexpected 的行为，比如 loss 会变成 b，本来 loss 应该是 P 的 :<

感觉你可以把计算loss的那段代码贴一下，应该更清楚点

[CPFLAME]

lazy graph 四卡ddp consistent + sbp 对齐情况 （弃用，请查看最新进展comment）

关于老版本lazy 输出的loss类型有疑问

现状

• 无法对齐的条件：在graph的loss输出为 broadcast之后的loss时，无法和lazy对齐

  • 代码片段为 loss = loss.to_consistent(sbp=flow.sbp.broadcast).to_local().numpy()
  • loss 曲线为
    

• 对齐的条件: 在graph的loss 输出为 0卡loss * world size时，可以和lazy对齐。

  • 代码片段为 loss = loss.to_local().numpy() * world_size
  • loss 曲线为
    

• 疑问：老版本lazy输出的loss，是broadcast之后的loss (sbp = b) ，还是0卡loss * world size(sbp = p)？

[Ldpe2G]

好奇 rank1, 2, 3 上的 loss 乘以 world size 得曲线是怎样的

[CPFLAME]

我再查一下，好像代码有问题

好奇 rank1, 2, 3 上的 loss 乘以 world size 得曲线是怎样的

[L1aoXingyu]

完成 gradient accumulation 的 loss 对齐，对齐单卡 batch32 和单卡 batch8-gradient-acc-4 的 loss 曲线
训练配置

batch size: 32
iter_num: 300
lr: 1e-3
optimizer: SGD, momentum=0.9
scheduler: None
oneflow commit: 1868c19f26d5

[CPFLAME]

更新：Bert多卡对齐 lazy与graph_consistent 进展

此问题已被pr6288修复，见下面进展comment。

oneflow版本：0.5.0.dev20210912+cu111

• 目前进展：当数据中data part为同一份时，loss可以对齐。当数据中data part为多份不同时，loss不对齐。
• 私人结论：graph_consistent的多卡训练是没有问题的，但是lazy的多卡训练数据读取可能有问题。
• 怀疑：lazy版本训练时，在把数据shuffle全部关掉的情况下，多卡读取的是同一份数据part0，而不是每张卡都会去读取parti
• 怀疑原因：lazy版本训练的代码，在只更改的gpu数量时，loss几乎完全一致，在1e-3量级。在所有的超参一致时，graph_consistent的单卡loss曲线，和lazy 单卡/多卡训练曲线几乎完全一致。

实验准备

由于lazy版本获取信息以及调试比较麻烦，为了证明猜想的正确性，我们用四卡的机器训练时，准备了三份数据集：

• 1.包含四个prat的数据，data_diff；
• 2.包含四个part的数据，但是顺序和上述的不一样，data_shuffle;
• 3.只包含一个part0的数据，但是复制了四分，data_repeat;

各个数据的分布如下图：

实验预期

• 按照正确的理解和预期来说，在这三个数据集上的表现应该是 data_diff == data_shuffle and data_diff != data_repeat, 在实验时，发现graph consistent符合我们正确的预期。
• 如果按照猜想的理解，lazy版本的多卡读取的是同一份数据part0，那么此时的表现应该是data_diff == data_repeat and data_diff != data_shuffle才符合我们对lazy猜想的预期。

实验结果

下面实验结果，可以证明我们猜想的正确性：graph consistent多卡没有问题，但是lazy的多卡训练会读取同一份数据part0

• 在graph上，data_diff == data_shuffle, 见图1；data_diff != data_repeat, 见图2；

• 在lazy上，data_diff != data_shuffle, 见图3；data_diff == data_repeat, 见图4 ；

• lazy和graph对齐上，两者的data_repeat可以对齐，其他的无法对齐，见图5

• 图1， graph的data_diff 和 data_shuffle的对比
  

• 图2， graph的data_diff 和 data_repeat的对比
  

• 图3， lazy的data_diff 和 data_shuffle的对比
  

• 图4， lazy的data_diff 和 data_repeat的对比
  

• 图5， lazy和graph的data_repeat和 data_repeat的对比
  

[L1aoXingyu]

在 oneflow/user/data/ofrecord_dataset.h 的 77 行下面代码上加上下面的 LOG

    in_stream_.reset(
        new PersistentInStream(DataFS(), local_file_paths, !shuffle_after_epoch_, false));
    for (std::string& fn : local_file_paths) {
     LOG(ERROR) << (void *)this << " " << fn << " " << parallel_num_ << " " << parallel_id_;

得到的结果

E0914 16:15:50.051710 42927 ofrecord_dataset.h:79] 0x7fd13800fd10 /dataset/bert/wiki_seq_len_128/part-0 1 0
E0914 16:15:50.050922 42925 ofrecord_dataset.h:79] 0x7fd11c0970e0 /dataset/bert/wiki_seq_len_128/part-0 1 0
E0914 16:15:50.052623 42926 ofrecord_dataset.h:79] 0x7fd15c1bd090 /dataset/bert/wiki_seq_len_128/part-0 1 0
E0914 16:15:50.052976 42927 ofrecord_dataset.h:79] 0x7fd13800fd10 /dataset/bert/wiki_seq_len_128/part-1 1 0
E0914 16:15:50.053155 42925 ofrecord_dataset.h:79] 0x7fd11c0970e0 /dataset/bert/wiki_seq_len_128/part-1 1 0
E0914 16:15:50.053894 42926 ofrecord_dataset.h:79] 0x7fd15c1bd090 /dataset/bert/wiki_seq_len_128/part-1 1 0
E0914 16:15:50.053983 42927 ofrecord_dataset.h:79] 0x7fd13800fd10 /dataset/bert/wiki_seq_len_128/part-2 1 0
E0914 16:15:50.054381 42925 ofrecord_dataset.h:79] 0x7fd11c0970e0 /dataset/bert/wiki_seq_len_128/part-2 1 0
E0914 16:15:50.055184 42926 ofrecord_dataset.h:79] 0x7fd15c1bd090 /dataset/bert/wiki_seq_len_128/part-2 1 0
E0914 16:15:50.055250 42927 ofrecord_dataset.h:79] 0x7fd13800fd10 /dataset/bert/wiki_seq_len_128/part-3 1 0
E0914 16:15:50.055503 42925 ofrecord_dataset.h:79] 0x7fd11c0970e0 /dataset/bert/wiki_seq_len_128/part-3 1 0
E0914 16:15:50.056347 42926 ofrecord_dataset.h:79] 0x7fd15c1bd090 /dataset/bert/wiki_seq_len_128/part-3 1 0
E0914 16:15:50.057664 42926 ofrecord_dataset.h:79] 0x7fd15c1bd090 /dataset/bert/wiki_seq_len_128/part-4 1 0
E0914 16:15:50.056713 42925 ofrecord_dataset.h:79] 0x7fd11c0970e0 /dataset/bert/wiki_seq_len_128/part-4 1 0
E0914 16:15:50.056460 42927 ofrecord_dataset.h:79] 0x7fd13800fd10 /dataset/bert/wiki_seq_len_128/part-4 1 0
...

看上去每张卡上都获得了相同的 data-part，如果不做 shuffle，相当于每次 iter 所有卡上拿到的数据是一样的

[CPFLAME]

再次更新：Bert多卡lazy与graph_consistent 已对齐

依赖于pr6288

用两个数据集分别跑了bert lazy和graph的四卡训练：

• 1.包含四个不同data part 数据，diffpart
• 2.包含同样的四个不同的data part数据, 但是顺序不一致，shuffle

数据的分布如下图：

以下为实验结果：loss基本对齐

[CPFLAME]

eager ddp 实验记录

已被pr6310修复

oneflow版本：0.5.0+cu111.git.b6ca28129

• 问题：当model.forward有多个返回值时，会报TensorTuple相关的错误。当model.forward只有一个返回值时，可以正常运行语句
• 条件：实验采用单卡跑eager的ddp，但是实际上单卡和多卡跑都会触发这个问题
• 错误信息：

Traceback (most recent call last):
  File "run_eager_pretraining.py", line 343, in <module>
    main()
  File "run_eager_pretraining.py", line 324, in main
    lr_scheduler,
  File "run_eager_pretraining.py", line 57, in pretrain
    prediction_scores, seq_relationship_scores = model(input_ids, segment_ids, input_mask)
  File "/home/chengpeng/oneflow_src/oneflow/python/oneflow/nn/module.py", line 84, in __call__
    result = hook(self, args, res)
  File "/home/chengpeng/oneflow_src/oneflow/python/oneflow/nn/parallel/ddp.py", line 68, in post_forward_hook
    convert_to_tensor_tuple([output, *ddp_state_for_reversed_params.keys()])
  File "/home/chengpeng/oneflow_src/oneflow/python/oneflow/framework/tensor_tuple_util.py", line 27, in convert_to_tensor_tuple
    return TensorTuple(args)
TypeError: __init__(): incompatible constructor arguments. The following argument types are supported:
    1. oneflow._oneflow_internal.TensorTuple()
    2. oneflow._oneflow_internal.TensorTuple(arg0: oneflow._oneflow_internal.TensorTuple)
    3. oneflow._oneflow_internal.TensorTuple(arg0: List[oneflow._oneflow_internal.Tensor])

Invoked with: [(tensor([[[0.8790, 0.6964, 0.5749,  ..., 0.1915, 2.0404, 0.9590],
         [0.6722, 0.8361, 0.7003,  ..., 0.8401, 1.4593, 0.4413],
         [1.5696, 1.0837, 1.1004,  ..., 0.9344, 1.3027, 0.8863],
         ...,
         [1.4227, 0.9753, 1.0381,  ..., 0.2575, 1.0045, 1.2830],
         [1.4711, 0.9875, 1.0947,  ..., 0.9168, 1.2916, 0.9532],
         [2.1704, 0.8288, 0.9799,  ..., 0.3927, 1.3398, 1.4371]],

        [[0.8613, 0.7739, 0.6275,  ..., 0.2011, 1.8372, 1.0068],
         [0.0153, 1.3177, 0.9809,  ..., 0.8060, 1.4150, 0.4955],
         [0.9502, 1.0506, 0.7459,  ..., 0.8688, 0.7254, 1.1827],
         ...,
         [1.2794, 1.2015, 0.5401,  ..., 0.9088, 0.8193, 0.7978],
         [2.0678, 1.0663, 1.5910,  ..., 0.9943, 0.8990, 1.1975],
         [2.1770, 0.9250, 1.0521,  ..., 0.3537, 1.2147, 1.4396]],

        [[0.8747, 0.7218, 0.5811,  ..., 0.1465, 2.0335, 0.9580],
         [0.3103, 0.9816, 0.9865,  ..., 0.3626, 1.5666, 0.7451],
         [1.2121, 1.3677, 0.5258,  ..., 0.6686, 1.9435, 0.7462],
         ...,
         [1.1546, 1.5225, 1.0765,  ..., 0.8886, 1.6548, 1.0114],
         [1.1560, 1.4781, 1.4344,  ..., 1.2447, 1.7169, 1.0197],
         [1.3914, 1.5772, 1.6321,  ..., 0.5102, 1.8968, 1.6388]],

        ...,

        [[0.9710, 0.9271, 0.6405,  ..., 0.3756, 1.5940, 0.9250],
         [0.5394, 1.5854, 0.5589,  ..., 0.8092, 1.2758, 0.5898],
         [0.7933, 1.7139, 0.7777,  ..., 0.7304, 0.5660, 1.3148],
         ...,
         [0.8432, 1.6537, 1.1371,  ..., 0.8997, 1.1822, 0.9852],
         [0.9176, 1.4704, 1.4282,  ..., 1.2452, 1.3234, 0.9216],
         [1.1842, 1.7367, 1.5877,  ..., 0.5134, 1.4082, 1.5590]],

        [[0.9004, 0.6829, 0.5809,  ..., 0.1883, 1.9453, 0.9190],
         [1.1199, 1.0601, 0.7142,  ..., 0.5268, 1.4391, 0.3941],
         [0.6110, 1.2423, 1.1521,  ..., 1.1103, 1.5401, 1.1477],
         ...,
         [1.0583, 1.4422, 1.0641,  ..., 0.9390, 1.5376, 1.0295],
         [1.1035, 1.3861, 1.4205,  ..., 1.2668, 1.6283, 0.9972],
         [1.3983, 1.5536, 1.5878,  ..., 0.5255, 1.7549, 1.6392]],

        [[0.9319, 0.7121, 0.5555,  ..., 0.1484, 2.0053, 1.0046],
         [0.1778, 1.2731, 0.9279,  ..., 0.7979, 1.5128, 0.4934],
         [1.5921, 1.0481, 0.6490,  ..., 0.9822, 1.0786, 1.2100],
         ...,
         [1.1461, 1.4780, 1.0662,  ..., 0.9224, 1.5571, 1.0558],
         [1.1612, 1.4581, 1.4172,  ..., 1.2474, 1.6505, 1.0881],
         [1.4717, 1.5679, 1.5753,  ..., 0.5098, 1.8399, 1.7004]]],
       device='cuda:0', dtype=oneflow.float32, grad_fn=<broadcast_add_backward>), tensor([[ 0.0237, -0.1002],
        [ 0.1160, -0.0664],
        [ 0.0434, -0.0974],
        [ 0.1227, -0.0644],
        [ 0.0458, -0.0984],
        [ 0.1560, -0.0432],
        [ 0.1940, -0.0382],
        [ 0.0647, -0.0679],
        [ 0.1095, -0.0644],
        [ 0.1752, -0.0356],
        [ 0.1375, -0.0656],
        [ 0.0194, -0.0750],
        [ 0.1485, -0.0556],
        [ 0.0556, -0.0874],
        [ 0.0475, -0.0712],
        [ 0.0043, -0.1024],
        [ 0.1213, -0.0789],
        [ 0.0871, -0.0788],
        [-0.0432, -0.1042],
        [ 0.0252, -0.0762],
        [ 0.1884, -0.0582],
        [ 0.0508, -0.0669],
        [ 0.1068, -0.0750],
        [ 0.0905, -0.0668],
        [ 0.1034, -0.0655],
        [ 0.1446, -0.0845],
        [ 0.1444, -0.0627],
        [ 0.0520, -0.0563],
        [ 0.1660, -0.0363],
        [ 0.1960, -0.0335],
        [ 0.0826, -0.0893],
        [ 0.0474, -0.0598]], device='cuda:0', dtype=oneflow.float32,
       grad_fn=<broadcast_add_backward>)), tensor([0., 0.], device='cuda:0', dtype=oneflow.float32,
       grad_fn=<accumulate_grad>), tensor([[-0.0159, -0.0007,  0.0055,  ...,  0.0018, -0.0217,  0.0340],
        [-0.0230,  0.0134, -0.0179,  ...,  0.0068,  0.0325,  0.0168]],

[CPFLAME]

bert实验进展更新：四卡 eager+ddp 已和 graph+consistent对齐

• 注意：在对齐loss时，需要注意loss的等价关系，由于eager+ddp是和pytorch进行对齐的，所以各个卡上的loss实际上是本卡的batch size上loss的平均值，如果要和graph consistent braodcast的loss进行严格对齐，换算关系应该是loss1 + loss2 + …. + lossn）/ n == graph consistent to broadcast loss n为卡的数量