Scaling Diffusion Transformers with Mixture of Experts Scaling Diffusion Transfo

Scaling Diffusion Transformers with Mixture of Experts

Scaling Diffusion Transformers with Mixture of Experts (DiT-MoE)

简介： Scaling Diffusion Transformers with Mixture of Experts (DiT-MoE) 通过将Diffusion Transformer与专家混合模型（MoE）结合，构建了一个稀疏版本的扩散Transformer，能够扩展到160亿参数。DiT-MoE在推理阶段具有高度优化的性能，且能够在性能上与稠密网络竞争。

主要内容：

1. 模型实现： 提供了DiT-MoE的PyTorch实现。
2. 预训练权重： 提供了论文中的预训练模型检查点。
3. 采样脚本： 提供了用于运行预训练DiT-MoE模型的采样脚本。
4. 训练脚本： 提供了使用PyTorch分布式数据并行(DDP)和DeepSpeed的训练脚本。

使用场景：

1. 图像生成： DiT-MoE可以用来生成高质量的图像，例如从ImageNet数据集生成的图像。
2. 大规模模型训练： 可以对大规模数据集进行训练，适用于需要高效推理的大模型。
3. 分布式训练与优化： 能够利用多节点分布式训练技术，加快模型训练速度，适用于深度学习研究机构和企业。
4. 特定任务的专家模型： 通过专家混合模型，可以训练特定任务的子模型，提高特定任务的性能。

具体使用说明：

1. 训练：

• 单节点示例：

torchrun –nnodes=1 –nproc_per_node=N train.py
–model DiT-S/2
–num_experts 8
–num_experts_per_tok 2
–data-path /path/to/imagenet/train
–image-size 256
–global-batch-size 256
–vae-path /path/to/vae

• 多节点示例：

torchrun –nnodes=8
–node_rank=0
–nproc_per_node=8
–master_addr=“10.0.0.0”
–master_port=1234
train.py
–model DiT-B/2
–num_experts 8
–num_experts_per_tok 2
–global-batch-size 1024
–data-path /path/to/imagenet/train
–vae-path /path/to/vae

2. 推理：

python sample.py
–model DiT-XL/2
–ckpt /path/to/model
–vae-path /path/to/vae
–image-size 256
–cfg-scale 1.5

未来计划：

1. 提供状态检查点给HuggingFace。
2. 发布合成数据。

引用： 如果使用此代码或模型，请引用以下文章：

@article{FeiDiTMoE2024, title={Scaling Diffusion Transformers to 16 Billion Parameters}, author={Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Jusnshi Huang}, year={2024}, journal={arXiv preprint}, }