Add fbgemm genai to the nightly docker

Author: xuzhao9

docker/gcp-a100-runner-dind.dockerfile

@@ -28,9 +28,9 @@ RUN sudo mkdir -p /workspace; sudo chown runner:runner /workspace
 # We assume that the host NVIDIA driver binaries and libraries are mapped to the docker filesystem
 
 # Use the CUDA installation scripts from pytorch/builder
+# Install CUDA 12.4 only to reduce docker size
 RUN cd /workspace; git clone https://github.com/pytorch/builder.git
-RUN sudo bash -c 'source /workspace/builder/common/install_cuda.sh; install_118; prune_118'
-RUN sudo bash -c 'source /workspace/builder/common/install_cuda.sh; install_121; prune_121'
+RUN sudo bash -c 'source /workspace/builder/common/install_cuda.sh; install_124; prune_124'
 
 # Install miniconda
 RUN wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /workspace/Miniconda3-latest-Linux-x86_64.sh

docker/torchbench-nightly.dockerfile

@@ -48,3 +48,8 @@ RUN cd /workspace/benchmark && \
 RUN cd /workspace/benchmark && \
     . ${SETUP_SCRIPT} && \
     python install.py
+
+# Install FBGEMM GENAI
+RUN cd /workspace/benchmark && \
+    . ${SETUP_SCRIPT} && \
+    python install.py --userbenchmark triton --fbgemm

install.py

@@ -46,7 +46,7 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
         choices=list_userbenchmarks(),
         help="Install requirements for optional components.",
     )
-    args = parser.parse_args()
+    args, extra_args = parser.parse_known_args()
 
     os.chdir(os.path.realpath(os.path.dirname(__file__)))
 
@@ -68,9 +68,11 @@ def pip_install_requirements(requirements_txt="requirements.txt"):
     if args.userbenchmark:
         # Install userbenchmark dependencies if exists
         userbenchmark_dir = REPO_ROOT.joinpath("userbenchmark", args.userbenchmark)
+        cmd = [sys.executable, "install.py"]
+        cmd.extend(extra_args)
         if userbenchmark_dir.joinpath("install.py").is_file():
             subprocess.check_call(
-                [sys.executable, "install.py"], cwd=userbenchmark_dir.absolute()
+                cmd, cwd=userbenchmark_dir.absolute()
             )
         sys.exit(0)
 

userbenchmark/triton/install.py

@@ -0,0 +1,22 @@
+import argparse
+import subprocess
+import sys
+
+from torchbenchmark import REPO_PATH
+FBGEMM_PATH = REPO_PATH.joinpath("submodules", "FBGEMM", "fbgemm_gpu")
+
+def install_fbgemm():
+    cmd = [sys.executable, "setup.py", "bdist_wheel", "--package_variant=genai"]
+    subprocess.check_call(cmd, cwd=FBGEMM_PATH)
+
+def test_fbgemm():
+    cmd = [sys.executable, "-c", '"import fbgemm_gpu.experimental.gen_ai"']
+    subprocess.check_call(cmd)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fbgemm", action="store_true", help="Install FBGEMM GPU")
+    args = parser.parse_args()
+    if args.fbgemm:
+        install_fbgemm()
+        test_fbgemm()