From 72eec0727b7f3ba0f8ccc8e51dcab33d7aefaa6b Mon Sep 17 00:00:00 2001
From: Benson Ma <bensonma415@meta.com>
Date: Tue, 7 Jan 2025 11:20:10 -0800
Subject: [PATCH] Fix versioning scheme for ROCm releases (#3554)

Summary:
- Fix versioning scheme for ROCm releases in setup.py


Reviewed By: spcyppt

Differential Revision: D67908722

Pulled By: q10
---
 .github/scripts/utils_pip.bash       |  2 +-
 .github/workflows/fbgemm_gpu_pip.yml |  2 +-
 fbgemm_gpu/setup.py                  | 10 +++++++++-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/.github/scripts/utils_pip.bash b/.github/scripts/utils_pip.bash
index cbe46a83a9..382f382ccd 100644
--- a/.github/scripts/utils_pip.bash
+++ b/.github/scripts/utils_pip.bash
@@ -85,7 +85,7 @@ __export_package_variant_info () {
       # e.g. rocm 6.2.4 => rocm6.2.4
       #
       # NOTE: Unlike CUDA-based releases, which ignores the minor patch version,
-      # ROCm-based releases use the full version string.
+      # ROCm-based releases may use the full version string.
       # See https://download.pytorch.org/whl/nightly/torch/ for examples.
       local variant_type="rocm"
       if [ "${rocm_version_arr[2]}" == "" ]; then
diff --git a/.github/workflows/fbgemm_gpu_pip.yml b/.github/workflows/fbgemm_gpu_pip.yml
index 65a0f4fd74..9b934555c2 100644
--- a/.github/workflows/fbgemm_gpu_pip.yml
+++ b/.github/workflows/fbgemm_gpu_pip.yml
@@ -197,7 +197,7 @@ jobs:
         ]
         # ROCm machines are limited, so we only test a subset of Python versions
         python-version: [ "3.9", "3.10", "3.11", "3.12", "3.13" ]
-        rocm-version: [ "6.1.2", "6.2.4" ]
+        rocm-version: [ "6.2.4", "6.3" ]
 
     steps:
     - name: Setup Build Container
diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py
index 7dc172cb11..723f07b535 100644
--- a/fbgemm_gpu/setup.py
+++ b/fbgemm_gpu/setup.py
@@ -166,7 +166,15 @@ def variant_version(self) -> str:
         elif self.args.package_variant == "rocm":
             if torch.version.hip is not None:
                 rocm_version = torch.version.hip.split(".")
-                pkg_vver = f"+rocm{rocm_version[0]}.{rocm_version[1]}"
+                # NOTE: Unlike CUDA-based releases, which ignores the minor patch version,
+                # ROCm-based releases may use the full version string.
+                # See https://download.pytorch.org/whl/nightly/torch/ for examples.
+                if len(rocm_version) > 2:
+                    pkg_vver = (
+                        f"+rocm{rocm_version[0]}.{rocm_version[1]}.{rocm_version[2]}"
+                    )
+                else:
+                    pkg_vver = f"+rocm{rocm_version[0]}.{rocm_version[1]}"
             else:
                 sys.exit(
                     "[SETUP.PY] The installed PyTorch variant is not ROCm; cannot determine the ROCm version!"