diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f6bebaaa0e..c2eb6c0083 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index bdd80ef4b4..ec2c4e14fe 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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(