fix: gate _grouped_mm dispatch patch on HIP < 7.13
AMD fixed the gfx1200 null HIP kernel in ROCm 7.13 (torch 2.11+).
Users on the new wheel now get the real GPU _grouped_mm kernel for
MoE workloads instead of the Python mm fallback.
Changes:
- worker.py: add _hip_ver_at_least() helper; wrap full _grouped_mm
patch in `if not _hip_ver_at_least(7, 13):` with else branch that
logs the skip reason; update section-1f comment to document the fix
- test_rocm_support.py: add 5 tests covering the helper definition,
the (7, 13) gate expression, the else branch, the skip log message,
and the AMD-format version string parsing (.split(".")[:2])
Verified: torch==2.11.0+rocm7.13.0 — 3D batch and grouped (offs)
variants both succeed; null crash only present on rocm7.12 and earlier.
This commit is contained in:
LeoBorcherding
parent
b313a4867c
commit
b33a90ee68
2 changed files with 135 additions and 72 deletions
|
|
@ -1291,7 +1291,7 @@ def run_training_process(
|
|||
pass
|
||||
|
||||
# ── 1f. Windows ROCm runtime patches ──
|
||||
# torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm 7.12 Windows),
|
||||
# torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows),
|
||||
# causing 0xC0000005 (access violation) during training.
|
||||
#
|
||||
# Root cause: the JitDecomp autograd decomposition system (NOT torch.compile)
|
||||
|
|
@ -1300,9 +1300,12 @@ def run_training_process(
|
|||
# JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm
|
||||
# with a safe Python fallback.
|
||||
#
|
||||
# Verified on torch==2.10.0+rocm7.12.0:
|
||||
# torch.library.Library("aten","IMPL").impl("_grouped_mm", fn, "CUDA")
|
||||
# correctly overrides the HIP kernel and the call succeeds.
|
||||
# Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped
|
||||
# (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200.
|
||||
# We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get
|
||||
# the real GPU kernel rather than our Python fallback.
|
||||
#
|
||||
# Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
|
||||
#
|
||||
# Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
|
||||
# Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor
|
||||
|
|
@ -1362,76 +1365,101 @@ def run_training_process(
|
|||
_bnb_rocm_ver,
|
||||
)
|
||||
|
||||
# Patch _grouped_mm CUDA dispatch with a safe Python mm fallback.
|
||||
try:
|
||||
import warnings as _warnings
|
||||
|
||||
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
|
||||
|
||||
def _grouped_mm_safe_impl(
|
||||
self, mat2, offs = None, bias = None, out_dtype = None
|
||||
):
|
||||
"""Safe fallback for _grouped_mm on gfx1200 (null HIP kernel)."""
|
||||
_t = _torch_for_rocm
|
||||
if offs is None:
|
||||
# Simple case: plain matrix multiply.
|
||||
result = _t.mm(self.contiguous(), mat2.contiguous())
|
||||
else:
|
||||
# Grouped case: offs[i] is the exclusive end-row of group i
|
||||
# in `self`; mat2 may be 3-D (num_groups, K, N) or 2-D.
|
||||
offs_list = offs.tolist()
|
||||
pieces = []
|
||||
prev = 0
|
||||
for idx, end in enumerate(offs_list):
|
||||
end = int(end)
|
||||
a_part = self[prev:end].contiguous()
|
||||
if mat2.dim() == 3:
|
||||
b_part = mat2[idx].contiguous()
|
||||
else:
|
||||
b_part = mat2.contiguous()
|
||||
pieces.append(_t.mm(a_part, b_part))
|
||||
prev = end
|
||||
# Include any trailing rows not covered by offs
|
||||
if prev < self.shape[0]:
|
||||
a_tail = self[prev:].contiguous()
|
||||
b_tail = (
|
||||
mat2[-1].contiguous()
|
||||
if mat2.dim() == 3
|
||||
else mat2.contiguous()
|
||||
)
|
||||
pieces.append(_t.mm(a_tail, b_tail))
|
||||
result = (
|
||||
_t.cat(pieces, dim = 0)
|
||||
if pieces
|
||||
else _t.zeros(
|
||||
0,
|
||||
mat2.shape[-1],
|
||||
device = self.device,
|
||||
dtype = self.dtype,
|
||||
)
|
||||
)
|
||||
if bias is not None:
|
||||
result = result + bias
|
||||
if out_dtype is not None:
|
||||
result = result.to(out_dtype)
|
||||
elif result.dtype != self.dtype:
|
||||
result = result.to(self.dtype)
|
||||
return result
|
||||
|
||||
with _warnings.catch_warnings():
|
||||
_warnings.simplefilter("ignore")
|
||||
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
|
||||
|
||||
_WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
|
||||
logger.info(
|
||||
"Windows ROCm: patched _grouped_mm CUDA dispatch "
|
||||
"(null HIP kernel on gfx1200 bypassed with safe mm fallback)"
|
||||
# Parse HIP version for the kernel-fix gate below.
|
||||
# torch.version.hip can be "7.13.99004", "7.2.0", etc.
|
||||
# We only need major.minor for the comparison.
|
||||
def _hip_ver_at_least(major: int, minor: int) -> bool:
|
||||
_hip_str = getattr(
|
||||
getattr(_torch_for_rocm, "version", None), "hip", None
|
||||
)
|
||||
except Exception as _patch_exc:
|
||||
logger.warning(
|
||||
"Windows ROCm: could not patch _grouped_mm — "
|
||||
"training may crash with 0xC0000005: %s",
|
||||
_patch_exc,
|
||||
if not _hip_str:
|
||||
return False
|
||||
try:
|
||||
_parts = [int(x) for x in str(_hip_str).split(".")[:2]]
|
||||
return (_parts[0], _parts[1]) >= (major, minor)
|
||||
except (ValueError, IndexError):
|
||||
return False
|
||||
|
||||
# _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12,
|
||||
# causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+).
|
||||
# Only install the Python fallback on the affected versions so users
|
||||
# on 7.13+ get the real GPU kernel for MoE workloads.
|
||||
if not _hip_ver_at_least(7, 13):
|
||||
try:
|
||||
import warnings as _warnings
|
||||
|
||||
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
|
||||
|
||||
def _grouped_mm_safe_impl(
|
||||
self, mat2, offs = None, bias = None, out_dtype = None
|
||||
):
|
||||
"""Python mm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
|
||||
_t = _torch_for_rocm
|
||||
if offs is None:
|
||||
# Simple case: plain matrix multiply.
|
||||
result = _t.mm(self.contiguous(), mat2.contiguous())
|
||||
else:
|
||||
# Grouped case: offs[i] is the exclusive end-row of
|
||||
# group i in `self`; mat2 may be 3-D or 2-D.
|
||||
offs_list = offs.tolist()
|
||||
pieces = []
|
||||
prev = 0
|
||||
for idx, end in enumerate(offs_list):
|
||||
end = int(end)
|
||||
a_part = self[prev:end].contiguous()
|
||||
if mat2.dim() == 3:
|
||||
b_part = mat2[idx].contiguous()
|
||||
else:
|
||||
b_part = mat2.contiguous()
|
||||
pieces.append(_t.mm(a_part, b_part))
|
||||
prev = end
|
||||
# Include any trailing rows not covered by offs
|
||||
if prev < self.shape[0]:
|
||||
a_tail = self[prev:].contiguous()
|
||||
b_tail = (
|
||||
mat2[-1].contiguous()
|
||||
if mat2.dim() == 3
|
||||
else mat2.contiguous()
|
||||
)
|
||||
pieces.append(_t.mm(a_tail, b_tail))
|
||||
result = (
|
||||
_t.cat(pieces, dim = 0)
|
||||
if pieces
|
||||
else _t.zeros(
|
||||
0,
|
||||
mat2.shape[-1],
|
||||
device = self.device,
|
||||
dtype = self.dtype,
|
||||
)
|
||||
)
|
||||
if bias is not None:
|
||||
result = result + bias
|
||||
if out_dtype is not None:
|
||||
result = result.to(out_dtype)
|
||||
elif result.dtype != self.dtype:
|
||||
result = result.to(self.dtype)
|
||||
return result
|
||||
|
||||
with _warnings.catch_warnings():
|
||||
_warnings.simplefilter("ignore")
|
||||
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
|
||||
|
||||
_WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
|
||||
logger.info(
|
||||
"Windows ROCm: patched _grouped_mm CUDA dispatch "
|
||||
"(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
|
||||
"bypassed with Python mm fallback)"
|
||||
)
|
||||
except Exception as _patch_exc:
|
||||
logger.warning(
|
||||
"Windows ROCm: could not patch _grouped_mm — "
|
||||
"training may crash with 0xC0000005: %s",
|
||||
_patch_exc,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, "
|
||||
"skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
|
||||
)
|
||||
|
||||
# ── 2. Now import ML libraries (fresh in this clean process) ──
|
||||
|
|
|
|||
|
|
@ -1884,6 +1884,41 @@ class TestWorkerWindowsRocmPatches:
|
|||
# Must gate on HIP version — code uses getattr chain: "version" and "hip"
|
||||
assert '"version"' in source and '"hip"' in source
|
||||
|
||||
def test_hip_ver_at_least_helper_defined(self):
|
||||
"""_hip_ver_at_least helper must be defined inside the Windows ROCm block."""
|
||||
source = _WORKER_PATH.read_text(encoding = "utf-8")
|
||||
assert "def _hip_ver_at_least(major: int, minor: int)" in source
|
||||
|
||||
def test_grouped_mm_patch_gated_on_hip_lt_713(self):
|
||||
"""_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13)."""
|
||||
source = _WORKER_PATH.read_text(encoding = "utf-8")
|
||||
# The guard must call _hip_ver_at_least with exactly (7, 13)
|
||||
assert "_hip_ver_at_least(7, 13)" in source
|
||||
# The patch must be inside the `if not` branch (negated guard)
|
||||
assert "if not _hip_ver_at_least(7, 13):" in source
|
||||
|
||||
def test_grouped_mm_hip_713_skip_message_present(self):
|
||||
"""worker.py must log a message when skipping the patch on HIP >= 7.13."""
|
||||
source = _WORKER_PATH.read_text(encoding = "utf-8")
|
||||
assert "HIP >= 7.13" in source
|
||||
assert "7.13" in source
|
||||
|
||||
def test_grouped_mm_patch_else_branch_present(self):
|
||||
"""An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+)."""
|
||||
source = _WORKER_PATH.read_text(encoding = "utf-8")
|
||||
# There must be an else: after the if not _hip_ver_at_least(7, 13): block
|
||||
gate_idx = source.find("if not _hip_ver_at_least(7, 13):")
|
||||
assert gate_idx != -1, "Version gate not found in worker.py"
|
||||
# The else: branch must appear after the gate
|
||||
else_idx = source.find("else:", gate_idx)
|
||||
assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found"
|
||||
|
||||
def test_hip_ver_at_least_handles_amd_version_format(self):
|
||||
"""_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004')."""
|
||||
source = _WORKER_PATH.read_text(encoding = "utf-8")
|
||||
# Must split the version string and take the first two parts
|
||||
assert 'split(".")[:2]' in source or ".split('.')[:2]" in source
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEST: install_python_stack.py -- _ROCM_TORCH_PKG_SPECS mapping
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue