Efficient Scaling of LLM Training with Flexible Context Parallelism¶

会议: NeurIPS 2026
arXiv: 2602.21788
代码: https://github.com/ivanniu/FCP
领域: LLM 其他
关键词: 上下文并行, 长文本训练, 动态并行调度, 异构序列, 负载均衡

一句话总结¶

针对真实长文本与多模态数据长度严重长尾分布导致的静态并行计算空转与通信冗余，提出灵活上下文并行框架 FCP，通过突破 2 的幂次限制的任意整数并行度和毫秒级二维动态规划调度算法，实现每微批次自适应并行重构，在 64 卡集群上最高取得 1.47× 平均吞吐提升和极端不平衡批次下 2.24× 加速。

研究背景与动机¶

大语言模型与多模态大模型正快速向超长上下文方向演化，然而自注意力机制关于序列长度 \(L\) 的平方计算复杂度 \(\mathcal{O}(L^2)\) 始终是长上下文分布式训练的算力与显存瓶颈。现有的主流分布式训练系统（如 Megatron-LM 和 DeepSpeed）普遍采用静态 4D 并行架构（数据并行 DP、张量并行 TP、流水线并行 PP、上下文/序列并行 CP/SP）。在面对同构定长序列时这类静态划分较为高效，但在真实世界的预训练与微调语料（如 GitHub 代码库、CommonCrawl 网页、多模态视频文本数据集）中，序列长度呈现出极度偏斜的长尾分布——绝大多数样本较短，仅有极少数长尾序列达到数十万甚至上百万 token。

静态并行策略在处理异构批次时暴露出严重的木桶效应与通信冗余。传统做法通常采用静态固定的 CP 度数或等 token 数量打包（packing），但这忽略了注意力计算的本质非线性特征：处理单个长序列 \(L\) 的 rank 需承担 \(\mathcal{O}(L^2)\) 计算量，而将多个短序列 \(\{l_i\}\) 简单打包到同一 rank 所需的计算量仅为 \(\sum \mathcal{O}(l_i^2)\)。由于 \(L^2 \gg \sum l_i^2\)，分配到超长文本或多模态视频帧序列的计算节点必然成为严重掉队者（straggler），迫使其他已完成短序列计算的节点在全局同步屏障处长期处于空转等待状态。此外，近期少数探索动态划分的工作（如 FlexSP）将通信组大小严格限制在 2 的幂次（powers-of-two），极大压缩了资源最优分配的搜索空间，且求解整数规划耗时高达数秒至十几秒，无法在训练中实时应用。

本文的切入角度是：如果抛弃 2 的幂次通信约束，借助 Ring-style CP 点对点环状通信天然支持任意正整数维度的灵活性，就能够将整集群计算单元细粒度解耦并动态编组，进而将复杂的微批次异构序列调度规约为极轻量的组合优化问题。核心 idea：基于 Ring-style 上下文并行构建支持任意正整数并行度的通信组重构机制，通过两阶段逼近算法（BFD 显存感知打包 + 二维动态规划资源分配）在毫秒级内完成每批次自适应调度，并在执行当前批次计算的同时异步重构下一批通信组，彻底消除长尾数据引起的负载失衡与算力空转。

方法详解¶

整体框架¶

FCP 作为非侵入式异步模块构建于 Megatron-LM 基础之上，将整个分布式系统的优化空间精准限定在上下文并行（CP）与数据并行（DP）维度，而将需要复杂权重重新切分的张量并行（TP）与流水线并行（PP）作为底层静态执行网格。FCP 的执行分为调度（Scheduling）与执行（Execution）两个阶段：全局批次（Global Batch）被切分为微批次后，CPU 端调度器在当前微批次正向/反向计算的同时异步求解下一微批次的最优分组策略；在进入新微批次前，调度器仅在 Megatron 并行状态管理器（MPU）中动态创建和调整 CP 通信进程组，保持其他并行拓扑不变。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    In["输入微批次 B：K 个异构长尾序列"] --> S1["阶段 1：显存感知序列打包 BFD<br/>降序显存排序与箱规初始化，聚合为 K' 个原子组"]
    S1 --> S2["阶段 2：二维动态规划资源分配 2D-DP<br/>放宽 2 的幂次约束，最小化组间最大执行时间 makespan"]
    subgraph S3["阶段 3：异步执行与通信组重构"]
        direction TB
        S3A["CPU 端异步重构下一批 CP 进程组<br/>Megatron MPU 毫秒级轻量切换"] --> S3B["NPU 集群重组 CP 通信组与 Ring Attention 并行训练"]
    end
    S2 --> S3
    S3 --> Out["输出：负载均衡计算与梯度规约"]

关键设计¶

1. 显存感知序列打包（BFD）：消除显存碎片与降维决策空间 针对微批次内序列数量较多且长度悬殊导致的决策变量激增与显存超限风险，FCP 首先基于最佳适应递减（Best-Fit Decreasing, BFD）策略将 \(K\) 个原始序列合并为 \(K' \le K\) 个原子序列组（Atomic Groups）。调度器将批次 \(\mathcal{B}\) 中的序列按显存开销降序排列，针对每个长序列根据单 rank 显存预算 \(E\) 计算满足显存约束所需的最小 CP 阶数 \(d_{\min, k} = \lceil \mathcal{M}(s_k) / E \rceil\)，以此建立具有容量为 \(d_{\min, k} \cdot E\) 的虚拟容器。随后，贪心地将较短的序列打包进现有容器的显存余量中。这一设计不仅严格保障了后续分配绝不触发 Out-of-Memory（OOM），更成功避免了将海量短序列单独放入高阶 CP 组造成的通信冗余，同时将动态规划的阶段数从 \(K\) 大幅压减至 \(K'\)。

2. 任意正整数并行度与二维动态规划调度：在多项式时间内逼近最优全局 makespan 针对传统序列并行方案（如 DeepSpeed-Ulysses）因 All-to-All 算子及注意力头数整除限制而被迫拘泥于 2 的幂次并行度这一核心瓶颈，FCP 采用支持点对点环状通信的 Ring-Attention，放宽至任意正整数并行度 \(d \in \mathbb{Z}^+\)，使得算力分配能以单卡（rank）为粒度贴合实际序列计算量。为了在毫秒级时间内求解该 NP-hard 的组划分与任务指派问题，FCP 设计了时间复杂度仅为 \(\mathcal{O}(K' N^2)\) 的二维动态规划（2D-DP）算法。令 \(DP[i][j]\) 表示前 \(i\) 个原子组分配 \(j\) 个 rank 时的最小最大执行时间（makespan），状态转移方程定义为：

\[DP[i][j] = \min_{d \in [d_{\min, i}, j - d']} \max(DP[i-1][j-d], \mathcal{T}(\mathcal{G}_i, d))\]

其中 \(d' = \sum_{m=1}^{i-1} d_{\min, m}\) 为保证前 \(i-1\) 个原子组满足最低显存约束所预留的最小 rank 数，\(\mathcal{T}(\mathcal{G}_i, d)\) 为第 \(i\) 个原子组在并行度 \(d\) 下的预测执行耗时。内层 \(\max\) 操作直接捕捉系统整体的木桶效应瓶颈，外层 \(\min\) 遍历当前组的所有可行并行度以寻找全局最优，最后自底向上回溯即可在数十毫秒内直接输出全局最优 CP 分组矩阵与度数。

3. 跨模态注意力感知耗时建模：精准覆盖因果与全自注意力的真实开销 为了让动态规划求解器获得高度保真的执行耗时预估，FCP 构建了能够统一覆盖纯文本 LLM 与多模态 MLLM 复杂注意力模式的分析型耗时评估模型。不同于纯文本普遍采用的因果注意力（Causal Attention），多模态模型（如 InternVL、Qwen-VL）中视觉 token 通常采用全双向自注意力（Full Attention）。模型引入掩码效率因子 \(\eta_k\)：对于纯因果自注意力取 \(\eta_k = 0\)，对于全双向注意力则赋予反映额外计算量的系数。针对包含 Multi-Head Attention（MHA）平方项与 Feed-Forward Network（FFN）线性项的计算耗时 \(\mathcal{T}_{cp}\)、以及环状点对点通信耗时 \(\mathcal{T}_{cm}\)，FCP 结合 Ring-Attention 通信与计算天然重叠的硬件特性，定义组总耗时模型：

\[\mathcal{T}(C_p) = \mathcal{T}_{cp}(C_p) + \mathcal{T}_{cm}(C_p) - \min(\mathcal{T}_{cpa}, \mathcal{T}_{cma})\]

式中通过扣减注意力计算耗时 \(\mathcal{T}_{cpa}\) 与注意力 KV 块传输耗时 \(\mathcal{T}_{cma}\) 的重叠交集，精准还原硬件真实执行时间。先验 Profiler 预先离线测定少量系数并常驻内存，运行时仅需代入多项式计算即可实现低于 8% 的端到端预估误差。

4. 计算与调度异步解耦引擎：将通信组更新开销完全隐藏在计算底层 针对动态并行引入的重配置开销是否会反噬训练吞吐的疑虑，FCP 将整个系统解耦为前台计算执行流与后台 CPU 规划流。当 NPU 集群正在执行当前批次的重度矩阵运算（耗时在数秒至十余秒）时，CPU 端调度器无缝启动下一批次的 BFD 打包、2D-DP 求解以及通信拓扑预生成，全流程调度仅耗时数百毫秒，被完全隐藏在计算流水线后方。在批次切换边界，调度器仅需对 PyTorch/Megatron 的通信组上下文执行局部增删，实测单次进程组重置耗时仅约 1 ms，做到了零气泡动态并行迁移。

实验关键数据¶

主实验¶

实验在 4 个节点（共 64 张华为昇腾 Ascend 910C NPU，每卡 64GB 显存，节点内 HCCS 互联）上开展，评估最长上下文达到 128k 的训练场景。对比基线包括 Megatron-LM（静态 4D 并行）与 DeepSpeed（Ulysses 序列并行），均在相同全局批次大小（GBS=512）下调优至各自最佳静态配置。

模型类别与评测模型	数据集 / 场景	FCP 迭代耗时 (s)	最佳基线耗时 (s)	加速比 (Gain)
Qwen3-8B (LLM)	GitHub (极度长尾代码)	12.85	17.60 (Megatron)	1.37×
Llama3-8B (LLM)	GitHub (极度长尾代码)	13.10	17.82 (Megatron)	1.36×
Qwen3-8B (LLM)	CommonCrawl (通用网页)	11.20	14.56 (DeepSpeed)	1.30×
Llama3-8B (LLM)	CommonCrawl (通用网页)	11.55	14.90 (DeepSpeed)	1.29×
Llama3.2-1B (LLM)	Wikipedia (密集短序列)	6.80	7.14 (DeepSpeed)	1.05×
InternVL3-8B (MLLM 全解冻)	MSRVTT (视频-文本)	14.20	19.60 (Megatron)	1.38×
Qwen3VL-8B (MLLM 全解冻)	MSRVTT (视频-文本)	13.90	19.04 (Megatron)	1.37×
InternVL3-8B (MLLM 全解冻)	OpenVid (高清视频)	15.10	20.39 (DeepSpeed)	1.35×
Qwen3VL-8B (MLLM 全解冻)	OpenVid (高清视频)	14.80	19.83 (DeepSpeed)	1.34×
极端不平衡合成批次	1×128k + 许多 2-8k 混合	-	-	2.24×

消融实验¶

消融与开销分析聚焦于 FCP 的调度开销、求解延迟与计算重叠能力（在 64 卡集群训练 Qwen3VL-2B 上的真实测试结果）：

实验配置 (GBS / 设备数)	计算耗时 Compute (s)	调度总耗时 Schedule (ms)	核心算法耗时 Solver (ms)	说明
GBS=128 (64 Devs)	6.34	448	23	调度仅占计算时间 7.0%，完全重叠掩盖
GBS=256 (64 Devs)	7.58	696	45	调度仅占计算时间 9.1%，完全重叠掩盖
GBS=512 (64 Devs)	11.65	905	86	求解仅 86ms，总调度 <1s，重叠率 100%
GBS=512 (32 Devs)	18.26	537	44	设备数减半，计算时间增加，更易掩盖开销
GBS=512 (16 Devs)	34.45	282	20	调度总耗时仅占计算时间 0.8%

关键发现¶

长尾分布越严重、模型规模越大，收益越显著：在 GitHub 和 OpenVid 等长尾效应显著的数据集上，FCP 相比基线取得 1.35×~1.38× 的显著加速；而在长度分布较为均匀密集的 Wikipedia 上，由于各 rank 负载天然接近，加速比收敛至 1.05×。这印证了动态并行正是解决长尾不平衡而生。
扩展性超越基线，通信负面效应被有效对冲：随着集群规模从 8 卡扩展至 64 卡，DeepSpeed 的单卡吞吐从 2.0 k tokens/s 降至 1.7 k tokens/s，Megatron-LM 从 1.8 降至 1.52 k tokens/s，而 FCP 相对于 DeepSpeed 的加速优势从 8 卡时的 1.02× 稳步放大至 64 卡时的 1.16×，展现出极强的抗通信衰减与近线性扩展能力。
对成本模型噪声具有极高容忍度：扰动实验表明，即使在成本预测系数中加入高达 20% 的高斯噪声，求解出的 CP 策略匹配度依然为 100%，单步执行耗时毫无增加；在 50% 极端噪声下，CP 分组匹配率仍超过 88%，耗时衰减被控制在 6% 以内。这是由于平方项 \(\mathcal{O}(L^2)\) 处于绝对主导地位，轻微扰动不改变候选方案的相对序次。

亮点与洞察¶

突破传统并行系统的 2 的幂次思维定势：通过选用无头数整除束缚的 Ring-Attention，将 CP 并行度泛化至任意正整数，从而把离散粗粒度的资源切分升级为细粒度弹性的连续谱系。
启发式装箱与动态规划的高效结合：先通过 BFD 显存打包消灭决策冗余并将 \(K\) 聚拢为 \(K'\)，再使用状态空间为 \(\mathcal{O}(K' N^2)\) 的 2D-DP 精确求解全局瓶颈，成功将原本指数级复杂度的全局调度压制到 86 毫秒以内。
极轻量的系统级异步工程闭环：充分利用深度学习训练迭代天然存在的流水线特征，将数据规划放在 CPU、矩阵计算留在 NPU，以仅 1 ms 的通信组切换成本换取整体 1.4×+ 的系统级吞吐提升。

局限与展望¶

依赖数据异构度：当训练语料为定长或经过精心等长填充截断的数据时，负载不均衡现象本就不存在，FCP 的调度收益将衰减至边际水平。
未将 TP/PP 纳入动态调度域：出于运行时权重跨卡重新切分（weight reshuffling）开销过大的考虑，FCP 将 TP 和 PP 视为静态预设网格，仅动态重排 CP 与 DP，尚未实现全维度 4D 动态弹性伸缩。
未来方向：探索与 MoE（专家混合模型）路由负载均衡的协同调度，以及在超大规模万卡集群上跨网络交换机拓扑（NUMA / Spine-Leaf）感知组构建。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 突破 2 的幂次限制并将动态上下文并行形式化为轻量 2D-DP 求解，构思巧妙。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖纯文本与全解冻/冻结多模态、多种长尾真实语料，且包含严谨的噪声抗扰测试。
写作质量: ⭐⭐⭐⭐⭐ 架构清晰、公式推导严谨、系统设计与软硬件权衡交代极为透彻。
价值: ⭐⭐⭐⭐⭐ 贴合超长上下文时代的训练刚需，代码开源且易于集成至主流 Megatron 生态。