Do not import torch.distributed in model_has_dtensor on a non-distributed build - #4250
Open
shoemoney wants to merge 1 commit into
Open
Conversation
…ild 🧩 Fixes huggingface#4249. model_has_dtensor() imports torch.distributed.tensor unconditionally. On a PyTorch build compiled without a distributed backend that import raises rather than returning False, and since v1.15.0 every prepare_model() call reaches it: accelerator.py:1802 device_placement = (... and not model_has_dtensor(model)) utils/other.py:243 from torch.distributed.tensor import DTensor ... torch/distributed/distributed_c10d.py:25 ModuleNotFoundError: No module named 'torch._C._distributed_c10d' AMD's official Windows ROCm wheels are exactly that build: torch.distributed imports, torch.distributed.is_available() is False, and torch._C._distributed_c10d does not exist. So all training on AMD GPUs on Windows dies at trainer start. In 1.14.0 the sole call site (accelerator.py:1878) sat behind `if self.multi_device and not (...)`, so a single-GPU run never reached it. The second call site added in huggingface#4181 is unguarded, which is what turned a multi-device-only path into one every run executes. Guarding model_has_dtensor itself rather than restoring the guard on the new call site, because the invariant belongs to the function: DTensor lives under torch.distributed, so a build without a distributed backend cannot hold DTensor parameters at all. That also covers any future call site. is_torch_distributed_available() is the existing idiom for this. It is already imported in this same file and already used as a guard 70 lines further down, and accelerator.py:2958 already computes `is_dtensor_available = torch.distributed.is_available() and ...`, so the "DTensor requires a distributed build" relationship is established. 2 tests added. The first reproduces the reported failure by making the import raise the way it does on those wheels, so it fails against unpatched sources with the bug report's own error: ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package The second pins that the guard does not short-circuit the real check on a normal build. tests/test_utils.py: 48 passed, 4 skipped. ruff check and ruff format clean.
danielhanchen
added a commit
to LeoBorcherding/unsloth
that referenced
this pull request
Sep 14, 2026
…io update`
The constraint is right, but only one install step passes -c, and on Windows that
step does not run.
install.ps1 resolves unsloth and unsloth-zoo itself with raw `uv pip install` and no
-c (6412, 6617, 6648), which is the resolve that pulls accelerate>=0.34.1 and lands
1.15.0. It then sets SKIP_STUDIO_BASE=1 at 7022, and install_python_stack.py:7712 is
`if skip_base: pass`, so the core-packages step at 7780 that would have applied the
cap is skipped. Nothing afterwards revisits accelerate: extras.txt, studio.txt,
data-designer-deps.txt and diffusers-pin.txt each resolve with 1.15.0 installed and
none of the four mentions it. pip documents this directly ("including a package in a
constraints file does not trigger installation of the package"), so with no step
requesting accelerate there is no resolution left to constrain.
That leaves install and Repair, the two paths in the report, still broken. Only
`unsloth studio update` benefits, because unsloth_cli/commands/studio.py:3728 pops
SKIP_STUDIO_BASE.
The anyio precedent this follows is two pieces, and only one was copied. The comment
at install_python_stack.py:932 says so: constraints.txt caps anyio, and
_repair_bad_anyio() at 951 handles the install that is already past the cap. It runs
as step 8b, outside the skip_base branch, so it reaches the installer handoff.
So add _repair_bad_accelerate() in the same shape, gated on Windows, as step 8c, and
generalise _installed_anyio_version() to _installed_version(package) rather than
copy its parser. The constraints line stays; it is what holds the version down on
the update path, and it was verified to downgrade an installed 1.15.0 there.
--no-deps is load bearing and is the one place this departs from the anyio step.
accelerate requires torch>=2.0.0, and `--force-reinstall accelerate<1.15.0` without
it resolves torch too: measured against a venv holding torch 2.11.0+rocm7.13.0, both
uv and pip replace it with torch 2.14.0 from PyPI. Copying the anyio flags verbatim
would have destroyed the ROCm build this exists to protect.
Verified
- The break and the fix, with no Windows or AMD hardware: an import hook that fails
torch.distributed.tensor the way AMD's Windows ROCm wheels do, then
Accelerator().prepare_model(nn.Linear(4, 4)). accelerate 1.15.0 raises the reported
ModuleNotFoundError, 1.14.0 returns normally.
- 240 simulated installs over [Windows, Linux, WSL, macOS arm, macOS Intel] x
[NVIDIA, AMD, CPU] x [installer handoff, studio update] x [accelerate absent,
1.14.0, 1.15.0, 2.0.0] x [torch, no-torch]. Every cell exits 0 and lands the
progress bar exactly on its total. The repair fires on 24 cells, all Windows with
accelerate at or above 1.15, and always as --no-deps.
- Markers in a -c file are honoured by pip 26.2.1 and uv 0.10.7, directly and
transitively, and the cap is inert on every non-Windows platform.
- `uv pip compile --python-platform` against the full constraints.txt: Windows x64
lands accelerate 1.14.0 next to transformers 5.5.0 and trl 0.23.1 with no conflict,
while Linux x64/arm64 and both macOS targets stay on 1.15.0. Windows ARM64 fails
identically with the cap, without it, and with no constraints file at all, on torch
wheels, so it is untouched by this.
- Removing the cap resolves 1.15.0 again on Windows, so lifting it once
huggingface/accelerate#4250 lands is a one-line revert plus this step.
- tests/studio/install plus test_torchao_select: 4115 passed, 6 skipped.
Both existing step-total guards needed the Windows count raised, which is the point
of having them: test_rocm_support.py's progress check and
test_windows_torch_flavor_invariant.py's _base_total, the latter re-executing the
real arithmetic out of the source.
New tests fail when the base_total bump is dropped (progress overshoots on Windows),
when --no-deps is swapped back for --force-reinstall, and when the Windows gate is
removed.
danielhanchen
added a commit
to unslothai/unsloth
that referenced
this pull request
Sep 14, 2026
…OCm training (#10819) * Cap accelerate below 1.15 on Windows, where 1.15 cannot start ROCm training accelerate 1.15.0 made Accelerator.prepare_model call model_has_dtensor() unconditionally, which imports torch.distributed.tensor and so torch._C._distributed_c10d. AMD's Windows ROCm wheels ship no c10d, so every single-GPU run fails at trainer start with ModuleNotFoundError. Reported upstream as huggingface/accelerate#4249. * Make the accelerate cap reach a fresh Windows install, not just `studio update` The constraint is right, but only one install step passes -c, and on Windows that step does not run. install.ps1 resolves unsloth and unsloth-zoo itself with raw `uv pip install` and no -c (6412, 6617, 6648), which is the resolve that pulls accelerate>=0.34.1 and lands 1.15.0. It then sets SKIP_STUDIO_BASE=1 at 7022, and install_python_stack.py:7712 is `if skip_base: pass`, so the core-packages step at 7780 that would have applied the cap is skipped. Nothing afterwards revisits accelerate: extras.txt, studio.txt, data-designer-deps.txt and diffusers-pin.txt each resolve with 1.15.0 installed and none of the four mentions it. pip documents this directly ("including a package in a constraints file does not trigger installation of the package"), so with no step requesting accelerate there is no resolution left to constrain. That leaves install and Repair, the two paths in the report, still broken. Only `unsloth studio update` benefits, because unsloth_cli/commands/studio.py:3728 pops SKIP_STUDIO_BASE. The anyio precedent this follows is two pieces, and only one was copied. The comment at install_python_stack.py:932 says so: constraints.txt caps anyio, and _repair_bad_anyio() at 951 handles the install that is already past the cap. It runs as step 8b, outside the skip_base branch, so it reaches the installer handoff. So add _repair_bad_accelerate() in the same shape, gated on Windows, as step 8c, and generalise _installed_anyio_version() to _installed_version(package) rather than copy its parser. The constraints line stays; it is what holds the version down on the update path, and it was verified to downgrade an installed 1.15.0 there. --no-deps is load bearing and is the one place this departs from the anyio step. accelerate requires torch>=2.0.0, and `--force-reinstall accelerate<1.15.0` without it resolves torch too: measured against a venv holding torch 2.11.0+rocm7.13.0, both uv and pip replace it with torch 2.14.0 from PyPI. Copying the anyio flags verbatim would have destroyed the ROCm build this exists to protect. Verified - The break and the fix, with no Windows or AMD hardware: an import hook that fails torch.distributed.tensor the way AMD's Windows ROCm wheels do, then Accelerator().prepare_model(nn.Linear(4, 4)). accelerate 1.15.0 raises the reported ModuleNotFoundError, 1.14.0 returns normally. - 240 simulated installs over [Windows, Linux, WSL, macOS arm, macOS Intel] x [NVIDIA, AMD, CPU] x [installer handoff, studio update] x [accelerate absent, 1.14.0, 1.15.0, 2.0.0] x [torch, no-torch]. Every cell exits 0 and lands the progress bar exactly on its total. The repair fires on 24 cells, all Windows with accelerate at or above 1.15, and always as --no-deps. - Markers in a -c file are honoured by pip 26.2.1 and uv 0.10.7, directly and transitively, and the cap is inert on every non-Windows platform. - `uv pip compile --python-platform` against the full constraints.txt: Windows x64 lands accelerate 1.14.0 next to transformers 5.5.0 and trl 0.23.1 with no conflict, while Linux x64/arm64 and both macOS targets stay on 1.15.0. Windows ARM64 fails identically with the cap, without it, and with no constraints file at all, on torch wheels, so it is untouched by this. - Removing the cap resolves 1.15.0 again on Windows, so lifting it once huggingface/accelerate#4250 lands is a one-line revert plus this step. - tests/studio/install plus test_torchao_select: 4115 passed, 6 skipped. Both existing step-total guards needed the Windows count raised, which is the point of having them: test_rocm_support.py's progress check and test_windows_torch_flavor_invariant.py's _base_total, the latter re-executing the real arithmetic out of the source. New tests fail when the base_total bump is dropped (progress overshoots on Windows), when --no-deps is swapped back for --force-reinstall, and when the Windows gate is removed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the accelerate repair nonfatal, so one unreachable index cannot fail an install `pip_install` does not return on failure. It reaches `_report_failed_command`, which prints and calls `sys.exit(result.returncode)`, so any step using it aborts the whole install. That is tolerable for the anyio repair, which only fires on a venv that is already past the cap. It is not tolerable here. On Windows this step fires on essentially every fresh install today, because install.ps1 resolves accelerate with no -c and lands 1.15.0, so a transient PyPI failure, a proxy, or an offline box would turn an install that used to finish into one that dies at step 8c. The bug this PR fixes leaves you with a Studio that cannot train; that regression would leave you with no Studio at all. So `pip_install_try`, which returns False instead of exiting, and a warning naming the consequence when it does. It is the same call `_ensure_rocm_torch` makes for the AMD index ("Nonfatal: a transient AMD-index failure must not abort the install"), and the same shape as the torchcodec step, which already reports "the rest of the install is unaffected". Found by running the real installer rather than by reading it: the end-to-end run on the Windows gfx1151 runner is what made the failure mode concrete. Verified - 720 simulated installs, adding a repair-fails dimension to the earlier sweep: [Windows, Linux, WSL, macOS arm, macOS Intel] x [NVIDIA, AMD, CPU] x [installer handoff, studio update] x [absent, 0.34.1, 1.14.0, 1.15.0, 1.15.0rc1, 2.0.0] x [torch, no-torch] x [repair succeeds, repair fails]. Every cell exits 0 with the progress bar on its total, including all 36 where the repair fails. The sweep fails the run if the installer exits. - Version parsing over the shapes importlib.metadata really returns: 1.15, 1.15.0.dev0, 1.15.0rc1 and 1.15.0.post1 all repair; 1.14.0, 1.9.0 and 0.34.1 do not; an epoch (1!1.0), an empty string, junk and a missing distribution all read as "leave it alone" rather than raising. - tests/studio/install plus test_torchao_select: 4117 passed, 6 skipped. ruff clean. Two new tests, both mutation checked against restoring `pip_install`: one asserts the installer survives a failing repair, the other reads the function's own source, since a future edit could swap the call back without changing anything a mock can observe. * Tighten the comments added by this PR Opening comment-reduction pass. Every surviving line names a specific mistake the next reader would otherwise make: that the constraints cap alone is enough, that --no-deps is optional, that pip_install is interchangeable with pip_install_try, and that step 8c could live inside the skip_base branch. * Final comment pass: drop the pointer line, keep every fact --------- Co-authored-by: danielhanchen <danielhanchen@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
model_has_dtensor()importstorch.distributed.tensorunconditionally. On a PyTorch build compiled without a distributed backend that import raises rather than returningFalse, and since v1.15.0 everyprepare_model()call reaches it:AMD's official Windows ROCm wheels are exactly that build:
torch.distributedimports,torch.distributed.is_available()isFalse, andtorch._C._distributed_c10ddoes not exist. So all training on AMD GPUs on Windows dies at trainer start.In 1.14.0 the sole call site (
accelerator.py:1878) sat behindif self.multi_device and not (...), so a single-GPU run never reached it. The second call site added in #4181 is unguarded, and that is what turned a multi-device-only path into one every run executes.Why guard the function rather than restore the guard on the new call site. The invariant belongs to
model_has_dtensoritself: DTensor lives undertorch.distributed, so a build without a distributed backend cannot hold DTensor parameters at all, and the honest answer there isFalserather than an exception. That also covers any future call site without anyone having to remember.is_torch_distributed_available()is the existing idiom for this. It is already imported in this same file and already used as a guard ~70 lines further down, andaccelerator.py:2958already computesis_dtensor_available = torch.distributed.is_available() and is_torch_version(">=", DTENSOR_PYTORCH_VERSION), so the "DTensor requires a distributed build" relationship is established in the codebase.Tests. Two added. The first reproduces the reported failure rather than stubbing it: it patches the import to raise the way it does on those wheels, so it fails against unpatched sources with the bug report's own error.
The second pins that the guard does not short-circuit the real check on a normal build.
tests/test_utils.py: 48 passed, 4 skipped.ruff checkandruff formatclean.I do not have one of the affected wheels, so the AMD-on-Windows end-to-end path is not something I verified directly. The unit test reproduces the exact import failure from the traceback, and the reporter has the hardware.
Fixes #4249
Before submitting
Pull Request section?
to it if that's the case. (prepare_model unconditionally imports DTensor, breaking single-GPU training on torch builds without a distributed backend #4249)
documentation guidelines, and
here are tips on formatting docstrings.