Refactor ZeRO, checkpoint and pipeline code

[zkh2016]

重构代码，各种策略解藕：

1. zero和checkpint独立控制，添加开关use_checkpoint，控制是否开启checkpoint
2. checkpoint开关放在CheckpointBlock中，可单模块开关是否使用checkpoint
3. 添加CustomLinear解决zero3参数无法释放问题
4. pipeline，zero2/3，checkpoint可以自由组合
5. 10B模型8卡A100 zero-3不开checkpoint训练性能提升15%，单卡显存增加35%

    bmt.init_distributed(
        zero_level=3, # 支持 Zero2 或 Zero3
        pipe_size=2 #配置流水线并行 stage 数
    )

 self.transformers = bmt.TransformerBlockList([
            bmt.CheckpointBlock(
                TransformerEncoder(
                    dim_model, dim_head, num_heads, dim_ff, bias, dtype
                ),
             use_checkpoint=False, #开启或关闭checkpoint，默认开启
            )
            for _ in range(num_layers)
        ])
 self.transformers(x, mask)

TODO：

1. 自动替换：torch.nn.Linear -> bmt.nn.Linear
2. 添加ZeRO-1的功能
3. 重构通信模块
4. 重构类名、文件名、目录结构

[Achazwl]

1. 找到 reentrant 报错的原因
2. False 可以省计算