【RFC】HyperParallel Trainer 新增 DiT 系列模型支持

1. 需求背景 & 价值

1.1 背景

DiT（Diffusion Transformer）属于扩散 Transformer 范式，是当前文生图、图生图领域的核心架构。与 LLM 的 Seq2Seq 训练流程不同，DiT 训练具有以下特点：

• 输入差异：图像经 VAE 编码后的 latent（4×32×32）而非 token 序列
• 条件机制：timestep 嵌入 + class 条件（或 text 条件）
• 损失定义：噪声预测 MSE（noise prediction），而非 next-token prediction
• 调度器依赖：DDPM/DDIM 等扩散调度器管理噪声添加过程

1.2 当前问题

现有 Trainer 基于 Seq2SeqTrainer 思路，直接迁移 DiT 会遇到：

• 输入/损失定义差异（latent vs token，MSE vs CrossEntropy）
• 扩散调度器（timestep 采样、q_sample 公式）无现成实现
• VAE 编码链路缺失（图像 → latent 空间转换）

1.3 核心价值

• 补全 hyper-parallel 在扩散模型领域的训练适配链路
• 验证 MindSpore + Ascend 后端对 DiT 范式的支持能力
• 为后续 SD3、Flux 等扩散 Transformer 提供最小可复现基线

---

2. 功能描述

2.1 DiT 最小适配器（batch 构造、loss 计算、checkpoint）

• Batch 构造：图像 latent + timestep + class_label 三元组
• Loss 计算：noise prediction MSE（模型预测噪声 vs 真实噪声）
• Checkpoint：支持 DiT 权重保存/加载（含 VAE 编码器可选）

2.2 扩散训练链路

• Timestep 采样：均匀/重要性采样策略
• q_sample：真实扩散公式 sqrt(alpha_t)*x + sqrt(1-alpha_t)*noise
• 噪声调度器：线性 beta 调度（可扩展至 cosine、sigmoid）

2.3 VAE 编码接入

• 支持 sd-vae-ft-mse 等标准 VAE 的编码/解码
• 缩放因子 0.18215 应用
• 训练时冻结 VAE，仅编码图像为 latent

2.4 并行策略验证

• 优先复用 fully_shard（FSDP）处理 DiT 的 Transformer Block 参数
• 逐步接入 TP（attention 头切分）和 PP（stage 切分）

---

3. 设计方案

3.1 整体架构

┌─────────────────────────────────────────────────────────┐
│  用户入口                                               │
│  DiTTrainer(config) → train() → save_checkpoint()       │
├─────────────────────────────────────────────────────────┤
│  数据流 (VAE + GeneratorDataset)                         │
│  ImageFolder → VAE.encode() → .npy → GeneratorDataset   │
│  → batch(latent, t, y)                                   │
├─────────────────────────────────────────────────────────┤
│  训练核心 (MindSpore TrainOneStepCell)                   │
│  q_sample(x, t, noise) → DiT(x_t, t, y) → MSE loss    │
│  → TrainOneStepCell(loss, optimizer)                     │
├─────────────────────────────────────────────────────────┤
│  模型注册 (hyper-parallel ModelSpec)                     │
│  hyper_parallel.models.dit.__init__ → register_spec()  │
└─────────────────────────────────────────────────────────┘

3.2 数据流设计

# 阶段 1：本地预处理（PyTorch）
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
latent = vae.encode(image).latent_dist.sample().mul_(0.18215)
np.save("img_000.npy", latent.numpy())

# 阶段 2：远程加载（MindSpore）
def latent_generator():
    for fname in os.listdir("./latents"):
        latent = np.load(f"./latents/{fname}")  # (4, 32, 32)
        yield latent.astype(np.float32), np.int32(0)

dataset = GeneratorDataset(source=latent_generator, column_names=["x", "y"]).batch(2)

3.3 训练 Step 设计

def construct(self, x, y):
    # 1. 随机采样 timestep
    t = ops.randint(0, self.num_timesteps, (x.shape[0],))

    # 2. 真实 q_sample（扩散公式）
    noise = ops.standard_normal(x.shape)
    x_t = self.diffusion.q_sample(x, t, noise)

    # 3. DiT 预测噪声
    model_output = self.model(x_t, t, y)

    # 4. MSE loss（noise prediction）
    loss = self.loss_fn(model_output, noise)
    return loss

3.4 ModelSpec 注册

# hyper_parallel/models/dit/__init__.py
def _build_dit(cfg):
    model_name = getattr(cfg.model, 'name', 'DiT-S/2')
    return DiT_models[model_name]()

register_spec(
    "DiT-S/2",
    ModelSpec(name="DiT-S/2", build_model_fn=_build_dit)
)

3.5 MindSpore 分布式接入

参考 examples/mindspore/llama3/fsdp_tp_example.py：

os.environ["HYPER_PARALLEL_PLATFORM"] = "mindspore"
from hyper_parallel import fully_shard
from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat

# 对 DiT 的 Transformer Block 应用 fully_shard
for block in model.blocks:
    fully_shard(block, mesh=mesh["dp"])

---

4. 对外 API

4.1 DiTTrainer

from dit_trainer import DiTTrainer

config = {
    'model_name': 'DiT-S/2',
    'weights_path': None,
    'lr': 1e-4,
    'weight_decay': 0,
    'num_timesteps': 1000,
}
trainer = DiTTrainer(config)
result = trainer.train_step(batch)  # {"loss": float}

4.2 数据生成器

from mindspore.dataset import GeneratorDataset

def latent_generator(latent_dir="./latents"):
    for fname in sorted(os.listdir(latent_dir)):
        latent = np.load(os.path.join(latent_dir, fname))
        yield latent.astype(np.float32), np.int32(0)

dataset = GeneratorDataset(source=latent_generator, column_names=["x", "y"]).batch(2)

4.3 ModelSpec 获取

from hyper_parallel.models.spec import get_spec
spec = get_spec("DiT-S/2")
model = spec.build_model_fn(cfg)

---

5. 使用约束

1. 框架限制：当前实现基于 MindSpore 2.9.0，Ascend 910B 验证通过
2. VAE 依赖：训练前需预处理图像为 latent（PyTorch diffusers 推荐）
3. Class 标签：当前使用 0 占位，后续接入 ImageNet class 或 text embedding
4. 并行策略：FSDP 已验证思路，TP/PP 需进一步测试
5. 静态图限制：当前使用动态图（PyNative），静态图需额外适配

---

6. 测试设计

6.1 单元测试

用例 ID	描述	期望
DIT-UT-01	DiT 模型 Forward 输出 shape	(B, 4, 32, 32)
DIT-UT-02	GaussianDiffusion.q_sample 公式正确性	与单卡 PyTorch 参考对齐
DIT-UT-03	TrainOneStepCell 1 个 step 无报错	loss 正常下降
DIT-UT-04	ModelSpec 注册与获取	get_spec("DiT-S/2") 成功
DIT-UT-05	VAE latent 加载与训练	10 个 .npy 文件跑通 5 step

6.2 精度对齐测试

用例 ID	描述	期望
DIT-AL-01	单卡 100 step 基准	平均 loss 稳定，保存 baseline
DIT-AL-02	多卡 DP 100 step	与单卡 loss 误差 < 5e-3
DIT-AL-03	不同 seed 复现性	同配置两次运行 loss 误差 < 1e-4

6.3 回归测试

• 现有 hyper-parallel 的 LLM 训练测试无回归
• MindSpore 后端示例（llama3）仍可正常运行

---

7. 规格 & 约束

7.1 规格

• 支持模型：DiT-S/2（起步），后续扩展 DiT-XL/2、DiT-B/2 等
• 支持后端：MindSpore + Ascend（优先），PyTorch 可扩展
• 支持数据集：ImageNet（class-conditional），后续支持 LAION（text-conditional）

7.2 约束

• VAE 编码需在训练前完成，不支持训练时实时编码（避免框架混用）
• 当前仅支持 class-conditional，text-conditional 需接入 T5/CLIP 编码器
• 多卡并行需申请额外算力券（当前单卡验证完成）

---

8. 参考

• 任务 Issue：#2100
• DiT 原始实现：facebookresearch/DiT
• MindSpore 示例：examples/mindspore/llama3/fsdp_tp_example.py
• VAE 权重：stabilityai/sd-vae-ft-mse
• 权重转换：292 参数映射（PyTorch → MindSpore）

---

附录：当前进度

里程碑	状态	关键数据
DiT MindSpore 模型改写	✅	dit_model_ms.py，33M 参数
权重转换	✅	292 参数映射，DiT-XL/2 ms.ckpt
GaussianDiffusion	✅	真实 q_sample 公式
独立 Trainer	✅	TrainOneStepCell，Ascend 910B
VAE 编码	✅	sd-vae-ft-mse → .npy
真实数据流	✅	GeneratorDataset + latent
ModelSpec 注册	✅	get_spec("DiT-S/2") 验证通过
100 step 单卡基准	✅	平均 loss 0.447393
MindSpore 分布式接入	🟡	参考 llama3 示例，待实现
精度对齐（多卡）	🟡	待实现