fix(rocm/win): restore _distributed_c10d + torchao stubs; fix BNB install
repo.amd.com torch wheels also omit torch._C._distributed_c10d on Windows (RCCL is not shipped on Windows). torch/distributed/__init__.py imports from it unconditionally at module level, so the stub must land in sys.modules before any torch.distributed import. torchao (pulled in by transformers.quantizers) walks torchao.float8.distributed_utils -> torch.distributed._functional_collectives -> distributed_c10d at import time. Stubbing torchao up-front short-circuits that chain. worker.py: - Restore _make_mod_stub / _StubSubpackageFinder / _StubSubpackageLoader - Restore _StubClassMeta for ProcessGroup.BackendType attribute access - Restore _distributed_c10d stub with __getattr__ (Windows only) - Restore torchao stubs (5 modules, Windows only) install_python_stack.py: - BNB AMD wheel install was inside the early-return branch that fires when torch is already a ROCm build (installed by install.ps1). Move BNB install outside that branch so it always runs on Windows ROCm — the PyPI bitsandbytes has only CUDA DLLs and fails to load on ROCm.
This commit is contained in:
LeoBorcherding
parent
d731c5fc63
commit
9c9d462ad6
2 changed files with 124 additions and 19 deletions
|
|
@ -1084,12 +1084,113 @@ def run_training_process(
|
|||
'Install for better performance: pip install "triton-windows<3.7"'
|
||||
)
|
||||
|
||||
# ── 1d. Ensure torch.distributed helper attrs are present ──
|
||||
# Single-GPU training never initialises the process group, so these helpers
|
||||
# are never called — but transformers/trl import them unconditionally at the
|
||||
# module level and crash when they're missing.
|
||||
# ── 1d. Pre-stub torch._C._distributed_c10d and torchao ──
|
||||
# Windows ROCm wheels (both repo.radeon.com and repo.amd.com) omit the
|
||||
# _distributed_c10d C++ extension — RCCL is not shipped on Windows.
|
||||
# torch/distributed/__init__.py and distributed_c10d.py both import from it
|
||||
# unconditionally at module level, so the stub must be in sys.modules BEFORE
|
||||
# any `import torch.distributed` call.
|
||||
#
|
||||
# torchao (pulled in by transformers.quantizers) imports
|
||||
# torch.distributed._functional_collectives → distributed_c10d at import
|
||||
# time. Stubbing the entire torchao package short-circuits that chain.
|
||||
import types as _types
|
||||
import importlib.machinery as _ilm
|
||||
import importlib.abc as _ilabc
|
||||
|
||||
_STUB_SENTINEL = object() # identity tag on every stub module
|
||||
|
||||
def _make_mod_stub(mod_name):
|
||||
m = _types.ModuleType(mod_name)
|
||||
m.__path__ = []
|
||||
m.__package__ = mod_name
|
||||
m._unsloth_stub = _STUB_SENTINEL
|
||||
m.__spec__ = _ilm.ModuleSpec(mod_name, loader=None, is_package=True)
|
||||
def _ga(attr, _m=m, _n=mod_name):
|
||||
if attr.startswith("__"):
|
||||
raise AttributeError(attr)
|
||||
child_name = f"{_n}.{attr}"
|
||||
child = _make_mod_stub(child_name)
|
||||
sys.modules.setdefault(child_name, child)
|
||||
setattr(_m, attr, child)
|
||||
return child
|
||||
m.__getattr__ = _ga
|
||||
return m
|
||||
|
||||
class _StubSubpackageLoader(_ilabc.Loader):
|
||||
def __init__(self, mod_name):
|
||||
self._mod_name = mod_name
|
||||
def create_module(self, spec):
|
||||
return _make_mod_stub(self._mod_name)
|
||||
def exec_module(self, module):
|
||||
pass
|
||||
|
||||
class _StubSubpackageFinder(_ilabc.MetaPathFinder):
|
||||
def find_spec(self, fullname, path, target=None):
|
||||
if "." not in fullname:
|
||||
return None
|
||||
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
|
||||
if parent is None:
|
||||
return None
|
||||
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
|
||||
return None
|
||||
return _ilm.ModuleSpec(fullname, _StubSubpackageLoader(fullname), is_package=True)
|
||||
|
||||
sys.meta_path.append(_StubSubpackageFinder())
|
||||
|
||||
# Metaclass so stub class attributes (e.g. ProcessGroup.BackendType.NCCL)
|
||||
# don't raise AttributeError.
|
||||
class _StubClassMeta(type):
|
||||
def __getattr__(cls, attr):
|
||||
if attr == "__members__":
|
||||
return {}
|
||||
if attr.startswith("__"):
|
||||
raise AttributeError(attr)
|
||||
child = _StubClassMeta(attr, (), {"__init__": lambda self, *a, **kw: None})
|
||||
setattr(cls, attr, child)
|
||||
return child
|
||||
|
||||
def _make_stub_class(name):
|
||||
return _StubClassMeta(name, (), {"__init__": lambda self, *a, **kw: None})
|
||||
|
||||
if sys.platform == "win32":
|
||||
# Stub torchao up-front so its import chain never reaches
|
||||
# torch.distributed._functional_collectives.
|
||||
for _tao_name in (
|
||||
"torchao",
|
||||
"torchao.quantization",
|
||||
"torchao.dtypes",
|
||||
"torchao.float8",
|
||||
"torchao.utils",
|
||||
):
|
||||
if _tao_name not in sys.modules:
|
||||
sys.modules[_tao_name] = _make_mod_stub(_tao_name)
|
||||
|
||||
# Stub torch._C._distributed_c10d so torch/distributed/__init__.py
|
||||
# and distributed_c10d.py can import from it without crashing.
|
||||
_c10d_key = "torch._C._distributed_c10d"
|
||||
if _c10d_key not in sys.modules:
|
||||
_c10d_stub = _types.ModuleType(_c10d_key)
|
||||
|
||||
def _c10d_stub_getattr(_attr):
|
||||
if _attr.startswith("__"):
|
||||
raise AttributeError(_attr)
|
||||
_cls = _make_stub_class(_attr)
|
||||
setattr(_c10d_stub, _attr, _cls)
|
||||
return _cls
|
||||
|
||||
_c10d_stub.__getattr__ = _c10d_stub_getattr
|
||||
sys.modules[_c10d_key] = _c10d_stub
|
||||
try:
|
||||
import torch._C as _torch_C_mod
|
||||
if not hasattr(_torch_C_mod, "_distributed_c10d"):
|
||||
_torch_C_mod._distributed_c10d = _c10d_stub
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── 1e. Ensure torch.distributed helper attrs are present ──
|
||||
# Single-GPU training never initialises the process group, so these helpers
|
||||
# are never called — but transformers/trl import them unconditionally.
|
||||
_td_stubs = {
|
||||
"is_initialized": lambda: False,
|
||||
"is_available": lambda: False,
|
||||
|
|
|
|||
|
|
@ -362,6 +362,8 @@ def _ensure_rocm_torch() -> None:
|
|||
gfx_arch = _detect_windows_gfx_arch()
|
||||
if not gfx_arch:
|
||||
return # no AMD GPU visible via hipinfo
|
||||
# Probe whether torch already links against HIP.
|
||||
_torch_already_rocm = False
|
||||
try:
|
||||
probe = subprocess.run(
|
||||
[
|
||||
|
|
@ -379,23 +381,25 @@ def _ensure_rocm_torch() -> None:
|
|||
timeout = 30,
|
||||
)
|
||||
if probe.returncode == 0 and probe.stdout.decode().strip() == "yes":
|
||||
_rocm_windows_torch_installed = True
|
||||
return # already ROCm torch
|
||||
_torch_already_rocm = True
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
index_url = _windows_rocm_index_url(gfx_arch)
|
||||
if index_url is None:
|
||||
print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping")
|
||||
return
|
||||
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
|
||||
pip_install(
|
||||
f"ROCm torch (Windows, {gfx_arch})",
|
||||
"--force-reinstall",
|
||||
"--index-url", index_url,
|
||||
"torch", "torchvision", "torchaudio",
|
||||
constrain = False,
|
||||
)
|
||||
# bitsandbytes Windows ROCm wheel.
|
||||
if not _torch_already_rocm:
|
||||
index_url = _windows_rocm_index_url(gfx_arch)
|
||||
if index_url is None:
|
||||
print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping")
|
||||
return
|
||||
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
|
||||
pip_install(
|
||||
f"ROCm torch (Windows, {gfx_arch})",
|
||||
"--force-reinstall",
|
||||
"--index-url", index_url,
|
||||
"torch", "torchvision", "torchaudio",
|
||||
constrain = False,
|
||||
)
|
||||
# Always install AMD Windows bitsandbytes — the PyPI wheel ships only
|
||||
# CUDA DLLs and will fail to load on ROCm. Install even when torch was
|
||||
# already a ROCm build so that `studio update` repairs a broken bnb.
|
||||
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
|
||||
if _bnb_win_url is not None:
|
||||
pip_install_try(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue