From 26f073d642cdd30ab8e285d99fe0539d0625f4f3 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Mon, 11 May 2026 05:53:18 -0500 Subject: [PATCH] fix: stub distributed tensor/functional_collectives to prevent missing C++ op crash on ROCm Windows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- studio/backend/core/training/worker.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a9589bb65e..a654b51f9b 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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