ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: stub distributed tensor/functional_collectives to prevent missing C++ op crash on ROCm Windows

Browse source 
torch._dynamo.trace_rules eagerly loads torch.distributed.tensor at import
time, which pulls in _functional_collectives.py. That file registers Meta
kernels for _c10d_functional C++ ops, but those ops are only registered
by torch._C._distributed_c10d — a C extension absent from ROCm Windows
wheels. Pre-stubbing the affected modules in sys.modules prevents the real
import chain from running and avoids the "operator does not exist" crash.
This commit is contained in:
LeoBorcherding 2026-05-11 05:53:18 -05:00
commit 26f073d642
1 changed files with 21 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

21
studio/backend/core/training/worker.py
View file

@ -1177,6 +1177,27 @@ def run_training_process(
if _fsdp_name not in sys.modules:
sys.modules[_fsdp_name] = _make_mod_stub(_fsdp_name)
# torch._dynamo.trace_rules.get_torch_obj_rule_map() eagerly loads
# torch.distributed.tensor, which in turn imports
# torch.distributed._functional_collectives. That module registers
# Meta kernels for ops in the _c10d_functional C++ namespace, but
# that namespace only exists when torch._C._distributed_c10d (the
# C extension absent from ROCm Windows wheels) has been loaded.
# Without it the impl() call raises "operator does not exist".
# Pre-stubbing these modules short-circuits the real import so
# torch._dynamo gets empty stub objects instead of crashing.
for _dist_name in (
"torch.distributed._functional_collectives",
"torch.distributed._functional_collectives_impl",
"torch.distributed.tensor",
"torch.distributed.tensor._ops",
"torch.distributed.tensor._ops._conv_ops",
"torch.distributed.tensor._dtensor_spec",
"torch.distributed.tensor.placement_types",
):
if _dist_name not in sys.modules:
sys.modules[_dist_name] = _make_mod_stub(_dist_name)
try:
import torch.distributed as _td