From e778e0e39152857a581f5f5b78c154f8eb9bc1bf Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Mon, 11 May 2026 06:09:51 -0500 Subject: [PATCH] fix: stub entire torchao package on ROCm Windows instead of individual ops torchao is not supported on ROCm Windows and its import chain transitively requires torch._C._distributed_c10d (absent from the ROCm Windows wheel). Rather than stub each missing op one by one, stub the whole torchao package upfront. Unsloth uses bitsandbytes for quantization, not torchao, so this has no functional impact. transformers gracefully handles an importable-but- empty torchao by disabling TorchAoHfQuantizer. --- studio/backend/core/training/worker.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index e8bdd59af3..8db2686f15 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1148,6 +1148,22 @@ def run_training_process( return _StubClassMeta(name, (), {"__init__": lambda self, *a, **kw: None}) if sys.platform == "win32": + # torchao is not supported on ROCm Windows and its import chain + # transitively pulls in torch._C._distributed_c10d (absent from the + # ROCm Windows wheel), causing cascading AttributeErrors. We don't + # use torchao quantization (unsloth uses bitsandbytes), so stub the + # entire package up-front. transformers falls back gracefully when + # torchao is importable but empty. + 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) + _c10d_key = "torch._C._distributed_c10d" if _c10d_key not in sys.modules: # guard: never overwrite real NVIDIA impl _c10d_stub = _types.ModuleType(_c10d_key)