已开启

PR1036 qwen3_vl_moe支持 Visual Encoder 独立 DP/CP 配置与自洽验证 #287

Aoalas创建于 7月13日

7月13日创建

关联PR:https://atomgit.com/mindspore/hyper-parallel/pull/1036

qwen3_vl_moe 支持 Visual Encoder 独立 DP/CP/异步 CP

背景

多模态训练中，Visual Encoder 与文本 Decoder 的输入形态和并行瓶颈并不一致。视觉侧需要处理图像或视频 patch embedding、变长视觉序列 attention、patch merge 以及 DeepStack 视觉特征注入；文本侧则主要沿用 LLM Decoder 的 FSDP/TP/CP/EP 组合。如果视觉塔只能跟随文本 Decoder 的全局并行策略，就难以单独控制视觉侧参数复制、视觉 attention 序列切分和通信开销。

本次任务面向 qwen3_vl_moe，在现有 HyperParallel 声明式并行能力之上，补齐 Visual Encoder 局部 DP/CP 配置、视觉 attention CP 边界、多模态 CP 数据流、精度对齐、性能对比、样例和文档。

任务目标

在 model.vision_parallel 下提供 Visual Encoder 局部并行配置入口；
支持 Visual Encoder 参数独立 DP/FSDP 策略，允许视觉塔参数复制而文本 Decoder 继续使用全局 DP/FSDP；
支持 Visual Encoder attention core 独立 CP，并覆盖 Pure Colossal 和 Pure Ulysses 两种 2 卡模式；
支持 Visual Encoder 异步 CP；
保持多模态 token 注入、mRoPE position ids、DeepStack 视觉特征在 CP 序列切分下正确对齐；
保持视觉 patch embedding 的 checkpoint 参数名和参数布局兼容；
提供精度对齐、2 卡性能对比、样例配置和使用文档。

最终差异

原始能力	最终能力
视觉塔主要跟随全局并行策略，无法单独表达视觉侧 DP/CP。	新增 `model.vision_parallel`，视觉塔可独立配置 `dp_shard`、`cp`、`ulysses_degree` 和 `async_cp`。
视觉 attention 是较完整的 block 内部逻辑，缺少可单独挂载 CP 的 core 边界。	拆出 `Qwen3VLMoeVisionSdpaCore`，仅在视觉 attention core 上挂载 `ContextParallel` 或 `AsyncContextParallel`。
文本 CP 会让输入序列变成本地 shard，多模态融合需要完整序列信息。	主模型在视觉 token 注入和 mRoPE 计算前 gather 完整文本序列，融合完成后再 slice 回本地 CP shard。
packed visual sequence 未按图片/帧边界做视觉 CP 切分与恢复。	视觉塔根据 `grid_thw -> cu_seqlens -> sequence_lengths` 对每个 packed 视觉序列切分，block 后再按原顺序 gather。
DeepStack 视觉特征在文本 CP 本地 shard 下缺少稳定映射。	DeepStack 特征按全局 visual mask 建立索引，再选择本 rank 文本序列切片对应的视觉特征。
patch embedding 依赖 `Conv3d` forward/backward 路径。	保留 `Conv3d` 参数名和布局，forward 使用等价 `F.linear` 投影，保证 checkpoint 兼容并规避 Ascend A2 `Conv3d backward` 不稳定路径。

配置入口

推荐在模型侧显式配置视觉并行策略：

model:
  name: qwen3_vl_moe
  vision_parallel:
    dp_shard: 1
    cp: 2
    ulysses_degree: 1
    async_cp: false
  config_overrides:
    vl: true
train:
  accelerator:
    dp_shard: 1
    cp: 2

配置项说明：

配置项	说明
`dp_shard`	Visual Encoder 的 DP shard size。`1` 表示视觉参数复制；未配置时跟随全局 `train.accelerator.dp_shard`。当前支持 `1` 或全局 DP shard size。
`cp`	Visual Encoder CP size。大于 `1` 时需要与 `train.accelerator.cp` 的 mesh size 匹配。
`ulysses_degree`	视觉 CP 模式选择。2 卡下 `1` 表示 Pure Colossal，`2` 表示 Pure Ulysses。
`async_cp`	是否在视觉 attention core 上启用异步 CP。默认关闭。

方案设计

1. 视觉并行配置解析

ModelConfig 提供 vision_parallel 字段，模型并行化阶段通过 get_vision_parallel_config() 统一读取视觉侧局部并行参数。这样视觉侧 DP/CP 不需要塞进文本 Decoder 的通用并行配置，也避免模型代码在多个位置重复解析 YAML。

2. Visual Encoder 独立 DP

视觉塔的 DP 逻辑由 hyper_parallel/models/qwen3_vl_moe/parallelize.py 处理：

_resolve_vision_dp_shard_size() 根据 model.vision_parallel.dp_shard 和全局 train.accelerator.dp_shard 得到视觉塔 DP shard size；
当 vision_parallel.dp_shard=1 时，_extend_visual_replicate_params() 将 model.visual 参数加入 FSDP replicate_params，视觉塔参数保持复制；
当视觉塔跟随全局 DP shard 时，_apply_vl_visual_tower() 将视觉 blocks、merger、DeepStack mergers 和视觉塔 root 作为单独 FSDP group 包装；
文本 Decoder 的 FSDP/TP/CP/EP 路径继续走原有全局并行策略。

3. Visual Encoder 独立 CP 与异步 CP

视觉 CP 作用在视觉 attention core 上：

ContextParallel(seq_dim=1, head_dim=2, ulysses_degree=...)
AsyncContextParallel(seq_dim=1, head_dim=2, ulysses_degree=...)

实现上将视觉 attention 拆为两层：

Qwen3VLMoeVisionAttention：负责 qkv 投影、RoPE、变长序列 split/cat 和输出投影；
Qwen3VLMoeVisionSdpaCore：只负责 core_attn(q, k, v)，输入输出采用 BSHD 布局，是 CP 注册边界。

普通 CP 通过 ContextParallel 挂载到 block.attn.sdpa_core；异步 CP 通过 AsyncContextParallel 挂载到同一个 core，并复用 q_cp_boundary、k_cp_boundary、v_cp_boundary 作为投影边界。这样 CP 通信只包住视觉 attention 的核心计算，不影响 patch embedding、MLP、merger 和文本 Decoder。

4. 多模态 CP 数据流

文本 CP 下，Trainer 会先按 CP rank 切分 input_ids、position_ids、attention_mask、labels，并同步切分 mm_token_type_ids，保证 loss shift 后的本地 label 与本地 logits 对齐。

进入 Qwen3VLMoeModel.forward() 后，模型会在多模态融合前 gather 完整文本序列，用完整序列完成视觉 token mask、image embedding scatter 和 mRoPE position ids 计算；融合完成后再按 CP rank 将 inputs_embeds、position_ids、attention_mask、visual_pos_masks 和 DeepStack 视觉特征切回本地 shard，交给文本 Decoder 的 CP 路径继续训练。

视觉塔内部则对 packed visual sequence 做局部 CP：根据 grid_thw 生成 cu_seqlens 和每个图片/帧序列长度，在进入视觉 blocks 前按 CP rank 切分 hidden_states 和 rotary_pos_emb，视觉 attention core 完成 CP 计算后，再将 last_hidden_state 和 deepstack_features gather 回完整视觉序列，供后续 image token 注入使用。

5. patch embedding 兼容性

Qwen3VLMoeVisionPatchEmbed 继续保留 self.proj = nn.Conv3d(...)，因此 checkpoint 中的参数名和参数布局不变；实际 forward 将输入 flatten 后使用 F.linear(hidden_states, self.proj.weight.view(...), self.proj.bias) 做等价投影。该方式不改变权重加载语义，同时避开 Ascend A2 上 Conv3d backward 的不稳定路径。

流程图

并行配置流程

flowchart TD
  A["train.yaml: model.vision_parallel"] --> B["get_vision_parallel_config"]
  B --> C{"dp_shard"}
  C -->|"1"| D["Visual Encoder params replicated"]
  C -->|"global dp_shard"| E["Visual Encoder uses visual FSDP mesh"]
  B --> F{"cp > 1"}
  F -->|"no"| G["Visual DP only"]
  F -->|"yes"| H["Resolve train.accelerator.cp mesh"]
  H --> I["visual.set_context_parallel_mesh(cp_mesh)"]
  I --> J{"async_cp"}
  J -->|"false"| K["ContextParallel on block.attn.sdpa_core"]
  J -->|"true"| L["AsyncContextParallel on sdpa_core with q/k/v boundaries"]

前向数据流

flowchart LR
  A["CP-local micro batch"] --> B["Gather full text sequence for VL fusion"]
  B --> C["Visual patch embedding<br/>Conv3d weights, linear compute"]
  C --> D["Shard packed visual sequence by CP rank"]
  D --> E["Visual blocks<br/>sdpa_core CP"]
  E --> F["Gather visual output and DeepStack features"]
  F --> G["Scatter image embeddings into text tokens"]
  G --> H["Build mRoPE position ids"]
  H --> I["Slice text sequence, masks and DeepStack features"]
  I --> J["Text Decoder CP/FSDP path"]
  J --> K["Logits and shifted-token loss"]

验证结果

功能与精度

验证项	结果
DeepStack CP 切片校验	3 个检查通过
Qwen3-VL-MoE 视觉塔 CPU parity	通过
1 卡 baseline smoke	通过
2 卡 baseline DP/FSDP smoke	通过
2 卡 Visual Encoder DP (`dp_shard=1`)	通过
2 卡 Visual Encoder CP Pure Colossal (`cp=2, ulysses_degree=1`)	通过
2 卡 Visual Encoder CP Pure Ulysses (`cp=2, ulysses_degree=2`)	通过
2 卡 Visual Encoder async CP Pure Colossal	通过

首步 loss 对齐结果：

模式	首步 loss
baseline DP	11.931214332580566
visual DP1	11.931214332580566
visual CP Pure Colossal	11.931214332580566
visual CP Pure Ulysses	11.931214332580566
visual async CP Pure Colossal	11.931214332580566

首步 loss 在 rel_tol=1e-6、abs_tol=1e-5 下完成对齐。100 step loss 对齐通过：最大绝对误差 0.00095845，平均绝对误差 0.00026391，首步误差 0，末步误差 0.00037098，满足 0.005 容差要求。

两卡性能对比

性能采样使用 tiny vl_dummy 配置，world_size=2，global_batch_size=2，max_seq_len=32，视觉 grid 为 2 x 2 x 2，warmup 2 step，统计 10 step。该结果用于同口径对比不同视觉并行路径。

模式	mean step time (s)	median step time (s)	tokens/s	相对 baseline DP	peak memory (MiB)
baseline DP	0.03308489	0.03243547	1934.418	1.0000	323.614
visual DP1	0.02392030	0.02345216	2675.551	1.3831	323.636
visual CP Pure Colossal	0.03314808	0.03304910	1930.730	0.9981	323.614
visual CP Pure Ulysses	0.03354174	0.03331385	1908.070	0.9864	323.614
visual async CP Pure Colossal	0.03493548	0.03402914	1831.948	0.9470	323.614

在该短序列 tiny 配置下，CP 通信开销占比较高，因此 CP 吞吐低于 visual DP1 属于预期范围；性能 PASS 的含义是五组数据完整、可比较，且对应训练链路均可正常执行。

使用入口

训练模板：

examples/qwen3_vl_30b_a3b_instruct/train.yaml

使用文档：

docs/guide/qwen3_vl_moe.md

视觉塔回归测试入口：

export HYPER_PARALLEL_PLATFORM=torch
python -m pytest -q tests/torch/qwen3_vl_moe/test_qwen3_vl_moe_vision_parity.py

2 卡 VL Trainer smoke 与 loss 对齐入口：

export HYPER_PARALLEL_PLATFORM=torch
export ASCEND_VISIBLE_DEVICES=0,1
export TORCH_DEVICE_BACKEND_AUTOLOAD=0
python -m pytest -q tests/torch/qwen3_vl_moe/test_qwen3_vl_moe_vl_trainer.py -s

涉及文件

核心代码：

hyper_parallel/models/qwen3_vl_vision/model.py
hyper_parallel/models/qwen3_vl_moe/model.py
hyper_parallel/models/qwen3_vl_moe/parallelize.py
hyper_parallel/trainer/base.py

文档与示例：

docs/guide/qwen3_vl_moe.md
examples/README.md
examples/qwen3_vl_30b_a3b_instruct/train.yaml

测试：

tests/torch/qwen3_vl_moe/_test_qwen3_vl_moe_vl_trainer.py
tests/torch/qwen3_vl_moe/test_qwen3_vl_moe_vl_trainer.py

范围边界

本次实现覆盖 qwen3_vl_moe 的 Visual Encoder 局部并行；
vision_parallel.cp > 1 依赖已有 train.accelerator.cp mesh，且视觉 packed sequence 长度需要满足 CP 切分要求；
vision_parallel.dp_shard 当前支持 1 或全局 DP shard size；
视觉 patch embedding 的参数名和 checkpoint 布局保持不变。