updated code to run both HSTU and DLRM on v6

Author: ajkv-google

recml/examples/dlrm_experiment.py

@@ -19,6 +19,12 @@
 import dataclasses
 from typing import Generic, Literal, TypeVar
 
+import sys
+import os
+# Add the RecML folder to the system path
+sys.path.append(os.path.join(os.getcwd(), "../../../RecML"))
+os.environ["KERAS_BACKEND"] = "jax"
+
 from etils import epy
 import fiddle as fdl
 import flax.linen as nn

requirements.txt

@@ -1,4 +1,6 @@
 absl-py==2.2.2
+aiofiles==25.1.0
+array-record==0.8.3
 astroid==3.3.9
 astunparse==1.6.3
 attrs==25.3.0
@@ -16,7 +18,7 @@ etils==1.12.2
 fiddle==0.3.0
 filelock==3.18.0
 flatbuffers==25.2.10
-flax==0.10.5
+flax==0.12.2
 fsspec==2025.3.2
 gast==0.6.0
 google-pasta==0.2.0
@@ -31,18 +33,22 @@ immutabledict==4.2.1
 importlib-resources==6.5.2
 iniconfig==2.1.0
 isort==6.0.1
-jax==0.6.0
-jaxlib==0.6.0
+jax==0.8.2
+jax-tpu-embedding==0.1.0.dev20251208
+jaxlib==0.8.2
 jaxtyping==0.3.1
-jinja2==3.1.6
+Jinja2==3.1.6
 kagglehub==0.3.11
 keras==3.9.2
 keras-hub==0.20.0
 libclang==18.1.1
 libcst==1.7.0
-markdown==3.8
+libtpu==0.0.32
+# libtpu-nightly is usually installed directly via URL, but pinning it helps tracking
+# libtpu-nightly==0.1.dev20240617+default 
+Markdown==3.8
 markdown-it-py==3.0.0
-markupsafe==3.0.2
+MarkupSafe==3.0.2
 mccabe==0.7.0
 mdurl==0.1.2
 ml-collections==1.1.0
@@ -54,23 +60,37 @@ nest-asyncio==1.6.0
 networkx==3.4.2
 nodeenv==1.9.1
 numpy==2.1.3
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cusparselt-cu12==0.6.2
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
 opt-einsum==3.4.0
 optax==0.2.4
 optree==0.15.0
-orbax-checkpoint==0.11.12
+orbax-checkpoint==0.11.31
 packaging==24.2
 platformdirs==4.3.7
 pluggy==1.5.0
+portpicker==1.6.0
 pre-commit==4.2.0
 promise==2.3
-# protobuf==5.29.4
+# protobuf==6.33.4
 psutil==7.0.0
 pyarrow==19.0.1
-pygments==2.19.1
+Pygments==2.19.1
 pylint==3.3.6
 pytest==8.3.5
 pytest-env==1.1.5
-pyyaml==6.0.2
+PyYAML==6.0.2
 regex==2024.11.6
 requests==2.32.3
 rich==14.0.0
@@ -84,21 +104,23 @@ tensorboard==2.19.0
 tensorboard-data-server==0.7.2
 tensorflow==2.19.0
 tensorflow-datasets==4.9.8
+tensorflow-io-gcs-filesystem==0.37.1
 tensorflow-metadata==1.17.1
 tensorflow-text==2.19.0
-tensorstore==0.1.73
+tensorstore==0.1.80
 termcolor==3.0.1
 toml==0.10.2
 tomlkit==0.13.2
 toolz==1.0.0
 torch==2.6.0
 tqdm==4.67.1
 treescope==0.1.9
-typing-extensions==4.13.2
+triton==3.2.0
+typing_extensions==4.13.2
 urllib3==2.4.0
 virtualenv==20.30.0
 wadler-lindig==0.1.5
-werkzeug==3.1.3
+Werkzeug==3.1.3
 wheel==0.45.1
 wrapt==1.17.2
-zipp==3.21.0
+zipp==3.21.0