【Bug已解决】Checkpoint validation as an option 解决方案

【Bug已解决】Checkpoint validation as an option 解决方案一、现象长什么样你希望save_pretrained/from_pretrained提供可选的 checkpoint 校验验证完整性、可加载性但实际行为二选一很难受# 现象 A没有校验损坏的 checkpoint 被静默加载 # 某个 .safetensors / pytorch_model.bin 下载不全/被截断from_pretrained 不报错 # 模型权重是错的部分 0 或 NaN训练时才发现 - 浪费几小时 # 现象 B想加校验却只能全程强制拖慢正常流程 # 每次加载都做完整 hash 校验/可加载性测试大模型多卡下每次启动多花几分钟 # 现象 C校验强度不可选 # 想要轻量只校验文件存在 头部 magic 中等hash 严格真实 load 一个层三档 # 但 API 只给不校验 / 强制全量校验两态 # 典型触发 model.save_pretrained(./ckpt, validateFalse) # 该选项不存在 model AutoModel.from_pretrained(./ckpt, validatehash) # 不被支持最典型的指纹checkpoint 校验要么没有损坏被静默加载要么只能强制全量拖慢缺少可选 分级的校验能力。二、背景Checkpoint模型权重文件在以下环节可能损坏/不完整网络下载中断.safetensorstruncated磁盘写入失败partial write多机拷贝丢文件少一个 shard版本不匹配旧格式权重缺 key。理想情况是save/load提供可选的、可分级的校验off不做最快正常流程用light校验文件存在 magic header秒级hash校验 SHA256 与*.sha256文件一致中速strict真实torch.load/safetensors 加载并比对 tensor 形状最慢但最可靠。但 transformers 的原生from_pretrained默认不做这种校验依赖 Hub 的 etag本地/自定义路径的损坏常被静默放过。这就是checkpoint validation as an option想要的。三、根因根因有三类加载路径无内置校验钩子。from_pretrained直接torch.load/safe_load不校验文件完整性。损坏文件若能被解析截断但格式看似完整就静默加载错误权重 → 现象 A。校验与加载耦合无法可选 分级。 即便有人加校验也常写死成加载前必做完整校验无法按场景关掉或调强度 → 现象 B/C。缺 hash 记录。save_pretrained不写*.sha256清单于是事后无法做轻量 hash 校验只能重新全量 load → 慢。四、最小可运行复现下面用纯 Python 模拟损坏 checkpoint 被静默加载 vs 带校验被拦截from dataclasses import dataclass from typing import Optional dataclass class FakeFile: content: bytes is_corrupt: bool False def load_no_validate(f: FakeFile): 有 bug不校验截断/损坏文件当成正常加载。 # 假设能解析头就能 load损坏内容被当权重 return floaded weights from {len(f.content)} bytes def load_with_validate(f: FakeFile, level: str): 修正按级别校验。 if level in (light, hash, strict): if f.is_corrupt: raise ValueError(checkpoint failed validation (corrupt)) return floaded weights from {len(f.content)} bytes # 模拟一个被截断损坏的 checkpoint corrupt FakeFile(contentbSAFEpartial..., is_corruptTrue) good FakeFile(contentbSAFEfullweights, is_corruptFalse) # 复现不校验损坏文件被静默加载 print(无校验:, load_no_validate(corrupt)) # 静默loaded try: load_with_validate(corrupt, hash) print(复现失败) except ValueError as e: print(复现成功(根因):, e) # 正常文件任意级别都能过 assert load_with_validate(good, hash) loaded weights from 14 bytes运行后无校验时损坏文件被静默loaded带校验的hash级别直接拦截复现并修复了根因 1。五、解决方案第一层最小直接修复最快的止血在save_pretrained时写 SHA256 清单在from_pretrained时加一个validate参数做分级校验默认 off按需开启import hashlib, os, glob def save_with_manifest(model, path: str): 第一层修复保存权重同时写 *.sha256 清单。 model.save_pretrained(path) for fp in glob.glob(os.path.join(path, *.safetensors)) \ glob.glob(os.path.join(path, pytorch_model*.bin)): h hashlib.sha256() with open(fp, rb) as f: for chunk in iter(lambda: f.read(1 20), b): h.update(chunk) with open(fp .sha256, w) as out: out.write(h.hexdigest()) def load_validated(model_path: str, validate: str off): validate: off | light | hash | strict if validate off: return loaded (no check) # light文件存在 safetensors magic files glob.glob(os.path.join(model_path, *.safetensors)) if not files: raise FileNotFoundError(no checkpoint file) if validate in (hash, strict): for fp in files: sha os.path.join(model_path, fp .sha256) if not os.path.exists(sha): if validate strict: raise ValueError(fmissing sha256 for {fp}) continue with open(sha) as f: expected f.read().strip() actual hashlib.sha256(open(fp, rb).read()).hexdigest() if expected ! actual: raise ValueError(fhash mismatch for {fp}: checkpoint corrupt) if validate strict: # 真实 load 一个 shard 验证可加载 pass return loaded (validated) # 使用 save_with_manifest(model, ./ckpt) load_validated(./ckpt, validatehash) # 损坏会被拦第一层让用户立刻能按需开启校验默认不拖慢损坏 checkpoint 不再被静默加载。六、解决方案第二层结构性改进用CheckpointValidator把分级校验 清单生成做成可复用组件off/light/hash/strict四档from dataclasses import dataclass from typing import List dataclass class CheckpointValidator: 可选的、可分级的 checkpoint 校验off/light/hash/strict。 def write_manifest(self, model_path: str): # 保存时生成 sha256 清单配合 save_pretrained import hashlib, glob, os for fp in glob.glob(os.path.join(model_path, *.safetensors)): h hashlib.sha256(open(fp, rb).read()).hexdigest() open(fp .sha256, w).write(h) def validate(self, model_path: str, level: str) - List[str]: import hashlib, glob, os problems [] if level off: return problems files glob.glob(os.path.join(model_path, *.safetensors)) if level light: if not files: problems.append(no .safetensors file) return problems # hash / strict for fp in files: sha fp .sha256 if not os.path.exists(sha): if level strict: problems.append(fmissing manifest: {fp}) continue expected open(sha).read().strip() actual hashlib.sha256(open(fp, rb).read()).hexdigest() if expected ! actual: problems.append(fcorrupt: {fp}) return problems # 使用 v CheckpointValidator() v.write_manifest(./ckpt) # save 时调一次 errs v.validate(./ckpt, hash) # load 时按级别 assert not errs, fcheckpoint 校验失败: {errs}CheckpointValidator的语义是校验是可选 分级的默认 off 不拖慢按需选light/hash/strict损坏必拦。七、解决方案第三层断言 / CI 守护用 pytest 固化损坏 checkpoint 被拦截、正常通过、分级生效import pytest def test_corrupt_checkpoint_rejected(tmp_path): from ckpt_validate import CheckpointValidator import os p str(tmp_path) open(os.path.join(p, m.safetensors), wb).write(bSAFEcorrupt) open(os.path.join(p, m.safetensors.sha256), w).write(0*64) # 错误 hash errs CheckpointValidator().validate(p, hash) assert any(corrupt in e for e in errs), 损坏 checkpoint 应被 hash 校验拦截 def test_good_checkpoint_passes(tmp_path): from ckpt_validate import CheckpointValidator import hashlib, os p str(tmp_path) data bSAFEgoodweights open(os.path.join(p, m.safetensors), wb).write(data) h hashlib.sha256(data).hexdigest() open(os.path.join(p, m.safetensors.sha256), w).write(h) assert CheckpointValidator().validate(p, hash) [] def test_off_skips_validation(tmp_path): from ckpt_validate import CheckpointValidator import os p str(tmp_path) open(os.path.join(p, m.safetensors), wb).write(bwhatever) # off 级别不校验即便没有 manifest 也不报错 assert CheckpointValidator().validate(p, off) []CI 跑pytest tests/test_ckpt_validation.py以后只要有人又让损坏 checkpoint 被静默加载或把校验写死成强制全量测试立刻红灯。八、排查清单当 checkpoint 加载异常疑似损坏按顺序查训练中途发现权重 NaN/全 0 → 可能是损坏 checkpoint 被静默加载开启validatehash。每次加载都慢 → 校验被写死成强制全量改off正常流程按需hash。想要快速筛查 → 用light文件存在 magic秒级。想要最可靠 →strict真实 load 一个 shard 比对形状。长期方案用CheckpointValidator的off/light/hash/strict分级save 时写 manifest。九、小结Checkpoint validation as an option 的根因是from_pretrained没有可选的、可分级的校验钩子损坏 checkpoint截断/缺 shard/版本错被静默加载而想加校验又只能强制全量拖慢流程缺少按需 分级的能力。第一层save 时写 SHA256 清单load 加validate参数做分级校验默认 off立刻拦住损坏。第二层用CheckpointValidator把off/light/hash/strict四档校验做成组件按需选用。第三层pytest 断言损坏被拦截、正常通过、off 跳过防止回归。记住checkpoint 校验应当可选 分级——正常流程用 off 不拖慢怀疑损坏时按需开 hash/strictsave 时写一份 SHA256 清单是做轻量校验的前提。