diff --git a/docs/source/feature/zch.md b/docs/source/feature/zch.md index 0aeb9316..08d7895f 100644 --- a/docs/source/feature/zch.md +++ b/docs/source/feature/zch.md @@ -21,6 +21,11 @@ feature_configs { - **zch_size**: 零冲突Hash的Bucket大小，Id数超过后会根据Id的驱逐策略进行淘汰 + > **注意**：`continue_train` / `fine_tune` 时**不要修改**已有 checkpoint 对应特征的 `zch_size`。 + > ZCH 的 Managed Collision 状态与 bucket 大小绑定；改大或改小都会在恢复 checkpoint 时失败 + > （见 [Issue #176](https://github.com/alibaba/TorchEasyRec/issues/176)）。 + > 若需要更大的动态容量，可考虑切换到 [DynamicEmbedding](./dynamicemb.md)（含 ZCH→dynamicemb 转换工具）。 + - **eviction_interval**: Id准入和驱逐策略执行的频率（训练步数间隔） - **eviction_policy**: 驱逐策略，可选`lfu`，`lru`，`distance_lfu`，详见下文驱逐策略 diff --git a/tzrec/tests/rank_integration_test.py b/tzrec/tests/rank_integration_test.py index 22f0e67e..68012e10 100644 --- a/tzrec/tests/rank_integration_test.py +++ b/tzrec/tests/rank_integration_test.py @@ -854,6 +854,52 @@ def test_multi_tower_din_zch_finetune_world_size_change(self): else: os.environ["TEST_NPROC_PER_NODE"] = prev_nproc + @unittest.skipIf(*gpu_unavailable) + @mark_ci_scope("gpu") + def test_multi_tower_din_zch_finetune_zch_size_change_fails(self): + # Changing zch_size on continue_train / fine_tune must fail fast with a + # clear error (issue #176), instead of MCH validate_state AssertionError. + self.success = utils.test_train_eval( + "tzrec/tests/configs/multi_tower_din_zch_fg_mock.config", + os.path.join(self.test_dir, "train"), + user_id="user_id", + item_id="item_id", + ) + self.assertTrue(self.success) + + # Bump one ZCH table size in the saved pipeline config, then fine-tune. + finetune_dir = os.path.join(self.test_dir, "finetune") + os.makedirs(finetune_dir, exist_ok=True) + pipeline_config = config_util.load_pipeline_config( + os.path.join(self.test_dir, "train/pipeline.config") + ) + changed = False + for feature in pipeline_config.feature_configs: + feature_type = feature.WhichOneof("feature") + feature_config = getattr(feature, feature_type) + if hasattr(feature_config, "zch") and feature_config.HasField("zch"): + feature_config.zch.zch_size = feature_config.zch.zch_size + 1 + changed = True + break + self.assertTrue(changed, "expected at least one zch feature in mock config") + finetune_config_path = os.path.join(finetune_dir, "pipeline.config") + config_util.save_message(pipeline_config, finetune_config_path) + + log_path = os.path.join(finetune_dir, "log_train_eval.txt") + self.success = utils.test_train_eval( + finetune_config_path, + finetune_dir, + user_id="user_id", + item_id="item_id", + args_str=( + f"--fine_tune_checkpoint {os.path.join(self.test_dir, 'train/train')}" + ), + ) + self.assertFalse(self.success) + with open(log_path, "r") as f: + log = f.read() + self.assertIn("ZCH (MCH) table size mismatch", log) + @mark_ci_scope("gpu") def test_multi_tower_din_with_fg_export_quant(self): self._test_rank_with_fg_quant( diff --git a/tzrec/utils/checkpoint_util.py b/tzrec/utils/checkpoint_util.py index 672ecd4c..ba9c43b9 100644 --- a/tzrec/utils/checkpoint_util.py +++ b/tzrec/utils/checkpoint_util.py @@ -681,6 +681,121 @@ def _needs_mch_redistribution(model_ckpt_path: str, cur_world_size: int) -> bool return saved_world_size != cur_world_size and saved_world_size % cur_world_size != 0 +_MCH_RAW_IDS_SUFFIX = "._mch_sorted_raw_ids" + + +def _mch_global_zch_size(local_zch_size: int, world_size: int) -> int: + """Global ZCH table size under uniform row-wise sharding.""" + return local_zch_size * world_size + + +def _is_mch_zch_size_compatible( + saved_buffer_size: int, + cur_local_zch_size: int, + saved_world_size: int, + cur_world_size: int, +) -> bool: + """Whether checkpoint MCH buffer size matches the current model capacity. + + DCP metadata may record either the global ShardedTensor size or a + per-rank local size. Accept any interpretation that yields the same + global ZCH capacity as the current model under uniform RW sharding. + """ + cur_global = _mch_global_zch_size(cur_local_zch_size, cur_world_size) + if saved_buffer_size == cur_global: + return True + if saved_world_size == cur_world_size and saved_buffer_size == cur_local_zch_size: + return True + if saved_buffer_size * saved_world_size == cur_global: + return True + return False + + +def _find_mch_raw_ids_metadata_key( + state_dict_metadata: Dict[str, Any], prefix: str +) -> Optional[str]: + """Locate ``_mch_sorted_raw_ids`` metadata key for an MC module prefix.""" + exact = f"{prefix}{_MCH_RAW_IDS_SUFFIX}" + if exact in state_dict_metadata: + return exact + suffix = f".{prefix}{_MCH_RAW_IDS_SUFFIX}" + for key in state_dict_metadata: + if key == exact or key.endswith(suffix) or key.endswith( + f"{prefix}{_MCH_RAW_IDS_SUFFIX}" + ): + return key + return None + + +def _check_mch_zch_size_compatible(model: nn.Module, model_ckpt_path: str) -> None: + """Fail fast when checkpoint ZCH table size disagrees with the model. + + Changing ``zch.zch_size`` across ``continue_train`` / fine-tune is not + supported. ``_output_segments_tensor`` is a fixed-shape (1025,) buffer, so + DCP loads the old partition boundaries even when ``zch_size`` changed; + ``MCHManagedCollisionModule.validate_state`` then asserts with an opaque + message (see https://github.com/alibaba/TorchEasyRec/issues/176). + """ + mc_modules = _find_mch_modules(model) + if not mc_modules or not os.path.exists(model_ckpt_path): + return + + try: + metadata = FileSystemReader(model_ckpt_path).read_metadata() + except Exception as e: + logger.warning( + f"Failed to read checkpoint metadata for ZCH size check " + f"from {model_ckpt_path}: {e}" + ) + return + + cur_world_size = dist.get_world_size() if dist.is_initialized() else 1 + try: + saved_world_size = _ckpt_world_size(model_ckpt_path) + except Exception as e: + logger.warning( + f"Failed to detect saved world size from {model_ckpt_path}: {e}" + ) + saved_world_size = cur_world_size + + # pyre-ignore [16] + state_dict_metadata = metadata.state_dict_metadata + mismatches: List[str] = [] + for prefix, m in mc_modules.items(): + key = _find_mch_raw_ids_metadata_key(state_dict_metadata, prefix) + if key is None: + continue + md = state_dict_metadata[key] + size = getattr(md, "size", None) + if size is None or len(size) < 1: + continue + saved_buffer_size = int(size[0]) + cur_local = int(m._zch_size) + if _is_mch_zch_size_compatible( + saved_buffer_size, cur_local, saved_world_size, cur_world_size + ): + continue + mismatches.append( + f"[{prefix}] saved_buffer_size={saved_buffer_size}, " + f"current_zch_size={cur_local}, " + f"saved_world_size={saved_world_size}, " + f"current_world_size={cur_world_size}" + ) + + if not mismatches: + return + + detail = "; ".join(mismatches) + raise RuntimeError( + "ZCH (MCH) table size mismatch when restoring checkpoint. " + "Changing feature zch.zch_size between training and continue_train / " + "fine_tune is not supported; keep the original zch_size, or start a " + "new train without restoring this checkpoint. " + f"Details: {detail}. " + "See https://github.com/alibaba/TorchEasyRec/issues/176" + ) + + def _strip_dmp_prefix(name: str) -> str: """Strip TorchRec DMP wrapper prefix from a module name.""" for prefix in ( @@ -901,6 +1016,8 @@ def restore_model( needs_mch_redistribution = _needs_mch_redistribution( model_ckpt_path, cur_world_size ) + if os.path.exists(model_ckpt_path): + _check_mch_zch_size_compatible(model, model_ckpt_path) meta = {} if os.path.exists(meta_path):