worker: remove _distributed_c10d stub; stub only torchao
The installed torch/distributed/__init__.py from repo.amd.com
(torch==2.10.0+rocm7.12.0) is now properly guarded with
`if is_available():`, so `import torch.distributed` alone is safe.
The crash only comes via torchao's import chain:
torchao.float8.distributed_utils
→ torch.distributed._functional_collectives (unguarded import)
→ torch.distributed.distributed_c10d
→ torch._C._distributed_c10d ← absent on Windows ROCm
Stubbing torchao short-circuits the chain entirely. No need to stub
_distributed_c10d. Remove _StubClassMeta and the _c10d stub block;
keep only _make_mod_stub + _StubSubpackageFinder + torchao seeds.
This commit is contained in:
parent
9c9d462ad6
commit
48406ad340
1 changed files with 10 additions and 49 deletions
|
|
@ -1084,21 +1084,20 @@ def run_training_process(
|
|||
'Install for better performance: pip install "triton-windows<3.7"'
|
||||
)
|
||||
|
||||
# ── 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.
|
||||
#
|
||||
# ── 1d. Stub torchao on Windows ROCm ──
|
||||
# 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.
|
||||
# torch.distributed._functional_collectives at module level, which imports
|
||||
# distributed_c10d.py unconditionally — that file crashes on Windows ROCm
|
||||
# because torch._C._distributed_c10d (the RCCL backend) is absent.
|
||||
# torch/distributed/__init__.py itself is guarded by `if is_available()`
|
||||
# so `import torch.distributed` alone is safe; the crash only comes via
|
||||
# torchao's import chain. Stubbing torchao short-circuits it entirely.
|
||||
# _StubSubpackageFinder handles any depth of torchao.xxx.yyy imports.
|
||||
import types as _types
|
||||
import importlib.machinery as _ilm
|
||||
import importlib.abc as _ilabc
|
||||
|
||||
_STUB_SENTINEL = object() # identity tag on every stub module
|
||||
_STUB_SENTINEL = object()
|
||||
|
||||
def _make_mod_stub(mod_name):
|
||||
m = _types.ModuleType(mod_name)
|
||||
|
|
@ -1138,24 +1137,8 @@ def run_training_process(
|
|||
|
||||
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.
|
||||
# Seed torchao top-level + key submodules; the finder handles the rest.
|
||||
for _tao_name in (
|
||||
"torchao",
|
||||
"torchao.quantization",
|
||||
|
|
@ -1166,28 +1149,6 @@ def run_training_process(
|
|||
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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue