Skip to content

[bugfix] 恢复 checkpoint 时若 ZCH zch_size 已变更,提前抛出明确错误#600

Open
fooSynaptic wants to merge 1 commit into
alibaba:masterfrom
fooSynaptic:bugfix/zch-size-mismatch-fail-fast
Open

[bugfix] 恢复 checkpoint 时若 ZCH zch_size 已变更,提前抛出明确错误#600
fooSynaptic wants to merge 1 commit into
alibaba:masterfrom
fooSynaptic:bugfix/zch-size-mismatch-fail-fast

Conversation

@fooSynaptic

Copy link
Copy Markdown

Hi team,

我在 continue_train / fine_tune 时也遇到了与 #176 相同的问题:修改特征的 zch.zch_size 后,checkpoint 仍能被 DCP 加载,但随后在 MCHManagedCollisionModule.validate_state 里以难懂的 AssertionError 失败。本 PR 针对这一场景,在 restore_model 的 DCP load 之前做兼容性检测,提前抛出可操作的 RuntimeError,并补充文档与 GPU 集成回归测。

Summary

  • 针对 #176 所述场景:在 restore_model 入口、DCP 加载前,比对 checkpoint metadata 中 _mch_sorted_raw_ids 的 buffer size 与当前模型各 MCH 模块的 _zch_size;不一致则 fail-fast,给出含模块前缀与 size 明细的错误信息。
  • 兼容已有 world size 变更逻辑:全局 ZCH 容量未变(仅 world size 变化、需 redistribution)时不误报。
  • metadata 读取失败时 warn 后放行(fail-open),避免阻断无关 restore。
  • 新增 GPU 集成测试 test_multi_tower_din_zch_finetune_zch_size_change_fails,覆盖与 issue 相同的 zch_size 变更 + fine_tune 路径;在 docs/source/feature/zch.md 注明 continue_train / fine_tune 不可修改 zch_size
  • 不做 MCH 状态 resize/remap;遇到相同需求时请保持原 zch_size 或换用 DynamicEmbedding。

根因

#176 描述一致:_output_segments_tensor 为固定形状 (1025,) 的 buffer。修改 zch_size 不改变该 tensor 形状,DCP 仍会加载旧分区边界,最终在 validate_state 中以晦涩 assert 失败。

Test plan

  • 我已在 自有 GPU 集成环境(单卡、torchrec 1.7.x 栈)复现 issue 相同路径(训练 → zch_size += 1 → fine_tune),RankIntegrationTest.test_multi_tower_din_zch_finetune_zch_size_change_fails 通过(约 77s);fine_tune 阶段日志含 ZCH (MCH) table size mismatch,符合预期

本地复现(需 GPU)

PYTHONPATH=. python -m tzrec.tests.rank_integration_test \
  RankIntegrationTest.test_multi_tower_din_zch_finetune_zch_size_change_fails -v

预期在 zch_size += 1 后的 fine_tune 阶段(与 issue 相同报错形态):

RuntimeError: ZCH (MCH) table size mismatch when restoring checkpoint. ...
Details: [...] saved_buffer_size=1000000, current_zch_size=1000001,
saved_world_size=1, current_world_size=1. See .../issues/176

Fixes #176

Detect MCH buffer size mismatch from DCP metadata before loading,
so continue_train/fine_tune with a changed zch_size raises a clear
RuntimeError instead of an opaque validate_state AssertionError.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

当我使用增量训练的时候,只改变修改一个feature的zch_size的大小,其他所有配置都不改变,会出错。

1 participant