Title 使用deepspeed训练Yi-34B 32K 以及200K上下文爆显存的问题

[lyccol]

Reminder

• [X] I have searched the Github Discussion and issues and have not found anything similar to this.

Environment

- OS:centos
- Python:3.10
- PyTorch:2.0.1+cu118
- CUDA:11.8

Current Behavior

在自己的数据上，对Yi-34B 做SFT，但是上下文调到32K及以上，会爆显存。

Expected Behavior

Deepspeed貌似只支持数据并行，是不是需要加入Megatron-LM的模型并行，通过拆分模型到多张卡上来进行训练。

官方训练过200K的Yi-34B模型，是否有现成的代码？

Steps to Reproduce

在自己的数据上，对Yi-34B 做SFT，但是上下文调到32K及以上，会爆显存。

Anything Else?

https://github.com/microsoft/DeepSpeed/discussions/1911

[Yimi81]

我们使用megatron在6nodes8GPUs/nodeA800 训练的200K-Yi-34B，不会提供现成的代码。

[Yimi81]

此外你可以尝试一下deepspeed zero3，把模型参数进行拆分到不同的GPU上，看看能不能跑起来

[puppet101]

@Yimi81 请问megatron直接全量训练200K的长度，只需要6*8张A800吗？ 有什么可以参考的开源实现或者技术介绍推荐一下吗？ 我这边单机8张A800，连4K长度都会OOM

[Yimi81]

可以参考一下我们的技术报告看看有没有什么细节，最低要求是6nodes8GPUs/nodeA800 ，3840G，具体的参数设置是保密的不会提供