[feat] remove patch_dynamicemb_eval_model (#286)

Author: tiankongdeguiji

docs/source/conf.py

@@ -29,7 +29,7 @@
 # -- Project information -----------------------------------------------------
 
 project = "tzrec"
-copyright = "2024, EasyRec Team"
+copyright = "2024-2025, EasyRec Team"
 author = "EasyRec Team"
 
 version_file = "../../tzrec/version.py"

docs/source/feature/dynamicemb.md

@@ -5,7 +5,7 @@ DynamicEmbedding 是特征零Hash冲突Id化的一种方式，它相比设置`ha
 注：目前使用DynamicEmbedding还处于实验阶段，配置和接口都可能调整，暂不包含在官方提供的镜像环境中，使用前需要额外安装如下whl包
 
 ```bash
-pip install https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/dynamicemb-0.0.1%2B20250929.6a3aeab-cp311-cp311-linux_x86_64.whl
+pip install https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/dynamicemb-0.0.1%2B20251013.60a6ddf-cp311-cp311-linux_x86_64.whl
 ```
 
 以id_feature的配置为例，DynamicEmbedding 只需在id_feature新增一个dynamicemb的配置字段

requirements/gpu.txt

@@ -1,4 +1,4 @@
-dynamicemb @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/dynamicemb-0.0.1%2B20250929.6a3aeab-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
+dynamicemb @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/dynamicemb-0.0.1%2B20251013.60a6ddf-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
 faiss_gpu_cu12 @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/faiss/faiss_gpu_cu12-1.11.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; python_version=="3.10"
 faiss_gpu_cu12 @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/faiss/faiss_gpu_cu12-1.11.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; python_version=="3.11"
 faiss_gpu_cu12 @ https://tzrec.oss-cn-beijing.aliyuncs.com/third_party/faiss/faiss_gpu_cu12-1.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; python_version=="3.12"

tzrec/main.py

@@ -74,7 +74,7 @@
 from tzrec.protos.model_pb2 import ModelConfig
 from tzrec.protos.pipeline_pb2 import EasyRecConfig
 from tzrec.protos.train_pb2 import TrainConfig
-from tzrec.utils import checkpoint_util, config_util, dynamicemb_util
+from tzrec.utils import checkpoint_util, config_util
 from tzrec.utils.dist_util import (
     DistributedModelParallel,
     create_train_pipeline,
@@ -835,8 +835,6 @@ def evaluate(
         model, device=device, mixed_precision=train_config.mixed_precision
     )
 
-    # TODO: remove it when DynamicEmbedding support eval without optimizer
-    dynamicemb_util._patch_dynamicemb_eval_model(model, pipeline_config.train_config)
     planner = create_planner(
         device=device,
         # pyre-ignore [16]

tzrec/utils/checkpoint_util.py

@@ -238,7 +238,7 @@ def restore_model(
             os.path.join(checkpoint_dir, "dynamicemb"),
             model,
             table_names=meta.get("dynamicemb_load_table_names", None),
-            optim=meta.get("dynamicemb_load_optim", True),
+            optim=meta.get("dynamicemb_load_optim", optimizer is not None),
         )
         logger.info(f"{os.environ.get('RANK', 0)} restore dynamic embedding finished.")
 

tzrec/utils/dynamicemb_util.py

@@ -44,12 +44,8 @@
     ShardMetadata,
 )
 from torchrec.modules.embedding_configs import BaseEmbeddingConfig, DataType
-from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
-from torchrec.optim.optimizers import SGD
 
-from tzrec.optim import optimizer_builder
 from tzrec.protos import feature_pb2
-from tzrec.protos.train_pb2 import TrainConfig
 
 has_dynamicemb = False
 try:
@@ -205,29 +201,6 @@ def build_dynamicemb_constraints(
     return constraints
 
 
-def _patch_dynamicemb_eval_model(model: nn.Module, train_config: TrainConfig) -> None:
-    """Patch model with optimizer when eval.
-
-    because DynamicEmbedding Eval need optimizer now.
-    """
-    if has_dynamicemb:
-        with_dynamicemb_feature = False
-        for feature in model.model._features:
-            if hasattr(feature.config, "dynamicemb") and feature.config.HasField(
-                "dynamicemb"
-            ):
-                with_dynamicemb_feature = True
-                break
-        if with_dynamicemb_feature:
-            sparse_optim_cls, _ = optimizer_builder.create_sparse_optimizer(
-                train_config.sparse_optimizer
-            )
-            trainable_params, frozen_params = model.model.sparse_parameters()
-            apply_optimizer_in_backward(sparse_optim_cls, trainable_params, {"lr": 0.0})
-            if len(frozen_params) > 0:
-                apply_optimizer_in_backward(SGD, frozen_params, {"lr": 0.0})
-
-
 if has_dynamicemb:
     enumerators.GUARDED_COMPUTE_KERNELS.add(EmbeddingComputeKernel.CUSTOMIZED_KERNEL)
 

tzrec/version.py

@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.9.2"
+__version__ = "0.9.3"