TorchTitan 静默数据损坏(SDC)检测实战:基于确定性重放的sdc_replayer全解析
【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan
导读
静默数据损坏(Silent Data Corruption, SDC)是指硬件或软件错误悄然篡改数据,却不产生任何显式报错,最终表现为训练 loss 曲线异常、模型质量劣化甚至隐性训练失败。TorchTitan 通过**确定性重放(deterministic replay)**机制检测此类问题:在受检的优化器步骤中,重放完全确定性的前向/反向计算,并比对所有可观测状态,一旦出现偏差即在所有 rank 上抛出SDCReplayMismatch异常。本文以 silent_data_corruption.md 为主线,结合 sdc_replayer.py 与 trainer.py 的源码实现,完整讲解该特性的启用方式、重放边界、签名构成、限制与失败上报机制,帮助你把它直接落地到自己的训练配置中。
SDC 重放检测的原理与定位
SDC 重放属于多种 SDC 检测策略之一。与之并列的替代方案还包括:
- 冗余硬件上的影子计算(shadow computation):在额外硬件上重复执行计算并交叉验证,成本高且依赖冗余资源;
- 算法级校验(algorithm-level checks):例如带校验和的矩阵乘法,属于算子内嵌检查。
TorchTitan 选择的重放方案则不同:它牺牲受检步骤的额外前向/反向时间,换取一个训练过程中的、与硬件无关的完整性检查。其核心前提是执行必须完全确定:同一输入、同一状态下的前向/反向必须产出完全一致的可观测结果,任何差异都只能由静默数据损坏引起。实现上,该特性默认关闭(config.sdc_replayer默认为None),不会给训练引入任何额外开销。
从 sdc_replayer.py 的模块文档可以看到其完整生命周期:
- 快照执行前状态(Python/CPU/加速器 RNG、注册的模块 buffer、调用方持有的标量),并记录哪些参数进入时没有梯度;
- 执行一次前向/反向,记录参考签名(reference signature):loss、梯度、buffer、RNG 推进、标量状态;
- 恢复快照并重放
num_replays次,将每次签名与参考签名比对; - 一旦出现任何分歧,在所有 rank 上抛出
SDCReplayMismatch;否则仅提交最后一次执行的结果。
配置方式:无 CLI 标志,纯编程式启用
SDC 重放没有 CLI 标志,必须在 config/recipe 中通过赋值配置对象来启用。config.sdc_replayer默认是None(即禁用重放)。这一点在 trainer.py 中有明确注释:该字段被tyro.conf.Suppress标注,从 tyro CLI 解析中隐藏,只能编程式赋值。
最小启用配置如下(来自原文档,可直接放入 config_registry 函数):
from torchtitan.observability.sdc_replayer import SDCReplayer config.debug.deterministic = True config.sdc_replayer = SDCReplayer.Config( num_steps=1, # optimizer steps checked after each (re)start; -1 checks every step num_replays=1, # re-executions compared against the reference )配置参数详解
| 参数 | 默认值 | 含义与取值范围 |
|---|---|---|
num_steps | 1 | 从 trainer 启动起、并在每次 checkpoint 加载后重新计数,需要检查多少个优化器步骤。-1表示检查每一步;默认值1只检查每次(重)启动后的第一步。校验规则:必须为-1或至少为1,否则抛出ValueError |
num_replays | 1 | 受检前向/反向相对初始参考执行被重放并比对的次数。必须至少为1;更大值可以捕捉单次重放容易漏掉的间歇性损坏 |
两个参数都定义在 sdc_replayer.py 的SDCReplayer.Configdataclass 中,并在__post_init__里做合法性校验(num_steps非-1时必须>= 1,num_replays必须>= 1)。
为什么默认只检查(重)启动后的第一步?因为训练刚恢复或刚初始化时,损坏最可能来自一次错误的 checkpoint 恢复或参数初始化。restart 语义通过reset_schedule()实现——trainer 在每次 checkpoint 加载后(load_state_dict)都会调用它重置计数(trainer.py),因此 "after each (re)start" 的语义被精确落实。
启用前置条件
重放要求同时满足以下条件,trainer.py 的_validate_sdc_replay会在配置解析阶段强制执行:
debug.deterministic = True:使用尽可能确定的算法(见 configs.py);debug.deterministic_warn_only = False:不能只对非确定性算子发出警告而不报错(configs.py);- 使用
torch.hash_tensor:SDCReplayer.__init__会在构造时通过_validate_hash_support探测当前设备上torch.hash_tensor的可用性(sdc_replayer.py)。
另外还有一条硬性组合限制:CUDA graphs 开启时num_replays必须为1,否则配置校验直接抛错:
ValueError: SDC replay supports at most one replay when CUDA graphs are enabled: set sdc_replayer.num_replays=1 or training.disable_cuda_graphs=True.一个真实的启用示例
graph_trainer/llama3/config_registry.py 展示了在实验性 recipe 中的实际用法——直接赋值config.sdc_replayer = SDCReplayer.Config()即可(使用两个参数的默认值)。
重放边界:什么被重放,什么不被重放
仅重放受检步骤的第一个前向/反向
一个受检优化器步骤中,只有第一个前向/反向调用(即一个梯度累积组)被重放;在流水线并行(PP)下,这个"一个梯度累积组"等价于一个完整的 pipeline schedule(包含所有 pipeline microbatch)。训练循环中的对应逻辑在 trainer.py:仅当fwd_bwd_index == 0时才走self.sdc_replayer.run_fwd_bwd(...),其余累积组直接执行、不检查。
为什么后面的累积组不检查?这是成本选择而非引擎限制:重放引擎可以检查任何前向/反向 callable,而后续累积组与第一组走的是相同的计算与通信路径,继续检查只会把受检步骤的开销成倍放大,却覆盖不到新的代码路径。因此当某步有多个累积组时,只有第一组受检。
状态恢复与提交语义
每次执行前都会恢复状态:参考执行与中间执行的结果均被丢弃,只提交最后一次执行产出的梯度、注册 buffer、RNG 推进、token 计数器和 loss。对应实现中,_run_checked先capture()基线,执行参考并clone()其签名,然后循环num_replays次:每次restore(baseline)后执行、比对签名,最后返回最后一次的 loss(sdc_replayer.py)。
梯度值永不快照:零梯度进入契约
这是本特性一个非常关键的设计细节:梯度值永远不会被快照。受检的前向/反向必须以"无待处理梯度"状态进入(None或全零,即zero_grad之后的状态),restore 直接依据该契约重建入口状态:
- 进入时无梯度的参数,恢复为
None; - 进入时是零张量的,原地清零(in-place zero),从而保留其存储地址供 CUDA graphs 使用,且受检步骤上不产生任何梯度大小的克隆或拷贝 kernel(sdc_replayer.py)。
违反该契约会产生假性不匹配(false mismatch)。源码注释给出了两类典型误用(sdc_replayer.py):
- 把非第一个梯度累积组路由进
run_fwd_bwd:restore 会抹掉部分累加值,导致重放与参考分歧; requires_grad=True的参数未纳入优化器:trainer 的zero_grad不会清它的梯度,后续受检步骤会带着累积梯度进入。这类参数应冻结为requires_grad=False。
执行后端与 scratch 状态的边界
以下执行后端都可以参与重放:eager 执行、torch.compile、CUDA graphs、symmetric-memory FSDP、distributed GEMM、async TP、DeepEP v2 与 HybridEP。GraphTrainer(实验性图训练器)也使用相同的重放边界(见 graph_trainer 目录)。
引擎对执行后端内部的 scratch 状态没有要求:只要一次前向/反向调用在返回前完成、后续调用在读取前覆盖这些状态即可。这类 scratch 状态既不被快照,也不进入重放签名。
签名构成与失败上报
签名里有什么
受检签名(signature)由_signature方法构建(sdc_replayer.py),包含:
- loss;
- 本地参数梯度(每个参数按其
named_parameters名称逐一哈希;None梯度记录为none模式占位); - 注册的模块 buffer(按模块名 + buffer 名逐一哈希;
Nonebuffer 同样有占位模式); - Python 与 torch RNG 状态(
random.getstate()、torch.get_rng_state()以及每设备的加速器 RNG 状态,如torch.cuda.get_rng_state_all(),签名项形如rng:cpu、rng:accelerator:<index>); - token 计数器(trainer 中唯一被重放前向/反向修改的标量
ntokens_seen,通过ScalarStateAccessor注册,见 trainer.py)。
每个张量项都会记录其 shape、dtype、device 组成的 schema,并以torch.hash_tensor生成设备端 digest,全部聚合到 loss 所在设备上做批量比对。
不匹配如何上报
一旦比对发现分歧,会在梯度裁剪、优化器、学习率调度器、checkpoint 保存之前,在所有 rank 上抛出SDCReplayMismatch(sdc_replayer.py)。异常信息精确标识:
step:全局优化器步骤;local_step:该步骤在当前检查调度中的 1 基位置(调度在 checkpoint 加载后重启,local_step也随之重置);replay:第几次重放发现的分歧;rank:最先发现分歧的 rank(通过all_reduce+all_gather_object在全体 rank 间汇总并选取第一个有分歧者,见 sdc_replayer.py);signature:第一个出现分歧的签名条目名(如loss、gradient:<i>:<name>、buffer:<i>:<name>、state:python_rng等)。
例如:
SDC replay mismatch: step=42, local_step=1, replay=1, rank=3, signature='gradient:0:layers.0.attn.q_proj.weight'成本模型
受检步骤的预期成本为1 + num_replays次前向/反向执行(1 次参考 +num_replays次重放)。未受检步骤不计算重放签名,零额外开销。
当前限制
CUDA graphs 与多次重放不兼容
CUDA graphs 目前要求sdc_replayer.num_replays=1。若需要更多次重放,则要恢复 graph 持有的梯度与可选 buffer 存储且不改变其被捕获的地址——当前实现做不到,因此在配置校验阶段直接禁止(trainer.py)。从 cudagraph.py 可以看到,CUDA graph 包装器会根据sdc_num_steps/sdc_num_replays计算预热与受检步骤数,将重放纳入 graph 捕获调度。
XOR 校验和的碰撞面
当前重放使用现有的基于 XOR 的torch.hash_tensor模式(sdc_replayer.py)。该校验和与顺序无关,因此:
- 置换(permutation)会碰撞:张量元素重排后哈希不变;
- 部分重复值损坏也可能碰撞。
源码中明确标注这一局限,并说明复杂张量(如 RoPE 缓存)按实部/虚部拆分后继承同样的限制;选用该哈希是因为它跨张量应用简单、对重放足够快。这意味着重放检测是"大概率发现损坏"的强信号,而非数学上完备的校验。
集成到训练循环:调用链速览
把以上机制串起来,SDC 重放在 trainer 中的完整调用链为:
- 配置校验:
JobConfig.__post_init__→_validate_sdc_replay(trainer.py); - 构造:
SDCReplayer(config.sdc_replayer, modules=self.model_parts, device=self.device, scalar_state={...})(trainer.py),同时校验torch.hash_tensor可用性; - CUDA graph 集成:若未禁用 CUDA graphs,
wrap_with_cuda_graph接收sdc_num_steps/sdc_num_replays(trainer.py); - 训练步:每个优化器步的
fwd_bwd_index == 0时调用run_fwd_bwd(fwd_bwd, step=self.step),调度判定num_steps == -1 or local_step < num_steps(sdc_replayer.py); - 调度重置:
load_state_dict加载 checkpoint 后调用reset_schedule()(trainer.py),使恢复后的第一步重新受检。
注意:SDC 重放状态是进程本地、不进 checkpoint的,其检查调度在每次 checkpoint 加载后自动重启(trainer.py 的注释明确了这一点)。
总结与推荐用法
- 日常训练:保持
config.sdc_replayer = None(默认),零开销; - 关键恢复点:在 checkpoint 恢复后启用默认配置
SDCReplayer.Config(),用最小成本验证"恢复是否正确"——这正是默认num_steps=1的设计意图; - 容错要求高、怀疑间歇性损坏:调大
num_replays(如2~3)以捕捉单次重放易漏的偶发损坏,同时注意 CUDA graphs 下必须保持num_replays=1; - 极致校验:
num_steps=-1让每一步的第一累积组都受检,代价是每步执行1 + num_replays次前向/反向。
SDC 重放并不能替代硬件层面的可靠性与算法级校验,但它以极低的接入成本(纯配置、无代码改动),为大规模分布式训练提供了训练过程中可随时开启的确定性完整性防线——尤其适合在跨节点、长周期训练中作为"体检"手段,把静默损坏暴露在梯度裁剪之前,避免污染优化器状态与最终 checkpoint。
【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考