Interleave 1F1B impl

https://github.com/NVIDIA/Megatron-LM/issues/169

原版不做Interleave的warmup 只需考虑offset为1

因此warmup_batch的计算公式为：

https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/megatron/schedules.py#L531-L533

然而Interleave版本的1F1B Warmup计算公式变成了：

https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/megatron/schedules.py#L222-L223

Untitled

原Issue第一点提到了使用这样的offset可以使用send_forward_backward_recv_forward_backward 的函数来优化P2P通信，通过打包通信加速。

但是还有一点是因为offset等于1时在interleave 中会遇到通信阻塞问题，比如按以下图分析

Untitled

Sequence-Level 1F1B

现在的LLM基本预训练任务都是GPT的Language Modeling任务，简单来说就是其Attention Mask是一个下三角的形状，每个Token只能看到其前面的Token，这个特性使得在序列维度做流水是可行的。

$$

P(\\mathbf{x}) = \\prod_{t=1}^{T} P(x_{t} \\mid x_{1}, x_{2}, \\ldots, x_{t-1})

$$

由于Transformer中，只有Attention模块会涉及序列间的依赖，只要解决这个依赖问题，实际上就可以在Transformer上进行序列并行。目前序列并行的主要方案：