已关闭

【实践文档】MagiHuman 15B 音视频扩散 DiT —— FSDP2 后端迁移与调优实践 #498

Ruiyu_Qiu创建于 7月9日关闭于 9月15日

7月9日创建

【实践文档】MagiHuman 15B 音视频扩散 DiT：FSDP2 后端迁移与调优实践

本文是 issue #357（基于 MindSpeed-MM FSDP2 开放源码补齐 MagiHuman 后端功能）的模型适配调优实践，
与迁移代码配套（PR：<https://gitcode.com/Ascend/MindSpeed-MM/merge_requests/2813>），代码引用对应该 PR 分支。
全部硬件 / 显存 / MFU 数据均来自 8×Ascend 910B3 实测。

全文分为两部分：

第一部分说明模型、数据、权重和训练入口如何接入 FSDP2；
第二部分给出性能调优过程、最终验证环境、MFU 计算口径和 5000 步结果。

第一部分：FSDP2 后端迁移

1. 模型简介与迁移目标

1.1 源模型结构

MagiHuman 是单流音视频扩散 DiT。

项目源码：GAIR-NLP/daVinci-MagiHuman
模型权重：GAIR/daVinci-MagiHuman

模型将 VIDEO、AUDIO 和 TEXT token 打包成一条序列，并预测 video/audio 的 flow-matching velocity。模型主体结构如下：

DiTModel
├── adapter                         # video/audio/text 输入投影和 RoPE
├── block
│   └── layers.0 ~ layers.39        # 40 层单流 Transformer layer
│       ├── attention               # GQA，40 个 query head、8 个 KV head
│       └── mlp                     # 首尾各 4 层为 modality-MoE
├── final_norm_video
├── final_norm_audio
├── final_linear_video
└── final_linear_audio

主要结构参数来自上游 inference/common/config.py：

参数	取值
总层数	40
hidden size	5120
head dim	128
query/KV head	40/8
modality-MoE 层	0～3、36～39
模态数	3（VIDEO、AUDIO、TEXT）
video/audio/text 输入宽度	192/64/3584

1.2 迁移目标

目标后端是 MindSpeed MM 插件式 FSDP2，统一入口为 mindspeed_mm/fsdp/train/trainer.py。本次迁移的落点为：

工作项	文件位置	实现方式
模型接入	`mindspeed_mm/fsdp/models/magihuman/magihuman_fsdp2/`	注册 `magihuman`，复用并适配上游 DiT
数据接入	`mindspeed_mm/fsdp/data/datasets/magihuman/`	注册 `magihuman_precomputed` 数据插件，并实现自定义 `collate_fn`
NPU 适配	`magihuman_fsdp2/npu_patch.py`	新增 MagiHuman NPU 适配
权重布局	`magihuman_fsdp2/weight_layout.py`	gated MLP 权重布局的一对互逆变换
checkpoint 转换	`checkpoint/fsdp/custom_model_converter/magihuman.py`	新增 `MagiHumanConverter`，在公开 HF 与内部 DCP 之间转换
FSDP2 配置	`examples/magihuman/magihuman_config_t2av.yaml`	提供全参数训练的顶层 YAML
启动脚本	`examples/magihuman/finetune_magihuman_t2av.sh`	通过 torchrun 启动训练器
性能计算	`mindspeed_mm/fsdp/tools/flops_tool/flops_magihuman.py`	新增 MagiHuman FLOPs/MFU 计算

2. 环境安装

2.1 环境准备

推荐使用如下的环境版本

组件	版本
Python	3.12.12
PyTorch	2.10.0
torch-npu	2.10.0 post4

2.2 环境搭建

# 1. 下载 MagiHuman base 权重。YAML 中的 model.model_name_or_path 和 model.checkpoint_path 指向 ${MODEL_ROOT}/base
MODEL_ROOT=/path/to/daVinci-MagiHuman
hf download GAIR/daVinci-MagiHuman \
  --include "base/*" \
  --local-dir "${MODEL_ROOT}"

# 2. 获取固定版本的上游代码
git clone https://github.com/GAIR-NLP/daVinci-MagiHuman.git
cd daVinci-MagiHuman
git checkout 209209b7086eba2020c5439265221495a8357322
cd ..

# 3. 复制上游 inference 模块到模型插件目录
cp -r daVinci-MagiHuman/inference \
  mindspeed_mm/fsdp/models/magihuman/inference

# 4. 删除 get_dit import，避免导入上游 CUDA checkpoint/distributed 初始化路径
sed -i '/dit_model import get_dit/d' \
  mindspeed_mm/fsdp/models/magihuman/inference/model/dit/__init__.py

3. 数据集准备及处理

3.1 数据格式（magihuman_precomputed）

离线将视频、音频编码为 latent，并将文本编码为 prompt embedding；训练阶段直接读取 .precomputed/* 下的 .pt 文件。

单条样本包含以下字段：

字段	形状	含义
`video_latent`	`[Nv, 192]`	经过 `1×2×2` patchify 的视频 latent
`video_grid_thw`	`(T, H, W)`	视频 token 网格，满足 `T×H×W=Nv`
`audio_latent`	`[Na, 64]`	音频 latent
`text_embed`	`[Nt, 3584]`	文本条件嵌入

3.2 预处理目录结构

.precomputed/
├── latents/*.pt
├── audio_latents/*.pt
└── conditions/*.pt

3.3 noising、packing 和 batch key

collate 对每条样本单独采样 sigma，构造：

x_t = (1 - sigma) * x0 + sigma * noise
target_velocity = noise - x0

video 和 audio 参与 noising，text embedding 保持不变。三种模态按 VIDEO、AUDIO、TEXT 的顺序拼成单流序列。输出给模型的字段为：

batch key	形状	用途
`x`	`[L, 3584]`	noised video/audio token 与 text condition
`coords_mapping`	`[L, 9]`	RoPE 坐标、实际尺寸和参考尺寸
`modality_mapping`	`[L]`	VIDEO=0、AUDIO=1、TEXT=2
`target_velocity`	`[L, 192]`	flow-matching target
`loss_mask`	`[L]`	video/audio 为 true，text 为 false
`cu_seqlens`	`[B+1]`	单流序列中各样本的累计边界
`max_seqlen`	标量	batch 内最长样本的 token 数

模型计算 loss 时会同时检查 modality_mapping 和 loss_mask，只有 mask 为 true 的 video/audio token 参与 flow-matching MSE；text token 和被显式屏蔽的 token 不计入 loss。

3.4 数据接入的配置方式

数据集注册和 collate_fn 就位后，在 YAML 中配置数据集类型、预计算目录和 dataloader：

data:
  dataset_param:
    dataset_type: magihuman_precomputed
    basic_parameters:
      dataset_dir: /path/to/.precomputed
      dataset: magihuman_precomputed
      max_samples: 10000
  dataloader_param:
    pin_memory: true
    shuffle: false
    dataloader_mode: sampler
    drop_last: true
    sampler_type: BaseRandomBatchSampler
    num_workers: 0
    collate_param:
      model_name: magihuman
      ignore_pad_token_for_loss: false

dataset_type 对应数据插件的注册名。dataset_dir 指向 3.2 节所示的 .precomputed目录，目录或必需子目录不存在时会直接报错。MagiHuman dataset 自带 collate_fn，dataloader 会优先使用它完成 noising、packing 和多样本边界组装。其余参数沿用样例 YAML中的默认配置即可。

4. 模型接入

4.1 模型插件组成

MagiHuman 的 FSDP2 模型插件位于 mindspeed_mm/fsdp/models/magihuman/magihuman_fsdp2/：

文件	职责
`__init__.py`	导入模型模块，触发 `magihuman` 注册
`modeling_magihuman.py`	构建上游 DiT、加载权重，并定义训练 forward 和 loss
`modified.py`	适配 DiT、attention 和 layer forward
`npu_patch.py`	替换上游 CUDA attention、编译器依赖和低效算子，并使自定义 linear 支持反向传播
`weight_layout.py`	gated MLP 权重布局的一对互逆变换

4.2 模型注册与训练接口

MagiHuman 不属于 Hugging Face PreTrainedModel 体系，采用自定义 BaseModel 接入。模型类需要实现以下接口：

接口	作用
`_from_config`	根据配置构建上游 `DiTModel` 结构
`from_pretrained`	在 CPU 上构建模型并加载本地 safetensors 权重
`forward`	接收 dataset 输出的 batch，并返回包含 `loss` 的结果

4.3 NPU 训练适配

为使模型能够在 NPU FSDP2 后端完成前向和反向训练，需要进行以下适配。表中的迁移实现文件均位于 mindspeed_mm/fsdp/models/magihuman/magihuman_fsdp2/。

适配项	代码位置	处理方式
CUDA 编译器依赖	`npu_patch.py::disable_magi_compile`	在导入 DiT 前提供最小兼容实现
CUDA attention	`npu_patch.py::npu_flash_attn_with_cp`、`modified.py::attention_forward`	改用 NPU 融合注意力；单样本使用 `BSND`，多样本使用 `TND` 变长注意力
linear 无 backward	`npu_patch.py::_DifferentiableBF16Linear`	改用 `F.linear`，由 autograd 生成 backward
DiT 前向与输出组装	`modified.py::dit_forward`	CP=1 时移除 Ulysses dispatch，并使用 fp32 输出 buffer
Transformer 层接口	`modified.py::block_forward`	接收框架传入的附加参数
RoPE、RMSNorm 逐元素算子	`npu_patch.py::patch_fused_elementwise_ops`	替换为 `npu_rotary_mul` 和 `npu_rms_norm`
gated MLP 的 stride-2 切片	`npu_patch.py::patch_swiglu7_deinterleaved`、`weight_layout.py`	权重改为连续半区布局，`swiglu7` 用 `chunk(2, -1)` 取值

前五项让模型能够在 NPU 上完成前向与反向，后两项来自第二部分的调优，替换上游低效的算子实现与内存访问方式，均由环境变量控制、默认开启。

4.4 权重布局与 checkpoint 转换

连续半区布局改变了训练进程内部的权重排列，而参数名称、shape 和 dtype 全部保持不变。漏掉任何一条加载或导出路径都不会报错，只会得到错误的结果，因此三条路径必须同时覆盖：

路径	入口	转换时机
加载 base 权重	`modeling_magihuman.py::from_pretrained`	权重加载完成、FSDP2 分片之前原地重排
在线加载 HF 权重	`MagiHumanWeightTransformPipeline`	由框架在分片前逐张量调用
checkpoint 导出	`MagiHumanConverter`	在 DCP 与公开 HF 格式之间做互逆转换

三条路径都以转换数量作为校验，期望值为 36 个 gated 层，数量不符直接报错；模型加载时还会在每个 rank 打印实际转换的层数。

训练 checkpoint 使用 DCP 格式，保存和恢复时无需手动转换。对外导出必须经过 MagiHumanConverter 恢复上游布局：

mm-convert MagiHumanConverter dcp_to_hf \
  --load_dir /path/to/dcp/iter_0005000 \
  --save_dir /path/to/exported-hf/base \
  --model_assets_dir /path/to/daVinci-MagiHuman/base

5. 训练 YAML 配置

下面按配置段说明仓库提供的全参训练样例 examples/magihuman/magihuman_config_t2av.yaml，第二部分的实测即使用该配置。

5.1 parallel 段：分片计划

parallel:
  tensor_parallel_size: 1
  fully_shard_parallel_size: auto
  fsdp_plan:
    apply_modules:
      - dit.block.layers.{*}
    param_dtype: bf16
    reduce_dtype: fp32
    num_to_forward_prefetch: 2
    num_to_backward_prefetch: 2
  ring_attention_size: 1
  ulysses_parallel_size: 1
  expert_parallel_size: 1
  expert_fully_shard_parallel_size: 1

tensor_parallel_size 为 1，因此 fully_shard_parallel_size: auto 使用总 world size作为 FSDP2 分片组。apply_modules 来自模型的 named_modules()，每个dit.block.layers.{*} 分别执行 fully_shard。当前配置未启用 TP、EP、Ring Attention或 Ulysses CP。

5.2 model 段：模型注册与权重路径

model:
  model_id: magihuman
  model_name_or_path: &MAGIHUMAN_CKPT /path/to/magihuman-base
  checkpoint_path: *MAGIHUMAN_CKPT
  enable_gradient_checkpointing: false

model_id 对应 @model_register.register("magihuman")。model_name_or_path 和checkpoint_path 指向同一份 MagiHuman base safetensors 目录；模型插件优先读取checkpoint_path，未配置时再使用 model_name_or_path。模型自身不启用 gradient checkpointing，重计算由顶层 features 段统一配置。

5.3 features 段：loss 与重计算

features:
  loss_cfg:
    loss_type: raw
  recompute: true
  recompute_plan:
    apply_modules:
      - dit.block.layers.{*}

loss_type: raw 表示训练引擎直接读取 MagiHumanForTraining.forward 返回的 .loss。当前配置对 40 个 dit.block.layers 全部启用重计算，以减少前向激活占用。

5.4 training 段：训练参数与插件

training:
  micro_batch_size: 1
  gradient_accumulation_steps: 1
  seed: 42
  lr: 1.0e-5
  lr_decay_style: cosine
  lr_warmup_ratio: 0.02
  lr_min: 1.0e-6
  weight_decay: 0.01
  train_iters: 5000
  clip_grad: 1.0
  init_model_with_meta_device: false
  optimizer: adamw
  adam_fused: false
  save: outputs/magihuman_t2av/checkpoints
  save_interval: 500
  use_deter_comp: false
  plugin:
    - mindspeed_mm/fsdp/models/magihuman/magihuman_fsdp2
    - mindspeed_mm/fsdp/data/datasets/magihuman

init_model_with_meta_device: false 时，模型插件在 CPU 上构建 DiT 并加载 base 权重，再进入FSDP2 分片。training.plugin 同时导入模型插件和数据插件，触发二者注册。单份 DCP 约 171 GB，包含模型、FP32 优化器和训练状态，需按磁盘容量设置 save_interval。

6. 启动训练

全参训练使用：

bash examples/magihuman/finetune_magihuman_t2av.sh

脚本通过 torchrun 启动 mindspeed_mm/fsdp/train/trainer.py。调整设备数时修改脚本中的NPUS_PER_NODE，YAML 的 fully_shard_parallel_size: auto 会根据总 world size 设置 FSDP2分片组。脚本中的 MAGIHUMAN_NPU_FUSED_OPS 和 MAGIHUMAN_DEINTERLEAVE_MLP 对应 4.3、4.4 节的两项适配，默认开启。

第二部分：性能调优实践

7. 验证环境

项目	配置
Python	3.12.12
PyTorch	2.10.0
torch-npu	2.10.0.post4
CANN	9.1.0
设备	8×Ascend 910B3，65536 MiB HBM/逻辑设备，单卡 20 个 AI core @ 1800 MHz
模型权重	MagiHuman base，7 个 safetensors 分片，331/331 个 key 命中
训练方式	全参数，非 LoRA

验收使用官方 5 秒、256p 代表性 shape 的合成 precomputed fixture。每条样本的输入为：

模态	形状	token 数
video latent	`[3840, 192]`，网格 `[32, 8, 15]`	3840
audio latent	`[126, 64]`	126
text embedding	`[640, 3584]`	640
packed sequence	`[4606, 3584]`	4606

fixture 由固定 seed 生成 clean tensor；训练时仍在 collate_fn 中随机采样 sigma/noise 并构造 flow-matching target。该输入用于验证官方代表性 tensor shape 下的训练链路与性能，不用于得出收敛性结论。

8. 调优过程

8.1 基线

8 卡全参数训练在未做算子与布局优化时可以稳定运行，显存仍有余量，但吞吐距离 24% MFU 的验收线较远：

指标	基线
稳态 step time	6926.65 ms
有效吞吐	48.0 TFLOP/s/卡
MFU（280 TFLOP/s 分母）	17.15%
峰值 HBM	约 51 GiB/卡

按 280 TFLOP/s 分母计算，达到 24% 需要把 step time 压到约 4.95 s 以内。显存不是瓶颈，问题集中在算子实现与内存访问方式上。

8.2 第一项优化：融合 RoPE 与 RMSNorm

上游的 RoPE 和 RMSNorm 都由多个 PyTorch 算子拼接而成，并在 40 层的 forward、backward 与重计算 forward 中反复执行，产生大量 Vector Core kernel 与 kernel launch：

RoPE：Slice、chunk、Neg、Cat、Mul、Add
RMSNorm：Pow、Mean、Add、Rsqrt、Mul

npu_patch.py 将二者替换为 torch_npu.npu_rotary_mul 和 torch_npu.npu_rms_norm。需要注意 MagiHuman 使用 partial RoPE：head dim 为 128，而 RoPE 宽度为 96，因此只对前 96 维做旋转，其余 32 维原样保留后再拼接。

数值验证结果：

项目	结果
RoPE forward（fp32）	bit-exact
RoPE gradient	bit-exact
RMSNorm（fp32）	最大误差约 1e-6
RMSNorm（bf16）	最大误差 1 个 bf16 ULP

性能变化：

指标	基线	融合算子后
step time	6926.65 ms	6517.50 ms
有效吞吐	48.0 TFLOP/s/卡	51.0 TFLOP/s/卡
MFU	17.15%	18.22%

该项提升约 5.9%，但仍未达到验收线，需要进一步定位。

8.3 Profiling 定位主要瓶颈

采集 8 卡全参数 Profiling，算子耗时分布如下：

算子	占比
ViewCopy	27.76%
MatMulV3 MIX_AIC	14.01%
MatMulV3 AI_CORE	12.01%
Mul	5.22%
StridedSlice	4.28%

ViewCopy 单项超过两路 MatMul 之和。其中 99.6% 来自 72 次大张量拷贝，每次输出形状 [4606, 13652]、耗时约 24.5 ms，合计约 1774 ms。72 次恰好对应 36 个 gated MLP 层各两次切片（层 0～3 使用 GELU7，不含 gate/linear 结构）。

根因在上游 SwiGLU 的取值方式：

x_glu = x[..., ::2]
x_linear = x[..., 1::2]

上游 up_gate_proj 的输出通道按 gate_0, linear_0, gate_1, linear_1, ... 交错排列，因此必须用 stride-2 切片分离两路输入。这样得到的视图在内存中不连续，NPU 无法直接用于后续计算，每次取值都会实例化成一次完整拷贝，有效带宽远低于 HBM 峰值。

8.4 第二项优化：gated MLP 连续半区布局

在加载权重时，把 up_gate_proj.weight 的输出通道由交错排列重排为两个连续半区：

gate_0, linear_0, gate_1, linear_1, ...   ->   gate_0, gate_1, ..., linear_0, linear_1, ...

swiglu7 相应改为 x.chunk(2, dim=-1)，两路输入本身连续，不再触发 ViewCopy。输出通道 i 仍对应同一组 (gate i, linear i)，down_proj 与计算结果都不受影响；重排在加载时原地完成，不额外占用显存。

转换范围为 36 层：

层范围	结构	处理
0～3	GELU7，非 gated	不转换
4～35	单模态 SwiGLU	整块转换
36～39	三模态 SwiGLU	权重含三个连续模态块，每块独立转换

最后 4 层必须按模态块独立处理，不能对整个矩阵一次重排。

数值验证按真实尺寸完成：单 expert SwiGLU 输出、三 expert 各自输出、权重梯度逆转换，以及 interleave→de-interleave→inverse 往返，均为 bit-exact。

调优期间出现过一次值得记录的失败：转换最初被插入到错误的模型构建路径，导致 chunk 已经启用但权重没有重排。训练照常运行、速度也符合预期，只有首步 loss 从预期的约 7.227 变为 5.165 才暴露问题。修正到真实的 from_pretrained 加载路径后，8 个 rank 均确认输出 converted 36 gated MLP layers。这类布局改动不改变参数名称与 shape，一旦漏掉某条路径不会报错，只会静默产生错误结果，因此三条加载/导出路径都加入了转换数量校验。

性能变化：

指标	融合算子后	布局优化后
稳态 step time	6517.50 ms	4638.10 ms
有效吞吐	51.0 TFLOP/s/卡	71.7 TFLOP/s/卡
MFU（280 TFLOP/s 分母）	18.22%	25.61%
峰值 HBM	基本不变	基本不变

8.5 优化路径汇总

阶段	step time	有效吞吐	MFU
8 卡全参基线	6926.65 ms	48.0 TFLOP/s/卡	17.15%
+ 融合 RoPE/RMSNorm	6517.50 ms	51.0 TFLOP/s/卡	18.22%
+ 连续半区 gated MLP	4638.10 ms	71.7 TFLOP/s/卡	25.61%

step time 累计下降 33.1%，吞吐由 48.0 提升至 71.7 TFLOP/s/卡，显存基本没有增加。跨过 24% 验收线的主要贡献来自消除 36 个 gated MLP 层产生的 72 次大规模 stride-2 ViewCopy，而非 FSDP2 参数的调整。

9. MFU 统计口径

统计排除首个初始化迭代，使用迭代 2～5000 的中位 step time。MagiHuman 的 modality-MoE 层按每个 token 实际激活的模态专家计算，因此工具读取到的 active linear parameters 为 10.934B。

python mindspeed_mm/fsdp/tools/flops_tool/flops_magihuman.py \
  --ckpt /path/to/daVinci-MagiHuman/base \
  --seqlens 4606 \
  --gbs 8 \
  --device_num 8 \
  --step_time 4.6348 \
  --hardware_flops 2.8e14

沿用仓库全参数等效的 6ND 统计口径，该配置得到 71.8 TFLOP/s/逻辑设备，对应 25.63% MFU。

硬件峰值取 Atlas 300T A2 训练卡产品规格的 280 TFLOPS，该规格标称的 20 个 AI core 与本机实测一致。若改用按本机 20 core × 1800 MHz 推导的 294.9 TFLOP/s 作为分母，MFU 为 24.33%，同样满足 ≥24% 的要求。MFU 数值与所取分母强相关，引用时需同时说明分母来源。

10. 5000 步结果

指标	结果
完成迭代	5000/5000
consumed samples	40000
稳态 step time 均值	4635.456 ms
稳态 step time 中位数	4634.8 ms
step time P95	4653.5 ms
末 500 步平均 step time	4635.914 ms
中位 step time 对应 MFU	25.63%
峰值 HBM	51384 / 65536 MiB per card
全程平均 loss	4.010511
末 500 步平均 loss	3.998362
最后一步 loss	3.997631
最后一步 grad norm（clip 前）	0.856
退出码	0

训练全程没有记录到 NaN、Inf、OOM、traceback 或分布式失败，gradient norm 保持有限。合成 fixture 中 clean latent 与 noise 独立采样，video 和 audio 两项 MSE 相加后的零预测统计基线约为 4；loss 在前期稳定到该基线附近，并在约 2000 步后持续下降。该结果用于说明迁移后的 MagiHuman FSDP2 全参数路径能够稳定完成 5000 步训练并达到 MFU 验收要求，不构成收敛性结论。

11. Checkpoint 保存、恢复与导出

项目	结果
DCP 体积	约 171 GB，含模型、FP32 AdamW 优化器与训练状态
恢复验证	从 `iter_0005000` 加载后继续执行第 5001 步，consumed samples 由 40000 续至 40008，退出码 0
HF 导出	7 个 safetensors 分片、约 29 GB、331/331 个 key
格式一致性	index、key、shape、dtype 均与上游公开格式一致
权重有效性	331 个 tensor 中 330 个相对 base 权重已发生变化

导出经由 MagiHumanConverter 完成，将内部连续半区布局逆转换回上游交错布局，并按参考文件的 dtype 写回公开格式。

loss 下降情况通过 TrainingLogParser 可视化如下，其中包含了开启CP前后的精度对比（减层然后对比开启CP前后的loss）：

对Ring Attention 的适配以及精度对齐结果如下：

young256成员

7月9日关联了里程碑：MindSpeed 26.2.0

yeqm成员

7月14日评论：

感谢您对 MindSpeed-MM 社区的贡献！MagiHuman 15B 音视频扩散 DiT 的 FSDP2 迁移、调优实践文档及关联 PR 已收到。后续相关代码将由社区 Maintainer 在 PR 中进行审核，请关注并根据评审意见及时更新。待代码审核及相关验证完成后，我们将进一步推进合入，感谢您的支持与配合！