[feat] distributed embedding export

.github/workflows/unittest_h20_ci.yml

@@ -1,6 +1,8 @@
 name: Unit Test H20 CI
 
 on:
+  # pull_request:
+  #   types: [opened, reopened, synchronize]
   workflow_dispatch:
 
 concurrency:

docs/source/usage/export.md

@@ -45,6 +45,8 @@ torchrun --master_addr=localhost --master_port=32555 \
   - **INPUT_TILE=3**：user侧fg和embedding计算仅一次，适用于user侧特征比较多的情况
 - INPUT_TILE_3_ONLINE: 配合INPUT_TILE=3使用，对User侧序列特征使用在线推理模式，序列特征在线模型服务中推理性能更好，但导出的模型无法用于离线预测
   - **INPUT_TILE_3_ONLINE=1**：启用序列特征的在线推理模式
+- USE_DISTRIBUTED_EMBEDDING: 开启分布式 embedding 导出模式，导出 dense graph 与分片 sparse embedding 参数；该模式会自动使用 `INPUT_TILE=3`
+  - **USE_DISTRIBUTED_EMBEDDING=1**：启用分布式 embedding 导出
 - ENABLE_AOT: 使用AOT(Ahead Of Time)编译优化导出的模型，可显著提升推理性能。**AOT 编译产物与导出机器的 GPU 架构强绑定，在线服务的 GPU 类型必须与导出时使用的 GPU 类型完全一致**，详见下文 [在 PAI 上导出 AOT 模型](#export-aot-on-pai) 章节
   - **ENABLE_AOT=1**: 使用AOT编译优化导出模型（sparse 部分用 JIT，dense 部分用 AOTI）
   - **ENABLE_AOT=2**: 使用统一 AOTI 模型编译优化 (sparse + dense 融合为单一 .pt2) [experimental]

tzrec/main.py

@@ -86,7 +86,10 @@
     create_train_pipeline,
     init_process_group,
 )
-from tzrec.utils.export_util import export_model
+from tzrec.utils.export_util import (
+    ensure_input_tile_for_distributed_embedding,
+    export_model,
+)
 from tzrec.utils.filesystem_util import url_to_fs
 from tzrec.utils.logging_util import ProgressLogger, logger
 from tzrec.utils.plan_util import create_planner, get_default_sharders
@@ -977,6 +980,8 @@ def export(
     """
     is_rank_zero = int(os.environ.get("RANK", 0)) == 0
 
+    ensure_input_tile_for_distributed_embedding()
+
     pipeline_config = config_util.load_pipeline_config(pipeline_config_path)
     ori_pipeline_config = copy.copy(pipeline_config)
 

tzrec/modules/embedding.py

@@ -12,7 +12,7 @@
 import os
 from collections import OrderedDict, defaultdict
 from functools import partial  # NOQA
-from typing import Dict, List, NamedTuple, Optional, Tuple, Union, cast
+from typing import Dict, FrozenSet, List, NamedTuple, Optional, Set, Tuple, Union, cast
 
 import torch
 from torch import nn
@@ -45,6 +45,7 @@
 from tzrec.utils.fx_util import (
     fx_int_item,
     fx_mark_keyed_tensor,
+    fx_mark_seq_ec_jt,
     fx_mark_seq_len,
     fx_mark_seq_tensor,
     fx_mark_tensor,
@@ -58,6 +59,7 @@
 torch.fx.wrap(fx_mark_tensor)
 torch.fx.wrap(fx_mark_seq_tensor)
 torch.fx.wrap(fx_mark_seq_len)
+torch.fx.wrap(fx_mark_seq_ec_jt)
 
 
 @torch.fx.wrap
@@ -901,7 +903,9 @@ def forward(
             kts.append(kt)
 
         if self.has_dense:
-            fx_mark_keyed_tensor(self._all_group_str + "__dense", dense_feature)
+            fx_mark_keyed_tensor(
+                self._all_group_str + "__dense", dense_feature, is_dense=True
+            )
             if self.has_dense_embedding:
                 kts.append(self.dense_ec(dense_feature))
             else:
@@ -977,6 +981,11 @@ def __init__(
         self._group_to_is_jagged = dict()
         self._group_to_sequence_length = OrderedDict()
 
+        seq_sparse_keys: Set[str] = set()
+        seq_sparse_keys_user: Set[str] = set()
+        query_sparse_keys: Set[str] = set()
+        query_sparse_keys_user: Set[str] = set()
+
         feat_to_group_to_emb_name = defaultdict(dict)
         group_name_to_suffix = {
             fg.group_name: fg.embedding_name_suffix for fg in feature_groups
@@ -1007,6 +1016,7 @@ def __init__(
             else:
                 shared_feature_flag[feature_name] = False
 
+        emb_name_to_feature_to_shared_name = defaultdict(dict)
         for feature_group in feature_groups:
             query_dim = 0
             sequence_dim = 0
@@ -1022,13 +1032,19 @@ def __init__(
             for name in feature_names:
                 shared_name = name
                 feature = name_to_feature[name]
+
                 if feature.is_sparse:
                     output_dim = feature.output_dim
                     emb_config = feature.emb_config
                     mc_module = feature.mc_module(device)
                     assert emb_config is not None
                     # we may/could modify ec name at feat_to_group_to_emb_name
                     emb_config.name = feat_to_group_to_emb_name[name][group_name]
+                    if shared_feature_flag[name]:
+                        shared_name = shared_name + "@" + emb_config.name
+                    emb_name_to_feature_to_shared_name[emb_config.name][name] = (
+                        shared_name
+                    )
                     embedding_dim = emb_config.embedding_dim
                     const = feature.parameter_constraints(emb_config)
 
@@ -1058,6 +1074,10 @@ def __init__(
                                 ] = const
                         if feature.is_sequence and feature.value_dim != 1:
                             self.has_mulval_seq_user = True
+                        if feature.is_sequence:
+                            seq_sparse_keys_user.add(shared_name)
+                        else:
+                            query_sparse_keys_user.add(shared_name)
                     else:
                         emb_configs = (
                             dim_to_mc_emb_configs[embedding_dim]
@@ -1084,9 +1104,10 @@ def __init__(
                                 ] = const
                         if feature.is_sequence and feature.value_dim != 1:
                             self.has_mulval_seq = True
-
-                    if shared_feature_flag[name]:
-                        shared_name = shared_name + "@" + emb_config.name
+                        if feature.is_sequence:
+                            seq_sparse_keys.add(shared_name)
+                        else:
+                            query_sparse_keys.add(shared_name)
                 else:
                     output_dim = feature.output_dim
                     if feature.is_sequence:
@@ -1124,12 +1145,28 @@ def __init__(
             self._group_output_dims[f"{group_name}.sequence"] = sequence_dims
             self._group_output_dims[group_name] = output_dims
 
+        self.seq_sparse_keys: FrozenSet[str] = frozenset(seq_sparse_keys)
+        self.seq_sparse_keys_user: FrozenSet[str] = frozenset(seq_sparse_keys_user)
+        self.query_sparse_keys: FrozenSet[str] = frozenset(query_sparse_keys)
+        self.query_sparse_keys_user: FrozenSet[str] = frozenset(query_sparse_keys_user)
+
         self.ec_dict = nn.ModuleDict()
         for k, emb_configs in dim_to_emb_configs.items():
             self.ec_dict[str(k)] = EmbeddingCollection(
                 list(emb_configs.values()), device=device
             )
 
+        self.ec_dict_features = OrderedDict()
+        for k, ec in self.ec_dict.items():
+            feat_list = []
+            for emb_config in ec._embedding_configs:
+                for feature_name in emb_config.feature_names:
+                    shared_name = emb_name_to_feature_to_shared_name[
+                        emb_config.name
+                    ].get(feature_name, feature_name)
+                    feat_list.append((feature_name, shared_name))
+            self.ec_dict_features[k] = feat_list
+
         self.mc_ec_dict = nn.ModuleDict()
         for k, emb_configs in dim_to_mc_emb_configs.items():
             self.mc_ec_dict[str(k)] = ManagedCollisionEmbeddingCollection(
@@ -1145,6 +1182,16 @@ def __init__(
                 self.ec_dict_user[str(k)] = EmbeddingCollection(
                     list(emb_configs.values()), device=device
                 )
+            self.ec_dict_features_user = OrderedDict()
+            for k, ec in self.ec_dict_user.items():
+                feat_list = []
+                for emb_config in ec._embedding_configs:
+                    for feature_name in emb_config.feature_names:
+                        shared_name = emb_name_to_feature_to_shared_name[
+                            emb_config.name
+                        ].get(feature_name, feature_name)
+                        feat_list.append((feature_name, shared_name))
+                self.ec_dict_features_user[k] = feat_list
 
             self.mc_ec_dict_user = nn.ModuleDict()
             for k, emb_configs in dim_to_mc_emb_configs_user.items():
@@ -1205,8 +1252,20 @@ def _forward_impl(
         dense_t_dict: Dict[str, torch.Tensor] = {}
 
         if self.has_sparse:
-            for ec in self.ec_dict.values():
-                sparse_jt_dict_list.append(ec(sparse_feature))
+            for feature_keys, ec in zip(
+                self.ec_dict_features.values(), self.ec_dict.values()
+            ):
+                d_jt = ec(sparse_feature)
+                new_d_jt = {}
+                for raw_key, shared_key in feature_keys:
+                    if (
+                        shared_key in self.seq_sparse_keys
+                        or shared_key in self.query_sparse_keys
+                    ):
+                        val = d_jt[raw_key]
+                        fx_mark_seq_ec_jt(f"{shared_key}", val)
+                        new_d_jt[shared_key] = val
+                sparse_jt_dict_list.append(new_d_jt)
 
         if self.has_mc_sparse:
             for ec in self.mc_ec_dict.values():
@@ -1216,8 +1275,20 @@ def _forward_impl(
             seq_mulval_length_jt_dict_list.append(sequence_mulval_lengths.to_dict())
 
         if self.has_sparse_user:
-            for ec in self.ec_dict_user.values():
-                sparse_jt_dict_list.append(ec(sparse_feature_user))
+            for feature_keys, ec in zip(
+                self.ec_dict_features_user.values(), self.ec_dict_user.values()
+            ):
+                d_jt = ec(sparse_feature_user)
+                new_d_jt = {}
+                for raw_key, shared_key in feature_keys:
+                    if (
+                        shared_key in self.seq_sparse_keys_user
+                        or shared_key in self.query_sparse_keys_user
+                    ):
+                        val = d_jt[raw_key]
+                        fx_mark_seq_ec_jt(f"{shared_key}", val)
+                        new_d_jt[shared_key] = val
+                sparse_jt_dict_list.append(new_d_jt)
 
         if self.has_mc_sparse_user:
             for ec in self.mc_ec_dict_user.values():

tzrec/utils/checkpoint_util.py

@@ -16,6 +16,7 @@
 import queue
 import re
 import shutil
+import tempfile
 import threading
 import weakref
 from dataclasses import replace
@@ -38,11 +39,26 @@
 )
 from torchrec.modules.mc_modules import MCHManagedCollisionModule
 
+from tzrec.acc.utils import is_input_tile_emb
 from tzrec.constant import TRAIN_EVAL_RESULT_FILENAME
 from tzrec.protos import export_pb2
 from tzrec.utils.dynamicemb_util import has_dynamicemb
 from tzrec.utils.logging_util import logger
 
+# INPUT_TILE=3 load-time fallback from user-side keys to non-user checkpoint keys.
+# Mirrors `tzrec/acc/utils.py:write_mapping_file_for_input_tile`.
+_INPUT_TILE_USER_REPLACEMENTS: Tuple[Tuple[str, str], ...] = (
+    (".ebc_user.embedding_bags.", ".ebc.embedding_bags."),
+    (".mc_ebc_user._embedding_module.", ".mc_ebc._embedding_module."),
+    (
+        ".mc_ebc_user._managed_collision_collection.",
+        ".mc_ebc._managed_collision_collection.",
+    ),
+    (".ec_list_user.", ".ec_list."),
+    (".mc_ec_list_user.", ".mc_ec_list."),
+    (".ec_dict_user.", ".ec_dict."),
+    (".mc_ec_dict_user.", ".mc_ec_dict."),
+)
 # queue token meaning "run a prune pass"; ``None`` means "stop the worker".
 _PRUNE_REQUEST = object()
 
@@ -124,6 +140,25 @@ def create_local_plan(self) -> LoadPlan:
                         "[{meta_fqn}], will be deprecated when tzrec version >= 1.0.0"
                     )
 
+            # INPUT_TILE=3 export adds user-side twin modules absent from
+            # training checkpoints. Remap to existing non-user keys, matching
+            # export_model_normal's emb_ckpt_mapping.txt fallback.
+            if (
+                is_input_tile_emb()
+                and meta_fqn not in self.metadata.state_dict_metadata
+            ):
+                for new_pat, old_pat in _INPUT_TILE_USER_REPLACEMENTS:
+                    if new_pat not in meta_fqn:
+                        continue
+                    candidate = meta_fqn.replace(new_pat, old_pat)
+                    if candidate in self.metadata.state_dict_metadata:
+                        logger.info(
+                            f"Remap INPUT_TILE=3 state [{fqn}] from [{candidate}]"
+                        )
+                        meta_fqn = candidate
+                        fqn_remap_set.add(fqn)
+                        break
+
             if meta_fqn in self.metadata.state_dict_metadata:
                 md = self.metadata.state_dict_metadata[meta_fqn]
             else:
@@ -794,6 +829,50 @@ def _a2a(
             m._buffers[name].copy_(new_meta)
 
 
+# Module-name segments that get a `_user` twin when INPUT_TILE=3 export
+# duplicates the embedding group into item/user halves. See
+# tzrec/modules/embedding.py:EmbeddingGroupImpl /
+# SequenceEmbeddingGroupImpl for the construction.
+_INPUT_TILE_USER_SEGMENTS = frozenset({"ebc", "mc_ebc", "ec_dict", "mc_ec_dict"})
+
+
+def _make_dynamicemb_input_tile_user_view(dynamicemb_path: str, view_path: str) -> str:
+    """Create a local symlink view for INPUT_TILE=3 dynamicemb loading.
+
+    The training checkpoint has non-user module paths, while INPUT_TILE=3
+    export adds user-side twins. Build aliases in a temporary local directory
+    instead of mutating the checkpoint directory.
+    """
+    entries = []
+    for entry in os.listdir(dynamicemb_path):
+        full_path = os.path.abspath(os.path.join(dynamicemb_path, entry))
+        if not os.path.isdir(full_path):
+            continue
+        entries.append((entry, full_path))
+        link_path = os.path.join(view_path, entry)
+        if not os.path.lexists(link_path):
+            os.symlink(full_path, link_path, target_is_directory=True)
+
+    for entry, full_path in entries:
+        segs = entry.split(".")
+        if any(seg.endswith("_user") for seg in segs):
+            continue
+        for i, seg in enumerate(segs):
+            if seg not in _INPUT_TILE_USER_SEGMENTS:
+                continue
+            user_segs = list(segs)
+            user_segs[i] = f"{seg}_user"
+            user_entry = ".".join(user_segs)
+            user_path = os.path.join(view_path, user_entry)
+            if os.path.lexists(user_path):
+                continue
+            os.symlink(full_path, user_path, target_is_directory=True)
+            logger.info(
+                f"created INPUT_TILE=3 dynamicemb alias {user_entry} -> {entry}"
+            )
+    return view_path
+
+
 def restore_model(
     checkpoint_dir: str,
     model: nn.Module,
@@ -872,16 +951,35 @@ def restore_model(
 
         dynamicemb_path = os.path.join(checkpoint_dir, "dynamicemb")
         if os.path.exists(dynamicemb_path):
+            # Training never sets INPUT_TILE, but exporting with
+            # INPUT_TILE=3 adds twin user-side modules (`ebc_user`,
+            # `mc_ebc_user`, `ec_dict_user`, `mc_ec_dict_user`) that
+            # share dynamic-embedding tables with their non-user counterparts.
+            # Build a local symlink view instead of mutating the checkpoint
+            # directory, which may be read-only or remote-mounted.
+            dynamicemb_load_path = dynamicemb_path
+            dynamicemb_view = None
+            if is_input_tile_emb():
+                dynamicemb_view = tempfile.TemporaryDirectory(
+                    prefix=f"tzrec_dynamicemb_rank{os.environ.get('RANK', '0')}_"
+                )
+                dynamicemb_load_path = _make_dynamicemb_input_tile_user_view(
+                    dynamicemb_path, dynamicemb_view.name
+                )
             logger.info(
                 f"RANK[{os.environ.get('RANK', 0)}] restoring dynamic embedding..."
             )
-            DynamicEmbLoad(
-                dynamicemb_path,
-                model,
-                table_names=meta.get("dynamicemb_load_table_names", None),
-                optim=meta.get("dynamicemb_load_optim", optimizer is not None),
-                counter=True,
-            )
+            try:
+                DynamicEmbLoad(
+                    dynamicemb_load_path,
+                    model,
+                    table_names=meta.get("dynamicemb_load_table_names", None),
+                    optim=meta.get("dynamicemb_load_optim", optimizer is not None),
+                    counter=True,
+                )
+            finally:
+                if dynamicemb_view is not None:
+                    dynamicemb_view.cleanup()
             logger.info(
                 f"RANK[{os.environ.get('RANK', 0)}] restore dynamic embedding finished."
             )

tzrec/utils/env_util.py

@@ -32,6 +32,11 @@ def use_rtp() -> bool:
     return flag
 
 
+def use_distributed_embedding() -> bool:
+    """Export model for distributed embedding mode of EAS processor."""
+    return os.environ.get("USE_DISTRIBUTED_EMBEDDING", "0") == "1"
+
+
 def enable_tma() -> bool:
     """Enable TMA (Tensor Memory Accelerator) for triton ops."""
     flag = os.environ.get("ENABLE_TMA", "0") == "1"