1. 问题背景

在利用 MindSpeed-MM 的 FSDP2 (Fully Sharded Data Parallel 2) 分布式并行策略微调或训练多模态大模型（如 Qwen2-VL、LLaVA-Next）时，动态分辨率（Dynamic Patching / Tiling）被广泛用于保护高分辨率图像中的局部细节。

然而，在实际的多机多卡压测中，我们观察到随着 Batch Size 与单样本图像 Tile 数量的变化，FSDP2 的多卡扩展效率（Scaling Efficiency）出现大幅退化。通过 CANN Profiling 进行深度性能下钻后，我们发现核心瓶颈在于：动态分辨率导致的分布式 Rank 间 Token 负载严重失衡（Straggler Effect）、计算 varlen FlashAttention 时 cu_seqlens 的高频 Host-to-Device (H2D) 同步阻塞，以及多模态 Packed Tensor 在 RoPE 阶段冗余的解包/打包（Unpacking/Packing）显存带宽开销。

2. 核心技术瓶颈剖析

2.1 FSDP2 间 Sample 级分发造成的 Token 负载严重失衡 (Straggler Effect)

在多模态场景下，不同样本对应的图像分辨率和切片数存在数倍乃至数十倍的差异：

• 分发机制缺陷：FSDP2 在对数据进行分片（Sharding）分发时，默认以样本数量（Sample-level）在各个 Rank 间进行等量划分。
• 负载倾斜：这导致在动态 Patching 场景下，有些 Rank 分配到了包含极长视频或多个高分辨率图像切片的样本，其分配到的 Token 数量极高（如 Rank 0 包含了 4096 个 tokens）；而邻近 Rank 仅分配到了低分辨率小图或纯文本样本，其 Token 数量极低（如 Rank 1 仅有 512 个 tokens）。
• 同步气泡：由于 Attention 的计算开销与 Token 数量的平方（或线性，在使用 FlashAttention 时）成正比，Rank 1 提前计算完毕并陷入长时间的 Idle 等待，在 FSDP2 的 all_gather 权重重组与梯度 reduce_scatter 边界处产生了巨大的同步空洞，导致多卡并行的整体效率被执行最慢的 Rank（木桶效应）锁死。

2.2 cu_seqlens 动态计算与 H2D 拷贝阻断指令发射 (Host-Device Sync Bubble)

在使用 npu_flash_attention_varlen 进行打包（Packed / Unpadded）无填充训练时，算子要求输入累加序列长度映射 cu_seqlens_q 与 cu_seqlens_k：

• 同步阻塞：由于每一批（Batch）中不同样本的图像 Tile 数量动态改变，cu_seqlens 必须在 Host 侧（CPU）根据当前 Batch 的 Token 拓扑进行动态计算，随后调用 .npu() 或 torch.from_numpy().to('npu') 将该一维 Tensor 下发到 Device。
• 流水线断裂：在 Profiler 中可以清晰地看到，即使该 Tensor 仅包含几个到几十个 int32 元素，该动态下发操作依然触发了隐式的流同步或 Host 同步，强行切断了 PyTorch 算子的发射流水线（Operator Launch Queue），使 NPU 在每个计算 Step 的起始阶段都出现了平均 80~200微秒 的指令空闲空洞。

2.3 Packed 1D 连续结构在多模态不规则 RoPE 计算下的拆装损耗 (Tensor Reshaping & Slice Overhead)

对于 Qwen2-VL 等模型，由于需要对 Vision 令牌应用 3D-RoPE，对 Text 令牌应用 1D-RoPE：

• 冗余搬运：当使用连续 Packed Tensor 格式时，框架在 Self-Attention 之前无法直接执行统一的 Fused RoPE。目前 MindSpeed-MM 被迫在 Python 层利用 index_select 或分片（Slicing）将打包的张量进行局部拆解（Unpacking），分别应用不同维度的位置编码后，再次通过 concat 重新拼接（Packing）为无填充的 1D 张量送入 FlashAttention。
• 显存风暴：这类频繁的数据切片与重拼接引入了大量冗余的内存拷贝指令，在每层 Attention 前后造成了高昂的 HBM 访存带宽消耗，并随 Sequence Length 增加呈现长尾化。

3. 复现说明与 Profiling 采样特征

我们在 FSDP2 + Sequence Parallel 下微调动态分辨率模型时，其 Rank 间的算力耗时严重不均衡：

# FSDP2 Rank 数据负载失衡模拟说明
# 设单机 2 卡 (Rank 0, Rank 1)，Batch Size = 2
# 样本 A (分发给 Rank 0): 包含一幅 1024x1024 图像，切片为 4 个 512x512 tile + 1 个全局缩略图，产生约 3120 个 vision tokens
# 样本 B (分发给 Rank 1): 包含一幅 256x256 图像，无切片，产生约 256 个 vision tokens

# 在 FSDP2 前向计算中：
# Rank 0 运行 npu_flash_attention_varlen 耗时约 t0
# Rank 1 运行 npu_flash_attention_varlen 耗时约 t1 (t1 << t0)
# Rank 1 将在 FSDP2 的权重重组 Barrier 处挂起，等待时间为 t0 - t1

CANN Profiler Timeline 采样视图特征：

1. FSDP2 Communication Stream 上的等待占比极高：在反向传播或前向反向交界处，HCCL 算子（如 all_gather）出现了由于 Rank 间到达时间（Arrival Time）不一致产生的大面积等待等待空洞（Waited state）。
2. 在每个 Epoch / Iteration 的开始处，均伴随着由于 Host 侧计算 cu_seqlens 并下发导致的 CPU 与 NPU Gap（调度空白期）。
3. 显存带宽曲线上，Attention 前后的 index_select / concat 拷贝算子占据了显著的非计算 HBM 读写周期。

4. 优化建议与方案技术探讨

针对多模态大模型在昇腾环境下的这一并行加速瓶颈，建议 MindSpeed-MM 团队考虑并探讨以下优化路径：

4.1 引入基于 Token 数量负载均衡的分布式采样器 (Token-based Load-Balancing Sampler)

在 FSDP2 的数据载入阶段（Dataloader / Sampler）：

• 建议设计一个自适应的 TokenLoadBalanceSampler。它不再死板地按样本数量（Sample-level）均分，而是提前扫描或估算当前 Batch 中每个多模态样本的实际 Token 数量。
• 通过启发式动态规划（Bin-packing 算法），合理组合并重排样本的分发顺序，确保分片到各个 Rank 上的 总 Token 数量（Sum of Tokens per Rank）尽可能对齐，从根本上消除 FSDP2 rank 间的木桶效应，使并行缩放效率重新逼近线性。

4.2 引入异步 Pinned Host-Device 内存区进行 cu_seqlens 下发

为了彻底消除 cu_seqlens 动态 H2D 同步产生的发射空白：

• 建议底层适配异步 H2D 传输通道。在 Host 侧使用页锁定内存（Paged-pinned Memory）预先创建 cu_seqlens 的影子缓冲区。
• Host 计算完偏移后，采用非阻塞的 rtMemcpyAsync 写入通信流，仅在实际进入 npu_flash_attention_varlen 执行前通过内部 Event 同步，使 Host 线程可以完全“放手”异步向后发射算子。

4.3 研发支持 1D Packed 混合多维位置编码的融合 RoPE 算子 (Fused Mixed-Dim RoPE)

• 建议通过 AscendNPU-IR 或直接结合 CANN 的高性能融合算子库，开发一个专为 Packed 1D Tensor 设计的 npu_fused_mixed_rope 算子。
• 允许通过传入一个轻量级的特征 Mask 或位置偏移向量，直接在 NPU 片上 L1/UB 内部对 Packed Tensor 的不同分段（Vision 或 Text）应用对应的 3D 或 1D 旋转位置编码，完全消除框架在 Eager 侧进行 Unpacking 和 Packing 的冗余搬运。

• 环境信息：NPU Ascend 910B (64GB), CANN 8.0.RC1, MindSpeed-MM (master), torch_npu (master)