From 43c0570e9f4bacc5d9b488afb6a1cb86e78766d7 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 21:39:01 +0000 Subject: [PATCH 01/50] Remove nvidia_wheel_versions --- jaxlib/jax.bzl | 5 ----- 1 file changed, 5 deletions(-) diff --git a/jaxlib/jax.bzl b/jaxlib/jax.bzl index d07287cf6684..fb9ff22ca06c 100644 --- a/jaxlib/jax.bzl +++ b/jaxlib/jax.bzl @@ -22,10 +22,6 @@ load("@jax_wheel//:wheel.bzl", "WHEEL_VERSION") load("@jax_wheel_version_suffix//:wheel_version_suffix.bzl", "WHEEL_VERSION_SUFFIX") load("@local_config_cuda//cuda:build_defs.bzl", _cuda_library = "cuda_library", _if_cuda_is_configured = "if_cuda_is_configured") load("@local_config_rocm//rocm:build_defs.bzl", _if_rocm_is_configured = "if_rocm_is_configured", _rocm_library = "rocm_library") - -# TODO(Intel-tf): Update `sycl` with `oneapi` when xla changes to use `oneapi`. -load("@local_config_sycl//sycl:build_defs.bzl", _if_oneapi_is_configured = "if_sycl_is_configured", _oneapi_library = "sycl_library") -load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS") load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION", "HERMETIC_PYTHON_VERSION_KIND") load("@rocm_external_test_deps//:external_deps.bzl", "EXTERNAL_DEPS") load("@rocm_prebuilt_test_deps//:external_deps.bzl", PREBUILT_EXTERNAL_DEPS = "EXTERNAL_DEPS") @@ -494,7 +490,6 @@ def _jax_wheel_impl(ctx): if ctx.attr.platform_version == "": fail("platform_version must be set to a valid cuda version for cuda wheels") args.add("--platform_version", ctx.attr.platform_version) # required for gpu wheels - args.add("--nvidia_wheel_versions_data", NVIDIA_WHEEL_VERSIONS) # required for gpu wheels if ctx.attr.enable_rocm: args.add("--enable-rocm", "True") if ctx.attr.platform_version == "": From bcef89c302e2fb199da3cf829bba498b362e6921 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 21:42:44 +0000 Subject: [PATCH 02/50] Make jaxlib targets visible --- jaxlib/BUILD | 2 +- jaxlib/rocm/BUILD | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/jaxlib/BUILD b/jaxlib/BUILD index 1dd304f5bbff..562cf538a514 100644 --- a/jaxlib/BUILD +++ b/jaxlib/BUILD @@ -41,7 +41,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//jax:internal"], + default_visibility = ["//visibility:public"], ) package_group( diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index a27c138bd42c..bfe9bbe9bfce 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -28,7 +28,7 @@ licenses(["notice"]) package( default_applicable_licenses = [], - default_visibility = ["//:__subpackages__"], + default_visibility = ["//visibility:public"], ) cc_library( From 733b7bf8b4a0b7fac6d83027791cb0999440fda2 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Wed, 12 Nov 2025 22:12:01 +0000 Subject: [PATCH 03/50] hipblas typedef fix --- jaxlib/gpu/solver_interface.cc | 12 ++++++------ jaxlib/gpu/vendor.h | 12 ++++-------- jaxlib/rocm/BUILD | 4 ++++ 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/jaxlib/gpu/solver_interface.cc b/jaxlib/gpu/solver_interface.cc index 7ebcf7483652..e530e5bc1f2f 100644 --- a/jaxlib/gpu/solver_interface.cc +++ b/jaxlib/gpu/solver_interface.cc @@ -63,8 +63,8 @@ JAX_GPU_DEFINE_GETRF(gpuDoubleComplex, gpusolverDnZgetrf); JAX_GPU_DEFINE_GETRF_BATCHED(float, gpublasSgetrfBatched); JAX_GPU_DEFINE_GETRF_BATCHED(double, gpublasDgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasComplex, gpublasCgetrfBatched); -JAX_GPU_DEFINE_GETRF_BATCHED(gpublasDoubleComplex, gpublasZgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuComplex, gpublasCgetrfBatched); +JAX_GPU_DEFINE_GETRF_BATCHED(gpuDoubleComplex, gpublasZgetrfBatched); #undef JAX_GPU_DEFINE_GETRF_BATCHED // QR decomposition: geqrf @@ -102,8 +102,8 @@ JAX_GPU_DEFINE_GEQRF(gpuDoubleComplex, gpusolverDnZgeqrf); JAX_GPU_DEFINE_GEQRF_BATCHED(float, gpublasSgeqrfBatched); JAX_GPU_DEFINE_GEQRF_BATCHED(double, gpublasDgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasComplex, gpublasCgeqrfBatched); -JAX_GPU_DEFINE_GEQRF_BATCHED(gpublasDoubleComplex, gpublasZgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuComplex, gpublasCgeqrfBatched); +JAX_GPU_DEFINE_GEQRF_BATCHED(gpuDoubleComplex, gpublasZgeqrfBatched); #undef JAX_GPU_DEFINE_GEQRF_BATCHED // Householder transformations: orgqr @@ -302,8 +302,8 @@ JAX_GPU_DEFINE_SYEVD(gpuDoubleComplex, gpusolverDnZheevd); JAX_GPU_DEFINE_SYRK(float, gpublasSsyrk); JAX_GPU_DEFINE_SYRK(double, gpublasDsyrk); -JAX_GPU_DEFINE_SYRK(gpublasComplex, gpublasCsyrk); -JAX_GPU_DEFINE_SYRK(gpublasDoubleComplex, gpublasZsyrk); +JAX_GPU_DEFINE_SYRK(gpuComplex, gpublasCsyrk); +JAX_GPU_DEFINE_SYRK(gpuDoubleComplex, gpublasZsyrk); #undef JAX_GPU_DEFINE_SYRK // Singular Value Decomposition: gesvd diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 9adb1f11776f..4818a70e6941 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -457,6 +457,8 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; #elif defined(JAX_GPU_HIP) +#define HIPBLAS_V2 1 + // IWYU pragma: begin_exports #include "rocm/include/hip/hip_cooperative_groups.h" #include "rocm/include/hip/hip_runtime_api.h" @@ -477,17 +479,11 @@ inline constexpr uint32_t kNumThreadsPerWarp = 32; // MIOpen lib. Remove when MIOpen support is complete. #define MIOPEN_STATUS_SUCCESS 0 -typedef hipFloatComplex gpuComplex; +typedef hipComplex gpuComplex; typedef hipDoubleComplex gpuDoubleComplex; -#if TF_ROCM_VERSION >= 70000 -typedef hipFloatComplex gpublasComplex; +typedef hipComplex gpublasComplex; typedef hipDoubleComplex gpublasDoubleComplex; -#else -typedef hipblasComplex gpublasComplex; -typedef hipblasDoubleComplex gpublasDoubleComplex; -#endif // TF_ROCM_VERSION >= 70000 - typedef struct hipsolverHandle_* gpusolverDnHandle_t; typedef hipblasFillMode_t gpublasFillMode_t; typedef hipsolverFillMode_t gpusolverFillMode_t; diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index bfe9bbe9bfce..2bb10751eb39 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -425,6 +425,10 @@ rocm_nanobind_extension( "@xla//xla/ffi/api:ffi", "@xla//xla/python:safe_static_init", ], + linkopts = [ + "-L/opt/rocm/lib", + "-lamdhip64", + ], ) cc_library( From 793d312782630007b3cb50eca13918b34b79158b Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Thu, 13 Nov 2025 21:50:02 +0000 Subject: [PATCH 04/50] No GPU fail --- jaxlib/rocm/BUILD | 1 + jaxlib/rocm/rocm_plugin_extension.cc | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD index 2bb10751eb39..5deb8c568ef1 100644 --- a/jaxlib/rocm/BUILD +++ b/jaxlib/rocm/BUILD @@ -562,6 +562,7 @@ rocm_nanobind_extension( enable_stub_generation = False, module_name = "rocm_plugin_extension", deps = [ + ":hip_gpu_kernel_helpers", ":py_client_gpu", "//jaxlib:kernel_nanobind_helpers", "//jaxlib/gpu:gpu_plugin_extension", diff --git a/jaxlib/rocm/rocm_plugin_extension.cc b/jaxlib/rocm/rocm_plugin_extension.cc index 7f3a6dd12b2e..525aea1011b6 100644 --- a/jaxlib/rocm/rocm_plugin_extension.cc +++ b/jaxlib/rocm/rocm_plugin_extension.cc @@ -23,6 +23,7 @@ limitations under the License. #include "jaxlib/gpu/gpu_plugin_extension.h" #include "jaxlib/gpu/py_client_gpu.h" #include "jaxlib/kernel_nanobind_helpers.h" +#include "jaxlib/gpu/gpu_kernel_helpers.h" namespace nb = nanobind; @@ -98,6 +99,13 @@ nb::dict FfiHandlers() { return dict; } +int ROCmDeviceCount() { + int device_count = -1; + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipInit(0))); + JAX_THROW_IF_ERROR(JAX_AS_STATUS(hipGetDeviceCount(&device_count))); + return device_count; +} + } // namespace NB_MODULE(rocm_plugin_extension, m) { @@ -124,5 +132,6 @@ NB_MODULE(rocm_plugin_extension, m) { return device_ordinal; }, nb::arg("data_value")); + m.def("get_device_count", &ROCmDeviceCount); } } // namespace jax From e3ad0ecbd6be5574f7e98679e4678e4b0ce98903 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 12 Feb 2026 14:06:16 -0800 Subject: [PATCH 05/50] Wrap HIP inline functions in anonymous namespaces in vendor.h --- jaxlib/gpu/vendor.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 4818a70e6941..31d554cf35b4 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -541,6 +541,7 @@ inline hipblasStatus_t gpublasCreate(gpublasHandle_t* handle) { return hipblasCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpublasCreate ::jax::hip::gpublasCreate #define gpublasSetStream hipblasSetStream #define gpublasSgeqrfBatched hipblasSgeqrfBatched @@ -597,6 +598,7 @@ inline hipsolverStatus_t gpusolverDnCreate(gpusolverDnHandle_t* handle) { return hipsolverCreate(reinterpret_cast(handle)); } } // namespace jax::hip + #define gpusolverDnCreate ::jax::hip::gpusolverDnCreate #define gpusolverDnSetStream hipsolverSetStream #define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo From a831ef2028e6b26227be617af9cd13ccffdd793c Mon Sep 17 00:00:00 2001 From: Dragoslav Sicarov Date: Tue, 10 Jun 2025 04:28:40 +0000 Subject: [PATCH 06/50] SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError --- jaxlib/gpu/vendor.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 31d554cf35b4..5c9098494886 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -794,7 +794,7 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPU_STREAM_NON_BLOCKING hipStreamNonBlocking #define gpuMalloc hipMalloc -#define gpuGetLastError hipGetLastError +#define gpuGetLastError hipExtGetLastError #define gpuGetErrorString hipGetErrorString #define gpuMemcpyAsync hipMemcpyAsync #define gpuMemcpyDeviceToDevice hipMemcpyDeviceToDevice From 58249a4a929be437d647e7b65a8c31b98097eac6 Mon Sep 17 00:00:00 2001 From: Charles Hofer Date: Fri, 14 Nov 2025 19:43:09 +0000 Subject: [PATCH 07/50] Add shared utility function get_rocm_version to test_util.py --- jax/_src/test_util.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py index 5d11761386bb..1906c42b5f4b 100644 --- a/jax/_src/test_util.py +++ b/jax/_src/test_util.py @@ -25,6 +25,7 @@ import math import os +from pathlib import Path import platform import re import sys @@ -389,6 +390,15 @@ def supported_dtypes() -> set[DTypeLike]: def is_device_rocm() -> bool: return 'rocm' in xla_bridge.get_backend().platform_version +def get_rocm_version(): + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + version_path = Path(rocm_path) / ".info" / "version" + if not version_path.exists(): + raise FileNotFoundError(f"Expected ROCm version file at {version_path}") + version_str = version_path.read_text().strip() + major, minor, *_ = version_str.split(".") + return int(major), int(minor) + def is_device_cuda() -> bool: return 'cuda' in xla_bridge.get_backend().platform_version From e587f9034f7f55f748002d0e2f8c280458f9b215 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Mon, 17 Nov 2025 19:34:29 +0000 Subject: [PATCH 08/50] Fix hipSparse CSR algorithm mappings for ROCm 7 --- jaxlib/gpu/vendor.h | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h index 5c9098494886..fb9c1f407145 100644 --- a/jaxlib/gpu/vendor.h +++ b/jaxlib/gpu/vendor.h @@ -778,10 +778,10 @@ inline hipsparseStatus_t gpusparseCreate(gpusparseHandle_t* handle) { #define GPUSPARSE_INDEX_32I HIPSPARSE_INDEX_32I #define GPUSPARSE_INDEX_64I HIPSPARSE_INDEX_64I #define GPUSPARSE_DENSETOSPARSE_ALG_DEFAULT HIPSPARSE_DENSETOSPARSE_ALG_DEFAULT -#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_MV_ALG_DEFAULT -#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_ALG_DEFAULT -#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_ALG_DEFAULT +#define GPUSPARSE_SPMV_COO_ALG HIPSPARSE_COOMV_ALG +#define GPUSPARSE_SPMV_CSR_ALG HIPSPARSE_CSRMV_ALG1 +#define GPUSPARSE_SPMM_COO_ALG HIPSPARSE_SPMM_COO_ALG1 +#define GPUSPARSE_SPMM_CSR_ALG HIPSPARSE_SPMM_CSR_ALG1 #define GPUSPARSE_INDEX_BASE_ZERO HIPSPARSE_INDEX_BASE_ZERO #define GPUSPARSE_OPERATION_NON_TRANSPOSE HIPSPARSE_OPERATION_NON_TRANSPOSE #define GPUSPARSE_OPERATION_TRANSPOSE HIPSPARSE_OPERATION_TRANSPOSE From 808994737bc6f80fcd4ed3e17743087bdf83e466 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Thu, 20 Nov 2025 01:30:24 +0200 Subject: [PATCH 09/50] =?UTF-8?q?Fix=20v=5Fpages=20quantization=20and=20ad?= =?UTF-8?q?just=20test=20params=20for=20ROCm=20compatibilit=E2=80=A6=20(#5?= =?UTF-8?q?60)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/pallas/gpu_paged_attention_test.py | 71 +++++++++++++++++++++++- 1 file changed, 70 insertions(+), 1 deletion(-) diff --git a/tests/pallas/gpu_paged_attention_test.py b/tests/pallas/gpu_paged_attention_test.py index 1b778c787a6d..6a1d8de22a78 100644 --- a/tests/pallas/gpu_paged_attention_test.py +++ b/tests/pallas/gpu_paged_attention_test.py @@ -112,6 +112,64 @@ class PagedAttentionKernelTest(PallasBaseTest): def setUp(self): super().setUp() + def _estimate_shared_memory_bytes(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Estimate shared memory usage for paged attention kernel.""" + dtype_size = jnp.dtype(dtype).itemsize + # Approximate calculation based on kernel's memory usage + # Q block: block_h * head_dim + # K/V blocks: pages_per_compute_block * page_size * head_dim + # Plus accumulators and intermediate values + block_k = pages_per_compute_block * page_size + estimated = dtype_size * ( + block_h * head_dim + # Q + 2 * block_k * head_dim + # K and V + block_h * block_k + # logits/attention weights + block_h * 8 # accumulators (m, l, etc.) in float32 + ) + return estimated + + def _adjust_params_for_shared_memory(self, block_h, pages_per_compute_block, + page_size, head_dim, dtype): + """Adjust parameters to fit within device shared memory limits. + + Uses XLA's DeviceDescription.shared_memory_per_block_optin() to query + the actual device capability rather than hardcoding values. + """ + try: + device = jax.local_devices()[0] + # Query XLA DeviceDescription for max shared memory per block + # This is exposed from stream_executor::DeviceDescription::shared_memory_per_block_optin() + max_smem = device.shared_memory_per_block_optin + except (AttributeError, IndexError): + # Fallback if XLA doesn't expose shared_memory_per_block_optin (older versions) + # or if no devices are available. Use conservative 48KB (safe for most GPUs). + max_smem = 48 * 1024 + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + # If within limits, no adjustment needed + if estimated <= max_smem: + return block_h, pages_per_compute_block, page_size + + # Try to reduce parameters to fit + while estimated > max_smem: + if pages_per_compute_block > 2: + pages_per_compute_block = pages_per_compute_block // 2 + elif page_size > 8: + page_size = page_size // 2 + elif block_h > 8: + block_h = block_h // 2 + else: + # Can't reduce further, will need to skip + return None, None, None + + estimated = self._estimate_shared_memory_bytes( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + return block_h, pages_per_compute_block, page_size + @jtu.sample_product( dtype=(jnp.float16,), page_size=(8, 16, 32), @@ -201,6 +259,17 @@ def test_quantized_paged_attention( if (quant_dtype == jnp.float8_e4m3fn and not jtu.is_cuda_compute_capability_at_least("8.9")): self.skipTest("Skipping since float8_e4m3fn is not supported on < sm89") + + # Check and adjust parameters if needed to fit device limits for ROCm + if jtu.is_device_rocm(): + adjusted = self._adjust_params_for_shared_memory( + block_h, pages_per_compute_block, page_size, head_dim, dtype) + + if adjusted == (None, None, None): + self.skipTest("Cannot adjust parameters to fit ROCm device shared memory limits") + + block_h, pages_per_compute_block, page_size = adjusted + max_kv_len = 2048 seq_lens = np.asarray([3, 256, 513, 1023, 2048], dtype=jnp.int32) q, k_pages, v_pages, block_tables = _generate_qkv( @@ -218,7 +287,7 @@ def test_quantized_paged_attention( k_, k_scales = (_quantize(k_pages, quant_dtype) if quantize_k else (k_pages, None)) - v_, v_scales = (_quantize(k_pages, quant_dtype) + v_, v_scales = (_quantize(v_pages, quant_dtype) if quantize_v else (v_pages, None)) o = paged_attention.paged_attention( From d9e7020ec5471a224a39f807a6157723b620402e Mon Sep 17 00:00:00 2001 From: Aleksei <208770786+Arech8@users.noreply.github.com> Date: Wed, 26 Nov 2025 17:19:56 +0200 Subject: [PATCH 10/50] Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563) When jaxlib was built in debug more, an assertion in LLVM code that lazy-loads VHLO dialect could fire, since the code path could execute in a multi-threaded environment, and LLVM dialect repositories aren't thread safe to modify. This patch applies the same changes that upstream makes to fix this: https://github.com/jax-ml/jax/commit/48c876227a67840d7bcd44d80d16edbc0e910335 (this includes disabling a call to `jax_mlir_ext.enter_multi_threaded_execution(context)` in `mlir.py`. Presumably, the whole functionality related to `enter_multi_threaded_execution()` multithreaded checks isn't ready yet, and it was prematurely rolled into the production code. Manual testing --- .gitignore | 4 +++ jax/_src/interpreters/mlir.py | 5 ++-- jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc | 35 +++++++++++++------------- 3 files changed, 23 insertions(+), 21 deletions(-) diff --git a/.gitignore b/.gitignore index d30c019b31e8..e55a9a5de101 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,7 @@ jax.iml /include/ /lib/ /share/ + +/compile_commands.json +/strace.txt +/external diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py index 1c348250dd7d..1433bb540ecf 100644 --- a/jax/_src/interpreters/mlir.py +++ b/jax/_src/interpreters/mlir.py @@ -634,9 +634,8 @@ def make_ir_context() -> ir.Context: # multi threaded execution aborts the process if we try to register a new # dialect after this point. The dialect registry in a context is not thread # safe, and a fatal error is much better than a data race. - # jax_mlir_ext.enter_multi_threaded_execution(context) - # TODO(phawkins): clean up users who add their own dialects to JAX's contexts - # and enable this. + # if jaxlib_version >= (0, 8): + # jax_mlir_ext.enter_multi_threaded_execution(context) return context diff --git a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc index 090a74e95056..243122889f14 100644 --- a/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc +++ b/jaxlib/mlir/_mlir_libs/jax_mlir_ext.cc @@ -310,24 +310,23 @@ NB_MODULE(_jax_mlir_ext, m) { MlirDialectRegistry c_registry = registry.get(); #define REGISTER_DIALECT(name) \ MlirDialectHandle name##_dialect = mlirGetDialectHandle__##name##__(); \ - mlirDialectHandleInsertDialect(name##_dialect, c_registry) - REGISTER_DIALECT(arith); - REGISTER_DIALECT(func); - REGISTER_DIALECT(math); - REGISTER_DIALECT(memref); - REGISTER_DIALECT(scf); - REGISTER_DIALECT(vector); - // TODO(jpienaar): these don't seem to have C API targets known to Bazel - unwrap(c_registry)->insert(); - unwrap(c_registry)->insert(); - unwrap(c_registry)->insert(); - - // For Mosaic GPU - REGISTER_DIALECT(cf); - REGISTER_DIALECT(gpu); - REGISTER_DIALECT(nvgpu); - REGISTER_DIALECT(nvvm); - REGISTER_DIALECT(llvm); + mlirDialectHandleInsertDialect(name##_dialect, registry) + REGISTER_DIALECT(arith); + REGISTER_DIALECT(func); + REGISTER_DIALECT(math); + REGISTER_DIALECT(memref); + REGISTER_DIALECT(scf); + REGISTER_DIALECT(vector); + // TODO(jpienaar): these don't seem to have C API targets known to Bazel + unwrap(registry)->insert(); + unwrap(registry)->insert(); + unwrap(registry)->insert(); + // For Mosaic GPU + REGISTER_DIALECT(cf); + REGISTER_DIALECT(gpu); + REGISTER_DIALECT(nvgpu); + REGISTER_DIALECT(nvvm); + REGISTER_DIALECT(llvm); #undef REGISTER_DIALECT mlirMosaicGpuRegisterSerdePass(); From 42a3be64c5da4430374827f0a7d99a7ed0344953 Mon Sep 17 00:00:00 2001 From: Aleksei <208770786+Arech8@users.noreply.github.com> Date: Wed, 26 Nov 2025 18:35:19 +0200 Subject: [PATCH 11/50] Add skip of test_is_finite() on Cuda (#565) (forgot this skip in the previous PR) --- tests/pallas/ops_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index 9fea8249e525..731684e334d6 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -1047,6 +1047,10 @@ def test_is_finite(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.test_device_matches(["cuda"]): + self.skipTest("Not tested on CUDA") # set this b/c this how the test was + # originally configured. Have no way to test cuda. + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( From 544c6d4a45c80a539c34eedc2b07a34874aa0f34 Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Mon, 15 Dec 2025 08:00:42 -0800 Subject: [PATCH 12/50] Add rocm test requirements file (#570) --- build/rocm-test-requirements.txt | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 build/rocm-test-requirements.txt diff --git a/build/rocm-test-requirements.txt b/build/rocm-test-requirements.txt new file mode 100644 index 000000000000..399237175957 --- /dev/null +++ b/build/rocm-test-requirements.txt @@ -0,0 +1,24 @@ +absl-py +build +cloudpickle +colorama>=0.4.4 +filelock +flatbuffers +hypothesis +mpmath>=1.3 +pillow>=10.4.0 +# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t +portpicker; python_version<"3.13" +pytest-xdist +pytest-json-report +pytest-html +pytest-csv +pytest-rerunfailures +pytest-html-merger +pytest-reportlog +wheel +rich +setuptools +matplotlib +opt-einsum +auditwheel From 4673584da681cea166c3fd6a437124511c171c97 Mon Sep 17 00:00:00 2001 From: charleshofer Date: Mon, 15 Dec 2025 11:03:51 -0600 Subject: [PATCH 13/50] Let the unit tests use build.py for setting up Bazel commands for unit tests (#582) --- build/build.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/build/build.py b/build/build.py index 38a11cedff46..d4368c444caa 100755 --- a/build/build.py +++ b/build/build.py @@ -667,6 +667,10 @@ async def main(): ) if "rocm" in args.wheels: + if not args.configure_only: + print("ERROR: This repo is not used for building the ROCm JAX plugins. Please use the new plugin repo: https://github.com/ROCm/rocm-jax") + exit(1) + wheel_build_command_base.append("--config=rocm_base") wheel_build_command_base.append("--config=rocm") if clang_local: From 1c79814a7a68202c6d4e15b50eca7e008a819477 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 13 Jan 2026 10:29:56 -0600 Subject: [PATCH 14/50] adding abort logic to rocm/jax (#590) --- conftest.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/conftest.py b/conftest.py index cd8414369b89..82d5273e2cdd 100644 --- a/conftest.py +++ b/conftest.py @@ -15,7 +15,10 @@ import os import pytest - +import json +import threading +import shutil +from datetime import datetime @pytest.fixture(autouse=True) def add_imports(doctest_namespace): From 9b5d70884640c6f58a359bc44c4971d4c8549325 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Wed, 14 Jan 2026 19:57:26 +0200 Subject: [PATCH 15/50] Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (#597) --- tests/pallas/ops_test.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index 731684e334d6..cb52c42e0012 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -1051,6 +1051,9 @@ def test_is_finite(self, dtype): self.skipTest("Not tested on CUDA") # set this b/c this how the test was # originally configured. Have no way to test cuda. + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( @@ -1103,6 +1106,9 @@ def test_is_finite_scalar(self, dtype): if jtu.is_device_rocm() and jtu.parse_version(jax.__version__) == (0, 8, 0): self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + if jtu.is_device_rocm(): + self.skipTest("is_finite not in Triton lowering for jax 0.8.0") + size = len(self.IS_FINITE_TEST_VALUES) @functools.partial( From 82bf13e1c47c11b7ade9ee4d5358027b459373a3 Mon Sep 17 00:00:00 2001 From: Pham Binh Date: Wed, 14 Jan 2026 19:57:53 +0200 Subject: [PATCH 16/50] Fix shared memory limit check for ROCm in test_dot (#596) --- tests/pallas/ops_test.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py index cb52c42e0012..a1af7d6f252a 100644 --- a/tests/pallas/ops_test.py +++ b/tests/pallas/ops_test.py @@ -2161,12 +2161,20 @@ def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y): self.skipTest("bfloat16 type are not supported on GPU") # Check shared memory limit: Triton loads lhs + rhs into shared memory if jtu.is_device_rocm(): + # ROCm: use correct formula with dynamic limit from rocminfo dtype_size = jnp.dtype(dtype).itemsize - if (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size > get_rocm_shared_memory_limit(): + shared_mem_bytes = (math.prod(lhs_shape) + math.prod(rhs_shape)) * dtype_size + shared_mem_limit = get_rocm_shared_memory_limit() + if shared_mem_bytes > shared_mem_limit: self.skipTest("Shared memory size limit exceeded") - elif math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) > (256 * 256) * 2: - self.skipTest("Shared memory size limit exceeded") - if (jax.local_devices()[0].shared_memory_per_block_optin == 99 * 1024 and + else: + # NVIDIA: keep original check + if ( + math.prod(lhs_shape) + math.prod(rhs_shape) + math.prod(out_shape) + > (256 * 256) * 2 + ): + self.skipTest("Shared memory size limit exceeded") + if (jax.local_devices()[0].device_kind == "NVIDIA L4" and dtype == jnp.float32 and lhs_and_rhs_shape in [ ((128, 16), (128, 256)), From ad47e1748b9abaa58a41b4f9c2180f2d49f29ad7 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Wed, 14 Jan 2026 12:14:02 -0600 Subject: [PATCH 17/50] Fix Numpy signatures test (#598) Co-authored-by: Daniel Suo Co-authored-by: Jake VanderPlas --- tests/lax_numpy_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index 3d3c9b9dab52..7252122fa548 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6425,6 +6425,10 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], +<<<<<<< HEAD +======= + 'var': ['mean'], +>>>>>>> a3f8af53f (Fix Numpy signatures test (#598)) 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 3b3b31cfdcf8195b140a84791e08913a08483e5d Mon Sep 17 00:00:00 2001 From: Ruturaj4 Date: Sun, 18 Jan 2026 10:18:54 -0600 Subject: [PATCH 18/50] fix merge arts --- tests/lax_numpy_test.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index 7252122fa548..033229c77c08 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6425,10 +6425,7 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], -<<<<<<< HEAD -======= 'var': ['mean'], ->>>>>>> a3f8af53f (Fix Numpy signatures test (#598)) 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 8a9adefd57c7d275654ec254924e0a63493d0ab2 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Thu, 22 Jan 2026 15:44:41 -0600 Subject: [PATCH 19/50] Enable RngShardingTests (#644) --- tests/array_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/array_test.py b/tests/array_test.py index cf08765c5893..890b701dd589 100644 --- a/tests/array_test.py +++ b/tests/array_test.py @@ -1550,7 +1550,7 @@ class RngShardingTest(jtu.JaxTestCase): # tests that the PRNGs are automatically sharded as expected @parameterized.named_parameters(("3", 3), ("4", 4), ("5", 5)) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_1d(self, num_devices): @jax.jit def f(x): @@ -1584,7 +1584,7 @@ def f(x): "mesh_shape": mesh_shape, "pspec": pspec} for mesh_shape in [(3, 2), (4, 2), (2, 3)] for pspec in [P('x', None), P(None, 'y'), P('x', 'y')]) - @jtu.skip_on_devices("gpu") + @jtu.skip_on_devices("cuda") def test_random_bits_is_pure_map_2d(self, mesh_shape, pspec): @jax.jit def f(x): From 4eb747358eb0d5d18a575ee996a8549a064c4812 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 12 Feb 2026 15:18:40 -0800 Subject: [PATCH 20/50] Enable test_variadic_reduce_window on ROCm (#647) --- tests/lax_vmap_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py index 81c7f5103f44..fa168d61b03c 100644 --- a/tests/lax_vmap_test.py +++ b/tests/lax_vmap_test.py @@ -759,6 +759,8 @@ def testSort(self, shape, dimension, arity, bdims, is_stable): # TODO Collapse # TODO Scatter + # b/183233858: variadic reduce-window not implemented on XLA:CUDA + @jtu.skip_on_devices("cuda") def test_variadic_reduce_window(self): # https://github.com/jax-ml/jax/discussions/9818 and # https://github.com/jax-ml/jax/issues/9837 From f360e13539c943a0b6e35fff4d1a37f107b93539 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 21/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 3823ca26d8f9..7332e4bab069 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -219,6 +220,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -592,6 +594,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. @@ -1054,6 +1057,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): # hipSPARSE segfault observed as of ROCm 7.2. From 81842f41c76489b913700e338f4f821310797143 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 22/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index 7332e4bab069..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,11 +137,8 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") csr_matmul = sparse_csr._csr_matvec if len(bshape) == 1 else sparse_csr._csr_matmat tol = {np.float32: 2E-5, np.float64: 1E-12, np.complex64: 1E-5, @@ -220,11 +217,8 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") op = lambda M: M.T if transpose else M @@ -594,11 +588,8 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") tol = {np.float32: 2E-5, np.float64: 2E-14} @@ -1057,11 +1048,8 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): - # hipSPARSE segfault observed as of ROCm 7.2. - # TODO(ROCm): Re-enable once hipSPARSE issue is fixed. self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") rng = sptu.rand_sparse(self.rng(), post=jnp.array) rng_b = jtu.rand_default(self.rng()) From 489fcf609d92cf98efedd36451e50ac2b6575036 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 23/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index a7bf9373c34c..df1a5a4e0286 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1048,6 +1051,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From c2ea7b4580e73c5c6e48c12551505b5e64af6bba Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 24/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index df1a5a4e0286..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1051,7 +1048,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 82a1e816b687b9d6f3d61fbfa11daba9d0ff9d8a Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 25/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index a7bf9373c34c..df1a5a4e0286 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1048,6 +1051,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 2c12a03cfc5ac0a8c2ac88a625ba3c4464639808 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 26/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index df1a5a4e0286..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1051,7 +1048,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 5c681b1018d4c1b254ef7e8efc607248fece26cd Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Wed, 28 Jan 2026 14:00:42 -0800 Subject: [PATCH 27/50] Enable testMultivariateNormalSingularCovariance on ROCm (#666) --- tests/random_lax_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py index 860f09a7f988..f7b68a483103 100644 --- a/tests/random_lax_test.py +++ b/tests/random_lax_test.py @@ -1024,7 +1024,7 @@ def testMultivariateNormalCovariance(self): check_dtypes=False) @jtu.sample_product(method=['cholesky', 'eigh', 'svd']) - @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. + @jtu.skip_on_devices('cuda', 'tpu') # Some NaNs on accelerators. ROCm supported def testMultivariateNormalSingularCovariance(self, method): # Singular covariance matrix https://github.com/jax-ml/jax/discussions/13293 mu = jnp.zeros((2,)) From 3757b64d80ba1fc595a9f12ffd91b7a35ac853ac Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Wed, 28 Jan 2026 16:38:17 -0600 Subject: [PATCH 28/50] Update Skip Reason Outputs (#663) --- jax/_src/test_util.py | 20 ++++++++++++++++---- tests/pallas/gpu_pallas_distributed_test.py | 15 +++++++++------ 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py index 1906c42b5f4b..ab0cc3849503 100644 --- a/jax/_src/test_util.py +++ b/jax/_src/test_util.py @@ -620,7 +620,14 @@ def skip_on_devices(*disabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = "Skipped on devices with tags: " + ", ".join(disabled_devices) + skip_messages = { + ("gpu",): "Skipped on all GPUs.", + ("cpu",): "Skipped on CPU.", + ("tpu",): "Skipped on TPU.", + ("cuda",): "Skipped on CUDA GPUs.", + ("rocm",): "Skipped on ROCm GPUs.", + } + skip_reason = skip_messages.get(disabled_devices) return _device_filter(lambda: not test_device_matches(disabled_devices), skip_reason) def run_on_devices(*enabled_devices, skip_reason=None): @@ -631,9 +638,14 @@ def run_on_devices(*enabled_devices, skip_reason=None): skip_reason: Optional custom skip message when test is skipped. """ if skip_reason is None: - skip_reason = ( - "Skipped unless running on devices with tags: " + ", ".join(enabled_devices) - ) + device_specific_skip_reasons = { + ("cpu",): "Skipped: CPU-only test.", + ("tpu",): "Skipped: TPU-only test.", + ("gpu",): "Skipped: GPU-only test.", + ("rocm",): "Skipped: ROCm-only test.", + ("cuda",): "Skipped: CUDA-only test.", + } + skip_reason = device_specific_skip_reasons.get(enabled_devices) return _device_filter(lambda: test_device_matches(enabled_devices), skip_reason) def device_supports_buffer_donation(): diff --git a/tests/pallas/gpu_pallas_distributed_test.py b/tests/pallas/gpu_pallas_distributed_test.py index f47e14e7ac21..169e24b2062b 100644 --- a/tests/pallas/gpu_pallas_distributed_test.py +++ b/tests/pallas/gpu_pallas_distributed_test.py @@ -133,14 +133,17 @@ def setUp(self): if jtu.is_device_rocm(): self.skipTest("Mosaic GPU is not supported on ROCm.") - if (not jtu.is_device_cuda() or + # Check mosaic support first (before GPU capability check) + if not mgpu.supports_cross_device_collectives(): + if jtu.test_device_matches(["rocm"]): + self.skipTest("Mosaic not supported on ROCm currently.") + else: + self.skipTest("NVSHMEM library unavailable.") + if (not jtu.test_device_matches(["cuda"]) or not jtu.is_cuda_compute_capability_at_least("9.0")): self.skipTest("Only works on GPU with capability >= sm90") - if not mgpu.supports_cross_device_collectives(): - self.skipTest( - "Skip test since cross-device collectives are not supported" - " (either NVSHMEM is not available in multi-process mode, or mixed" - " mode is used).") + if jax.process_count() == 1: + self.skipTest("Test requires multiple processes.") if os.environ.get("XLA_PYTHON_CLIENT_ALLOCATOR", "") == "platform": self.skipTest("NVSHMEM doesn't work with the platform allocator.") From 7ec9fe0f5e884597d0bb0873bee144c90470d201 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 29/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index a7bf9373c34c..df1a5a4e0286 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1048,6 +1051,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 411d4faab4a0b00a89f166c0f472be7c774797cc Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 30/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index df1a5a4e0286..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1051,7 +1048,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 130ca42212c15befce6bcfe88f5ccaf64ee0d36d Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Thu, 29 Jan 2026 13:12:07 -0600 Subject: [PATCH 31/50] Skip testCudaArrayInterfaceOnNonCudaFails on ROCm platform (#677) --- tests/array_interoperability_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py index d033a490f547..b14e9171f952 100644 --- a/tests/array_interoperability_test.py +++ b/tests/array_interoperability_test.py @@ -312,7 +312,7 @@ def testCudaArrayInterfaceOnNonCudaFails(self): self.assertFalse(hasattr(x, "__cuda_array_interface__")) with self.assertRaisesRegex( AttributeError, - "__cuda_array_interface__ is only defined for .*GPU buffers.", + "__cuda_array_interface__ is only defined for GPU buffers.", ): _ = x.__cuda_array_interface__ From 837654d25aa3ef414a362b4dc09c60c5838e7bf2 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 32/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index a7bf9373c34c..df1a5a4e0286 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1048,6 +1051,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 2b8c7fe3d04bca4b51007526668b81180c6b7a9f Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 33/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index df1a5a4e0286..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1051,7 +1048,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From d0a11b30d9ac63838f02157f1be262329e58f48d Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 13:38:16 -0600 Subject: [PATCH 34/50] Skip sparse tests on ROCm due to hipSPARSE issue (#652) --- tests/sparse_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index a7bf9373c34c..df1a5a4e0286 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,6 +137,7 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -217,6 +218,7 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -588,6 +590,7 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1048,6 +1051,7 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") + @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 3e79165b6442c01f7c20c3a2eca1e8c73d8f49e0 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 23 Jan 2026 15:35:23 -0600 Subject: [PATCH 35/50] Update sparse test skip messages in v0.8.2 (#653) --- tests/sparse_test.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/sparse_test.py b/tests/sparse_test.py index df1a5a4e0286..a7bf9373c34c 100644 --- a/tests/sparse_test.py +++ b/tests/sparse_test.py @@ -137,7 +137,6 @@ def test_csr_fromdense_ad(self, shape, dtype): dtype=jtu.dtypes.floating + jtu.dtypes.complex, ) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matmul_ad(self, shape, dtype, bshape): if jtu.is_device_rocm(): self.skipTest("test_csr_matmul_ad not supported on ROCm due to hipSPARSE issue") @@ -218,7 +217,6 @@ def test_csr_fromdense(self, shape, dtype): dtype=all_dtypes, transpose=[True, False], ) - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_matvec(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_matvec not supported on ROCm due to hipSPARSE issue") @@ -590,7 +588,6 @@ def test_coo_spmm(self, shape, dtype, transpose): transpose=[True, False], ) @jtu.run_on_devices("gpu") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_csr_spmv(self, shape, dtype, transpose): if jtu.is_device_rocm(): self.skipTest("test_csr_spmv not supported on ROCm due to hipSPARSE issue") @@ -1051,7 +1048,6 @@ def test_transpose(self, shape, dtype, Obj): ) for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])) @jax.default_matmul_precision("float32") - @jtu.skip_on_devices("rocm") # skipping on ROCm due to known issue in hipSPARSE def test_matmul(self, shape, dtype, Obj, bshape): if jtu.is_device_rocm(): self.skipTest("test_matmul not supported on ROCm due to hipSPARSE issue") From 9d4fce1a71eb18e78aa2235553636d19381acb88 Mon Sep 17 00:00:00 2001 From: Manjunath Gaonkar Date: Fri, 6 Feb 2026 12:48:48 -0600 Subject: [PATCH 36/50] Remove 'mean' from unsupported params for jnp.var (#689) --- tests/lax_numpy_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py index 033229c77c08..3d3c9b9dab52 100644 --- a/tests/lax_numpy_test.py +++ b/tests/lax_numpy_test.py @@ -6425,7 +6425,6 @@ def testWrappedSignaturesMatch(self): 'stack': ['casting'], 'tri': ['like'], 'unravel_index': ['order'], - 'var': ['mean'], 'vstack': ['casting'], 'zeros': ['order', 'like'], 'zeros_like': ['subok', 'order'] From 7601b8681a63150267ec56bba0191478025df282 Mon Sep 17 00:00:00 2001 From: AratiGanesh Date: Mon, 9 Feb 2026 05:08:48 -0800 Subject: [PATCH 37/50] Skipping testEighTinyNorm due to hipSolver issues (#697) --- tests/linalg_test.py | 241 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 241 insertions(+) diff --git a/tests/linalg_test.py b/tests/linalg_test.py index 867168b1c300..b99492724fd5 100644 --- a/tests/linalg_test.py +++ b/tests/linalg_test.py @@ -497,6 +497,247 @@ def testEigBatching(self, shape, dtype): self.assertTrue(np.all(np.linalg.norm( np.matmul(args, vs) - ws[..., None, :] * vs) < 1e-3)) + @jtu.sample_product( + n=[0, 4, 5, 50, 512], + dtype=float_types + complex_types, + lower=[True, False], + ) + def testEigh(self, n, dtype, lower): + rng = jtu.rand_default(self.rng()) + eps = np.finfo(dtype).eps + args_maker = lambda: [rng((n, n), dtype)] + + uplo = "L" if lower else "U" + + a, = args_maker() + a = (a + np.conj(a.T)) / 2 + w, v = jnp.linalg.eigh(np.tril(a) if lower else np.triu(a), + UPLO=uplo, symmetrize_input=False) + w = w.astype(v.dtype) + tol = 2 * n * eps + self.assertAllClose( + np.eye(n, dtype=v.dtype), + np.matmul(np.conj(T(v)), v), + atol=tol, + rtol=tol, + ) + + with jax.numpy_rank_promotion('allow'): + tol = 100 * eps + self.assertLessEqual( + np.linalg.norm(np.matmul(a, v) - w * v), tol * np.linalg.norm(a) + ) + + self._CompileAndCheck( + partial(jnp.linalg.eigh, UPLO=uplo), args_maker, rtol=eps + ) + + # Compare eigenvalues against Numpy using double precision. We do not compare + # eigenvectors because they are not uniquely defined, but the two checks above + # guarantee that that they satisfy the conditions for being eigenvectors. + double_type = dtype + if dtype == np.float32: + double_type = np.float64 + if dtype == np.complex64: + double_type = np.complex128 + w_np = np.linalg.eigvalsh(a.astype(double_type)) + tol = 8 * eps + self.assertAllClose( + w_np.astype(w.dtype), w, atol=tol * np.linalg.norm(a), rtol=tol + ) + + @jax._src.config.explicit_x64_dtypes("allow") + @jtu.run_on_devices("gpu") + @unittest.skip("Needs a large amount of GPU memory, doesn't work in CI") + def testEighLargeMatrix(self): + # https://github.com/jax-ml/jax/issues/33062 + n = 16384 + A = jnp.eye(n, dtype=jnp.float64) + jax.block_until_ready(jax.lax.linalg.eigh(A)) + + @jtu.sample_product( + start=[0, 1, 63, 64, 65, 255], + end=[1, 63, 64, 65, 256], + ) + @jtu.run_on_devices("tpu") # TODO(rmlarsen: enable on other devices) + def testEighSubsetByIndex(self, start, end): + if start >= end: + return + dtype = np.float32 + n = 256 + rng = jtu.rand_default(self.rng()) + eps = np.finfo(dtype).eps + args_maker = lambda: [rng((n, n), dtype)] + subset_by_index = (start, end) + k = end - start + (a,) = args_maker() + a = (a + np.conj(a.T)) / 2 + + v, w = lax.linalg.eigh( + a, symmetrize_input=False, subset_by_index=subset_by_index + ) + w = w.astype(v.dtype) + + self.assertEqual(v.shape, (n, k)) + self.assertEqual(w.shape, (k,)) + with jax.numpy_rank_promotion("allow"): + tol = 200 * eps + self.assertLessEqual( + np.linalg.norm(np.matmul(a, v) - w * v), tol * np.linalg.norm(a) + ) + tol = 3 * n * eps + self.assertAllClose( + np.eye(k, dtype=v.dtype), + np.matmul(np.conj(T(v)), v), + atol=tol, + rtol=tol, + ) + + self._CompileAndCheck(partial(jnp.linalg.eigh), args_maker, rtol=eps) + + # Compare eigenvalues against Numpy. We do not compare eigenvectors because + # they are not uniquely defined, but the two checks above guarantee that + # that they satisfy the conditions for being eigenvectors. + double_type = dtype + if dtype == np.float32: + double_type = np.float64 + if dtype == np.complex64: + double_type = np.complex128 + w_np = np.linalg.eigvalsh(a.astype(double_type))[ + subset_by_index[0] : subset_by_index[1] + ] + tol = 20 * eps + self.assertAllClose( + w_np.astype(w.dtype), w, atol=tol * np.linalg.norm(a), rtol=tol + ) + + def testEighZeroDiagonal(self): + a = np.array([[0., -1., -1., 1.], + [-1., 0., 1., -1.], + [-1., 1., 0., -1.], + [1., -1., -1., 0.]], dtype=np.float32) + w, v = jnp.linalg.eigh(a) + w = w.astype(v.dtype) + eps = jnp.finfo(a.dtype).eps + with jax.numpy_rank_promotion('allow'): + self.assertLessEqual( + np.linalg.norm(np.matmul(a, v) - w * v), 2.5 * eps * np.linalg.norm(a) + ) + + + def testEighTinyNorm(self): + if jtu.is_device_rocm(): + # numerical errors seen as of ROCm 7.2 due to hipSolver issue + # TODO: re-enable the test once the hipSolver issue is fixed + self.skipTest("testEighNorm not supported on ROCm due to hipSOLVER issue") + rng = jtu.rand_default(self.rng()) + a = rng((300, 300), dtype=np.float32) + eps = jnp.finfo(a.dtype).eps + a = eps * (a + np.conj(a.T)) + w, v = jnp.linalg.eigh(a) + w = w.astype(v.dtype) + with jax.numpy_rank_promotion("allow"): + self.assertLessEqual( + np.linalg.norm(np.matmul(a, v) - w * v), 80 * eps * np.linalg.norm(a) + ) + + @jtu.sample_product( + rank=[1, 3, 299], + ) + def testEighRankDeficient(self, rank): + rng = jtu.rand_default(self.rng()) + eps = jnp.finfo(np.float32).eps + a = rng((300, rank), dtype=np.float32) + a = a @ np.conj(a.T) + w, v = jnp.linalg.eigh(a) + w = w.astype(v.dtype) + with jax.numpy_rank_promotion("allow"): + self.assertLessEqual( + np.linalg.norm(np.matmul(a, v) - w * v), + 85 * eps * np.linalg.norm(a), + ) + + @jtu.sample_product( + n=[0, 4, 5, 50, 512], + dtype=float_types + complex_types, + lower=[True, False], + ) + def testEighIdentity(self, n, dtype, lower): + tol = np.finfo(dtype).eps + uplo = "L" if lower else "U" + + a = jnp.eye(n, dtype=dtype) + w, v = jnp.linalg.eigh(a, UPLO=uplo, symmetrize_input=False) + w = w.astype(v.dtype) + self.assertLessEqual( + np.linalg.norm(np.eye(n) - np.matmul(np.conj(T(v)), v)), tol + ) + with jax.numpy_rank_promotion('allow'): + self.assertLessEqual(np.linalg.norm(np.matmul(a, v) - w * v), + tol * np.linalg.norm(a)) + + @jtu.sample_product( + shape=[(4, 4), (5, 5), (50, 50)], + dtype=float_types + complex_types, + ) + def testEigvalsh(self, shape, dtype): + rng = jtu.rand_default(self.rng()) + n = shape[-1] + def args_maker(): + a = rng((n, n), dtype) + a = (a + np.conj(a.T)) / 2 + return [a] + self._CheckAgainstNumpy( + np.linalg.eigvalsh, jnp.linalg.eigvalsh, args_maker, tol=2e-5 + ) + + @jtu.sample_product( + shape=[(1, 1), (4, 4), (5, 5), (25, 25), (2, 10, 10)], + dtype=float_types + complex_types, + lower=[True, False], + ) + def testEighGrad(self, shape, dtype, lower): + if platform.system() == "Windows": + self.skipTest("Skip on Windows due to tolerance issues.") + rng = jtu.rand_default(self.rng()) + a = rng(shape, dtype) + a = (a + np.conj(T(a))) / 2 + ones = np.ones((a.shape[-1], a.shape[-1]), dtype=dtype) + a *= np.tril(ones) if lower else np.triu(ones) + # Gradient checks will fail without symmetrization as the eigh jvp rule + # is only correct for tangents in the symmetric subspace, whereas the + # checker checks against unconstrained (co)tangents. + f = partial(_normalizing_eigh, lower=lower, symmetrize_input=True) + norm_a = jnp.linalg.norm(a) + eps = 2e-5 * norm_a + atol = 5e-3 * norm_a + rtol = 0.025 + jtu.check_grads(f, (a,), 2, atol=atol, rtol=rtol, eps=eps) + + def testEighGradPrecision(self): + rng = jtu.rand_default(self.rng()) + a = rng((3, 3), np.float32) + jtu.assert_dot_precision( + lax.Precision.HIGHEST, partial(jvp, jnp.linalg.eigh), (a,), (a,)) + + def testEighGradRankPromotion(self): + rng = jtu.rand_default(self.rng()) + a = rng((10, 3, 3), np.float32) + jvp(jnp.linalg.eigh, (a,), (a,)) # doesn't crash + + @jtu.sample_product( + shape=[(1, 1), (4, 4), (5, 5), (300, 300)], + dtype=float_types + complex_types, + ) + def testEighBatching(self, shape, dtype): + rng = jtu.rand_default(self.rng()) + shape = (10,) + shape + args = rng(shape, dtype) + args = (args + np.conj(T(args))) / 2 + ws, vs = vmap(jsp.linalg.eigh)(args) + ws = ws.astype(vs.dtype) + norm = np.max(np.linalg.norm(np.matmul(args, vs) - ws[..., None, :] * vs)) + self.assertLess(norm, 1.4e-2) @jtu.sample_product( shape=[(1,), (4,), (5,)], From c21c3b8dc5cddb02d77665bbdb6dd4200f6a86c7 Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Fri, 20 Feb 2026 11:17:09 -0600 Subject: [PATCH 38/50] Abort detection CI workflow (#688) --- .github/workflows/pytest_rocm_abort.yml | 162 ++++++++++++++++ .../wheel_tests_nightly_release_abort.yml | 53 ++++++ ci/run_pytest_rocm_abort.sh | 179 ++++++++++++++++++ conftest.py | 4 - 4 files changed, 394 insertions(+), 4 deletions(-) create mode 100644 .github/workflows/pytest_rocm_abort.yml create mode 100644 .github/workflows/wheel_tests_nightly_release_abort.yml create mode 100755 ci/run_pytest_rocm_abort.sh diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml new file mode 100644 index 000000000000..fcfdb06855d5 --- /dev/null +++ b/.github/workflows/pytest_rocm_abort.yml @@ -0,0 +1,162 @@ +# CI - Pytest ROCm (Abort Support) +# +# This workflow runs the ROCm tests with Pytest in ROCm GHCR containers, +# using the ROCm `pytest-abort` retry wrapper to detect/retry aborts/crashes. +# +# It can be triggered manually via workflow_dispatch or called by other workflows +# via workflow_call. +# +# It consists of the following job: +# run-tests: +# - Runs in ROCm container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest) +# - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from latest release. +# - Executes the `run_pytest_rocm_abort.sh` script, which installs wheel artifacts and +# runs the ROCm tests with Pytest under `pytest-abort-retry`. +name: CI - Pytest ROCm (Abort Support) + +on: + workflow_dispatch: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: choice + default: "linux-x86-64-4gpu-amd" + options: + - "linux-x86-64-1gpu-amd" + - "linux-x86-64-4gpu-amd" + - "linux-x86-64-8gpu-amd" + python: + description: "Which Python version to use?" + type: choice + default: "3.11" + options: + - "3.11" + - "3.12" + rocm-version: + description: "Which ROCm version to test?" + type: choice + default: "7.2.0" + options: + - "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: choice + default: "head" + options: + - "head" + - "pypi_latest" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + type: choice + default: '0' + options: + - '0' + - '1' + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + type: string + default: 'gs://jax-nightly-artifacts/latest' + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + workflow_call: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: string + default: "linux-x86-64-4gpu-amd" + python: + description: "Which Python version to use?" + type: string + default: "3.11" + rocm-version: + description: "Which ROCm version to test?" + type: string + default: "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: string + default: "head" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + default: '0' + type: string + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + default: 'gs://jax-nightly-artifacts/latest' + type: string + +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-tests: + defaults: + run: + # Set the shell to bash as GitHub actions run with /bin/sh by default + shell: bash + runs-on: ${{ inputs.runner }} + continue-on-error: true + # Run in ROCm GHCR container with GPU access + container: + image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings + name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') || + (contains(inputs.runner, '4gpu') && '4gpu') || + (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}" + + env: + JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" + JAXCI_PYTHON: "python${{ inputs.python }}" + JAXCI_ENABLE_X64: "0" + + steps: + - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0 + with: + persist-credentials: false + - name: Download JAX ROCm wheels + uses: ./.github/actions/download-jax-rocm-wheels + with: + python: ${{ inputs.python }} + rocm-version: ${{ inputs.rocm-version }} + jaxlib-version: ${{ inputs.jaxlib-version }} + skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} + gcs_download_uri: ${{ inputs.gcs_download_uri }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Install Python dependencies + run: | + $JAXCI_PYTHON -m pip install uv~=0.5.30 + $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt + # Halt for testing + - name: Wait For Connection + uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c + with: + halt-dispatch-input: ${{ inputs.halt-for-connection }} + - name: Run Pytest ROCm tests (abort support) + timeout-minutes: 180 + run: ./ci/run_pytest_rocm_abort.sh + - name: Upload pytest results to artifact + if: always() + uses: actions/upload-artifact@v4 + with: + name: logs_abort + path: | + logs_abort/ + if-no-files-found: warn + retention-days: 2 + overwrite: true diff --git a/.github/workflows/wheel_tests_nightly_release_abort.yml b/.github/workflows/wheel_tests_nightly_release_abort.yml new file mode 100644 index 000000000000..518306956eff --- /dev/null +++ b/.github/workflows/wheel_tests_nightly_release_abort.yml @@ -0,0 +1,53 @@ +# CI - Wheel Tests (Nightly/Release) (ROCm abort only) +# +# This workflow runs only the ROCm wheel tests using the abort/retry wrapper workflow. +name: CI - Wheel Tests (Nightly/Release) (ROCm abort only) + +on: + workflow_dispatch: + inputs: + gcs_download_uri: + description: "GCS location URI from where the artifacts should be downloaded" + required: true + default: 'gs://jax-nightly-artifacts/latest' + type: string + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)" + required: true + default: '0' + type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection? (yes/no)' + required: false + default: 'no' + type: string + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: true +permissions: {} + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-pytest-rocm: + uses: ./.github/workflows/pytest_rocm_abort.yml + strategy: + fail-fast: false # don't cancel all jobs on failure + matrix: + runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] + python: ["3.11", "3.12", "3.13", "3.14"] + rocm: [ + {version: "7.2.0", tag: "rocm720"}, + ] + name: "Pytest ROCm abort (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" + with: + runner: ${{ matrix.runner }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.version }} + rocm-tag: ${{ matrix.rocm.tag }} + jaxlib-version: "head" + skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} + gcs_download_uri: ${{inputs.gcs_download_uri}} + halt-for-connection: ${{inputs.halt-for-connection}} diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh new file mode 100755 index 000000000000..f17e7184b3bf --- /dev/null +++ b/ci/run_pytest_rocm_abort.sh @@ -0,0 +1,179 @@ +#!/bin/bash +# Copyright 2024 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Runs Pytest ROCm tests (with ROCm pytest-abort retry wrapper). +# Requires the jaxlib and ROCm plugin wheels to be present inside $JAXCI_OUTPUT_DIR (../dist) +# +# -e: abort script if one command fails +# -u: error if undefined variable used +# -x: log all commands +# -o history: record shell history +# -o allexport: export all functions and variables to be available to subscripts +set -exu -o history -o allexport + +source ci/envs/default.env + +# Install jaxlib and ROCm plugin wheels inside the $JAXCI_OUTPUT_DIR directory +echo "Installing wheels locally..." +source ./ci/utilities/install_wheels_locally.sh + +# Print all the installed packages +echo "Installed packages:" +"$JAXCI_PYTHON" -m uv pip freeze + +"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))" + +rocm-smi + +# ============================================================================== +# Set up the generic test environment variables +# ============================================================================== +export PY_COLORS=1 +export JAX_SKIP_SLOW_TESTS=true +export NCCL_DEBUG=WARN +export TF_CPP_MIN_LOG_LEVEL=0 +export JAX_ENABLE_X64="$JAXCI_ENABLE_X64" + +# ============================================================================== +# Calculate the optimal number of parallel processes for pytest +# This will be the minimum of: GPU capacity, CPU core count, and a system RAM limit. +# ============================================================================== + +export gpu_count=$(rocminfo | egrep -c "Device Type:\\s+GPU") +echo "Number of GPUs detected: $gpu_count" + +# Query GPU 0 memory using rocm-smi +export memory_per_gpu_mib=$(rocm-smi -d 0 --showmeminfo vram | grep -i "vram total" | awk '{print int($NF/1024/1024)}' | head -1) +echo "Reported memory per GPU: $memory_per_gpu_mib MiB" + +# Convert effective memory from MiB to GiB. +export memory_per_gpu_gib=$((memory_per_gpu_mib / 1024)) +echo "Effective memory per GPU: $memory_per_gpu_gib GiB" + +# Allow 2 GiB of GPU RAM per test. +export max_tests_per_gpu=$((memory_per_gpu_gib / 2)) +echo "Max tests per GPU (assuming 2GiB/test): $max_tests_per_gpu" + +export num_processes=$((gpu_count * max_tests_per_gpu)) +echo "Initial number of processes based on GPU capacity: $num_processes" + +export num_cpu_cores=$(nproc) +echo "Number of CPU cores available: $num_cpu_cores" + +# Reads total memory from /proc/meminfo (in KiB) and converts to GiB. +export total_ram_gib=$(awk '/MemTotal/ {printf \"%.0f\", $2/1048576}' /proc/meminfo) +echo "Total system RAM: $total_ram_gib GiB" + +# Set a safety limit for system RAM usage, e.g., 1/6th of total. +export host_memory_limit=$((total_ram_gib / 6)) +echo "Host memory process limit (1/6th of total RAM): $host_memory_limit" + +if [[ $num_cpu_cores -lt $num_processes ]]; then + num_processes=$num_cpu_cores + echo "Adjusting num_processes to match CPU core count: $num_processes" +fi + +if [[ $host_memory_limit -lt $num_processes ]]; then + num_processes=$host_memory_limit + echo "Adjusting num_processes to match host memory limit: $num_processes" +fi + +if [[ 16 -lt $num_processes ]]; then + num_processes=16 + echo "Reducing num_processes to $num_processes" +fi + +echo "Final number of processes to run: $num_processes" + +export JAX_ENABLE_ROCM_XDIST="$gpu_count" +export XLA_PYTHON_CLIENT_ALLOCATOR=platform +export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1 --xla_gpu_enable_nccl_comm_splitting=false --xla_gpu_enable_command_buffer=" + +# Disable core dumps just in case +ulimit -c 0 + +# Keep deselected tests in one place for the abort wrapper. +ROCM_PYTEST_DESELECT_ARGS=( + --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data + --deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices + --deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric +) + +# --max-runs: retry the entire pytest run up to N times on abort/crash. +# --max-worker-restart: restart crashed xdist workers up to N times. +# --maxfail: stop the run after N test failures. +rocm_test_cmd() { + local abort_flag="${1:-0}" + shift + if [[ "$abort_flag" == "1" ]]; then + pytest-abort-retry --max-runs 3 --clear-crash-log -- "$@" + else + "$@" + fi +} + +rocm_log_tail_on_failure() { + local logfile="$1" + local status="$2" + if [[ "$status" -ne 0 ]]; then + echo "Pytest failed (exit=$status). Showing last 200 lines of $logfile:" + tail -n 200 "$logfile" || true + else + echo "Pytest output saved to $logfile (uploaded as artifact)." + fi +} + +rocm_install_extra_requirements() { + if [[ -n "${GITHUB_WORKSPACE:-}" ]]; then + cd "$GITHUB_WORKSPACE" + fi + + # Install extra requirements. + "$JAXCI_PYTHON" -m uv pip install pytest-timeout pytest-html pytest-csv pytest-json-report pytest-abort +} + +rocm_install_extra_requirements + +echo "Running ROCm tests (with abort/retry wrapper)..." +mkdir -p logs_abort +logfile="logs_abort/jax_ToT_UT_abort.log" + +# pytest-abort output directories (must be set before running pytest). +export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running" +export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl" +mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR" + +set +e +rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart=200 --tb=short --timeout=1200 --timeout-method=thread tests \ + "${ROCM_PYTEST_DESELECT_ARGS[@]}" \ + --json-report \ + --json-report-file=logs_abort/tests-report-abort.json \ + --csv=logs_abort/tests-report-abort.csv \ + --html=logs_abort/tests-report-abort.html \ + --self-contained-html \ + >"$logfile" 2>&1 +pytest_status=$? +set -e +rocm_log_tail_on_failure "$logfile" "$pytest_status" + +echo "Postprocessing reports with crashed tests..." +pytest-abort-postprocess \ + --crash-log "$PYTEST_ABORT_CRASHED_TESTS_LOG" \ + --json-report logs_abort/tests-report-abort.json \ + --html-report logs_abort/tests-report-abort.html \ + --csv-report logs_abort/tests-report-abort.csv \ + >>"$logfile" 2>&1 + +exit "$pytest_status" diff --git a/conftest.py b/conftest.py index 82d5273e2cdd..8f905bf53af8 100644 --- a/conftest.py +++ b/conftest.py @@ -15,10 +15,6 @@ import os import pytest -import json -import threading -import shutil -from datetime import datetime @pytest.fixture(autouse=True) def add_imports(doctest_namespace): From 3f5828e9543f114a192adf8000f2e311086df0cf Mon Sep 17 00:00:00 2001 From: Gulsum Gudukbay Akbulut Date: Tue, 24 Feb 2026 09:59:51 -0600 Subject: [PATCH 39/50] Abort-Detection: Fix halt-for-connection input (#712) --- .github/workflows/pytest_rocm_abort.yml | 13 +++++++++++++ ci/run_pytest_rocm_abort.sh | 7 +++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pytest_rocm_abort.yml b/.github/workflows/pytest_rocm_abort.yml index fcfdb06855d5..854b6681978b 100644 --- a/.github/workflows/pytest_rocm_abort.yml +++ b/.github/workflows/pytest_rocm_abort.yml @@ -64,6 +64,10 @@ on: description: 'Should this workflow run wait for a remote connection?' type: string default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' workflow_call: inputs: runner: @@ -94,6 +98,14 @@ on: description: "GCS location prefix from where the artifacts should be downloaded" default: 'gs://jax-nightly-artifacts/latest' type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + max-worker-restart: + description: "Max xdist worker restarts (passed to pytest --max-worker-restart)" + type: string + default: '50' permissions: {} @@ -123,6 +135,7 @@ jobs: JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" JAXCI_PYTHON: "python${{ inputs.python }}" JAXCI_ENABLE_X64: "0" + MAX_WORKER_RESTART: "${{ inputs['max-worker-restart'] }}" steps: - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0 diff --git a/ci/run_pytest_rocm_abort.sh b/ci/run_pytest_rocm_abort.sh index f17e7184b3bf..6c6ab6d18281 100755 --- a/ci/run_pytest_rocm_abort.sh +++ b/ci/run_pytest_rocm_abort.sh @@ -73,7 +73,7 @@ export num_cpu_cores=$(nproc) echo "Number of CPU cores available: $num_cpu_cores" # Reads total memory from /proc/meminfo (in KiB) and converts to GiB. -export total_ram_gib=$(awk '/MemTotal/ {printf \"%.0f\", $2/1048576}' /proc/meminfo) +export total_ram_gib=$(awk '/MemTotal/ {printf "%.0f", $2/1048576}' /proc/meminfo) echo "Total system RAM: $total_ram_gib GiB" # Set a safety limit for system RAM usage, e.g., 1/6th of total. @@ -150,13 +150,16 @@ echo "Running ROCm tests (with abort/retry wrapper)..." mkdir -p logs_abort logfile="logs_abort/jax_ToT_UT_abort.log" +# Allow the workflow to override worker restart limit. +max_worker_restart="${MAX_WORKER_RESTART:-50}" + # pytest-abort output directories (must be set before running pytest). export PYTEST_ABORT_LAST_RUNNING_DIR="logs_abort/last_running" export PYTEST_ABORT_CRASHED_TESTS_LOG="logs_abort/crashed_tests.jsonl" mkdir -p "$PYTEST_ABORT_LAST_RUNNING_DIR" set +e -rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart=200 --tb=short --timeout=1200 --timeout-method=thread tests \ +rocm_test_cmd 1 "$JAXCI_PYTHON" -m pytest -n "$num_processes" --max-worker-restart="$max_worker_restart" --tb=short --timeout=1200 --timeout-method=thread tests \ "${ROCM_PYTEST_DESELECT_ARGS[@]}" \ --json-report \ --json-report-file=logs_abort/tests-report-abort.json \ From 6283bbcf25f76377592987e99744eeec9c872f20 Mon Sep 17 00:00:00 2001 From: Yanyao Wang Date: Tue, 17 Mar 2026 17:40:53 -0500 Subject: [PATCH 40/50] fix: add rocm_sysdeps/lib to wheel RUNPATH (#737) Signed-off-by: Yanyao Wang --- jaxlib/rocm/rocm_rpath.bzl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl index 354082c0d200..1da2d640ca42 100644 --- a/jaxlib/rocm/rocm_rpath.bzl +++ b/jaxlib/rocm/rocm_rpath.bzl @@ -35,8 +35,11 @@ _ROCM_LINK_ONLY = "@local_config_rocm//rocm:link_only" _WHEEL_RPATHS = [ "-Wl,-rpath,$$ORIGIN/../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,$$ORIGIN/../../rocm/lib", + "-Wl,-rpath,$$ORIGIN/../../rocm/lib/rocm_sysdeps/lib", "-Wl,-rpath,/opt/rocm/lib", + "-Wl,-rpath,/opt/rocm/lib/rocm_sysdeps/lib", ] def _wheel_features(): From 2422c280b386c9ea4e4308d42c1332cee69c9798 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Wed, 8 Apr 2026 08:58:12 -0500 Subject: [PATCH 41/50] Temporarily disable the cron trigger for the cont. wheel tests workflow --- .github/workflows/wheel_tests_continuous.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index ef5fb02f7ced..c043a8c80eb3 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -34,8 +34,8 @@ permissions: actions: read on: - schedule: - - cron: "0 */3 * * *" # Run once every 3 hours + # schedule: + # - cron: "0 */3 * * *" # Run once every 3 hours workflow_dispatch: # allows triggering the workflow run manually concurrency: From 8d4fbef1799c41bf491f59d0fb71e52381e066de Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Sun, 10 May 2026 14:26:51 -0500 Subject: [PATCH 42/50] Add placeholder for nightly benchmark workflow (#768) --- .../wheel_benchmarks_nightly_release.yml | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 .github/workflows/wheel_benchmarks_nightly_release.yml diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml new file mode 100644 index 000000000000..d0a31a7cefe5 --- /dev/null +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -0,0 +1,23 @@ +name: CI - Wheel Tests (Nightly/Release) + +on: + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test-runner: + runs-on: linux-x86-64-8gpu-amd + + strategy: + matrix: + python-version: ["3.12"] + rocm-version: ["7.2.0"] + + steps: + - name: Test runner label + run: | + echo "python=${{ matrix.python-version }}" + echo "rocm=${{ matrix.rocm-version }}" From eb58ba290c8a71730e8259ef41e291d4f709247a Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Sun, 10 May 2026 14:40:43 -0500 Subject: [PATCH 43/50] Rename the benchmarks workflow consistently (#770) --- .github/workflows/wheel_benchmarks_nightly_release.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml index d0a31a7cefe5..95e6e3ce19cc 100644 --- a/.github/workflows/wheel_benchmarks_nightly_release.yml +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -1,4 +1,4 @@ -name: CI - Wheel Tests (Nightly/Release) +name: CI - Wheel Benchmarks (Nightly/Release) on: workflow_dispatch: From 2483b469b9a59e73e8d928860f36f9f911819585 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Thu, 23 Apr 2026 17:28:24 +0000 Subject: [PATCH 44/50] Add ROCm benchmark workflow for MaxText --- .github/workflows/benchmark_rocm.yml | 203 ++++++++++++++++++ .../wheel_benchmarks_nightly_release.yml | 63 ++++-- .../maxtext_rocm/cmp_maxtext_rocm.py | 85 ++++++++ .../maxtext_rocm/exp_maxtext_rocm.yml | 4 + .../maxtext_rocm/run_maxtext_rocm.sh | 94 ++++++++ ci/collect_ci_manifest_rocm.py | 141 ++++++++++++ ci/upload_rocm_artifacts.sh | 29 +++ 7 files changed, 606 insertions(+), 13 deletions(-) create mode 100644 .github/workflows/benchmark_rocm.yml create mode 100644 ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py create mode 100644 ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml create mode 100755 ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh create mode 100644 ci/collect_ci_manifest_rocm.py create mode 100644 ci/upload_rocm_artifacts.sh diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml new file mode 100644 index 000000000000..0458fe2a2b89 --- /dev/null +++ b/.github/workflows/benchmark_rocm.yml @@ -0,0 +1,203 @@ +# CI - Benchmark ROCm +# +# This workflow runs ROCm benchmarks in ROCm team's GHCR containers. +# It can be triggered manually via workflow_dispatch or called by other workflows +# via workflow_call. +# +# It consists of the following job: +# run-benchmark: +# - Runs in ROCm team's container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest) +# - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from S3. +# - Executes the target benchmark runner script at `targets//run.sh`. +name: CI - Benchmark ROCm +on: + workflow_dispatch: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: choice + default: "linux-x86-64-8gpu-amd" + options: + - "linux-x86-64-1gpu-amd" + - "linux-x86-64-4gpu-amd" + - "linux-x86-64-8gpu-amd" + python: + description: "Which Python version to use?" + type: choice + default: "3.12" + options: + - "3.11" + - "3.12" + rocm-version: + description: "Which ROCm version to benchmark?" + type: choice + default: "7.2.0" + options: + - "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + target: + description: "Benchmark target" + type: choice + default: "maxtext" + options: + - "maxtext" + workload: + description: "Benchmark workload" + type: string + default: "llama3_8b" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: choice + default: "head" + options: + - "head" + - "pypi_latest" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + type: choice + default: '0' + options: + - '0' + - '1' + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + type: string + default: 'gs://jax-nightly-artifacts/latest' + s3_download_uri: + description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" + type: string + default: 'latest' + use-te: + description: "Install Transformer Engine" + type: choice + default: '1' + options: + - '0' + - '1' + halt-for-connection: + description: 'Should this workflow run wait for a remote connection?' + type: string + default: 'no' + workflow_call: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: string + default: "linux-x86-64-8gpu-amd" + python: + description: "Which Python version to use?" + type: string + default: "3.12" + rocm-version: + description: "Which ROCm version to benchmark?" + type: string + default: "7.2.0" + rocm-tag: + description: "ROCm tag for container image (e.g., rocm720)" + type: string + default: "rocm720" + target: + description: "Benchmark target" + type: string + default: "maxtext" + workload: + description: "Benchmark workload" + type: string + default: "llama3_8b" + jaxlib-version: + description: "Which jaxlib version to use? (head/pypi_latest)" + type: string + default: "head" + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to skip downloading the jaxlib and plugins from GCS (e.g for testing a jax only release)" + type: string + default: '0' + gcs_download_uri: + description: "GCS location prefix from where the artifacts should be downloaded" + type: string + default: 'gs://jax-nightly-artifacts/latest' + s3_download_uri: + description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" + type: string + default: 'latest' + use-te: + description: "Install Transformer Engine" + type: string + default: '1' +permissions: + id-token: write + contents: read + +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" + +jobs: + run-benchmark: + defaults: + run: + # Set the shell to bash as GitHub actions run with /bin/sh by default + shell: bash + runs-on: ${{ inputs.runner }} + continue-on-error: true + # Run in ROCm team's GHCR container with GPU access + container: + image: ghcr.io/rocm/jax-base-ubu24.rocm722:2d65281b00de2bcafc811247d563b6e5e7c887af + credentials: + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --shm-size 64G --env-file /etc/podinfo/gha-gpu-isolation-settings + name: "${{ (contains(inputs.runner, '1gpu') && '1gpu') || + (contains(inputs.runner, '4gpu') && '4gpu') || + (contains(inputs.runner, '8gpu') && '8gpu') }}, ROCm ${{ inputs.rocm-version }}, py${{ inputs.python }}" + + env: + JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" + JAXCI_PYTHON: "python${{ inputs.python }}" + JAXCI_ENABLE_X64: "0" + USE_TE: "${{ inputs.use-te }}" + INPUT_ROCM_VERSION: "${{ inputs.rocm-version }}" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - name: Download JAX ROCm wheels + uses: ./.github/actions/download-jax-rocm-wheels + with: + python: ${{ inputs.python }} + rocm-version: ${{ inputs.rocm-version }} + jaxlib-version: ${{ inputs.jaxlib-version }} + skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }} + gcs_download_uri: ${{ inputs.gcs_download_uri }} + s3_download_uri: ${{ inputs.s3_download_uri }} + - name: Install Python dependencies + run: | + $JAXCI_PYTHON -m pip install uv~=0.11.2 + $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt + $JAXCI_PYTHON -m pip install pytest-json-report + # Halt for testing + - name: Wait For Connection + uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c + with: + halt-dispatch-input: ${{ inputs.halt-for-connection }} + - name: Run ROCm benchmarks + env: + TARGET: ${{ inputs.target }} + WORKLOAD: ${{ inputs.workload }} + timeout-minutes: 120 + run: ./ci/benchmark_targets/${TARGET}/run_${TARGET}_rocm.sh" --workload "${WORKLOAD}" + - name: Upload GitHub artifacts + if: always() + continue-on-error: true + uses: actions/upload-artifact@v4 + with: + name: benchmark-artifacts-${{ inputs.target }}-${{ inputs.workload }}-${{ inputs.runner }}-py${{ inputs.python }}-rocm${{ inputs.rocm-version }} + path: run_artifacts/result.json + if-no-files-found: warn + - name: Upload CI artifacts to S3 + if: always() + env: + S3_BUCKET_NAME: jax-ci-amd + run: ./ci/upload_artifacts.sh run_artifacts \ No newline at end of file diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml index 95e6e3ce19cc..069881761964 100644 --- a/.github/workflows/wheel_benchmarks_nightly_release.yml +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -2,22 +2,59 @@ name: CI - Wheel Benchmarks (Nightly/Release) on: workflow_dispatch: + inputs: + gcs_download_uri: + description: "GCS location URI from where the artifacts should be downloaded" + required: true + default: 'gs://jax-nightly-artifacts/latest' + type: string + skip-download-jaxlib-and-plugins-from-gcs: + description: "Whether to download only the jax wheel from GCS (e.g for testing a jax only release)" + required: true + default: '0' + type: string + halt-for-connection: + description: 'Should this workflow run wait for a remote connection? (yes/no)' + required: false + default: 'no' + type: string concurrency: - group: ${{ github.workflow }}-${{ github.ref }} + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: true +permissions: {} -jobs: - test-runner: - runs-on: linux-x86-64-8gpu-amd +env: + UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" +jobs: + run-benchmark-rocm: + permissions: + id-token: write + contents: read + checks: write + uses: ./.github/workflows/benchmark_rocm.yml strategy: - matrix: - python-version: ["3.12"] - rocm-version: ["7.2.0"] - - steps: - - name: Test runner label - run: | - echo "python=${{ matrix.python-version }}" - echo "rocm=${{ matrix.rocm-version }}" + fail-fast: false # don't cancel all jobs on failure + matrix: + runner: ["linux-x86-64-8gpu-amd"] + python: ["3.12"] + rocm: [ + {version: "7.2.0", tag: "rocm720"}, + ] + target: ["maxtext"] + workload: ["gemma_4b"] + name: "Benchmark ROCm (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" + with: + runner: ${{ matrix.runner }} + python: ${{ matrix.python }} + rocm-version: ${{ matrix.rocm.version }} + rocm-tag: ${{ matrix.rocm.tag }} + target: ${{ matrix.target }} + workload: ${{ matrix.workload }} + jaxlib-version: "head" + skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} + gcs_download_uri: ${{ inputs.gcs_download_uri }} + s3_download_uri: ${{ inputs.s3_download_uri }} + use-te: "1" + te-wheel-url: "https://github.com/ROCm/maxtext/releases/download/te-rocm-wheels-2026-05-04-86438dc3d04e/transformer_engine-2.12.0.dev0+86438dc3-1.mi355-cp312-cp312-linux_x86_64.whl" diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py new file mode 100644 index 000000000000..e2d6a40ad75e --- /dev/null +++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +import argparse +import json +import statistics +from pathlib import Path + +METRIC_FIELD_INDEX = 3 +WARMUP_STEPS = 2 + + +def read(path): + return ( + Path(path).read_text(errors="replace") if path and Path(path).exists() else "" + ) + + +def parse_metric_values(log_text): + values = [] + for line in log_text.splitlines(): + if "completed step:" not in line: + continue + fields = line.split(",") + if len(fields) <= METRIC_FIELD_INDEX: + continue + try: + values.append(float(fields[METRIC_FIELD_INDEX].split(":", 1)[1].strip())) + except Exception: + pass + return values + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--log", required=True) + p.add_argument("--expected", required=True) + p.add_argument("--config", required=True) + p.add_argument("--env-file", default="") + p.add_argument("--requirements", required=True) + p.add_argument("--target", required=True) + p.add_argument("--workload", required=True) + p.add_argument("--run-code", required=True) + p.add_argument("--model-run-started-at", required=True) + p.add_argument("--model-run-completed-at", required=True) + p.add_argument("--out", required=True) + args = p.parse_args() + + expected = json.loads(Path(args.expected).read_text()) + values = parse_metric_values(read(args.log)) + + samples = values[WARMUP_STEPS:] + observed = statistics.median(samples) if samples else None + + baseline = float(expected["baseline"]) + threshold = float(expected["threshold_percent"]) + + distance = None + cmp_code = 1 + + if int(args.run_code) == 0 and observed is not None and baseline != 0: + raw = ((observed - baseline) / baseline) * 100.0 + distance = abs(raw) + cmp_code = 0 if raw >= -threshold else 1 + + result = { + "benchmark_schema_version": 1, + "kind": "benchmark", + "target": args.target, + "workload": args.workload, + "run_code": int(args.run_code), + "cmp_code": cmp_code, + "distance_percent": round(distance, 4) if distance is not None else None, + "model_run_started_at": args.model_run_started_at, + "model_run_completed_at": args.model_run_completed_at, + "workload_config_raw": read(args.config), + "workload_env_raw": read(args.env_file), + "requirements_raw": read(args.requirements), + "expected_config_raw": read(args.expected), + } + + Path(args.out).write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + return cmp_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml new file mode 100644 index 000000000000..ca1bacc5e890 --- /dev/null +++ b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml @@ -0,0 +1,4 @@ +{ + "baseline_ms": 2600.0, + "threshold_percent": 5.0 +} \ No newline at end of file diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh new file mode 100755 index 000000000000..f650fdd0eabf --- /dev/null +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +set -euo pipefail + +WORKLOAD="gemma3-4b" +[[ $# -eq 0 || ( $# -eq 2 && "$1" == "--workload" ) ]] || { + echo "usage: $0 [--workload NAME]" >&2 + exit 2 +} +[[ $# -eq 2 ]] && WORKLOAD="$2" + +ROOT="${ROOT:-$PWD}" +TARGET_DIR="${ROOT}/benchmark_targets/rocm_maxtext" +RUN_DIR="${ROOT}/run_artifacts" +REPO_DIR="${ROOT}/maxtext" +WORK_DIR="${REPO_DIR}/src" + +CFG_FILE="${REPO_DIR}/configs/models/${WORKLOAD}.yml" +ENV_FILE="${REPO_DIR}/configs/models/${WORKLOAD}.env.sh" +REQ_FILE="${WORK_DIR}/dependencies/requirements/requirements_rocm_jax_0.8.2.txt" +EXPECTED_FILE="${TARGET_DIR}/expected.json" + +RUN_LOG="${RUN_DIR}/maxtext.log" +BENCHMARK_JSON="${RUN_DIR}/benchmark.json" +RESULT_JSON="${RUN_DIR}/result.json" + +PYTHON_BIN="${JAXCI_PYTHON:-python3}" + +mkdir -p "${RUN_DIR}" + +source ci/envs/default.env +source ./ci/utilities/install_wheels_locally.sh + +[[ -d "${REPO_DIR}/.git" ]] || \ + git clone --depth 1 --branch main https://github.com/ROCm/maxtext.git "${REPO_DIR}" + +[[ -f "${CFG_FILE}" ]] || { echo "missing config file: ${CFG_FILE}" >&2; exit 2; } +[[ -f "${EXPECTED_FILE}" ]] || { echo "missing expected file: ${EXPECTED_FILE}" >&2; exit 2; } + +[[ -f "${REQ_FILE}" ]] && { + echo "[setup] installing MaxText requirements" + "${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" +} + +[[ -f "${ENV_FILE}" ]] && source "${ENV_FILE}" + +[[ "${USE_TE:-0}" == "1" ]] && { + : "${TE_WHEEL_URL:?TE_WHEEL_URL must be set when USE_TE=1}" + echo "[setup] installing Transformer Engine" + "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}" +} + +export PY_COLORS=1 +export NCCL_DEBUG=WARN +export TF_CPP_MIN_LOG_LEVEL=0 +export JAX_ENABLE_X64="${JAXCI_ENABLE_X64}" +export XLA_PYTHON_CLIENT_ALLOCATOR=platform +export XLA_PYTHON_CLIENT_PREALLOCATE=false + +MODEL_RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +set +e +pushd "${WORK_DIR}" >/dev/null +"${PYTHON_BIN}" -m maxtext.trainers.pre_train.train \ + "$(realpath "${CFG_FILE}")" \ +> "${RUN_LOG}" 2>&1 +RUN_CODE=$? +popd >/dev/null +set -e + +MODEL_RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +CMP_CODE=1 +"${PYTHON_BIN}" "${TARGET_DIR}/benchmark_result.py" \ + --log "${RUN_LOG}" \ + --expected "${EXPECTED_FILE}" \ + --config "${CFG_FILE}" \ + --env-file "${ENV_FILE}" \ + --requirements "${REQ_FILE}" \ + --target rocm_maxtext \ + --workload "${WORKLOAD}" \ + --run-code "${RUN_CODE}" \ + --model-run-started-at "${MODEL_RUN_STARTED_AT}" \ + --model-run-completed-at "${MODEL_RUN_COMPLETED_AT}" \ + --out "${BENCHMARK_JSON}" || CMP_CODE=$? + +"${PYTHON_BIN}" ci/make_ci_manifest.py \ + --extra "${BENCHMARK_JSON}" \ + --out "${RESULT_JSON}" + +rm -f "${RUN_LOG}" "${BENCHMARK_JSON}" + +[[ -s "${RESULT_JSON}" ]] && touch "${RUN_DIR}/_SUCCESS" + +exit $(( RUN_CODE != 0 || CMP_CODE != 0 )) \ No newline at end of file diff --git a/ci/collect_ci_manifest_rocm.py b/ci/collect_ci_manifest_rocm.py new file mode 100644 index 000000000000..77b010ece756 --- /dev/null +++ b/ci/collect_ci_manifest_rocm.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +import re +import subprocess +from pathlib import Path + + +def env(name, default="unknown"): + return os.environ.get(name, default) + + +def cmd(args): + try: + return subprocess.check_output( + args, text=True, stderr=subprocess.DEVNULL + ).strip() + except Exception: + return "" + + +def sh(command): + return cmd(["bash", "-lc", command]) + + +def github_run_started_at(): + repo = env("GITHUB_REPOSITORY", "") + run_id = env("GITHUB_RUN_ID", "") + raw = cmd( + [ + "curl", + "-fsSL", + "-H", + "Accept: application/vnd.github+json", + f"https://api.github.com/repos/{repo}/actions/runs/{run_id}", + ] + ) + try: + return json.loads(raw).get("run_started_at", "") + except Exception: + return "" + + +def gpu_count(runner): + match = re.search(r"([0-9]+)gpu", runner or "") + return int(match.group(1)) if match else None + + +def image_digest(rocm_tag): + repo = f"rocm/jax-base-ubu24.{rocm_tag}" + token_raw = cmd( + [ + "curl", + "-fsSL", + f"https://ghcr.io/token?service=ghcr.io&scope=repository:{repo}:pull", + ] + ) + try: + token = json.loads(token_raw).get("token", "") + except Exception: + return "" + + headers = sh( + "curl -fsSL -D - " + f"-H 'Authorization: Bearer {token}' " + "-H 'Accept: application/vnd.docker.distribution.manifest.v2+json' " + f"https://ghcr.io/v2/{repo}/manifests/latest " + "-o /dev/null" + ) + for line in headers.splitlines(): + if line.lower().startswith("docker-content-digest:"): + return line.split(":", 1)[1].strip() + return "" + + +def package_snapshot(): + python = env("JAXCI_PYTHON", "python3") + pkgs = cmd([python, "-m", "pip", "list", "--format=freeze"]) + return "|".join( + line + for line in pkgs.splitlines() + if re.search(r"^(jax|jaxlib)==|pjrt|plugin|transformer[-_]engine", line) + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--extra", action="append", default=[]) + parser.add_argument("--out", required=True) + args = parser.parse_args() + + started = github_run_started_at() + date = started.split("T", 1)[0] if started else sh("date -u +%F") + + repo = env("GITHUB_REPOSITORY") + run_id = env("GITHUB_RUN_ID") + attempt = env("GITHUB_RUN_ATTEMPT") + workflow = env("GITHUB_WORKFLOW", "") + runner = env("INPUT_RUNNER") + rocm_tag = env("INPUT_ROCM_TAG") + + manifest = { + "schema_version": 1, + "run_key": f"{date}_{run_id}_{attempt}", + "run_started_at": started, + "run_completed_at": sh("date -u +%Y-%m-%dT%H:%M:%SZ"), + "github_run_url": f"https://github.com/{repo}/actions/runs/{run_id}", + "github_repository": repo, + "github_ref_name": env("GITHUB_REF_NAME"), + "github_ref": env("GITHUB_REF"), + "github_sha": env("GITHUB_SHA"), + "github_event_name": env("GITHUB_EVENT_NAME"), + "github_run_id": run_id, + "github_run_attempt": attempt, + "github_run_number": env("GITHUB_RUN_NUMBER"), + "github_workflow": workflow, + "is_nightly": "nightly" if "nightly" in workflow.lower() else "continuous", + "github_job": env("GITHUB_JOB"), + "python_version": env("INPUT_PYTHON", env("JAXCI_HERMETIC_PYTHON_VERSION")), + "rocm_version": env("INPUT_ROCM_VERSION"), + "rocm_tag": rocm_tag, + "gpu_count": gpu_count(runner), + "runner": runner, + "base_image_name": f"ghcr.io/rocm/jax-base-ubu24.{rocm_tag}:latest", + "base_image_digest": image_digest(rocm_tag), + "jax_packages_raw": package_snapshot(), + "wheels_sha_raw": sh("sha256sum dist/*.whl 2>/dev/null || true").replace( + "\n", "|" + ), + } + + for extra in args.extra: + if extra and Path(extra).exists(): + manifest.update(json.loads(Path(extra).read_text())) + + Path(args.out).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + + +if __name__ == "__main__": + main() diff --git a/ci/upload_rocm_artifacts.sh b/ci/upload_rocm_artifacts.sh new file mode 100644 index 000000000000..0778697cdc3f --- /dev/null +++ b/ci/upload_rocm_artifacts.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +: "${S3_BUCKET_NAME:?}" +ARTIFACT_DIR="${1:-run_artifacts}" +RESULT_FILE="${ARTIFACT_DIR}/result.json" + +[[ -f "${RESULT_FILE}" ]] || { + echo "missing result.json" >&2 + exit 2 +} + +RUN_KEY="${DATE}_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" +COMBO="py$(norm "${INPUT_PYTHON}")-rocm$(norm "${INPUT_ROCM_VERSION}")-${GPU_PART}" +PREFIX="${GITHUB_REPOSITORY}/${GITHUB_REF_NAME}/${IS_NIGHTLY}/${RUN_KEY}/${COMBO}" + +DEST="s3://${S3_BUCKET_NAME}/${TEST_LOGS_ROOT}/${PREFIX}" + +echo "[upload] ${DEST}" + +aws s3 cp --only-show-errors "${RESULT_FILE}" "${DEST}/result.json" + +if [[ -f "${ARTIFACT_DIR}/logs.tar.gz" ]]; then + aws s3 cp --only-show-errors "${ARTIFACT_DIR}/logs.tar.gz" "${DEST}/logs.tar.gz" +fi + +printf '' | aws s3 cp --only-show-errors - "${DEST}/_SUCCESS" + +echo "[done]" \ No newline at end of file From 1d25b87d02a174009ffd153a900dc8c025087b4e Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Sun, 10 May 2026 21:49:57 +0000 Subject: [PATCH 45/50] Resolve latest MaxText Transformer Engine wheel from ROCm MaxText releases --- .github/workflows/benchmark_rocm.yml | 12 +++++----- .../wheel_benchmarks_nightly_release.yml | 8 +++++-- .../maxtext_rocm/run_maxtext_rocm.sh | 24 +++++++++++++++---- 3 files changed, 32 insertions(+), 12 deletions(-) diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml index 0458fe2a2b89..bc3c648e0c26 100644 --- a/.github/workflows/benchmark_rocm.yml +++ b/.github/workflows/benchmark_rocm.yml @@ -1,14 +1,14 @@ # CI - Benchmark ROCm # -# This workflow runs ROCm benchmarks in ROCm team's GHCR containers. +# This workflow runs the ROCm benchmarks in ROCm team's GHCR containers. # It can be triggered manually via workflow_dispatch or called by other workflows # via workflow_call. # # It consists of the following job: -# run-benchmark: +# run-benchmarks: # - Runs in ROCm team's container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest) # - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from S3. -# - Executes the target benchmark runner script at `targets//run.sh`. +# - Executes the target benchmark scripts at `targets//run.sh`. name: CI - Benchmark ROCm on: workflow_dispatch: @@ -135,7 +135,7 @@ env: UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple" jobs: - run-benchmark: + run-benchmarks: defaults: run: # Set the shell to bash as GitHub actions run with /bin/sh by default @@ -144,7 +144,7 @@ jobs: continue-on-error: true # Run in ROCm team's GHCR container with GPU access container: - image: ghcr.io/rocm/jax-base-ubu24.rocm722:2d65281b00de2bcafc811247d563b6e5e7c887af + image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest # zizmor: ignore[unpinned-images] credentials: username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} @@ -187,7 +187,7 @@ jobs: TARGET: ${{ inputs.target }} WORKLOAD: ${{ inputs.workload }} timeout-minutes: 120 - run: ./ci/benchmark_targets/${TARGET}/run_${TARGET}_rocm.sh" --workload "${WORKLOAD}" + run: ./ci/benchmark_targets/${TARGET}_rocm/run_${TARGET}_rocm.sh --workload "${WORKLOAD}" - name: Upload GitHub artifacts if: always() continue-on-error: true diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml index 069881761964..9b4989229bd2 100644 --- a/.github/workflows/wheel_benchmarks_nightly_release.yml +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -1,3 +1,9 @@ +# CI - Wheel Tests (Nightly/Release) +# +# This workflow is used to benchmark the JAX wheels that were built by internal CI jobs. +# +# 1. run-benchmark-rocm: Calls the `benchmark_rocm.yml` workflow which downloads the JAX wheels and +# ROCm plugins from S3 and runs ROCm benchmarks. name: CI - Wheel Benchmarks (Nightly/Release) on: @@ -55,6 +61,4 @@ jobs: jaxlib-version: "head" skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} gcs_download_uri: ${{ inputs.gcs_download_uri }} - s3_download_uri: ${{ inputs.s3_download_uri }} use-te: "1" - te-wheel-url: "https://github.com/ROCm/maxtext/releases/download/te-rocm-wheels-2026-05-04-86438dc3d04e/transformer_engine-2.12.0.dev0+86438dc3-1.mi355-cp312-cp312-linux_x86_64.whl" diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index f650fdd0eabf..e68b987e97af 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -43,11 +43,27 @@ source ./ci/utilities/install_wheels_locally.sh [[ -f "${ENV_FILE}" ]] && source "${ENV_FILE}" -[[ "${USE_TE:-0}" == "1" ]] && { - : "${TE_WHEEL_URL:?TE_WHEEL_URL must be set when USE_TE=1}" - echo "[setup] installing Transformer Engine" +if [[ "${USE_TE:-0}" == "1" ]]; then + echo "[setup] resolving latest Transformer Engine wheel" + PY_TAG="cp$(echo "${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" | tr -d '.')" + TE_WHEEL_URL="$( + curl -fsSL https://api.github.com/repos/ROCm/maxtext/releases \ + | grep "browser_download_url:" \ + | grep "te-rocm-wheels-" \ + | grep "${PY_TAG}" \ + | grep linux_x86_64.whl \ + | head -n1 \ + | cut -d '"' -f4 + )" + + [[ -n "${TE_WHEEL_URL}" ]] || { + echo "failed to resolve Transformer Engine wheel for ${PY_TAG}" >&2 + exit 1 + } + + echo "[setup] installing Transformer Engine from ${TE_WHEEL_URL}" "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}" -} +fi export PY_COLORS=1 export NCCL_DEBUG=WARN From 24fb105b7bfd39a5783ac9bb9a22a1df53c0c1bb Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Mon, 11 May 2026 01:03:25 +0000 Subject: [PATCH 46/50] Load MaxText ROCm benchmark configs and requirements from ROCm/maxtext --- .../wheel_benchmarks_nightly_release.yml | 2 +- .../maxtext_rocm/run_maxtext_rocm.sh | 31 ++++++++----------- ...t_rocm.py => collect_run_manifest_rocm.py} | 0 3 files changed, 14 insertions(+), 19 deletions(-) rename ci/{collect_ci_manifest_rocm.py => collect_run_manifest_rocm.py} (100%) diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml index 9b4989229bd2..3507b830af7b 100644 --- a/.github/workflows/wheel_benchmarks_nightly_release.yml +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -49,7 +49,7 @@ jobs: {version: "7.2.0", tag: "rocm720"}, ] target: ["maxtext"] - workload: ["gemma_4b"] + workload: ["gemma3-4b"] name: "Benchmark ROCm (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" with: runner: ${{ matrix.runner }} diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index e68b987e97af..1c895c1ee149 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -8,17 +8,12 @@ WORKLOAD="gemma3-4b" } [[ $# -eq 2 ]] && WORKLOAD="$2" -ROOT="${ROOT:-$PWD}" -TARGET_DIR="${ROOT}/benchmark_targets/rocm_maxtext" -RUN_DIR="${ROOT}/run_artifacts" -REPO_DIR="${ROOT}/maxtext" +JAX_DIR="${JAX_DIR:-$PWD}" +TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/maxtext_rocm" +RUN_DIR="${TARGET_DIR}/run_artifacts" +REPO_DIR="${TARGET_DIR}/maxtext" WORK_DIR="${REPO_DIR}/src" -CFG_FILE="${REPO_DIR}/configs/models/${WORKLOAD}.yml" -ENV_FILE="${REPO_DIR}/configs/models/${WORKLOAD}.env.sh" -REQ_FILE="${WORK_DIR}/dependencies/requirements/requirements_rocm_jax_0.8.2.txt" -EXPECTED_FILE="${TARGET_DIR}/expected.json" - RUN_LOG="${RUN_DIR}/maxtext.log" BENCHMARK_JSON="${RUN_DIR}/benchmark.json" RESULT_JSON="${RUN_DIR}/result.json" @@ -31,27 +26,28 @@ source ci/envs/default.env source ./ci/utilities/install_wheels_locally.sh [[ -d "${REPO_DIR}/.git" ]] || \ - git clone --depth 1 --branch main https://github.com/ROCm/maxtext.git "${REPO_DIR}" + git clone --depth 1 --branch add-rocm-benchmark-configs https://github.com/ROCm/maxtext.git "${REPO_DIR}" + +CFG_FILE="${REPO_DIR}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" +REQ_FILE="${REPO_DIR}/src/dependencies/requirements/requirements_rocm_benchmark.txt" +EXP_FILE="${TARGET_DIR}/exp_maxtext_rocm.yml" [[ -f "${CFG_FILE}" ]] || { echo "missing config file: ${CFG_FILE}" >&2; exit 2; } -[[ -f "${EXPECTED_FILE}" ]] || { echo "missing expected file: ${EXPECTED_FILE}" >&2; exit 2; } +[[ -f "${EXP_FILE}" ]] || { echo "missing expected file: ${EXP_FILE}" >&2; exit 2; } [[ -f "${REQ_FILE}" ]] && { echo "[setup] installing MaxText requirements" "${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" } -[[ -f "${ENV_FILE}" ]] && source "${ENV_FILE}" - if [[ "${USE_TE:-0}" == "1" ]]; then echo "[setup] resolving latest Transformer Engine wheel" PY_TAG="cp$(echo "${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" | tr -d '.')" TE_WHEEL_URL="$( curl -fsSL https://api.github.com/repos/ROCm/maxtext/releases \ - | grep "browser_download_url:" \ + | grep "browser_download_url" \ | grep "te-rocm-wheels-" \ | grep "${PY_TAG}" \ - | grep linux_x86_64.whl \ | head -n1 \ | cut -d '"' -f4 )" @@ -60,7 +56,7 @@ if [[ "${USE_TE:-0}" == "1" ]]; then echo "failed to resolve Transformer Engine wheel for ${PY_TAG}" >&2 exit 1 } - + echo "[setup] installing Transformer Engine from ${TE_WHEEL_URL}" "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}" fi @@ -90,9 +86,8 @@ CMP_CODE=1 --log "${RUN_LOG}" \ --expected "${EXPECTED_FILE}" \ --config "${CFG_FILE}" \ - --env-file "${ENV_FILE}" \ --requirements "${REQ_FILE}" \ - --target rocm_maxtext \ + --target maxtext_rocm \ --workload "${WORKLOAD}" \ --run-code "${RUN_CODE}" \ --model-run-started-at "${MODEL_RUN_STARTED_AT}" \ diff --git a/ci/collect_ci_manifest_rocm.py b/ci/collect_run_manifest_rocm.py similarity index 100% rename from ci/collect_ci_manifest_rocm.py rename to ci/collect_run_manifest_rocm.py From 34b2d0f8bea5e21f4098cab76a8728a56e3394a7 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Mon, 11 May 2026 11:55:49 +0000 Subject: [PATCH 47/50] Revisit ROCm benchmark results and run-manifest collection --- .github/workflows/benchmark_rocm.yml | 11 +- .../maxtext_rocm/cmp_maxtext_rocm.py | 51 +++---- .../maxtext_rocm/exp_maxtext_rocm.yml | 2 +- .../maxtext_rocm/run_maxtext_rocm.sh | 122 ++++++++++------ ci/collect_run_manifest_rocm.py | 130 +++++++++++------- 5 files changed, 196 insertions(+), 120 deletions(-) diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml index bc3c648e0c26..8c5af7ae4488 100644 --- a/.github/workflows/benchmark_rocm.yml +++ b/.github/workflows/benchmark_rocm.yml @@ -157,8 +157,11 @@ jobs: JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" JAXCI_PYTHON: "python${{ inputs.python }}" JAXCI_ENABLE_X64: "0" - USE_TE: "${{ inputs.use-te }}" INPUT_ROCM_VERSION: "${{ inputs.rocm-version }}" + INPUT_RUNNER: "${{ inputs.runner }}" + INPUT_ROCM_TAG: "${{ inputs.rocm-tag }}" + USE_TE: "${{ inputs.use-te }}" + steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: @@ -175,8 +178,6 @@ jobs: - name: Install Python dependencies run: | $JAXCI_PYTHON -m pip install uv~=0.11.2 - $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt - $JAXCI_PYTHON -m pip install pytest-json-report # Halt for testing - name: Wait For Connection uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c @@ -187,14 +188,14 @@ jobs: TARGET: ${{ inputs.target }} WORKLOAD: ${{ inputs.workload }} timeout-minutes: 120 - run: ./ci/benchmark_targets/${TARGET}_rocm/run_${TARGET}_rocm.sh --workload "${WORKLOAD}" + run: ./ci/benchmark_targets/${TARGET}_rocm/run_${TARGET}_rocm.sh "${WORKLOAD}" - name: Upload GitHub artifacts if: always() continue-on-error: true uses: actions/upload-artifact@v4 with: name: benchmark-artifacts-${{ inputs.target }}-${{ inputs.workload }}-${{ inputs.runner }}-py${{ inputs.python }}-rocm${{ inputs.rocm-version }} - path: run_artifacts/result.json + path: ci/benchmark_targets/${{ inputs.target }}_rocm/run_artifacts/${{ inputs.workload }}/result.json if-no-files-found: warn - name: Upload CI artifacts to S3 if: always() diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py index e2d6a40ad75e..e0ba94a53860 100644 --- a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py +++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py @@ -1,17 +1,23 @@ #!/usr/bin/env python3 + +# Evaluates benchmark results against expected thresholds and writes +# benchmark-specific result metadata for inclusion in the final ROCm +# CI run manifest. + import argparse import json import statistics from pathlib import Path METRIC_FIELD_INDEX = 3 -WARMUP_STEPS = 2 +WARMUP_STEPS = 4 def read(path): - return ( - Path(path).read_text(errors="replace") if path and Path(path).exists() else "" - ) + path = Path(path) + if not path.exists(): + return "" + return path.read_text(errors="replace") def parse_metric_values(log_text): @@ -30,27 +36,26 @@ def parse_metric_values(log_text): def main(): - p = argparse.ArgumentParser() - p.add_argument("--log", required=True) - p.add_argument("--expected", required=True) - p.add_argument("--config", required=True) - p.add_argument("--env-file", default="") - p.add_argument("--requirements", required=True) - p.add_argument("--target", required=True) - p.add_argument("--workload", required=True) - p.add_argument("--run-code", required=True) - p.add_argument("--model-run-started-at", required=True) - p.add_argument("--model-run-completed-at", required=True) - p.add_argument("--out", required=True) - args = p.parse_args() - - expected = json.loads(Path(args.expected).read_text()) + parser = argparse.ArgumentParser() + parser.add_argument("--log", required=True) + parser.add_argument("--expected", required=True) + parser.add_argument("--config", required=True) + parser.add_argument("--requirements", required=True) + parser.add_argument("--target", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--run-code", required=True) + parser.add_argument("--model-run-started-at", required=True) + parser.add_argument("--model-run-completed-at", required=True) + parser.add_argument("--out", required=True) + args = parser.parse_args() + + expected = json.loads(read(args.expected)) values = parse_metric_values(read(args.log)) samples = values[WARMUP_STEPS:] observed = statistics.median(samples) if samples else None - baseline = float(expected["baseline"]) + baseline = float(expected["baseline_ms"]) threshold = float(expected["threshold_percent"]) distance = None @@ -59,20 +64,18 @@ def main(): if int(args.run_code) == 0 and observed is not None and baseline != 0: raw = ((observed - baseline) / baseline) * 100.0 distance = abs(raw) - cmp_code = 0 if raw >= -threshold else 1 + cmp_code = 0 if raw <= threshold else 1 result = { "benchmark_schema_version": 1, - "kind": "benchmark", "target": args.target, "workload": args.workload, "run_code": int(args.run_code), "cmp_code": cmp_code, - "distance_percent": round(distance, 4) if distance is not None else None, + "distance_percent": (round(distance, 4) if distance is not None else None), "model_run_started_at": args.model_run_started_at, "model_run_completed_at": args.model_run_completed_at, "workload_config_raw": read(args.config), - "workload_env_raw": read(args.env_file), "requirements_raw": read(args.requirements), "expected_config_raw": read(args.expected), } diff --git a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml index ca1bacc5e890..c0ccb9c6f48a 100644 --- a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml +++ b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml @@ -1,4 +1,4 @@ { - "baseline_ms": 2600.0, + "baseline_ms": 2250.00, "threshold_percent": 5.0 } \ No newline at end of file diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index 1c895c1ee149..1f4ff4fb9d1a 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -1,48 +1,79 @@ #!/usr/bin/env bash +# Runs a MaxText ROCm benchmark workload and produces benchmark-specific +# result metadata for inclusion in the final ROCm CI run manifest. +# +# The benchmark result payload is written to benchmark.json and later +# merged into the final result.json by collect_rocm_run_metadata.py. +# +# This script: +# - installs MaxText benchmark dependencies +# - optionally installs Transformer Engine +# - runs the benchmark workload +# - evaluates benchmark results against expected thresholds +# - writes benchmark-specific metadata to benchmark.json +# +# The final ROCm CI result manifest is produced separately by +# collect_rocm_run_metadata.py. set -euo pipefail -WORKLOAD="gemma3-4b" -[[ $# -eq 0 || ( $# -eq 2 && "$1" == "--workload" ) ]] || { - echo "usage: $0 [--workload NAME]" >&2 - exit 2 -} -[[ $# -eq 2 ]] && WORKLOAD="$2" +WORKLOAD="${1:-gemma3-4b}" JAX_DIR="${JAX_DIR:-$PWD}" + +PYTHON_BIN="${JAXCI_PYTHON:-python3}" +PYTHON_VERSION="${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" +JAX_ENABLE_X64="${JAXCI_ENABLE_X64:-0}" +USE_TE="${USE_TE:-0}" + TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/maxtext_rocm" -RUN_DIR="${TARGET_DIR}/run_artifacts" -REPO_DIR="${TARGET_DIR}/maxtext" -WORK_DIR="${REPO_DIR}/src" +MAXTEXT_DIR="${TARGET_DIR}/maxtext" +MAXTEXT_SRC_DIR="${MAXTEXT_DIR}/src" -RUN_LOG="${RUN_DIR}/maxtext.log" -BENCHMARK_JSON="${RUN_DIR}/benchmark.json" +RUN_DIR="${TARGET_DIR}/run_artifacts/${WORKLOAD}" +RUN_LOG="${RUN_DIR}/model_run.log" +BENCH_JSON="${RUN_DIR}/benchmark.json" RESULT_JSON="${RUN_DIR}/result.json" -PYTHON_BIN="${JAXCI_PYTHON:-python3}" +CFG_FILE="${MAXTEXT_DIR}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" +REQ_FILE="${MAXTEXT_DIR}/src/dependencies/requirements/requirements_rocm_benchmark.txt" +EXP_FILE="${TARGET_DIR}/exp_maxtext_rocm.yml" mkdir -p "${RUN_DIR}" -source ci/envs/default.env -source ./ci/utilities/install_wheels_locally.sh +source "${JAX_DIR}/ci/envs/default.env" +source "${JAX_DIR}/ci/utilities/install_wheels_locally.sh" -[[ -d "${REPO_DIR}/.git" ]] || \ - git clone --depth 1 --branch add-rocm-benchmark-configs https://github.com/ROCm/maxtext.git "${REPO_DIR}" +if [[ ! -d "${MAXTEXT_DIR}/.git" ]]; then + git clone \ + --depth 1 \ + --branch add-rocm-benchmark-configs \ + https://github.com/ROCm/maxtext.git \ + "${MAXTEXT_DIR}" +fi -CFG_FILE="${REPO_DIR}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" -REQ_FILE="${REPO_DIR}/src/dependencies/requirements/requirements_rocm_benchmark.txt" -EXP_FILE="${TARGET_DIR}/exp_maxtext_rocm.yml" +[[ -f "${CFG_FILE}" ]] || { + echo "missing config file: ${CFG_FILE}" >&2 + exit 2 +} -[[ -f "${CFG_FILE}" ]] || { echo "missing config file: ${CFG_FILE}" >&2; exit 2; } -[[ -f "${EXP_FILE}" ]] || { echo "missing expected file: ${EXP_FILE}" >&2; exit 2; } +[[ -f "${REQ_FILE}" ]] || { + echo "missing requirements file: ${REQ_FILE}" >&2 + exit 2 +} -[[ -f "${REQ_FILE}" ]] && { - echo "[setup] installing MaxText requirements" - "${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" +[[ -f "${EXP_FILE}" ]] || { + echo "missing expected file: ${EXP_FILE}" >&2 + exit 2 } -if [[ "${USE_TE:-0}" == "1" ]]; then - echo "[setup] resolving latest Transformer Engine wheel" - PY_TAG="cp$(echo "${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" | tr -d '.')" +echo "Installing MaxText ROCm benchmark requirements" +"${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" + +if [[ "${USE_TE}" == "1" ]]; then + echo "Resolving latest Transformer Engine wheel" + + PY_TAG="cp$(echo "${PYTHON_VERSION}" | tr -d '.')" + TE_WHEEL_URL="$( curl -fsSL https://api.github.com/repos/ROCm/maxtext/releases \ | grep "browser_download_url" \ @@ -53,11 +84,11 @@ if [[ "${USE_TE:-0}" == "1" ]]; then )" [[ -n "${TE_WHEEL_URL}" ]] || { - echo "failed to resolve Transformer Engine wheel for ${PY_TAG}" >&2 + echo "Failed to resolve Transformer Engine wheel" >&2 exit 1 } - echo "[setup] installing Transformer Engine from ${TE_WHEEL_URL}" + echo "Installing Transformer Engine from ${TE_WHEEL_URL}" "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}" fi @@ -70,21 +101,29 @@ export XLA_PYTHON_CLIENT_PREALLOCATE=false MODEL_RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" +echo "Starting MaxText workload: ${WORKLOAD}" + set +e -pushd "${WORK_DIR}" >/dev/null +pushd "${MAXTEXT_SRC_DIR}" >/dev/null + "${PYTHON_BIN}" -m maxtext.trainers.pre_train.train \ "$(realpath "${CFG_FILE}")" \ > "${RUN_LOG}" 2>&1 + RUN_CODE=$? + popd >/dev/null set -e +echo "Completed MaxText workload: ${WORKLOAD}" + MODEL_RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" -CMP_CODE=1 -"${PYTHON_BIN}" "${TARGET_DIR}/benchmark_result.py" \ +CMP_CODE=0 + +"${PYTHON_BIN}" "${TARGET_DIR}/cmp_maxtext_rocm.py" \ --log "${RUN_LOG}" \ - --expected "${EXPECTED_FILE}" \ + --expected "${EXP_FILE}" \ --config "${CFG_FILE}" \ --requirements "${REQ_FILE}" \ --target maxtext_rocm \ @@ -92,14 +131,19 @@ CMP_CODE=1 --run-code "${RUN_CODE}" \ --model-run-started-at "${MODEL_RUN_STARTED_AT}" \ --model-run-completed-at "${MODEL_RUN_COMPLETED_AT}" \ - --out "${BENCHMARK_JSON}" || CMP_CODE=$? - -"${PYTHON_BIN}" ci/make_ci_manifest.py \ - --extra "${BENCHMARK_JSON}" \ + --out "${BENCH_JSON}" || CMP_CODE=$? + +"${PYTHON_BIN}" "${JAX_DIR}/ci/collect_run_manifest_rocm.py" \ + --runner "${INPUT_RUNNER}" \ + --python-version "${PYTHON_VERSION}" \ + --python-bin "${PYTHON_BIN}" \ + --rocm-version "${INPUT_ROCM_VERSION}" \ + --rocm-tag "${INPUT_ROCM_TAG}" \ + --extra "${BENCH_JSON}" \ --out "${RESULT_JSON}" -rm -f "${RUN_LOG}" "${BENCHMARK_JSON}" +rm -f "${RUN_LOG}" "${BENCH_JSON}" [[ -s "${RESULT_JSON}" ]] && touch "${RUN_DIR}/_SUCCESS" -exit $(( RUN_CODE != 0 || CMP_CODE != 0 )) \ No newline at end of file +exit $(( RUN_CODE != 0 || CMP_CODE != 0 )) diff --git a/ci/collect_run_manifest_rocm.py b/ci/collect_run_manifest_rocm.py index 77b010ece756..d2acf1d1f306 100644 --- a/ci/collect_run_manifest_rocm.py +++ b/ci/collect_run_manifest_rocm.py @@ -1,9 +1,32 @@ #!/usr/bin/env python3 + +# Collect ROCm CI run metadata and optionally merges additional JSON +# payloads into the final result manifest. +# +# Extra JSON files are merged using dict.update() semantics. Fields in +# extra payloads may override existing manifest fields if names collide. +# +# Benchmark/test-specific payloads should avoid using reserved top-level +# run metadata keys such as: +# - github_* +# - run_* +# - python_version +# - rocm_ + +#!/usr/bin/env python3 +# Collects ROCm CI run metadata and writes the final result manifest. +# +# Optional --extra JSON payloads are merged into the manifest using +# dict.update() semantics. Extra payloads should avoid reserved run metadata +# keys such as github_*, run_*, python_version, rocm_*, runner, and gpu_count. + import argparse import json import os import re import subprocess +import urllib.request +from datetime import datetime, timezone from pathlib import Path @@ -11,7 +34,7 @@ def env(name, default="unknown"): return os.environ.get(name, default) -def cmd(args): +def run(args): try: return subprocess.check_output( args, text=True, stderr=subprocess.DEVNULL @@ -21,23 +44,34 @@ def cmd(args): def sh(command): - return cmd(["bash", "-lc", command]) + return run(["bash", "-lc", command]) + + +def utc_now(): + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def get_json(url, headers=None): + request = urllib.request.Request(url, headers=headers or {}) + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode()) + + +def get_response_headers(url, headers=None): + request = urllib.request.Request(url, headers=headers or {}) + with urllib.request.urlopen(request, timeout=30) as response: + return response.headers def github_run_started_at(): repo = env("GITHUB_REPOSITORY", "") run_id = env("GITHUB_RUN_ID", "") - raw = cmd( - [ - "curl", - "-fsSL", - "-H", - "Accept: application/vnd.github+json", - f"https://api.github.com/repos/{repo}/actions/runs/{run_id}", - ] - ) try: - return json.loads(raw).get("run_started_at", "") + data = get_json( + f"https://api.github.com/repos/{repo}/actions/runs/{run_id}", + headers={"Accept": "application/vnd.github+json"}, + ) + return data.get("run_started_at", "") except Exception: return "" @@ -49,34 +83,27 @@ def gpu_count(runner): def image_digest(rocm_tag): repo = f"rocm/jax-base-ubu24.{rocm_tag}" - token_raw = cmd( - [ - "curl", - "-fsSL", - f"https://ghcr.io/token?service=ghcr.io&scope=repository:{repo}:pull", - ] - ) try: - token = json.loads(token_raw).get("token", "") + token_data = get_json( + f"https://ghcr.io/token?service=ghcr.io&scope=repository:{repo}:pull" + ) + token = token_data.get("token", "") + if not token: + return "" + headers = get_response_headers( + f"https://ghcr.io/v2/{repo}/manifests/latest", + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/vnd.docker.distribution.manifest.v2+json", + }, + ) + return headers.get("Docker-Content-Digest", "") except Exception: return "" - headers = sh( - "curl -fsSL -D - " - f"-H 'Authorization: Bearer {token}' " - "-H 'Accept: application/vnd.docker.distribution.manifest.v2+json' " - f"https://ghcr.io/v2/{repo}/manifests/latest " - "-o /dev/null" - ) - for line in headers.splitlines(): - if line.lower().startswith("docker-content-digest:"): - return line.split(":", 1)[1].strip() - return "" - -def package_snapshot(): - python = env("JAXCI_PYTHON", "python3") - pkgs = cmd([python, "-m", "pip", "list", "--format=freeze"]) +def package_snapshot(python_bin): + pkgs = run([python_bin, "-m", "pip", "list", "--format=freeze"]) return "|".join( line for line in pkgs.splitlines() @@ -86,25 +113,25 @@ def package_snapshot(): def main(): parser = argparse.ArgumentParser() + parser.add_argument("--runner", required=True) + parser.add_argument("--python-version", required=True) + parser.add_argument("--python-bin", default="python3") + parser.add_argument("--rocm-version", required=True) + parser.add_argument("--rocm-tag", required=True) parser.add_argument("--extra", action="append", default=[]) parser.add_argument("--out", required=True) args = parser.parse_args() - started = github_run_started_at() date = started.split("T", 1)[0] if started else sh("date -u +%F") - repo = env("GITHUB_REPOSITORY") run_id = env("GITHUB_RUN_ID") attempt = env("GITHUB_RUN_ATTEMPT") workflow = env("GITHUB_WORKFLOW", "") - runner = env("INPUT_RUNNER") - rocm_tag = env("INPUT_ROCM_TAG") - manifest = { "schema_version": 1, "run_key": f"{date}_{run_id}_{attempt}", "run_started_at": started, - "run_completed_at": sh("date -u +%Y-%m-%dT%H:%M:%SZ"), + "run_completed_at": utc_now(), "github_run_url": f"https://github.com/{repo}/actions/runs/{run_id}", "github_repository": repo, "github_ref_name": env("GITHUB_REF_NAME"), @@ -117,22 +144,23 @@ def main(): "github_workflow": workflow, "is_nightly": "nightly" if "nightly" in workflow.lower() else "continuous", "github_job": env("GITHUB_JOB"), - "python_version": env("INPUT_PYTHON", env("JAXCI_HERMETIC_PYTHON_VERSION")), - "rocm_version": env("INPUT_ROCM_VERSION"), - "rocm_tag": rocm_tag, - "gpu_count": gpu_count(runner), - "runner": runner, - "base_image_name": f"ghcr.io/rocm/jax-base-ubu24.{rocm_tag}:latest", - "base_image_digest": image_digest(rocm_tag), - "jax_packages_raw": package_snapshot(), + "python_version": args.python_version, + "rocm_version": args.rocm_version, + "rocm_tag": args.rocm_tag, + "gpu_count": gpu_count(args.runner), + "runner": args.runner, + "base_image_name": f"ghcr.io/rocm/jax-base-ubu24.{args.rocm_tag}:latest", + "base_image_digest": image_digest(args.rocm_tag), + "jax_packages_raw": package_snapshot(args.python_bin), "wheels_sha_raw": sh("sha256sum dist/*.whl 2>/dev/null || true").replace( "\n", "|" ), } for extra in args.extra: - if extra and Path(extra).exists(): - manifest.update(json.loads(Path(extra).read_text())) + path = Path(extra) + if path.exists(): + manifest.update(json.loads(path.read_text())) Path(args.out).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") From 37fe3f7ad54e44ffa684c4b36055e2db091e0f2c Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Mon, 11 May 2026 18:26:02 +0000 Subject: [PATCH 48/50] Revisit ROCm artifact upload to S3 for reusability --- .github/workflows/benchmark_rocm.yml | 16 ++++--- .../maxtext_rocm/cmp_maxtext_rocm.py | 2 +- .../maxtext_rocm/run_maxtext_rocm.sh | 4 +- ci/collect_run_manifest_rocm.py | 2 +- ci/upload_rocm_artifacts.sh | 42 +++++++++++-------- 5 files changed, 40 insertions(+), 26 deletions(-) diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml index 8c5af7ae4488..28b7c62eed70 100644 --- a/.github/workflows/benchmark_rocm.yml +++ b/.github/workflows/benchmark_rocm.yml @@ -197,8 +197,14 @@ jobs: name: benchmark-artifacts-${{ inputs.target }}-${{ inputs.workload }}-${{ inputs.runner }}-py${{ inputs.python }}-rocm${{ inputs.rocm-version }} path: ci/benchmark_targets/${{ inputs.target }}_rocm/run_artifacts/${{ inputs.workload }}/result.json if-no-files-found: warn - - name: Upload CI artifacts to S3 - if: always() - env: - S3_BUCKET_NAME: jax-ci-amd - run: ./ci/upload_artifacts.sh run_artifacts \ No newline at end of file + # - name: Upload CI artifacts to S3 + # if: always() + # env: + # S3_BUCKET_NAME: jax-ci-amd + # run: | + # RUN_KEY="$(date -u +%F)_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" + # PREFIX="jax-ci-bench-logs/${GITHUB_REPOSITORY}/${GITHUB_REF_NAME}/${RUN_KEY}" + # ./ci/upload_rocm_artifacts.sh \ + # "ci/benchmark_targets/${{ inputs.target }}_rocm/run_artifacts/${{ inputs.workload }}" \ + # "${PREFIX}" \ + # success \ No newline at end of file diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py index e0ba94a53860..ad74615d8a3a 100644 --- a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py +++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py @@ -80,7 +80,7 @@ def main(): "expected_config_raw": read(args.expected), } - Path(args.out).write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + Path(args.out).write_text(json.dumps(result, indent=2) + "\n") return cmp_code diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index 1f4ff4fb9d1a..1190398d1614 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -101,7 +101,7 @@ export XLA_PYTHON_CLIENT_PREALLOCATE=false MODEL_RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" -echo "Starting MaxText workload: ${WORKLOAD}" +echo "Running MaxText workload: ${WORKLOAD}.." set +e pushd "${MAXTEXT_SRC_DIR}" >/dev/null @@ -115,7 +115,7 @@ RUN_CODE=$? popd >/dev/null set -e -echo "Completed MaxText workload: ${WORKLOAD}" +echo "..Completed MaxText workload: ${WORKLOAD}" MODEL_RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" diff --git a/ci/collect_run_manifest_rocm.py b/ci/collect_run_manifest_rocm.py index d2acf1d1f306..592195ee3fff 100644 --- a/ci/collect_run_manifest_rocm.py +++ b/ci/collect_run_manifest_rocm.py @@ -162,7 +162,7 @@ def main(): if path.exists(): manifest.update(json.loads(path.read_text())) - Path(args.out).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + Path(args.out).write_text(json.dumps(manifest, indent=2) + "\n") if __name__ == "__main__": diff --git a/ci/upload_rocm_artifacts.sh b/ci/upload_rocm_artifacts.sh index 0778697cdc3f..bdae41f03703 100644 --- a/ci/upload_rocm_artifacts.sh +++ b/ci/upload_rocm_artifacts.sh @@ -1,29 +1,37 @@ #!/usr/bin/env bash set -euo pipefail -: "${S3_BUCKET_NAME:?}" -ARTIFACT_DIR="${1:-run_artifacts}" -RESULT_FILE="${ARTIFACT_DIR}/result.json" +# Uploads a file or directory to S3. +# +# Usage: +# upload_rocm_artifacts.sh [success] +# +# If the optional third argument is "success", an empty _SUCCESS marker is +# written after the upload completes. -[[ -f "${RESULT_FILE}" ]] || { - echo "missing result.json" >&2 - exit 2 -} +: "${S3_BUCKET_NAME:?}" -RUN_KEY="${DATE}_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" -COMBO="py$(norm "${INPUT_PYTHON}")-rocm$(norm "${INPUT_ROCM_VERSION}")-${GPU_PART}" -PREFIX="${GITHUB_REPOSITORY}/${GITHUB_REF_NAME}/${IS_NIGHTLY}/${RUN_KEY}/${COMBO}" +SRC="${1:?missing source path}" +DEST_PREFIX="${2:?missing S3 destination prefix}" +WRITE_SUCCESS="${3:-}" -DEST="s3://${S3_BUCKET_NAME}/${TEST_LOGS_ROOT}/${PREFIX}" +[[ -e "${SRC}" ]] || { + echo "missing source path: ${SRC}" >&2 + exit 2 +} -echo "[upload] ${DEST}" +DEST="s3://${S3_BUCKET_NAME}/${DEST_PREFIX}" -aws s3 cp --only-show-errors "${RESULT_FILE}" "${DEST}/result.json" +echo "[upload] ${SRC} -> ${DEST}" -if [[ -f "${ARTIFACT_DIR}/logs.tar.gz" ]]; then - aws s3 cp --only-show-errors "${ARTIFACT_DIR}/logs.tar.gz" "${DEST}/logs.tar.gz" +if [[ -d "${SRC}" ]]; then + aws s3 cp --only-show-errors "${SRC}" "${DEST}" --recursive +else + aws s3 cp --only-show-errors "${SRC}" "${DEST}/$(basename "${SRC}")" fi -printf '' | aws s3 cp --only-show-errors - "${DEST}/_SUCCESS" +if [[ "${WRITE_SUCCESS}" == "success" ]]; then + printf '' | aws s3 cp --only-show-errors - "${DEST}/_SUCCESS" +fi -echo "[done]" \ No newline at end of file +echo "[done]" From 9183f204d5a9cbcf275a474889af7f6ebff345c4 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Wed, 13 May 2026 18:46:21 +0000 Subject: [PATCH 49/50] Remove TE installation to keep the model lightweight --- .github/workflows/benchmark_rocm.yml | 12 --------- .../wheel_benchmarks_nightly_release.yml | 1 - .../maxtext_rocm/cmp_maxtext_rocm.py | 2 +- .../maxtext_rocm/run_maxtext_rocm.sh | 25 ------------------- ci/collect_run_manifest_rocm.py | 2 +- 5 files changed, 2 insertions(+), 40 deletions(-) diff --git a/.github/workflows/benchmark_rocm.yml b/.github/workflows/benchmark_rocm.yml index 28b7c62eed70..74b51589e753 100644 --- a/.github/workflows/benchmark_rocm.yml +++ b/.github/workflows/benchmark_rocm.yml @@ -70,13 +70,6 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" type: string default: 'latest' - use-te: - description: "Install Transformer Engine" - type: choice - default: '1' - options: - - '0' - - '1' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -123,10 +116,6 @@ on: description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)" type: string default: 'latest' - use-te: - description: "Install Transformer Engine" - type: string - default: '1' permissions: id-token: write contents: read @@ -160,7 +149,6 @@ jobs: INPUT_ROCM_VERSION: "${{ inputs.rocm-version }}" INPUT_RUNNER: "${{ inputs.runner }}" INPUT_ROCM_TAG: "${{ inputs.rocm-tag }}" - USE_TE: "${{ inputs.use-te }}" steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/.github/workflows/wheel_benchmarks_nightly_release.yml b/.github/workflows/wheel_benchmarks_nightly_release.yml index 3507b830af7b..76a7b0584552 100644 --- a/.github/workflows/wheel_benchmarks_nightly_release.yml +++ b/.github/workflows/wheel_benchmarks_nightly_release.yml @@ -61,4 +61,3 @@ jobs: jaxlib-version: "head" skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}} gcs_download_uri: ${{ inputs.gcs_download_uri }} - use-te: "1" diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py index ad74615d8a3a..98de57e92c89 100644 --- a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py +++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py @@ -62,7 +62,7 @@ def main(): cmp_code = 1 if int(args.run_code) == 0 and observed is not None and baseline != 0: - raw = ((observed - baseline) / baseline) * 100.0 + raw = ((baseline - observed) / baseline) * 100.0 distance = abs(raw) cmp_code = 0 if raw <= threshold else 1 diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index 1190398d1614..50ee8de170f4 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -7,7 +7,6 @@ # # This script: # - installs MaxText benchmark dependencies -# - optionally installs Transformer Engine # - runs the benchmark workload # - evaluates benchmark results against expected thresholds # - writes benchmark-specific metadata to benchmark.json @@ -23,7 +22,6 @@ JAX_DIR="${JAX_DIR:-$PWD}" PYTHON_BIN="${JAXCI_PYTHON:-python3}" PYTHON_VERSION="${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" JAX_ENABLE_X64="${JAXCI_ENABLE_X64:-0}" -USE_TE="${USE_TE:-0}" TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/maxtext_rocm" MAXTEXT_DIR="${TARGET_DIR}/maxtext" @@ -69,29 +67,6 @@ fi echo "Installing MaxText ROCm benchmark requirements" "${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" -if [[ "${USE_TE}" == "1" ]]; then - echo "Resolving latest Transformer Engine wheel" - - PY_TAG="cp$(echo "${PYTHON_VERSION}" | tr -d '.')" - - TE_WHEEL_URL="$( - curl -fsSL https://api.github.com/repos/ROCm/maxtext/releases \ - | grep "browser_download_url" \ - | grep "te-rocm-wheels-" \ - | grep "${PY_TAG}" \ - | head -n1 \ - | cut -d '"' -f4 - )" - - [[ -n "${TE_WHEEL_URL}" ]] || { - echo "Failed to resolve Transformer Engine wheel" >&2 - exit 1 - } - - echo "Installing Transformer Engine from ${TE_WHEEL_URL}" - "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}" -fi - export PY_COLORS=1 export NCCL_DEBUG=WARN export TF_CPP_MIN_LOG_LEVEL=0 diff --git a/ci/collect_run_manifest_rocm.py b/ci/collect_run_manifest_rocm.py index 592195ee3fff..3175e41fd06e 100644 --- a/ci/collect_run_manifest_rocm.py +++ b/ci/collect_run_manifest_rocm.py @@ -107,7 +107,7 @@ def package_snapshot(python_bin): return "|".join( line for line in pkgs.splitlines() - if re.search(r"^(jax|jaxlib)==|pjrt|plugin|transformer[-_]engine", line) + if re.search(r"^(jax|jaxlib)==|pjrt|plugin", line) ) From 182dd2a873c3bd27f176ba9e7f67689a33bd5840 Mon Sep 17 00:00:00 2001 From: Pakize Sanal Date: Fri, 15 May 2026 20:17:09 +0000 Subject: [PATCH 50/50] Update benchmark target scripts for more generic use cases --- .../collect_bench_manifest_rocm.py | 52 ++++++ .../maxtext_rocm/cmp_maxtext_rocm.py | 154 +++++++++++------- .../maxtext_rocm/exp_maxtext_rocm.yml | 28 +++- .../maxtext_rocm/run_maxtext_rocm.sh | 143 ++++++++-------- 4 files changed, 244 insertions(+), 133 deletions(-) create mode 100644 ci/benchmark_targets/collect_bench_manifest_rocm.py diff --git a/ci/benchmark_targets/collect_bench_manifest_rocm.py b/ci/benchmark_targets/collect_bench_manifest_rocm.py new file mode 100644 index 000000000000..f9844333a48c --- /dev/null +++ b/ci/benchmark_targets/collect_bench_manifest_rocm.py @@ -0,0 +1,52 @@ +# Copyright 2026 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import argparse +import json +from pathlib import Path + + +def read(path): + return Path(path).read_text(errors="replace") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--target", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--run-code", required=True, type=int) + parser.add_argument("--run-started-at", required=True) + parser.add_argument("--run-completed-at", required=True) + parser.add_argument("--raw", action="append", default=[]) + parser.add_argument("--out", required=True) + args = parser.parse_args() + + benchmark = { + "target": args.target, + "workload": args.workload, + "run_code": args.run_code, + "run_started_at": args.run_started_at, + "run_completed_at": args.run_completed_at, + } + + for raw in args.raw: + name, path = raw.split("=", 1) + benchmark[f"{name}_raw"] = read(path) + + Path(args.out).write_text(json.dumps({"benchmark": benchmark}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py index 98de57e92c89..ce425a31773a 100644 --- a/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py +++ b/ci/benchmark_targets/maxtext_rocm/cmp_maxtext_rocm.py @@ -1,87 +1,121 @@ -#!/usr/bin/env python3 - -# Evaluates benchmark results against expected thresholds and writes -# benchmark-specific result metadata for inclusion in the final ROCm -# CI run manifest. - +# Copyright 2026 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Compares workload benchmark metrics against configured baselines. +# +# Each workload defines: +# - which metrics are evaluated +# - how metric values are extracted from logs +# - whether lower or higher values are better +# - acceptable regression thresholds +# +# The output contains one comparison result per metric. import argparse import json import statistics from pathlib import Path - -METRIC_FIELD_INDEX = 3 -WARMUP_STEPS = 4 +import yaml def read(path): - path = Path(path) - if not path.exists(): - return "" - return path.read_text(errors="replace") + return Path(path).read_text(errors="replace") + +def load_workload_config(path, workload): + config = yaml.safe_load(read(path)) or {} + if workload not in config: + raise KeyError(f"Missing workload in baseline config: {workload}") + return config[workload] -def parse_metric_values(log_text): + +def parse_metric_values(log, metric_config): values = [] - for line in log_text.splitlines(): - if "completed step:" not in line: - continue - fields = line.split(",") - if len(fields) <= METRIC_FIELD_INDEX: + pattern = metric_config["log_pattern"] + field_index = int(metric_config["field_index"]) + warmup_steps = int(metric_config.get("warmup_steps", 0)) + for line in log.splitlines(): + if pattern not in line: continue try: - values.append(float(fields[METRIC_FIELD_INDEX].split(":", 1)[1].strip())) - except Exception: - pass - return values + value = line.split(",")[field_index] + values.append(float(value.split(":", 1)[1].strip())) + except (IndexError, ValueError): + continue + + return values[warmup_steps:] + + +def regression_percent(value, baseline, direction): + if direction == "lower_is_better": + regression = ((value - baseline) / baseline) * 100.0 + elif direction == "higher_is_better": + regression((baseline - value) / baseline) * 100.0 + else: + raise ValueError(f"Unknown comparison direction: {direction}") + + return abs(regression) + + +def evaluate_metric(log, name, metric_config): + values = parse_metric_values(log, metric_config) + value = statistics.median(values) if values else None + + baseline = float(metric_config["baseline"]) + threshold = float(metric_config["threshold_percent"]) + direction = metric_config["direction"] + + regression = None + cmp_code = 1 + if value is not None and baseline != 0: + regression = regression_percent(value, baseline, direction) + cmp_code = int(regression > threshold) + + return { + "name": name, + "value": round(value, 2) if value is not None else None, + "baseline": baseline, + "threshold_percent": threshold, + "direction": direction, + "regression_percent": ( + round(regression, 4) if regression is not None else None + ), + "cmp_code": cmp_code, + } def main(): parser = argparse.ArgumentParser() parser.add_argument("--log", required=True) - parser.add_argument("--expected", required=True) - parser.add_argument("--config", required=True) - parser.add_argument("--requirements", required=True) - parser.add_argument("--target", required=True) + parser.add_argument("--baseline", required=True) parser.add_argument("--workload", required=True) - parser.add_argument("--run-code", required=True) - parser.add_argument("--model-run-started-at", required=True) - parser.add_argument("--model-run-completed-at", required=True) parser.add_argument("--out", required=True) args = parser.parse_args() - expected = json.loads(read(args.expected)) - values = parse_metric_values(read(args.log)) - - samples = values[WARMUP_STEPS:] - observed = statistics.median(samples) if samples else None + log = read(args.log) + workload = load_workload_config(args.baseline, args.workload) - baseline = float(expected["baseline_ms"]) - threshold = float(expected["threshold_percent"]) - - distance = None - cmp_code = 1 - - if int(args.run_code) == 0 and observed is not None and baseline != 0: - raw = ((baseline - observed) / baseline) * 100.0 - distance = abs(raw) - cmp_code = 0 if raw <= threshold else 1 - - result = { - "benchmark_schema_version": 1, - "target": args.target, - "workload": args.workload, - "run_code": int(args.run_code), - "cmp_code": cmp_code, - "distance_percent": (round(distance, 4) if distance is not None else None), - "model_run_started_at": args.model_run_started_at, - "model_run_completed_at": args.model_run_completed_at, - "workload_config_raw": read(args.config), - "requirements_raw": read(args.requirements), - "expected_config_raw": read(args.expected), + metrics = { + evaluate_metric(log, name, config) + for name, config in workload["metrics"].items() } - Path(args.out).write_text(json.dumps(result, indent=2) + "\n") - return cmp_code + Path(args.out).write_text( + json.dumps({"benchmark": {"metrics": metrics}}, indent=2) + "\n" + ) + + return max(metric["cmp_code"] for metric in metrics) if __name__ == "__main__": diff --git a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml index c0ccb9c6f48a..5e887517a525 100644 --- a/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml +++ b/ci/benchmark_targets/maxtext_rocm/exp_maxtext_rocm.yml @@ -1,4 +1,24 @@ -{ - "baseline_ms": 2250.00, - "threshold_percent": 5.0 -} \ No newline at end of file +# Copyright 2026 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Baseline for JAX benchmarks +gemma3_4b: + benchmark_metric: + baseline: 2500 + threshold: 0.10 # regression allowance (%) + direction: higher_is_better + log_pattern: completed_step + field_index: 3 + warmup_step: 4 \ No newline at end of file diff --git a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh index 50ee8de170f4..dfb65578ec98 100755 --- a/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh +++ b/ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh @@ -1,124 +1,129 @@ -#!/usr/bin/env bash -# Runs a MaxText ROCm benchmark workload and produces benchmark-specific -# result metadata for inclusion in the final ROCm CI run manifest. +#!/bin/bash +# Copyright 2026 The JAX Authors. # -# The benchmark result payload is written to benchmark.json and later -# merged into the final result.json by collect_rocm_run_metadata.py. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at # -# This script: -# - installs MaxText benchmark dependencies -# - runs the benchmark workload -# - evaluates benchmark results against expected thresholds -# - writes benchmark-specific metadata to benchmark.json +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Runs a MaxText ROCm benchmark workload. # -# The final ROCm CI result manifest is produced separately by -# collect_rocm_run_metadata.py. +# This script: +# - prepares the benchmark environment +# - executes the workload +# - collects benchmark run metadata +# - evaluates benchmark metrics +# - generates the final ROCm run manifest set -euo pipefail -WORKLOAD="${1:-gemma3-4b}" +WORKLOAD="${1:-gemma3_4b}" JAX_DIR="${JAX_DIR:-$PWD}" -PYTHON_BIN="${JAXCI_PYTHON:-python3}" +PYTHON="${JAXCI_PYTHON:-python3}" PYTHON_VERSION="${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" -JAX_ENABLE_X64="${JAXCI_ENABLE_X64:-0}" -TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/maxtext_rocm" -MAXTEXT_DIR="${TARGET_DIR}/maxtext" -MAXTEXT_SRC_DIR="${MAXTEXT_DIR}/src" +TARGET="maxtext_rocm" +TARGET_DIR="${JAX_DIR}/ci/benchmark_targets/${TARGET}" +MAXTEXT="${TARGET_DIR}/maxtext" RUN_DIR="${TARGET_DIR}/run_artifacts/${WORKLOAD}" -RUN_LOG="${RUN_DIR}/model_run.log" -BENCH_JSON="${RUN_DIR}/benchmark.json" -RESULT_JSON="${RUN_DIR}/result.json" - -CFG_FILE="${MAXTEXT_DIR}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" -REQ_FILE="${MAXTEXT_DIR}/src/dependencies/requirements/requirements_rocm_benchmark.txt" -EXP_FILE="${TARGET_DIR}/exp_maxtext_rocm.yml" mkdir -p "${RUN_DIR}" source "${JAX_DIR}/ci/envs/default.env" source "${JAX_DIR}/ci/utilities/install_wheels_locally.sh" -if [[ ! -d "${MAXTEXT_DIR}/.git" ]]; then +if [[ ! -d "${MAXTEXT}/.git" ]]; then git clone \ --depth 1 \ --branch add-rocm-benchmark-configs \ https://github.com/ROCm/maxtext.git \ - "${MAXTEXT_DIR}" + "${MAXTEXT}" fi -[[ -f "${CFG_FILE}" ]] || { - echo "missing config file: ${CFG_FILE}" >&2 - exit 2 -} +for file in \ + "${MAXTEXT}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" \ + "${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt" \ + "${TARGET_DIR}/baseline.yml"; do + [[ -f "${file}" ]] || { + echo "missing required file: ${file}" >&2 + exit 2 + } +done -[[ -f "${REQ_FILE}" ]] || { - echo "missing requirements file: ${REQ_FILE}" >&2 - exit 2 -} - -[[ -f "${EXP_FILE}" ]] || { - echo "missing expected file: ${EXP_FILE}" >&2 - exit 2 -} - -echo "Installing MaxText ROCm benchmark requirements" -"${PYTHON_BIN}" -m pip install -r "${REQ_FILE}" +"${PYTHON}" -m pip install -r \ + "${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt" export PY_COLORS=1 export NCCL_DEBUG=WARN export TF_CPP_MIN_LOG_LEVEL=0 -export JAX_ENABLE_X64="${JAXCI_ENABLE_X64}" + +export JAX_ENABLE_X64="${JAXCI_ENABLE_X64:-0}" + export XLA_PYTHON_CLIENT_ALLOCATOR=platform export XLA_PYTHON_CLIENT_PREALLOCATE=false -MODEL_RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" - -echo "Running MaxText workload: ${WORKLOAD}.." +RUN_STARTED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" set +e -pushd "${MAXTEXT_SRC_DIR}" >/dev/null -"${PYTHON_BIN}" -m maxtext.trainers.pre_train.train \ - "$(realpath "${CFG_FILE}")" \ -> "${RUN_LOG}" 2>&1 +pushd "${MAXTEXT}/src" >/dev/null + +"${PYTHON}" -m maxtext.trainers.pre_train.train \ + "$(realpath "maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml")" \ +> "${RUN_DIR}/run.log" 2>&1 RUN_CODE=$? popd >/dev/null + set -e -echo "..Completed MaxText workload: ${WORKLOAD}" +RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" -MODEL_RUN_COMPLETED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" +"${PYTHON}" "${JAX_DIR}/ci/collect_bench_manifest_rocm.py" \ + --target "${TARGET}" \ + --workload "${WORKLOAD}" \ + --run-code "${RUN_CODE}" \ + --run-started-at "${RUN_STARTED_AT}" \ + --run-completed-at "${RUN_COMPLETED_AT}" \ + --raw workload_config="${MAXTEXT}/src/maxtext/configs/gpu/models/${WORKLOAD}-rocm.yml" \ + --raw requirements="${MAXTEXT}/src/dependencies/requirements/requirements_rocm_benchmark.txt" \ + --raw baseline_config="${TARGET_DIR}/baseline.yml" \ + --out "${RUN_DIR}/benchmark_meta.json" CMP_CODE=0 -"${PYTHON_BIN}" "${TARGET_DIR}/cmp_maxtext_rocm.py" \ - --log "${RUN_LOG}" \ - --expected "${EXP_FILE}" \ - --config "${CFG_FILE}" \ - --requirements "${REQ_FILE}" \ - --target maxtext_rocm \ +"${PYTHON}" "${TARGET_DIR}/cmp_maxtext_rocm.py" \ + --log "${RUN_DIR}/run.log" \ + --baseline "${TARGET_DIR}/baseline.yml" \ --workload "${WORKLOAD}" \ - --run-code "${RUN_CODE}" \ - --model-run-started-at "${MODEL_RUN_STARTED_AT}" \ - --model-run-completed-at "${MODEL_RUN_COMPLETED_AT}" \ - --out "${BENCH_JSON}" || CMP_CODE=$? + --out "${RUN_DIR}/benchmark_metrics.json" || CMP_CODE=$? -"${PYTHON_BIN}" "${JAX_DIR}/ci/collect_run_manifest_rocm.py" \ +"${PYTHON}" "${JAX_DIR}/ci/collect_run_manifest_rocm.py" \ --runner "${INPUT_RUNNER}" \ --python-version "${PYTHON_VERSION}" \ - --python-bin "${PYTHON_BIN}" \ + --python-bin "${PYTHON}" \ --rocm-version "${INPUT_ROCM_VERSION}" \ --rocm-tag "${INPUT_ROCM_TAG}" \ - --extra "${BENCH_JSON}" \ - --out "${RESULT_JSON}" + --extra "${RUN_DIR}/benchmark_meta.json" \ + --extra "${RUN_DIR}/benchmark_metrics.json" \ + --out "${RUN_DIR}/result.json" -rm -f "${RUN_LOG}" "${BENCH_JSON}" +rm -f \ + "${RUN_DIR}/run.log" \ + "${RUN_DIR}/benchmark_meta.json" \ + "${RUN_DIR}/benchmark_metrics.json" -[[ -s "${RESULT_JSON}" ]] && touch "${RUN_DIR}/_SUCCESS" +[[ -s "${RUN_DIR}/result.json" ]] && touch "${RUN_DIR}/_SUCCESS" exit $(( RUN_CODE != 0 || CMP_CODE != 0 ))