diff --git a/.github/workflows/amd-ci.yml b/.github/workflows/amd-ci.yml index 643172d7a..26e53ed18 100644 --- a/.github/workflows/amd-ci.yml +++ b/.github/workflows/amd-ci.yml @@ -12,10 +12,10 @@ jobs: runs-on: "amd-gfx942-mi325" timeout-minutes: 45 container: - image: rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview + image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1 options: >- --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root - --volume /home/runner/.triton:/github/home/.triton + --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB env: TRITON_BUILD_WITH_CLANG_LLD: "TRUE" TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE" @@ -26,8 +26,8 @@ jobs: - name: Install dependencies run: | rm -f /usr/local/bin/cmake - apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake - pip3 install -i https://test.pypi.org/simple hip-python>=6.3.0 # (or whatever Rocm version you have) + apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake miopen-hip autoconf libtool flex ninja-build clang lld + pip3 install -i https://test.pypi.org/simple hip-python>=7.1 # (or whatever Rocm version you have) pip3 install pybind11 - name: Checkout uses: actions/checkout@v4 @@ -58,10 +58,10 @@ jobs: runs-on: "amd-gfx942-mi325" timeout-minutes: 60 container: - image: rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview + image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1 options: >- --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root - --volume /home/runner/.triton:/github/home/.triton + --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB env: TRITON_BUILD_WITH_CLANG_LLD: "TRUE" TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE" @@ -72,9 +72,9 @@ jobs: steps: - name: Install dependencies run: | - apt-get update -y && apt install -y libopenmpi-dev - pip3 install -i https://test.pypi.org/simple hip-python>=6.3.0 # (or whatever Rocm version you have) - pip3 install pybind11 + apt-get update -y && apt install -y libopenmpi-dev cmake miopen-hip autoconf libtool flex ninja-build clang lld + pip3 install -i https://test.pypi.org/simple hip-python>=7.1 # (or whatever Rocm version you have) + pip3 install pybind11 psutil - name: Checkout uses: actions/checkout@v4 with: @@ -103,10 +103,10 @@ jobs: runs-on: "amd-gfx942-mi325" timeout-minutes: 45 container: - image: rocm/pytorch:rocm6.4.3_ubuntu22.04_py3.10_pytorch_release_2.5.1 + image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1 options: >- --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root - --volume /home/runner/.triton:/github/home/.triton + --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB env: TRITON_BUILD_WITH_CLANG_LLD: "TRUE" TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE" @@ -118,8 +118,8 @@ jobs: - name: Install dependencies run: | rm -f /usr/local/bin/cmake - apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake - pip3 install -i https://test.pypi.org/simple hip-python==6.4.3.555.40 + apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake ninja-build clang lld + pip3 install -i https://test.pypi.org/simple hip-python>=7.1 pip3 install pybind11 - name: Checkout uses: actions/checkout@v4 diff --git a/.gitignore b/.gitignore index ffb38f58a..a840320d1 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,9 @@ python/triton_dist/_C/*nvshmem hf_configs* *.whl +3rdparty/rocm-systems +3rdparty/rocshmem + # conflicts temp_conflicts temp_merge_repo @@ -82,4 +85,4 @@ tmp* *zst* .cursor -*.csv \ No newline at end of file +*.csv diff --git a/.gitmodules b/.gitmodules index 11e11c8cc..c61680680 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,3 @@ -[submodule "3rdparty/rocshmem"] - path = 3rdparty/rocshmem - url = https://github.com/ROCm/rocSHMEM.git [submodule "3rdparty/triton"] path = 3rdparty/triton url = https://github.com/ByteDance-Seed/triton.git diff --git a/3rdparty/mori b/3rdparty/mori index f07a21c92..cef7ecb92 160000 --- a/3rdparty/mori +++ b/3rdparty/mori @@ -1 +1 @@ -Subproject commit f07a21c926e827bf1a5c24bb5338b14be88b6702 +Subproject commit cef7ecb921e8abfbfd74e376e15d2fd283ab0004 diff --git a/3rdparty/rocshmem b/3rdparty/rocshmem deleted file mode 160000 index 7b5765ec0..000000000 --- a/3rdparty/rocshmem +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 7b5765ec0e137db04b50d657ed83ea239be8dd58 diff --git a/python/triton_dist/test/amd/test_rocshmem_api.py b/python/triton_dist/test/amd/test_rocshmem_api.py index 20c7c7d95..9c747dbe9 100644 --- a/python/triton_dist/test/amd/test_rocshmem_api.py +++ b/python/triton_dist/test/amd/test_rocshmem_api.py @@ -209,9 +209,8 @@ def test_rocshmem_memcpy(): HIP_CHECK(cp_res) - pyrocshmem.rocshmem_barrier_all_on_stream(cur_stream.cuda_stream) - torch.cuda.synchronize() + pyrocshmem.rocshmem_barrier_all_on_stream(cur_stream.cuda_stream) try: torch.testing.assert_close(comm_buffs[peer], one) diff --git a/shmem/rocshmem_bind/build.sh b/shmem/rocshmem_bind/build.sh index f6497dab8..f90944a48 100755 --- a/shmem/rocshmem_bind/build.sh +++ b/shmem/rocshmem_bind/build.sh @@ -6,11 +6,6 @@ set -e SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" PROJECT_ROOT=$(realpath ${SCRIPT_DIR}) -function apt_install_deps() { - apt update -y - apt-get install -y miopen-hip -} - function build_pyrocshmem_cmake() { pushd ${PROJECT_ROOT}/pyrocshmem mkdir -p build @@ -58,8 +53,6 @@ export ROCM_INSTALL_DIR="/opt/rocm" export PATH="${OPENMPI_UCX_INSTALL_DIR}/bin:$PATH" export LD_LIBRARY_PATH="${OPENMPI_UCX_INSTALL_DIR}/lib:$LD_LIBRARY_PATH" -apt_install_deps - # build rocshmem bash -x ${PROJECT_ROOT}/build_rocshmem.sh # build rocshmem bitcode diff --git a/shmem/rocshmem_bind/build_rocshmem.sh b/shmem/rocshmem_bind/build_rocshmem.sh index 9bb464288..91becc603 100755 --- a/shmem/rocshmem_bind/build_rocshmem.sh +++ b/shmem/rocshmem_bind/build_rocshmem.sh @@ -7,6 +7,49 @@ SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" PROJECT_ROOT=$(realpath ${SCRIPT_DIR}) ROCSHMEM_SRC_DIR=${PROJECT_ROOT}/../../3rdparty/rocshmem +sys_path="${PROJECT_ROOT}/../../3rdparty/rocm-systems" + +if [ -d "${ROCSHMEM_SRC_DIR}" ]; then + pushd "${PROJECT_ROOT}/../.." + active=$(git config submodule.3rdparty/rocshmem.active || echo "nil") + if [ "${active}" = "true" ]; then + echo "Error: Rocshmem submodule still active, please delete it" + quit=1 + fi + popd + pushd "${ROCSHMEM_SRC_DIR}" + url=$(git remote get-url origin || echo "nil") + if [ "${url}" = "https://github.com/ROCm/rocSHMEM.git" ]; then + echo "Error: Old rocshmem checkout found, please delete it" + quit=1 + fi + popd + if ! [ -z "${quit}" ]; then + exit $quit + fi + + if ! [ "$(ls -A "${ROCSHMEM_SRC_DIR}")" ]; then + rmdir "${ROCSHMEM_SRC_DIR}" + fi +fi + +rocm_systems_tag=hip-version_7.12.60610 + +if ! [ -d "${ROCSHMEM_SRC_DIR}" ]; then + echo "Creating sparse checkout" + pushd "${PROJECT_ROOT}/../.." + git clone "https://github.com/ROCm/rocm-systems.git" -b "${rocm_systems_tag}" --depth 1 --sparse "${sys_path}" + popd + pushd "${sys_path}" + git config core.sparseCheckoutCone true + git sparse-checkout set projects/rocshmem + ln -s rocm-systems/projects/rocshmem ../rocshmem + popd +fi + +pushd "${sys_path}" +git checkout "${rocm_systems_tag}" +popd pushd ${ROCSHMEM_SRC_DIR} @@ -17,7 +60,6 @@ OMPI_INSTALL_DIR="${OMPI_INSTALL_DIR:-/opt/ompi_build}" # build ompi, ucx if [ ! -e "${OMPI_INSTALL_DIR}" ]; then # prepare for building ompi, ucx - apt-get install autoconf libtool flex -y BUILD_DIR=${OMPI_INSTALL_DIR} bash ${ROCSHMEM_SRC_DIR}/scripts/install_dependencies.sh else echo "ompi exists, skip building ompi and ucx" diff --git a/shmem/rocshmem_bind/pyrocshmem/setup.py b/shmem/rocshmem_bind/pyrocshmem/setup.py index 250a30566..1f242a224 100644 --- a/shmem/rocshmem_bind/pyrocshmem/setup.py +++ b/shmem/rocshmem_bind/pyrocshmem/setup.py @@ -108,7 +108,7 @@ def setup_pytorch_extension() -> setuptools.Extension: """Setup CppExtension for PyTorch support""" include_dirs, library_dirs, libraries = [], [], [] - deps = [hip_deps(), mpi_deps(), rocshmem_deps()] + deps = [rocshmem_deps(), hip_deps(), mpi_deps()] for include_dir, library_dir, library in deps: include_dirs += include_dir diff --git a/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh b/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh index d1797c564..1dc3b56cf 100755 --- a/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh +++ b/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh @@ -39,12 +39,11 @@ SOURCE_MAP=( ["${ROCSHMEM_SRC}/src/ipc/context_ipc_device_coll.cpp"]="rocshmem_context_ipc_device_coll.bc" ["${ROCSHMEM_SRC}/src/ipc_policy.cpp"]="rocshmem_ipc_policy.bc" ["${ROCSHMEM_SRC}/src/gda/context_gda_device.cpp"]="rocshmem_context_gda_device.bc" + ["${ROCSHMEM_SRC}/src/gda/context_gda_device_coll.cpp"]="rocshmem_context_gda_device_coll.bc" ["${ROCSHMEM_SRC}/src/gda/backend_gda.cpp"]="rocshmem_backend_gda.bc" ["${ROCSHMEM_SRC}/src/gda/queue_pair.cpp"]="rocshmem_queue_pair.bc" ["${ROCSHMEM_SRC}/src/gda/ionic/queue_pair_ionic.cpp"]="rocshmem_queue_pair_ionic.bc" ["${ROCSHMEM_SRC}/src/gda/mlx5/queue_pair_mlx5.cpp"]="rocshmem_queue_pair_mlx5.bc" - ["${ROCSHMEM_SRC}/src/gda/mlx5/segment_builder.cpp"]="rocshmem_segment_builder.bc" - ["${ROCSHMEM_SRC}/src/gda/endian.cpp"]="rocshmem_endian.bc" ["${ROCSHMEM_SRC}/src/team.cpp"]="rocshmem_team.bc" ["${ROCSHMEM_SRC}/src/sync/abql_block_mutex.cpp"]="rocshmem_abql_block_mutex.bc" ["${ROCSHMEM_SRC}/src/util.cpp"]="rocshmem_util.bc"