fix(studio/worker): inject --gcc-install-dir for HIP source builds on Ubuntu 24.04 (#5517)

* fix(studio/worker): inject --gcc-install-dir for HIP source builds on Ubuntu 24.04

On Ubuntu 24.04 + ROCm clang-20, the HIP source-build fallback in
`_install_package_wheel_first` (causal-conv1d, mamba-ssm source fallback,
flash-attn source fallback) dies at:

  /opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10:
    fatal error: 'cstdlib' file not found

Root cause: clang-20 picks the highest-numbered /usr/lib/gcc/x86_64-linux-gnu/<N>
runtime dir by default. On 24.04 that's gcc-14, whose runtime objects ship in
the gcc-14 package but whose C++ headers (/usr/include/c++/14) come from
libstdc++-14-dev — NOT in the default apt set. libstdc++-13-dev IS in the
default set, so /usr/include/c++/13 exists. clang has no way to discover
that asymmetry and the build fails.

Fix: new `_hipcc_gcc_install_dir()` helper iterates gcc 14 → 11 and returns
the first /usr/lib/gcc/x86_64-linux-gnu/<N> dir where BOTH the runtime AND
/usr/include/c++/<N> exist. The HIP branch of `_install_package_wheel_first`
appends `--gcc-install-dir=<that path>` to HIPCC_COMPILE_FLAGS_APPEND before
invoking pip. Respects an existing `--gcc-install-dir` in the env var
(user-set takes precedence); preserves any other flags the user has set
(appends to the end rather than overwriting). No-op on non-HIP, non-Linux,
non-x86_64.

Mirrors the same fix bbf004c added to studio/setup.sh for the llama.cpp HIP
build branch (#5301), but via env var since pip-driven source builds can't
take CMake flags directly.

Verified on Ryzen AI MAX+ 395 / Radeon 8060S (gfx1151) / Ubuntu 24.04 /
ROCm 7.13 nightly: `_hipcc_gcc_install_dir()` returns
`/usr/lib/gcc/x86_64-linux-gnu/13`, which matches the manual workaround
that already lets `pip install causal-conv1d` succeed on this hardware.

Tests added (8 new in test_training_worker_flash_attn.py):
- test_hipcc_gcc_install_dir_picks_highest_with_headers
- test_hipcc_gcc_install_dir_picks_14_when_headers_exist
- test_hipcc_gcc_install_dir_returns_none_when_no_match
- test_hipcc_gcc_install_dir_returns_none_on_non_linux
- test_hipcc_gcc_install_dir_returns_none_on_non_x86_64
- test_install_injects_gcc_install_dir_on_hip_source_build
- test_install_appends_to_existing_hipcc_compile_flags
- test_install_respects_user_gcc_install_dir
- test_install_does_not_inject_env_on_cuda

Per @danielhanchen's suggestion in
https://github.com/unslothai/unsloth/pull/5434#issuecomment-4469980122

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* review: apply gemini-code-assist suggestion on _run_kwargs env handling

Use _run_kwargs.get("env", os.environ).copy() + key-mutation instead of
rebuilding env from os.environ directly. Today both forms are equivalent
(no earlier code in _install_package_wheel_first sets _run_kwargs["env"]),
but the .get().copy() pattern survives any future env modification added
upstream of this block without silently throwing it away.

No behavioural change; tests already assert the final HIPCC_COMPILE_FLAGS_APPEND
value, not the env-construction pattern.

Per https://github.com/unslothai/unsloth/pull/5517#discussion_r... (gemini-code-assist[bot])

---------

Co-authored-by: h34v3nzc0dex <h34v3nzc0dex@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Paul Durkin 2026-05-18 02:05:30 -06:00 committed by GitHub
commit 388ade4c84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 335 additions and 0 deletions

View file

@ -77,6 +77,38 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
)
def _hipcc_gcc_install_dir() -> str | None:
"""Return the highest-numbered ``/usr/lib/gcc/x86_64-linux-gnu/<N>`` that has
BOTH the gcc runtime dir AND the corresponding ``/usr/include/c++/<N>`` C++
headers, or ``None`` if no match (or non-Linux / non-x86_64).
Ubuntu 24.04 ships ``/usr/lib/gcc/x86_64-linux-gnu/14/`` (gcc-14 runtime
objects) but does NOT ship ``/usr/include/c++/14`` in its default apt set;
libstdc++ headers come from ``libstdc++-13-dev``. ROCm clang-20 picks the
highest-numbered runtime dir by default, finds no ``<cstdlib>``, and the
HIP source build fails with::
/opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10:
fatal error: 'cstdlib' file not found
Returning a path lets the caller pass ``--gcc-install-dir=<path>`` to clang
via ``HIPCC_COMPILE_FLAGS_APPEND``. Mirrors the same loop ``bbf004c`` added
to ``studio/setup.sh`` for the llama.cpp HIP build branch (PR #5301).
"""
if not sys.platform.startswith("linux"):
return None
import platform as _platform
if _platform.machine().lower() != "x86_64":
return None
for _ver in (14, 13, 12, 11):
_runtime = f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}/include"
_headers = f"/usr/include/c++/{_ver}"
if os.path.isdir(_runtime) and os.path.isdir(_headers):
return f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}"
return None
def _install_package_wheel_first(
*,
event_queue: Any,
@ -212,6 +244,30 @@ def _install_package_wheel_first(
}
if is_hip:
_run_kwargs["timeout"] = 1800
# On Ubuntu 24.04 + ROCm clang-20, the HIP source build (causal-conv1d,
# mamba-ssm source fallback, flash-attn source fallback) defaults to
# /usr/lib/gcc/x86_64-linux-gnu/14/ which has the runtime dir but no
# /usr/include/c++/14 headers, and dies at:
# __clang_hip_runtime_wrapper.h:112:10:
# fatal error: 'cstdlib' file not found
# Inject --gcc-install-dir for a gcc whose C++ headers actually exist.
# Respect any pre-existing --gcc-install-dir in HIPCC_COMPILE_FLAGS_APPEND
# (user knows best); otherwise append. Mirrors the same fix bbf004c
# added to studio/setup.sh for the llama.cpp HIP build (PR #5301).
_existing_flags = os.environ.get("HIPCC_COMPILE_FLAGS_APPEND", "")
if "--gcc-install-dir" not in _existing_flags:
_gcc_dir = _hipcc_gcc_install_dir()
if _gcc_dir is not None:
_appended = (f"{_existing_flags} --gcc-install-dir={_gcc_dir}").strip()
_env = _run_kwargs.get("env", os.environ).copy()
_env["HIPCC_COMPILE_FLAGS_APPEND"] = _appended
_run_kwargs["env"] = _env
logger.info(
"HIP source build for %s: appended "
"--gcc-install-dir=%s to HIPCC_COMPILE_FLAGS_APPEND",
display_name,
_gcc_dir,
)
try:
result = _sp.run(pypi_cmd, **_run_kwargs)

View file

@ -6,6 +6,7 @@ from __future__ import annotations
import builtins
import subprocess
import sys
from typing import Any
from unittest import mock
from core.training import worker
@ -22,6 +23,17 @@ def _missing_flash_attn_import():
return fake_import
def _missing_module_import(missing: str):
real_import = builtins.__import__
def fake_import(name, globals = None, locals = None, fromlist = (), level = 0):
if name == missing:
raise ImportError
return real_import(name, globals, locals, fromlist, level)
return fake_import
def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch):
monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False)
assert worker._should_try_runtime_flash_attn_install(32767) is False
@ -193,3 +205,270 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch):
release_tag = worker._MAMBA_SSM_RELEASE_TAG,
release_base_url = "https://github.com/state-spaces/mamba/releases/download",
)
# ────────────────────────────────────────────────────────────────────
# HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo).
# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14,
# so ROCm clang-20 picks it and fails with 'cstdlib' file not found
# when building causal-conv1d (or any other HIP source fallback).
# _hipcc_gcc_install_dir() finds a gcc dir that has both halves; the
# _install_package_wheel_first HIP branch passes it to clang via
# HIPCC_COMPILE_FLAGS_APPEND. Parallel to bbf004c's setup.sh fix for
# the llama.cpp HIP build (PR #5301).
# ────────────────────────────────────────────────────────────────────
def _isdir_for_layout(*existing: str):
"""Return an os.path.isdir replacement that only treats the given
absolute paths as directories. Lets a test simulate exactly which
gcc runtime dirs and C++ header dirs exist on the host."""
valid = set(existing)
def fake_isdir(path: str) -> bool:
return path in valid
return fake_isdir
def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch):
"""gcc-14 has runtime but no /usr/include/c++/14; loop falls through
to gcc-13 which has both. This is the exact Ubuntu 24.04 layout."""
monkeypatch.setattr(sys, "platform", "linux")
import platform as _platform
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
monkeypatch.setattr(
worker.os.path,
"isdir",
_isdir_for_layout(
"/usr/lib/gcc/x86_64-linux-gnu/14/include", # runtime present
# but no /usr/include/c++/14 — typical Ubuntu 24.04 default
"/usr/lib/gcc/x86_64-linux-gnu/13/include",
"/usr/include/c++/13", # libstdc++-13-dev installed
),
)
assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/13"
def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch):
"""If the user has libstdc++-14-dev installed, prefer gcc-14."""
monkeypatch.setattr(sys, "platform", "linux")
import platform as _platform
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
monkeypatch.setattr(
worker.os.path,
"isdir",
_isdir_for_layout(
"/usr/lib/gcc/x86_64-linux-gnu/14/include",
"/usr/include/c++/14",
),
)
assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/14"
def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch):
"""No gcc dir has both halves → return None and skip the env injection
rather than guessing wrong and surfacing a confusing build failure."""
monkeypatch.setattr(sys, "platform", "linux")
import platform as _platform
monkeypatch.setattr(_platform, "machine", lambda: "x86_64")
monkeypatch.setattr(worker.os.path, "isdir", lambda path: False)
assert worker._hipcc_gcc_install_dir() is None
def test_hipcc_gcc_install_dir_returns_none_on_non_linux(monkeypatch):
"""Don't probe gcc layout on macOS / Windows — early-return."""
monkeypatch.setattr(sys, "platform", "darwin")
def _isdir_should_not_be_called(_path):
raise AssertionError("isdir should not be called on non-Linux")
monkeypatch.setattr(worker.os.path, "isdir", _isdir_should_not_be_called)
assert worker._hipcc_gcc_install_dir() is None
def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch):
"""ROCm clang-20 on aarch64 has a different libstdc++ layout."""
monkeypatch.setattr(sys, "platform", "linux")
import platform as _platform
monkeypatch.setattr(_platform, "machine", lambda: "aarch64")
assert worker._hipcc_gcc_install_dir() is None
def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None):
"""Common scaffolding for tests that exercise the HIP source-build
branch of _install_package_wheel_first end-to-end. The package isn't
installed yet, no prebuilt wheel exists, hipcc is on PATH, and the
fake env reports an HIP torch."""
monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d"))
monkeypatch.setattr(
worker,
"probe_torch_wheel_env",
lambda timeout = 30: {
"hip_version": "7.13.26176",
"python_tag": "cp312",
"torch_mm": "2.11",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
},
)
monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None)
monkeypatch.setattr(
worker.shutil,
"which",
lambda name: "/opt/rocm/bin/hipcc" if name == "hipcc" else None,
)
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
monkeypatch.setattr(worker, "_hipcc_gcc_install_dir", lambda: gcc_dir)
def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch):
"""HIP source-build with no user-set HIPCC_COMPILE_FLAGS_APPEND →
subprocess env carries --gcc-install-dir=<detected path>."""
monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False)
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
captured: dict[str, str] = {}
def fake_run(cmd, **kwargs):
captured.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(cmd, 0, "")
monkeypatch.setattr(worker._sp, "run", fake_run)
worker._install_package_wheel_first(
event_queue = [],
import_name = "causal_conv1d",
display_name = "causal-conv1d",
pypi_name = "causal-conv1d",
pypi_version = "1.6.2.post1",
filename_prefix = "causal_conv1d",
release_tag = "v1.6.2.post1",
release_base_url = "https://example.com",
)
assert (
captured.get("HIPCC_COMPILE_FLAGS_APPEND")
== "--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13"
)
def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch):
"""User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value
keeps the user's flags AND adds --gcc-install-dir at the end."""
monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO")
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
captured: dict[str, str] = {}
def fake_run(cmd, **kwargs):
captured.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(cmd, 0, "")
monkeypatch.setattr(worker._sp, "run", fake_run)
worker._install_package_wheel_first(
event_queue = [],
import_name = "causal_conv1d",
display_name = "causal-conv1d",
pypi_name = "causal-conv1d",
pypi_version = "1.6.2.post1",
filename_prefix = "causal_conv1d",
release_tag = "v1.6.2.post1",
release_base_url = "https://example.com",
)
assert captured.get("HIPCC_COMPILE_FLAGS_APPEND") == (
"-O3 -DFOO --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13"
)
def test_install_respects_user_gcc_install_dir(monkeypatch):
"""User explicitly set --gcc-install-dir=… already → don't touch it.
Avoids two competing --gcc-install-dir flags on the clang command line."""
monkeypatch.setenv(
"HIPCC_COMPILE_FLAGS_APPEND",
"--gcc-install-dir=/opt/custom/gcc-13",
)
_make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
captured: dict[str, str] | None = {"_called": "no"}
def fake_run(cmd, **kwargs):
env = kwargs.get("env")
if env is not None:
captured.clear()
captured.update(env)
else:
captured["_called"] = "yes_no_env"
return subprocess.CompletedProcess(cmd, 0, "")
monkeypatch.setattr(worker._sp, "run", fake_run)
worker._install_package_wheel_first(
event_queue = [],
import_name = "causal_conv1d",
display_name = "causal-conv1d",
pypi_name = "causal-conv1d",
pypi_version = "1.6.2.post1",
filename_prefix = "causal_conv1d",
release_tag = "v1.6.2.post1",
release_base_url = "https://example.com",
)
# subprocess.run was invoked without env override (the user already
# set HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left
# the env alone — the existing value is inherited normally).
assert captured == {"_called": "yes_no_env"}
def test_install_does_not_inject_env_on_cuda(monkeypatch):
"""CUDA path (no hip_version in env) → no env override at all."""
monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False)
monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d"))
monkeypatch.setattr(
worker,
"probe_torch_wheel_env",
lambda timeout = 30: {
"python_tag": "cp312",
"torch_mm": "2.11",
"cuda_major": "12",
"cxx11abi": "TRUE",
"platform_tag": "linux_x86_64",
},
)
monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None)
monkeypatch.setattr(worker.shutil, "which", lambda name: None)
monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
# If _hipcc_gcc_install_dir were called on CUDA we'd want to know.
monkeypatch.setattr(
worker,
"_hipcc_gcc_install_dir",
lambda: (_ for _ in ()).throw(AssertionError("must not run on CUDA")),
)
captured: dict[str, Any] = {}
def fake_run(cmd, **kwargs):
captured["env_in_kwargs"] = "env" in kwargs
return subprocess.CompletedProcess(cmd, 0, "")
monkeypatch.setattr(worker._sp, "run", fake_run)
worker._install_package_wheel_first(
event_queue = [],
import_name = "causal_conv1d",
display_name = "causal-conv1d",
pypi_name = "causal-conv1d",
pypi_version = "1.6.2.post1",
filename_prefix = "causal_conv1d",
release_tag = "v1.6.2.post1",
release_base_url = "https://example.com",
)
# CUDA branch never sets the env, never invokes the gcc helper.
assert captured.get("env_in_kwargs") is False