diff --git a/.github/scripts/Dockerfile.ci.deps b/.github/scripts/Dockerfile.ci.deps index 722c8f620..14aa8f76e 100644 --- a/.github/scripts/Dockerfile.ci.deps +++ b/.github/scripts/Dockerfile.ci.deps @@ -3,120 +3,88 @@ # See LICENSE for license information. # # TE CI deps image -# -# ROCm installer: https://raw.githubusercontent.com/ROCm/TheRock/release/therock-7.13/dockerfiles/install_rocm_tarball.sh FROM ubuntu:24.04 ARG DEBIAN_FRONTEND=noninteractive SHELL ["/bin/bash", "-euo", "pipefail", "-c"] -ARG ROCK_RELEASE_VERSION=7.13.0 -ARG GPU_ARCH=gfx942 +ARG ROCK_RELEASE_VERSION=7.14.0 ARG PYTHON_VERSION=3.12 ARG PYTHON_ABI=cp312 -ARG TORCH_VERSION=2.10.0 -ARG TORCHVISION_VERSION=0.25.0 -ARG TORCHAUDIO_VERSION=2.10.0 -ARG TRITON_VERSION=3.6.0 -ARG JAX_VERSION=0.10.2 -ARG JAX_BRANCH=rocm-jax-v0.10.2 -ARG XLA_BRANCH=rocm-jaxlib-v0.10.2 -ARG FA_VERSION=v2.8.1 +ARG TORCH_VERSION=2.12.0 +ARG TORCHVISION_VERSION=0.27.0 +# triton 3.7.1 ships as a git snapshot: triton-3.7.1+.rocm +ARG TRITON_VERSION=3.7.1 +ARG TRITON_LOCAL=git0263a6a6 +ARG JAX_VERSION=0.11.0 +ARG FA_VERSION=v2.8.3 ARG AITER_COMMIT=77455e3ecf4f0d28756afc452e914940c45b944b -ARG INSTALL_ROCM_TARBALL_SH_URL=https://raw.githubusercontent.com/ROCm/TheRock/release/therock-7.13/dockerfiles/install_rocm_tarball.sh +ARG GPU_ARCH=gfx950;gfx942 # Base OS packages RUN apt-get update && apt-get install -y --no-install-recommends \ ca-certificates curl \ git vim \ build-essential cmake ninja-build pkg-config liblzma-dev \ - openjdk-17-jdk-headless \ + libnuma1 libnuma-dev \ python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev python3-pip \ && rm -rf /var/lib/apt/lists/* -# Native ROCm tarball → /opt/rocm -RUN case "${GPU_ARCH}" in \ - gfx942) _amd_gpu_family=gfx94X-dcgpu ;; \ - gfx950) _amd_gpu_family=gfx950-dcgpu ;; \ - esac \ - && curl -fsSL -o /tmp/install_rocm_tarball.sh "${INSTALL_ROCM_TARBALL_SH_URL}" \ - && chmod +x /tmp/install_rocm_tarball.sh \ - && /tmp/install_rocm_tarball.sh "${ROCK_RELEASE_VERSION}" "${_amd_gpu_family}" stable \ - && rm -f /tmp/install_rocm_tarball.sh - # Isolated Python env for pip packages RUN python${PYTHON_VERSION} -m venv /opt/venv -# Default container env: venv on PATH; /opt/rocm tarball for flash-attention / aiter. -ENV GPU_ARCH=${GPU_ARCH} \ - VIRTUAL_ENV=/opt/venv \ - PATH=/opt/venv/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +# Default container env: venv on PATH. ROCm is provided by the pip rocm-sdk and +# resolved on demand via `rocm-sdk path --root` (no /opt/rocm install). +ENV VIRTUAL_ENV=/opt/venv \ + PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -RUN python -m pip install --upgrade pip setuptools wheel \ +RUN python -m pip install --upgrade pip setuptools wheel pyyaml \ && pip install --no-cache-dir ipython pytest fire pydantic pybind11 ninja pandas expecttest onnxscript -# Python ROCm SDK + torch from https://repo.amd.com/rocm/whl// -RUN case "${GPU_ARCH}" in \ - gfx942) _amd_gpu_family=gfx94X-dcgpu ;; \ - gfx950) _amd_gpu_family=gfx950-dcgpu ;; \ - esac \ - && W="https://repo.amd.com/rocm/whl/${_amd_gpu_family}" \ - && LIBS_PKG="rocm-sdk-libraries-$(echo "${_amd_gpu_family}" | tr '[:upper:]' '[:lower:]')" \ +# Python ROCm SDK + torch from https://repo.amd.com/rocm/whl-multi-arch/ +# Multi-arch layout: arch-agnostic wheels (rocm-sdk-libraries, torch, torchvision, +# triton) installed once, then a per-GPU device wheel (rocm-sdk-device-, +# amd_torch_device_, amd_torchvision_device_) for every arch in +# GPU_ARCH (semicolon-separated, overridable via --build-arg). +RUN W="https://repo.amd.com/rocm/whl-multi-arch" \ && ROCM_WHEEL_TAG="rocm${ROCK_RELEASE_VERSION}" \ && pip install --no-cache-dir \ --extra-index-url "${W}" \ "rocm-sdk-core==${ROCK_RELEASE_VERSION}" \ "rocm-sdk-devel==${ROCK_RELEASE_VERSION}" \ - "${LIBS_PKG}==${ROCK_RELEASE_VERSION}" \ + "rocm-sdk-libraries==${ROCK_RELEASE_VERSION}" \ "${W}/rocm-${ROCK_RELEASE_VERSION}.tar.gz" \ "${W}/torch-${TORCH_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ "${W}/torchvision-${TORCHVISION_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ - "${W}/torchaudio-${TORCHAUDIO_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ - "${W}/triton-${TRITON_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" + "${W}/triton-${TRITON_VERSION}%2B${TRITON_LOCAL}.${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ + && for arch in ${GPU_ARCH//;/ }; do \ + pip install --no-cache-dir --extra-index-url "${W}" \ + "rocm-sdk-device-${arch}==${ROCK_RELEASE_VERSION}" \ + "${W}/amd_torch_device_${arch}-${TORCH_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" \ + "${W}/amd_torchvision_device_${arch}-${TORCHVISION_VERSION}%2B${ROCM_WHEEL_TAG}-${PYTHON_ABI}-${PYTHON_ABI}-linux_x86_64.whl" ; \ + done # Install Jax -RUN pip install --no-cache-dir patchelf auditwheel numpy scipy \ - && curl -fsSL -o /usr/local/bin/bazel \ - https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-amd64 \ - && chmod +x /usr/local/bin/bazel \ - && git clone --branch "${JAX_BRANCH}" --depth 1 https://github.com/ROCm/jax.git /tmp/jax-src \ - && git clone --branch "${XLA_BRANCH}" --depth 1 https://github.com/ROCm/xla.git /tmp/xla-src \ - && cd /tmp/jax-src \ - && JAX_DIST=/tmp/jax-dist \ - && mkdir -p "${JAX_DIST}" \ - && ROCM_SDK_ROOT="$(rocm-sdk path --root)" \ - && python build/build.py build \ - --wheels=jax-rocm-plugin,jax-rocm-pjrt \ - --bazel_startup_options="--bazelrc=build/rocm/rocm.bazelrc" \ - --bazel_startup_options="--output_user_root=/tmp/bazel-jax" \ - --bazel_options=--config=rocm_release_wheel \ - --bazel_options=--override_repository=xla=/tmp/xla-src \ - --bazel_options=--override_module=xla=/tmp/xla-src \ - --bazel_options=--repo_env=ML_WHEEL_TYPE=release \ - --bazel_options=--//jaxlib/tools:jaxlib_git_hash="$(git rev-parse HEAD)" \ - --bazel_options=--repo_env=ROCM_PATH="${ROCM_SDK_ROOT}" \ - --python_version="${PYTHON_VERSION}" \ - --rocm_path="${ROCM_SDK_ROOT}" \ - --rocm_version=7 \ - --rocm_amdgpu_targets="${GPU_ARCH}" \ - --verbose \ - --output_path="${JAX_DIST}" \ - && pip wheel . --no-deps -w "${JAX_DIST}" \ - && pip install --no-cache-dir "jaxlib==${JAX_VERSION}" "${JAX_DIST}"/*.whl \ - && pip install --no-cache-dir "jax==${JAX_VERSION}" \ - && rm -rf /tmp/jax-src /tmp/xla-src /tmp/jax-dist /tmp/bazel-jax +RUN pip install --no-cache-dir \ + "jax==${JAX_VERSION}" \ + "jax-rocm7-pjrt==${JAX_VERSION}" \ + "jax-rocm7-plugin==${JAX_VERSION}" -RUN git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ +# Install Flash Attention +RUN export ROCM_PATH="$(rocm-sdk path --root)" \ + && git clone --branch "${FA_VERSION}" --depth 1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention \ && cd /tmp/flash-attention \ - && ROCM_PATH=/opt/rocm GPU_ARCHS="${GPU_ARCH}" python setup.py install \ + && GPU_ARCHS="${GPU_ARCH}" python setup.py install \ && rm -rf /tmp/flash-attention -RUN git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ +# Install Aiter +RUN export ROCM_PATH="$(rocm-sdk path --root)" \ + && git clone --no-checkout https://github.com/ROCm/aiter.git /tmp/aiter \ && cd /tmp/aiter \ && git checkout "${AITER_COMMIT}" \ && git submodule update --init --recursive \ - && ROCM_PATH=/opt/rocm GPU_ARCHS="${GPU_ARCH}" pip install --no-build-isolation --no-cache-dir . \ + && pip install --no-build-isolation --no-cache-dir . \ && rm -rf /tmp/aiter WORKDIR /workspace/ diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml index 5591f1f39..606b1e835 100644 --- a/.github/workflows/rocm-ci.yml +++ b/.github/workflows/rocm-ci.yml @@ -109,7 +109,7 @@ jobs: timeout-minutes: 270 runs-on: ${{ matrix.arch_label == 'mi30x' && 'linux-te-mi30x-4' || matrix.arch_label == 'mi35x' && 'linux-te-mi35x-4' }} env: - DOCKER_IMAGE: ${{ inputs.docker_image_override || format('{0}{1}', needs.select_image.outputs.image-tag, matrix.arch_label == 'mi30x' && '_gfx942' || matrix.arch_label == 'mi35x' && '_gfx950') }} + DOCKER_IMAGE: ${{ inputs.docker_image_override || needs.select_image.outputs.image-tag }} strategy: fail-fast: false matrix: @@ -262,7 +262,7 @@ jobs: timeout-minutes: 210 runs-on: ${{ matrix.arch_label == 'mi30x' && 'linux-te-mi30x-8' || matrix.arch_label == 'mi35x' && 'linux-te-mi35x-8' }} env: - DOCKER_IMAGE: ${{ inputs.docker_image_override || format('{0}{1}', needs.select_image.outputs.image-tag, matrix.arch_label == 'mi30x' && '_gfx942' || matrix.arch_label == 'mi35x' && '_gfx950') }} + DOCKER_IMAGE: ${{ inputs.docker_image_override || needs.select_image.outputs.image-tag }} strategy: fail-fast: false matrix: diff --git a/benchmarks/attention/benchmark_attention_rocm.py b/benchmarks/attention/benchmark_attention_rocm.py index b234d374f..60fcaa9aa 100644 --- a/benchmarks/attention/benchmark_attention_rocm.py +++ b/benchmarks/attention/benchmark_attention_rocm.py @@ -138,7 +138,6 @@ def setup_backend_env(backend_name, use_ck_bwd_v3=True, use_ck_fwd_v3=True, use_ ROCPROF_STATS_CSV = "results.stats.csv" - def _profiler_python_code(model, attention, column_name, benchmark_dir): return ( f"import sys; sys.path.insert(0, {benchmark_dir!r}); " @@ -181,7 +180,9 @@ def _run_attention_profiler(model, attention, column_name, dirname): py_code, ] try: - result = subprocess.run(cmd, capture_output=True, text=True) + # Pass os.environ explicitly: C code may set ROCPROFILER_REGISTER_LIBRARY + # via setenv() without updating Python's environ cache, which breaks rocprofv3. + result = subprocess.run(cmd, capture_output=True, text=True, env=os.environ.copy()) except FileNotFoundError: print( "WARNING: rocprofv3 not found on PATH; kernel timing columns may be empty.", diff --git a/ci/_utils.sh b/ci/_utils.sh index f2478fc76..d09776b93 100644 --- a/ci/_utils.sh +++ b/ci/_utils.sh @@ -313,8 +313,8 @@ check_test_filter() { start_message() { echo "Started with TEST_LEVEL=$TEST_LEVEL sGPU='$TEST_SGPU' mGPU='$TEST_MGPU' at `date`" - _rocm_path=$(resolve_rocm_path) - _rocm_path=`$REALPATH "$_rocm_path" 2>/dev/null || echo "$_rocm_path"` + export ROCM_PATH=$(resolve_rocm_path) + _rocm_path=`$REALPATH "$ROCM_PATH" 2>/dev/null || echo "$ROCM_PATH"` echo "ROCM PATH: $_rocm_path" python3 --version } diff --git a/ci/ci_config.json b/ci/ci_config.json index d2f01bfac..db70c7455 100644 --- a/ci/ci_config.json +++ b/ci/ci_config.json @@ -1,5 +1,5 @@ { "docker_images": { - "default": "registry-sc-harbor.amd.com/framework/te-ci:therock_7.13.0_ubuntu24.04_py3.12_pytorch_2.10.0_triton_3.6.0_jax_0.10.2_fa_2.8.1_aiter_77455e3ecf" + "default": "registry-sc-harbor.amd.com/framework/te-ci:therock_7.14.0_ubuntu24.04_py3.12_pytorch_2.12.0_triton_3.7.1_git0263a6a6_jax_0.11.0_fa_2.8.3_aiter_77455e3ecf" } } diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 92f73d588..48b713ad6 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -35,6 +35,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_rope.xml $ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PATH/tests/pytorch/nvfp4 || test_fail "test_nvfp4" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8.xml $TE_PATH/tests/pytorch/mxfp8 || test_fail "test_mxfp8" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor.xml $TE_PATH/tests/pytorch/test_quantized_tensor.py || test_fail "test_quantized_tensor.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xml $TE_PATH/tests/pytorch/test_torch_compile.py || test_fail "test_torch_compile.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 3a4405533..ec9cfcc22 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2126,8 +2126,25 @@ def test_gpt_cuda_graph(dtype, bs, model): for param1, param2 in zip(block.parameters(), graphed_block.parameters()): param2.copy_(param1) - out, grads = _test_gpt_e2e_cuda_graph(block, bs, dtype, config, False) - graphed_out, graphed_grads = _test_gpt_e2e_cuda_graph(graphed_block, bs, dtype, config, True) + # WAR (ROCm): torch>=2.12 (pytorch/pytorch#179053) makes hipBLASLt handles + # per-(device, stream). The graph-capture stream's handle is created lazily + # during capture, and hipblasLtCreate performs an internal hipMalloc that is + # illegal mid-capture, failing with HIP error 900 ("operation not permitted + # when stream is capturing"). The capture_begin pre-init + # (pytorch/pytorch#180692) only covers the calling thread, not the autograd + # backward thread, so torch's own bmm in the captured backward still trips it. + # Route torch's matmul/bmm off hipBLASLt for this test until PyTorch + # extends the pre-init to cover it. + _prev_blas_library = None + if IS_HIP_EXTENSION: + _prev_blas_library = torch.backends.cuda.preferred_blas_library() + torch.backends.cuda.preferred_blas_library("cublas") + try: + out, grads = _test_gpt_e2e_cuda_graph(block, bs, dtype, config, False) + graphed_out, graphed_grads = _test_gpt_e2e_cuda_graph(graphed_block, bs, dtype, config, True) + finally: + if _prev_blas_library is not None: + torch.backends.cuda.preferred_blas_library(_prev_blas_library) params = list(block.parameters()) graphed_params = list(graphed_block.parameters()) diff --git a/tests/pytorch/test_sanity_import.py b/tests/pytorch/test_sanity_import.py index 84d2e8aa0..418a34e5f 100644 --- a/tests/pytorch/test_sanity_import.py +++ b/tests/pytorch/test_sanity_import.py @@ -9,20 +9,29 @@ if __name__ == "__main__": print("OK") -AMDSMI_SRC = "/opt/rocm/share/amd_smi" - - def _amdsmi_pythonpath(): """Return ROCm amdsmi source dir if the package is not already installed. torch counts devices through amdsmi whenever it is importable, reading HIP_VISIBLE_DEVICES at call time rather than asking an already initialized HIP runtime -- which is what makes a visible-devices change after import observable. + + With TheRock pip wheels, bindings live under the rocm-sdk-core tree + (``/share/amd_smi``). """ - import importlib.util, os - if importlib.util.find_spec("amdsmi") is not None or not os.path.isdir(AMDSMI_SRC): + import importlib.util + + if importlib.util.find_spec("amdsmi") is not None: return None - return AMDSMI_SRC + try: + from rocm_sdk_core._cli import _get_core_module_path + + candidate = _get_core_module_path() / "share/amd_smi" + if candidate.is_dir(): + return str(candidate) + except (ImportError, ModuleNotFoundError, OSError): + pass + return None def _env_with_amdsmi_pythonpath(env, amdsmi_path): diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 309a5d124..1286492a6 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -26,6 +26,7 @@ from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer +from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch import ( is_fp8_available, is_mxfp8_available, @@ -123,21 +124,25 @@ def __fx_repr__(self): _Q = get_opaque_type_name(ToyQuantizer) def _make_qfactory(tag: str): - """Return a qfactory that produces ToyQuantizer instances tagged with *tag*.""" + """Return a qfactory that produces ToyQuantizer instances tagged with *tag*. + + The factory dispatches on ``QuantizerRole.tensor_type``; the roles are + supplied by :meth:`ToyLinear.get_quantizer_roles`. + """ quantizers = { - role: ToyQuantizer(tag=f"{tag}:{role}") - for role in ( - "linear_input", - "linear_weight", - "linear_output", - "linear_grad_output", - "linear_grad_input", + tensor_type: ToyQuantizer(tag=f"{tag}:{tensor_type}") + for tensor_type in ( + "input", + "weight", + "output", + "grad_output", + "grad_input", ) } - def qfactory(role: str): - return quantizers[role] + def qfactory(role: QuantizerRole): + return quantizers[role.tensor_type] return qfactory @@ -163,6 +168,22 @@ def __init__( ) torch.nn.init.normal_(self.weight) + def get_quantizer_roles(self, *, fwd: bool, num_quantizers: int): + # Supplying explicit roles keeps CustomRecipeState from emitting a + # warning (which would graph-break under fullgraph=True) and lets the + # qfactory dispatch per tensor slot. Order must match the module's + # quantizer array (FP8FwdTensorIdx / FP8BwdTensorIdx). + if fwd: + return [ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ] + return [ + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ] + def _get_weight_tensors(self): return [self.weight] diff --git a/tests/pytorch/triton_kernels/conftest.py b/tests/pytorch/triton_kernels/conftest.py new file mode 100644 index 000000000..9cbf0bffd --- /dev/null +++ b/tests/pytorch/triton_kernels/conftest.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""Work around a ROCm HSA-runtime teardown segfault after the triton_kernels tests. + +On ROCm 7.14 the HIP runtime's atexit handler calls ``hsa_shut_down()``, which +destroys the GPU agents and their AQL queues; ``AqlQueue::~AqlQueue()`` then +writes to an already-freed HSA doorbell signal and the process dies with SIGSEGV. +Observed backtrace (``libamdhip64.so.7`` / ``libhsa-runtime64.so.1``):: + + exit() -> -> rocr::HSA::hsa_shut_down() + -> Runtime::Release() -> Unload() -> DestroyAgents() + -> GpuAgent::~GpuAgent() -> AqlQueue::~AqlQueue() + -> hsa_signal_store_screlease <-- SIGSEGV + +The crash happens *after* every test has passed and pytest has already written +its JUnit XML report, so the run is functionally green, but the process exits +139 and ``ci/_utils.sh``'s exit-code gate records a suite error. ``test_norms.py`` +reliably trips it because it allocates the most Triton streams/queues of the +triton_kernels suites; the fault itself is purely in the ROCm libraries' shutdown +path, not in TE. + +Bypass the buggy C-level atexit handler by hard-exiting with pytest's real exit +status once the session (and its report writing) has finished. Real failures +still propagate through ``exitstatus``, and a crash *during* the run (before +``pytest_sessionfinish``) is unaffected and still surfaces as an error. Remove +once the ROCm HSA shutdown crash is fixed. +""" + +import os +import sys + +import pytest +import torch + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session, exitstatus): + # Only ROCm hits the hsa_shut_down teardown segfault; leave CUDA/CPU exit + # semantics (and their normal atexit cleanup) untouched. + if getattr(torch.version, "hip", None) is None: + return + # trylast ensures the junitxml plugin and te_ci_result_sink have already + # written their reports in this same hook before we hard-exit. + sys.stdout.flush() + sys.stderr.flush() + os._exit(0 if exitstatus == 0 else int(exitstatus)) diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 2d29f7e06..edb865027 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -26,11 +26,11 @@ namespace cuda { #ifndef __HIP_PLATFORM_AMD__ namespace { -// String with build-time CUDA include path +// Build-time CUDA include path #include "string_path_cuda_include.h" } // namespace -#endif // #ifndef __HIP_PLATFORM_AMD__ +#endif // __HIP_PLATFORM_AMD__ int num_devices() { auto query_num_devices = []() -> int { diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f65052f9b..7ca09f701 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -595,6 +595,12 @@ def get_ub(name: str, use_fp8: bool): return _ub_communicators[key] +@torch.compiler.assume_constant_result +def get_ub_is_fp8(name: str, use_fp8: bool) -> bool: + """Query is_fp8_ubuf for a named UB communicator; treated as compile-time constant.""" + return get_ub(name, use_fp8).is_fp8_ubuf() + + def destroy_ub(): """Destroy all allocated userbuffer communicators.""" global _ub_communicators, _ub_with_cublasmp, _ub_initialized @@ -603,6 +609,9 @@ def destroy_ub(): _ub_initialized = False global layers_atomic_ring_exchange layers_atomic_ring_exchange = [] + # Compiled graphs may have baked is_fp8_ubuf() via assume_constant_result; + # reset so re-init with different settings doesn't read stale constants. + torch.compiler.reset() def fill_userbuffers_buffer_for_all_gather( @@ -1094,7 +1103,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): - return + if recipe_state.recipe is recipe: + return # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and # 2 (grad_output and grad_input) for bwd diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index e9342c064..988dab9fa 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -25,6 +25,7 @@ from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -1094,8 +1095,10 @@ def wgrad_gemm( if ctx.ln_out_needs_gather: # Gathered input is internal clear_tensor_data(ln_out_total) - if ctx.parallel_mode == "row" and ctx.sequence_parallel: - # Gathered grad output tensor is internal + if ctx.sequence_parallel and ( + ctx.parallel_mode == "row" or (ctx.parallel_mode == "column" and ctx.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal clear_tensor_data(grad_output) # Update grad input if overlapping reduce-scatter with wgrad GEMM @@ -1736,14 +1739,10 @@ def forward( is_first_microbatch = False if self.ub_overlap_rs_fprop: - if get_ub( - self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True if self.ub_overlap_rs_dgrad: - if get_ub( - self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True inp = self.prepare_forward( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index cea110f67..e5a9675fc 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -26,6 +26,7 @@ fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -2367,7 +2368,7 @@ def forward( fp8_output = False if self.ub_overlap_rs: - if get_ub("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()).is_fp8_ubuf(): + if get_ub_is_fp8("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True inp = self.prepare_forward(inp, num_gemms=2) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b91a3e8fd..556348fde 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -23,6 +23,7 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, + get_ub_is_fp8, is_ub_initialized, using_cublasmp_backend, quantize_weight, @@ -122,6 +123,8 @@ class LinearFwdArgs: fp8_output: bool save_original_input: bool backward_override: Optional[str] + dgrad_use_split_accumulator: bool + wgrad_use_split_accumulator: bool custom: bool debug: bool @@ -192,7 +195,8 @@ class LinearBwdArgs: # --- Numerical / dtype config --- activation_dtype: Optional[torch.dtype] = None fp8: bool = False - fp8_recipe: Optional[Recipe] = None + dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD + wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD backward_override: Optional[str] = None is_weight_param_quantized: bool = False custom: bool = False @@ -695,7 +699,8 @@ def _linear_setup_ctx( # Numerical / dtype config bwd_args.activation_dtype = fwd_args.activation_dtype bwd_args.fp8 = fp8 - bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator + bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator bwd_args.backward_override = backward_override bwd_args.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) bwd_args.custom = fwd_args.custom @@ -1001,11 +1006,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. weight_fp8.update_usage(columnwise_usage=True) # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_DGRAD - if bwd_args.fp8: - recipe = bwd_args.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator + use_split_accumulator = bwd_args.dgrad_use_split_accumulator # Update grad input quantizer if grad_input_quantizer is not None: @@ -1178,11 +1179,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. grad_output = grad_output_quantizer(grad_output) # Figure out whether to use split accumulator - use_split_accumulator = _2X_ACC_WGRAD - if bwd_args.fp8: - recipe = bwd_args.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator + use_split_accumulator = bwd_args.wgrad_use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad if bwd_args.is_first_microbatch is not None: @@ -1273,8 +1270,11 @@ def wgrad_gemm( elif bwd_args.backward_input_needs_gather: # Gathered input tensor is internal clear_tensor_data(inputmat_total) - if bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: - # Gathered grad output tensor is internal + if bwd_args.sequence_parallel and ( + bwd_args.parallel_mode == "row" + or (bwd_args.parallel_mode == "column" and bwd_args.fp8) + ): + # Gathered (row-SP) or quantized (column-SP FP8) grad_output is internal clear_tensor_data(grad_output) # Update grad input if overlapping reduce-scatter with wgrad GEMM @@ -1888,14 +1888,10 @@ def forward( is_first_microbatch = False if self.ub_overlap_rs_fprop: - if get_ub( - self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_fprop", FP8GlobalStateManager.is_fp8_enabled()): fp8_output = True if self.ub_overlap_rs_dgrad: - if get_ub( - self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled() - ).is_fp8_ubuf(): + if get_ub_is_fp8(self.ub_name + "_dgrad", FP8GlobalStateManager.is_fp8_enabled()): fp8_grad = True inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor)) @@ -1933,8 +1929,15 @@ def forward( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) + dgrad_use_split_accumulator = _2X_ACC_DGRAD + wgrad_use_split_accumulator = _2X_ACC_WGRAD if self.fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + _recipe = FP8GlobalStateManager.get_fp8_recipe() + backward_override = _recipe.backward_override + if hasattr(_recipe, "fp8_gemm_dgrad"): + dgrad_use_split_accumulator = _recipe.fp8_gemm_dgrad.use_split_accumulator + if hasattr(_recipe, "fp8_gemm_wgrad"): + wgrad_use_split_accumulator = _recipe.fp8_gemm_wgrad.use_split_accumulator else: backward_override = None custom = is_custom(input_quantizer) or is_custom(weight_quantizer) @@ -1989,6 +1992,8 @@ def forward( fp8_output=fp8_output, save_original_input=self.save_original_input, backward_override=backward_override, + dgrad_use_split_accumulator=dgrad_use_split_accumulator, + wgrad_use_split_accumulator=wgrad_use_split_accumulator, custom=custom, debug=debug, # weight-workspace caching