[Feature] 希望增加对 32B 模型的多机多卡训练支持

[xushengG]

希望能官方支持并提供 32B Dense 模型的多机多卡分布式训练方案。

具体期望包括：

1. 集成或提供基于主流分布式框架（如 PyTorch FSDP、DeepSpeed ZeRO-3 ）的多机多卡适配代码。
2. 支持跨节点级别的张量并行（TP）或流水线并行（PP）。
3. 提供相关的启动脚本示例（例如基于 torchrun）以及多机网络通信配置的参考文档。

其他补充信息

• **我们目前的硬件环境：8个计算节点，每个节点 8 张 NPU 64GB

[QingGo]

你好！非常感谢你提出这个实用的功能需求～

这个需求涉及的分布式训练、大模型适配、NPU 兼容等工作内容非常多，工程量很大。
因为本项目是个人独立开发维护，目前仅在 GPU 单机环境 下完成了功能验证和调试，客观上要完整落地这个需求存在很大难度。

我梳理了满足你需求的技术实现路径，按实际开发优先级排序如下：

1. DDP 基础分布式训练
2. NPU 环境适配
3. DeepSpeed ZeRO 1/2
4. DeepSpeed ZeRO-3
5. FSDP
6. TP
7. PP

目前我手头只有单卡硬件环境，需要租用 GPU/NPU 才能进行单机多卡、多机多卡的调试和测试，这也是一个瓶颈。

我的规划是：
DDP/NPU适配/DeepSpeed ZeRO 1/2 后续会支持，到时候应该可以支持 13B 左右模型的全参数微调，32B 模型的 PEFT 模式微调（冻结 Backbone，仅训练 Engram/LoRA）
ZeRO-3 / FSDP / TP/PP 对我来说 开发优先级非常低，个人现阶段没有精力和条件推进；
但我非常欢迎社区有志者一起参与，通过提交 PR / 贡献代码的方式，共同完善这块功能！

[aivolcano]

There is a clash if execute multiple GPUs. The embedding is sparse type (the corresponding code is nn.Embedding(sparse=True) in this version of implementing), but the Adam is dense. there is an error in this pair nn.Embedding(sparse=True) and torch.optim.Adam(), Claude code recommended torch.optim.SparseAdam(), which may fix this error. My question is that the original paper from deepseek using torch.optim.Adam() as the optimizer for training. How to fix this clash? I think we can fix it together if you are interested

[QingGo]

There is a clash if execute multiple GPUs. The embedding is sparse type (the corresponding code is nn.Embedding(sparse=True) in this version of implementing), but the Adam is dense. there is an error in this pair nn.Embedding(sparse=True) and torch.optim.Adam(), Claude code recommended torch.optim.SparseAdam(), which may fix this error. My question is that the original paper from deepseek using torch.optim.Adam() as the optimizer for training. How to fix this clash? I think we can fix it together if you are interested

Actually, I already use SparseAdam for training with nn.Embedding(sparse=True). I believe that in DeepSeek's paper, when they mention using "Adam" for the Engram Optimizer, they are referring to a broader category that includes not just torch.optim.Adam, but also variants like AdamW and SparseAdam. I suspect they specifically used SparseAdam with sparse embeddings.

That said, torch.optim.SparseAdam() doesn't have great support for multiple GPUs, so I'm afraid we'll need to find another approach to enable that. Alternatively, you can try setting EngramConfig(use_sparse_embeddings=False), though this will come with the cost of higher memory usage.

Also, contributions to the codebase are very welcome! If you'd like to contribute, please pull the code from the dev branch for development. I've done quite a bit of refactoring there, so it might be quite far from the master branch at this point.

[QingGo]

The newest dev branch now supports multiple GPUs (deepspeed zero 1/2 or DDP) @aivolcano, and also supports NPU @xushengG . I haven't tried it in a real environment though. Welcome to have a try. Any feedback will be appreciated.
You can refer to the core documents on the dev branch for testing:

• Full step-by-step tutorial: tutorial.md
• Dedicated distributed training guide: distributed.md