资讯详情

资讯详情

长上下文训练 256k 不难:Halo 上下文并行(CP + Ulysses 注意力)原理与实测

【免费下载链接】haloHalo is an open-source framework built by White Circle for training large language and multimodal models项目地址https://gitcode.com/gh_mirrors/halo34/halo点击查看免费下载长上下文训练一直是多卡训练里的硬骨头序列一长激活显存和注意力的平方复杂度就会把单卡撑爆。开源大模型训练框架Halo通过**上下文并行Context Parallelism, CP**加Ulysses 注意力把 256k 长度的序列切到多张 GPU 上精确训练实测峰值显存可降到单卡方案的约一半。本文讲清它的原理、使用方式和真实基准数据。为什么长上下文训练这么难多卡并行里模型太大可以用张量并行TPMoE 专家太多可以用专家并行EP但序列太长是另一类瓶颈——它吃的是激活显存不是参数显存。EP 切的是专家对激活没帮助所以 32K 以上的长序列训练需要一条专门的路径把序列本身切到多张卡上。这就是 Halo 上下文并行CP存在的意义每张 GPU 只持有seq_len / cp_size个 token 的激活其余计算照常分摊注意力层则靠通信临时拼回完整序列。官方说明见 agent-docs/parallelism/context-parallelism.md。Ulysses 注意力三次 all-to-all 换一次精确计算朴素的每卡切一段序列有个致命问题注意力要看全序列切断了跨块依赖。Ulysses 的思路很聪明——在注意力核函数周围重新分配张量每个 rank 持有S/cp_size个 token、但全部注意力头两次 all-to-all 把张量换成全量S个 token、但只有H/cp_size个头供 Flash Attention 计算Q 一次、融合后的 K|V 一次、输出再一次反向传播按同样三步反向执行。多头注意力在头维度上是相互独立的所以这种交换是精确的不是近似。RoPE 在本地块上完成优化路径省掉了位置索引的全局搬运原生 GQA 支持还省掉了repeat_kv的拷贝。入口实现是 UlyssesCPModelWrapper各类模型的注意力包装器在 src/distributed/context_parallel/layers/ 下注册表 registry.py 直接从子类树派生支持列表。图中展示cp 4、64 个 Q 头 / 8 个 KV 头的例子交换前每卡只有对角块本 rank 的一个序列块 × 全部头交换后每卡拿到全部 S 个 token 的本 rank 16 个头注意力计算即覆盖完整序列。一次 CP 训练步的完整数据流理解了注意力层再看整步训练入口split_sequence_for_cp把 batch 收窄为每 rank 一个连续序列块S % cp_size 0collator 自动补齐Embedding 和所有 decoder 层MLP / MoE 专家都只在本块上跑CP 对此无感知logits 保持分块从不 all-gatherLoss 做两处关键修正边界 token因果标签偏移会让每块少一个预测非末块用最后一个 logit 预测下一块首 token全局归一化的 sum lossloss_i cp_size * local_sum / global_tokensglobal_tokens是 CP 组内 all-reduce 出的非 mask token 总数保证每个 token 对梯度的贡献完全相等。一句话总结只有注意力层看得到完整序列其余所有层和 loss 都只处理本 rank 的块。GRPO 等训练目标的分块求和路径在 autograd.py 中提供 autograd 感知的cp_sum_rows。上手两条命令即可开启纯 CP任意支持的模型长序列 SFTtorchrun --nproc_per_node4 scripts/training/sft.py \ examples/sft/qwen3/qwen3-4b-ultrachat.yaml --context_parallel_size4 --packingfalseMoE 模型的 EPCP节点内 EPep_group_size等于 NVLink 域大小torchrun --nproc_per_node8 scripts/training/sft.py \ examples/sft/gptoss/gptoss-20b-multinode-ep.yaml \ --expert_parallel_size8 --ep_scopenode --context_parallel_size2 --packingfalse开跑前的硬性前提 前提说明头数整除cp_size必须同时整除 Q 头和 KV 头数GPT-OSS 64Q/8KVCP8 ✅CP3 ❌序列长度必须是cp_size的倍数sft.py 自动设置pad_to_multiple_of无需手动 collator拓扑cp_size不能超过 NVLink 域all-to-all 必须留在 NVLink 内注意力实现必须是 Flash AttentionFA2/FA3/FA4eager/sdpa会被拒绝仅右 paddingUlysses 路径走稠密因果核左 padding 批次会在 forward 时直接报错与 packing 互斥packing/padding_free在 CP 下被拒绝支持的模型包括 Qwen3 dense/MoE、Qwen3-VL 文本分支、GPT-OSS、GLM-4 MoE Lite、Mistral4MLA 路径、Bailing/Ling 2.0、Cohere2 MoE线性注意力、序列轴卷积等结构如 Qwen3.5/3.6 完整检查点、LFM-2、Gemma 4暂不支持 CP。实测256k 不是数字游戏以下数据来自官方基准gpt-oss-20b8× B300b1GC 开启256k 序列实测对比tok/s/GPU · 峰值显存 GiB配置吞吐峰值显存stock TRL (z3)4,463148Halo EP1 z3稠密5,728112Halo EP8CP8 (DP1)2,39577Halo dense CP-only (z3)1,83656三个结论稠密 EP1 z3 赢吞吐256k 下比 stock TRL 快 28%显存还少 36 GiB64k 之后 CP 越大越好128k 处 EP8CP8 在吞吐和显存上同时优于 EP8CP4——注意力是二次复杂度序列减半带来的收益远大于多一次 all-gather 的开销dense CP-only 是全榜最省显存128k→256k 仅 32–56 GiB且不依赖 DeepEP。再看 EP8CP8 在 16k–64k 的表现SeqLentok/s/GPUTFLOPS峰值显存16,3847,14924924.6 GiB32,7689,03140229.9 GiB65,5368,86256641.9 GiB注意TFLOPS 随序列变长而上升——更长的序列摊薄了 Ulysses all-to-all 的固定成本显存则只从 25 GiB 缓慢涨到 42 GiB。CP 的本质是一笔划算的交易用少量 per-GPU 吞吐换取便宜的长上下文。完整表格见 throughput-benchmarks.md 与 halo-vs-stock-trl.md。CP 与 EP/TP 怎么搭配data_parallel_size world_size / cp_sizeCP 直接缩减 DPEP 则正交不减 DPEPCP 是合法组合CP 在注意力层做 all-to-allEP 在 MoE 块做 DeepEP dispatch/combine位置不同互不冲突。8 卡节点上的合法形状如ep8/cp2 → DP4、ep8/cp4 → DP2、ep8/cp8 → DP1TPCP 被拒绝两者对同一组连续 rank 做两次不兼容的划分配置期即报错所有组合由 parallelism_config.py 中的白名单在配置期校验跑错的组合会在加载模型前就被明确拒绝。注意事项与限制 ⚠️训练器仅 SFT、SMPO、offline GRPOoffline GRPO 要求全量微调声明支持 CP其余训练器会在构造时拒绝不支持在线生成generate()会直接报错每 rank 只有一个块Ulysses 注意力没有 KV-cache 路径评测请用保存的检查点另起进程文本 only多模态输入pixel_values在 CP 下被拒绝torch.compile对 CP 无收益all-to-all 在每个注意力层打断计算图保持默认关闭即可。小结Halo 把长上下文训练抽象成一个干净的问题序列切块 Ulysses 精确注意力交换 全局归一化的分块 loss。前提校验严格配置期拒绝一切不合法组合、实现精确非近似、实测扎实256k 显存减半、TFLOPS 随序列增长。如果你正被 128k/256k 的 OOM 卡住--context_parallel_size加一个参数可能就够用了。更多细节可查阅 上下文并行指南 和 并行策略总览。赞分享【免费下载链接】haloHalo is an open-source framework built by White Circle for training large language and multimodal models项目地址https://gitcode.com/gh_mirrors/halo34/halo点击查看免费下载相关推荐长上下文训练不再爆显存NeMo AutoModel 序列打包与上下文并行 CP 优化完整指南长上下文训练不再爆显存NeMo AutoModel 序列打包与上下文并行 CP 优化完整指南 ! NeMo AutoModel 项目架构图 https://rMindSpeed-LLM Mamba 长序列训练Mamba-CP 上下文并行方案实战指南MindSpeed LLM Mamba 长序列训练Mamba CP 上下文并行方案实战指南 导读 本指南面向在昇腾平台上使用 MindSpeed LLM 训练人工智能大模型预训练分布式训练微调强化学习模型推理服务模型评测Ascend256K超长上下文如何实现MiniCPM-V-4.7-35B-A3B-4bit线性注意力全注意力混合机制原理256K超长上下文如何实现MiniCPM V 4.7 35B A3B 4bit线性注意力全注意力混合机制原理 MiniCPM V 4.7 35B A3B 4上一篇Arco Design React 表单实战使用 Form 实现带二次密码校验的注册表单下一篇Windows窗口美化终极指南打造个性化毛玻璃特效完整教程创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →