DDP损失归一化:各卡token数不同怎么办

黎 浩然/ 18 10 月, 2026/ 大语言模型/LARGELANGUAGEMODEL/LLM, 教程/TUTORIAL, 计算机/COMPUTER/ 0 comments

DDP会同步梯度,但不会自动把“各卡平均loss的平均”改成“所有有效token的平均”。当各卡有效token数不同,先在每卡取mean,再平均梯度,会给每张卡相同权重,而不是给每个token相同权重。

文章目录
  1. 先把想优化的量写出来
  2. 一张卡一个token,另一张卡三个
  3. 为什么本地损失要乘W
  4. 九项检查比较两种梯度
  5. 某卡没有有效标签,不等于可以单独退出
  6. 梯度累积时,分母应覆盖整个更新单元
  7. 日志标量和反向标量分开记录
  8. 用分片不变性检查目标
  9. 来源与补发说明

若目标是整个更新单元的token平均损失,记有效token总数为N、参与梯度平均的进程数为W,每卡先求损失和Sᵣ,再用W·Sᵣ/N反向,才能在本文假设下抵消DDP的跨卡平均。下面推导这个系数,并用精确分数检查不等长、等长和某卡没有有效token的情况。Read in English。

先把想优化的量写出来

设第r个进程有nᵣ个有效token,Sᵣ是这些token损失之和。希望优化的目标是L=(ΣᵣSᵣ)/(Σᵣnᵣ),分母N=Σᵣnᵣ。有效计数由任务定义,例如只统计有监督标签的位置;补齐后的张量长度不一定等于有效计数。

DDP文档说明跨进程梯度被平均。本文只讨论参与进程数固定、使用通常梯度平均、各进程协调执行一次更新的情形,不假定自定义通信hook、join下的动态除数或其他并行封装也采用同一系数。

这里的N是与参数无关的计数。若目标有类别或样本权重,应先依据实际损失公式确定归一化质量;它可能不是简单token数。不能因为变量叫num_tokens,就忽略损失函数内部已经使用了另一种mean分母。

分片对照还要求:重新分配数据后,每个token对应的损失函数与数值保持一致。批次依赖层或随机性变化可能破坏这个前提。参考测试应先控制这些因素,再把差异归因于归一化。

一张卡一个token,另一张卡三个

用单个参数θ和损失ℓ(θ,y)=½(θ−y)²构造教学例子。θ=0时,第一张卡只有目标0,其损失导数为0;第二张卡有三个目标2,每项导数为−2,导数和为−6。

每卡分别取mean得到0和−2,DDP再平均后是−1。但四个有效token合并后的平均导数为(0−2−2−2)/4=−1.5。前者把两张卡各当作一半数据,后者给四个token各四分之一权重。两者优化的量不同,差异不是通信浮点误差。

如果把四个目标重新划成两张卡各两个,局部mean平均就能在本例中得到−1.5。只改变分片方式,错误实现的梯度就变化,这给了一个明确的回归测试方向。它比观察多卡训练loss是否“看起来正常”更能定位归一化问题。

rank0:1个token,局部mean权重½
rank1:3个token,局部mean权重½
全局token目标:两卡权重¼与¾
原创权重图:rank等权与token等权只有在有效数量相等时自然一致。

为什么本地损失要乘W

让每卡反向的标量为Lᵣ=W·Sᵣ/N。通常DDP平均得到(1/W)Σᵣ∇Lᵣ=(1/W)Σᵣ(W/N)∇Sᵣ=(Σᵣ∇Sᵣ)/N,正是全局token平均梯度。关键是先按共同分母归一化,再补偿已知的平均除数。

若先得到的是本地mean,即Sᵣ/nᵣ,在nᵣ>0时也可以乘W·nᵣ/N。但显式损失和更容易表达零有效token的分片,不必先生成一个未定义的局部mean,再尝试乘零修复。

相反,若通信实现返回梯度和而不是平均,补偿W会多乘一次。采用自定义hook、改变进程组或使用其他封装时,应重新写出实际归约式。W指参与本次归约的组大小,不应靠机器数或配置文件里的某个旧值猜测。

每卡损失和Sᵣ
共同分母N,补偿W:W·Sᵣ/N
DDP平均1/W → 全局token平均梯度
原创归约图:补偿系数依赖本文假定的固定跨卡平均。

九项检查比较两种梯度

代码直接计算平方损失的解析导数θ−y,用Fraction避免舍入干扰。它没有启动进程、调用all_reduce或运行PyTorch;证明的是归一化代数,不是通信正确性或GPU性能。

# Exact gradient algebra; no distributed framework is executed.
from fractions import Fraction as F

def global_gradient(groups, theta=F(0)):
    values=[v for group in groups for v in group]
    if not values:
        raise ValueError("no valid tokens")
    return sum((theta-y for y in values),F(0))/len(values)

def averaged_local_means(groups, theta=F(0)):
    return sum((sum((theta-y for y in g),F(0))/len(g)
                for g in groups),F(0))/len(groups)

def scaled_sum_gradient(groups, theta=F(0)):
    total=sum(map(len,groups))
    if total==0:
        raise ValueError("no valid tokens")
    world=len(groups)
    local=[F(world,total)*sum((theta-y for y in g),F(0)) for g in groups]
    return sum(local,F(0))/world

groups=[[F(0)],[F(2),F(2),F(2)]]
assert global_gradient(groups)==F(-3,2)
assert averaged_local_means(groups)==-1
assert scaled_sum_gradient(groups)==F(-3,2)
assert scaled_sum_gradient([[F(0),F(2)],[F(2),F(2)]])==F(-3,2)
assert averaged_local_means([[F(0),F(2)],[F(2),F(2)]])==F(-3,2)
assert scaled_sum_gradient([[],[F(2),F(2)]])==-2
assert scaled_sum_gradient(groups,F(1))==global_gradient(groups,F(1))
assert F(1,2)!=F(1,4)
try:
    scaled_sum_gradient([[],[]])
except ValueError:
    pass
else:
    raise AssertionError("expected zero-count rejection")
print("9 DDP normalization checks passed")

本地九项检查通过。不等长分片暴露局部mean平均的偏差,等长分片作为对照;某一分片为空时,全局和方式仍有数学定义;所有分片都为空时显式拒绝。改变θ后再次对照,避免测试只碰巧在θ=0成立。

某卡没有有效标签,不等于可以单独退出

数学上,若nᵣ=0而N>0,该卡的Sᵣ应贡献零。但分布式执行还要求所有进程按兼容顺序参加需要的同步。不能看到本卡没有标签,就让它跳过反向或某次collective,而其他卡继续等待。

实际实现需要保留与训练图及所用DDP配置兼容的零贡献路径,并在真实环境测试。一个与模型完全无关的新建常数零,不自动保证触发预期的梯度同步。本文标准库里的空列表只表达数学贡献为零,没有提供可直接替代该路径的训练代码。

若全局N=0,应由全部进程一致采取事先约定的策略,例如共同跳过本更新单元。是否推进优化器、调度器与统计步数也要一致定义;不能每卡根据本地计数分别决定。数据不足和通信步序是两个必须同时解决的问题。

梯度累积时,分母应覆盖整个更新单元

若一次optimizer.step包含多个微批,目标中的N通常应覆盖本次更新涉及的全部有效token,而不是每个微批各自的N。否则仍可能变成微批平均的平均,只是把错误重复到了多卡环境。

可以在更新开始前掌握该单元计数,再按统一目标缩放各次损失和;或者采用经过验证、等价的梯度缩放方案。具体方式要与no_sync、AMP、梯度裁剪和最终同步位置配合,不能把本篇单步代数未经验证地拼接成完整训练循环。

此前梯度累积的损失归约讨论了跨微批权重,本篇增加跨rank梯度平均这一层。两层都要追踪分母,才能把一次更新对应到同一个明确目标。

日志标量和反向标量分开记录

为了反向使用W·Sᵣ/N,并不意味着要把这些本地标量简单当作全局训练指标。日志可分别汇总已脱离反向图的损失和与有效计数,再计算ΣSᵣ/N。记录中同时保存分片计数,能解释某次更新为什么与以往规模不同。

不要为了汇总日志,随意把一个跨卡聚合后的标量重新接回损失图。可微分collective与普通指标聚合的自动求导语义不同,应按真实API验证;仅仅“每卡显示同一个数”不能证明反向正确。

最终对照应固定数据与初始参数,用单进程合并批作为参考,再改变分片大小比较梯度。允许浮点归约顺序导致的合理容差,但应先排除整数倍或token权重差异。本篇没有执行这项多进程验收,因此不把算术通过写成完整DDP实现已验证。

用分片不变性检查目标

若同一组有效token只改变如何分配到rank,期望的全局token平均目标就不应改变。把这个要求写进测试,比凭卡数调整学习率更直接:学习率变化可能掩盖幅度差异,却不能恢复每个token应有的相对权重。

若这套推导与分片对照方法帮你定位了多卡归一化问题,欢迎通过文末支持区自愿支持。可复用的结果,是能明确说明每次更新究竟平均了谁的梯度。

来源与补发说明

DistributedDataParallel

依据原计划时段前已有的PyTorch 2.8版本定义核对。推导、算术代码和两幅图为原创教学内容。计划文章日期为2026年10月18日23:30,本次遗漏于2026年10月19日补发,实际公开时间另行核验报告。未运行PyTorch训练、分布式进程或GPU性能测试。

支持

如果文中的解释或示例帮你解决了问题、节省了时间,欢迎自愿打赏支持本站。

微信支持二维码,点击查看大图
微信
支付宝支持二维码,点击查看大图
支付宝
Buy Me a Coffee,点击支持本站
Buy Me a Coffee

二维码可点击放大。更多支持方式见支持页面。

Share this Post

Leave a Comment

您的邮箱地址不会被公开。 必填项已用 * 标注

*
*