Added Dockerfile for CI images & Upgrate CI to ROCm 7.2 by VeeraRajasekhar · Pull Request #195 · ROCm/TransformerEngine

VeeraRajasekhar · 2025-05-28T22:36:25Z

Description

Added the dockerfile, which can be used to create the ci-artifactory images.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Added a new file docker/Dockerfile

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

docker/Dockerfile

wenchenvincent

Please address the comments.

docker/Dockerfile

ipanfilo

Why conversations are marked as resolved w/o any actual action?

VeeraRajasekhar · 2025-05-29T15:50:54Z

Why conversations are marked as resolved w/o any actual action?

Some of them, I have resolved, some I have currently resolved in my local, just to keep track I will mark them resolved.

docker/Dockerfile.ci.deps

wenchenvincent · 2025-08-27T14:49:04Z

@VeeraRajasekhar Is this PR still needed?

wenchenvincent · 2026-01-14T03:41:43Z

@VeeraRajasekhar Could you remind me of what we had decided on this PR? It seemed that it is no longer relevant and we should close it.

VeeraRajasekhar · 2026-02-11T03:06:55Z

Hi @ipanfilo, @wangye805

I have updated this PR with latest 7.2 docker file and moved to .github/scripts.

Let me know if I need to add an action to automate docker build and upload to our artifactory?

Thanks.

ipanfilo · 2026-02-11T04:36:13Z

.github/scripts/Dockerfile.ci.deps

+RUN pip install ipython pytest fire pydantic pybind11 ninja pandas
+RUN apt-get update && apt-get install -y vim
+
+ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151


Why is this arg needed?

Not needed, when I manually installed torch I mentioned this, I can remove this now.

ipanfilo · 2026-02-11T04:36:51Z

.github/scripts/Dockerfile.ci.deps

@@ -0,0 +1,27 @@
+# TE CI Dockerfile


Put copyright here

ipanfilo · 2026-02-11T04:37:14Z

.github/scripts/Dockerfile.ci.deps

+
+ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151
+
+# Install flash-attention v2.8.1


Better make FA version arg

ipanfilo · 2026-02-11T04:54:07Z

.github/scripts/Dockerfile.ci.deps

+
+
+RUN pip install \
+    https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/jax_rocm7_pjrt-0.8.0%2Brocm7.2.0-py3-none-manylinux_2_28_x86_64.whl \


Since base image is ARG, JAX wheels should also be built based on ARG cause their ROCm and Python versions should match. So in general there should be ARGs for ROCm version, base image, JAX version, JAX wheels

VeeraRajasekhar · 2026-02-12T01:54:26Z

I had to force push to include new FA 2.8.3 support commit and my changes for 7.2 support to run the CI.

Thanks.

ROCm's jax.nn.scaled_matmul kernels require the contracting dimension (K) to be at least 64. Without this validation, backward pass GEMMs with K < 64 cause segmentation faults. Added K >= 64 check in _check_mxfp8_gemm_support() for JAX GEMM on ROCm. Fixes: test_dense_grad_fp8[MXFP8_1D_SCALING-with_jax_gemm_True-64-32-64]

VeeraRajasekhar · 2026-02-14T20:55:27Z

@Micky774, please review the following,

Analysis on testing on Jax & xla 0.8.2

(Not Supported) jax.nn.scaled_matmul (MXFP8) on ROCm crashes with a segmentation fault when the contracting dimension (K) is less than 64.

import functools
import jax
import jax.numpy as jnp
from jax import nn

key = jax.random.PRNGKey(0)
key_a, key_b = jax.random.split(key)
B, M, N, K = 1, 128, 128, 32


lhs = jax.random.normal(key_a, (B, M, K), dtype=jnp.float32)
rhs = jax.random.normal(key_b, (B, N, K), dtype=jnp.float32)

# 1. high-precision matmul
ref = jnp.einsum("bmk,bnk->bmn", lhs, rhs)

# 2. mxfp8 matmul
configs = [nn.get_scaled_dot_general_config("mxfp8")] * 3
scaled_dot = functools.partial(
    nn.scaled_dot_general,
    configs=configs,
    preferred_element_type=jnp.float32,
)

out = scaled_dot(lhs, rhs, (((2,), (2,)), ((0,), (0,))))

# compare results
print("high-precision ref: ")
print(ref)

print("mxfp8 out: ")
print(out)

max_abs = jnp.max(jnp.abs(out - ref))
max_rel = max_abs / jnp.max(jnp.abs(ref))

print("max abs error:", max_abs)
print("max rel error:", max_rel)

[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x32x64_UR_2: K must be a multiple of workgroupTile.k=64 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_16x16x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x16x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_16x32x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x32x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x16x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_16x64x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x32x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x64x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_16x16x256_UR_2: K must be a multiple of workgroupTile.k=256 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x64x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_128x32x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x128x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_128x64x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x128x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_16x64x256_UR_2: K must be a multiple of workgroupTile.k=256 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_192x32x128_UR_2: M must be a multiple of workgroupTile.m=192
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_192x32x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_128x128x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x192x128_UR_2: N must be a multiple of workgroupTile.n=192
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_32x192x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_192x64x128_UR_2: M must be a multiple of workgroupTile.m=192
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_192x64x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x192x128_UR_2: N must be a multiple of workgroupTile.n=192
[rr/error] Predicate mismatch for RR_GEMM_TN_FP8_FP8_Float_Float_Float_SA_B_SB_B_WGT_64x192x128_UR_2: K must be a multiple of workgroupTile.k=128 * unrollK=1
Segmentation fault (core dumped)

VeeraRajasekhar · 2026-02-15T07:12:16Z

https://github.com/ROCm/TransformerEngine/actions/runs/22024090773/job/63637830333

Level=3 testing had no issues.

Micky774 · 2026-02-16T16:31:07Z

@Micky774, please review the following,

Analysis on testing on Jax & xla 0.8.2

(Not Supported) jax.nn.scaled_matmul (MXFP8) on ROCm crashes with a segmentation fault when the contracting dimension (K) is less than 64.
...

This is a failure on certain configs for hipblaslt, which we already have tickets open for. We don't support these configs in TE anyways, so it's a known issue and not a blocker. Shouldn't be a problem for this PR.

tests/jax/utils.py

ipanfilo · 2026-02-17T16:15:04Z

.github/scripts/Dockerfile.ci.deps

+    https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/jaxlib-${JAX_VERSION}%2Brocm${ROCM_VERSION}.0-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-manylinux_2_27_x86_64.whl
+
+WORKDIR /workspace/
+CMD ["/bin/bash"]


ipanfilo · 2026-02-17T16:17:26Z

.github/scripts/Dockerfile.ci.deps

+RUN apt-get update && apt-get install -y vim
+
+# Install flash-attention
+ENV GPU_ARCHS=gfx90a;gfx950;gfx942


For CI purposes gfx90a is not needed. It is OK to keep if does not affect build time

.github/scripts/Dockerfile.ci.deps

VeeraRajasekhar requested review from ipanfilo, wangye805 and wenchenvincent May 28, 2025 22:37

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

wenchenvincent requested changes May 29, 2025

View reviewed changes

ipanfilo reviewed May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

wangye805 requested changes May 29, 2025

View reviewed changes

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

docker/Dockerfile Outdated Show resolved Hide resolved

ipanfilo requested changes May 29, 2025

View reviewed changes

wangye805 reviewed May 29, 2025

View reviewed changes

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

ipanfilo requested changes May 29, 2025

View reviewed changes

docker/Dockerfile.ci.deps Outdated Show resolved Hide resolved

VeeraRajasekhar force-pushed the dockerfile branch from c4913b2 to e49b365 Compare February 11, 2026 03:04

VeeraRajasekhar self-assigned this Feb 11, 2026

VeeraRajasekhar requested review from ipanfilo and wangye805 February 11, 2026 03:07

ipanfilo reviewed Feb 11, 2026

View reviewed changes

VeeraRajasekhar added 7 commits February 12, 2026 01:35

Added Dockerfile for CI images

30833e7

Addressed reviews

342eb9a

Addressed reviews

7ba78c3

Addressed flash attention commit to tag

4fff8ff

Moved Dockerfile.ci.deps to .github/scripts, Updated dockerfile

41c9e42

Addressed comments

d1ecdd3

Adjusted tolerance for gfx942 for ROCm7.2 support

bdc75a2

VeeraRajasekhar force-pushed the dockerfile branch from e49b365 to bdc75a2 Compare February 12, 2026 01:53

VeeraRajasekhar added 3 commits February 12, 2026 01:56

Updated CI docker image

dde1a02

[CI][Jax] Update minimum JAX version for MXFP8 JAX GEMM to 0.8.2

9662dfb

VeeraRajasekhar requested a review from ipanfilo February 14, 2026 21:00

VeeraRajasekhar changed the title ~~Added Dockerfile for CI images~~ Added Dockerfile for CI images & Upgrate CI to ROCm 7.2 Feb 16, 2026

ipanfilo requested changes Feb 17, 2026

View reviewed changes

Addressed Comments

869c260

VeeraRajasekhar requested a review from ipanfilo February 17, 2026 17:00


		ARG PYTORCH_ROCM_ARCH=gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx950;gfx1151

		# Install flash-attention v2.8.1



		RUN pip install \
		https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/jax_rocm7_pjrt-0.8.0%2Brocm7.2.0-py3-none-manylinux_2_28_x86_64.whl \

Conversation

VeeraRajasekhar commented May 28, 2025

Description

Type of change

Changes

Checklist:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wenchenvincent left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ipanfilo left a comment

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented May 29, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wenchenvincent commented Aug 27, 2025

Uh oh!

wenchenvincent commented Jan 14, 2026

Uh oh!

VeeraRajasekhar commented Feb 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

VeeraRajasekhar commented Feb 12, 2026

Uh oh!

VeeraRajasekhar commented Feb 14, 2026

Analysis on testing on Jax & xla 0.8.2

Uh oh!

VeeraRajasekhar commented Feb 15, 2026

Uh oh!

Micky774 commented Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Analysis on testing on Jax & xla 0.8.2

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Micky774 commented Feb 16, 2026 •

edited

Loading