[Distributed] Fix wgrad bug in Sparse Operation Kit. (#918)

Use new_git_repository to manage sok dependency,update sok for fixing localized mode wgrad bug. Signed-off-by: aleliu <aleliu@nvidia.com>
DeepRec-AI · Aug 9, 2023 · 4983e02 · 4983e02
1 parent 4191627
commit 4983e02
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/tensorflow/tools/pip_package/build_sok.sh b/tensorflow/tools/pip_package/build_sok.sh
@@ -16,5 +16,4 @@ export MAKEFLAGS=-j$(nproc)
 export SOK_COMPILE_GPU_SM="70;75;80"
 cd ./bazel-DeepRec/external/hugectr/sparse_operation_kit
 
-"${PYTHON_BIN_PATH:-python}" setup.py bdist_wheel
-pip install ./dist/merlin_sok-1.2.0-cp38-cp38-linux_x86_64.whl
+"${PYTHON_BIN_PATH:-python}" setup.py install
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
@@ -1369,13 +1369,12 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
         ],
     )
 
-    http_archive(
+    new_git_repository(
         name = "hugectr",                                     # Apache License 2.0
-        build_file = "//third_party:hugectr.BUILD",
-        strip_prefix = "HugeCTR-23.06.00",
-        urls = [
-            "https://github.com/NVIDIA-Merlin/HugeCTR/archive/refs/tags/v23.06.00.tar.gz",
-        ],
+	build_file = "//third_party:hugectr.BUILD",
+	commit = "869028c1c32bdcda2f18efc88d54f0527ed28d6d",
+	init_submodules = True,
+	remote = "https://github.com/NVIDIA-Merlin/HugeCTR.git",
     )
 
 def tf_bind():