Fix allocator fragmentation on Windows/WSL and correct the torch 2.10 alloc conf boundary#933
Fix allocator fragmentation on Windows/WSL and correct the torch 2.10 alloc conf boundary#933danielhanchen merged 2 commits intomainunslothai/unsloth-zoo:mainfrom fix/alloc-conf-windows-wsl-torch210unslothai/unsloth-zoo:fix/alloc-conf-windows-wsl-torch210Copy head branch name to clipboard
Conversation
…oc conf boundary Long context QLoRA training on Windows/WSL had no allocator fragmentation mitigation: expandable_segments is unsupported there and the platform branches deleted all allocator config, leaving cudaMalloc retry storms and OOM despite low steady state memory. Add a roundup_power2_divisions fallback for NVIDIA CUDA on Windows/WSL, keeping any existing user config. Also correct an off-by-one: the unified PYTORCH_ALLOC_CONF is only read by the caching allocator from torch 2.10, so torch 2.9.x silently dropped expandable_segments on every platform. Gate on IS_TORCH_2_10_OR_NEWER and write the legacy PYTORCH_CUDA_ALLOC_CONF on torch <= 2.9.x. Adds tests/test_alloc_conf_platform_matrix.py covering the torch version boundary, the Windows/WSL fallback, user precedence, opt-out, vLLM standby, and AMD/Intel isolation.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
| # detection + fallback wiring against drift (same source-inspection style | ||
| # as tests/test_upstream_import_fixes_drift.py). | ||
| src = os.path.join(_REPO_ROOT, "unsloth_zoo", "__init__.py") | ||
| text = open(src).read() |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8499274d0b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _dt.DEVICE_COUNT = 1 | ||
| _dt.ALLOW_PREQUANTIZED_MODELS = True | ||
| _sys.modules["unsloth_zoo.device_type"] = _dt | ||
| import unsloth_zoo # runs the import-time allocator block |
There was a problem hiding this comment.
Make the subprocess allocator tests self-contained
When unsloth is not installed or its optional CI installation fails, every behavioral case exits at this import because unsloth_zoo.__init__ requires find_spec("unsloth"); the parent conftest's attempted import is swallowed and its sys.modules setup cannot cross into this fresh subprocess. As a result, the new test file produces no RESULT: and fails all matrix cases instead of testing the allocator logic. Provide the child with the required package/marker stub or explicitly skip these tests when that dependency is unavailable.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d0764f9: the test subprocesses now inject a minimal unsloth package marker so the find_spec("unsloth") guard passes without unsloth installed, and set UNSLOTH_ALLOW_CPU so unsloth_zoo imports without a GPU. The cases now run self-contained.
| if os.environ.get(_alloc_key, "").strip() == "": | ||
| os.environ[_alloc_key] = "roundup_power2_divisions:[32:256,64:128,256:64,>:32]" |
There was a problem hiding this comment.
Preserve explicit empty allocator settings
On Windows/WSL, a user who explicitly sets the selected allocator variable to "" to disable allocator options is now overwritten with the roundup fallback, because this condition treats an absent key and an explicit empty value identically. That breaks the established empty-string workaround used by gradient_checkpointing.py and prevents users from opting out of allocator configuration without discovering the new separate environment variable; only install the fallback when _alloc_key is absent, not when it was intentionally supplied empty.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d0764f9: the fallback now only fills when the allocator var is absent, so an explicit empty value (the opt-out advised in gradient_checkpointing.py) is preserved instead of overwritten. Added regression tests for both the unified and legacy vars.
…lf-contained Address review feedback on #933: - Apply the Windows/WSL roundup fallback only when the allocator var is absent, so an explicit empty value (the opt-out advised in gradient_checkpointing.py) is preserved instead of overwritten. - Stub the unsloth package marker inside the test subprocesses so they run self-contained even when unsloth is not installed, and set UNSLOTH_ALLOW_CPU. - Add regression tests for the explicit-empty opt-out.
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…unslothai#931/unslothai#933/others) Merge origin/main into the streaming branch. Only unsloth_zoo/mlx/trainer.py conflicted, in MLXTrainingConfig.__init__'s wholesale-copy detection: unslothai#922 (true grad-norm reporting + compiled global-norm clipping) replaced this branch's _MLX_CONFIG_OPTIONAL_COPY_FIELDS-based check with its own _appended_fields set {compile_max_variants, report_grad_norm}, and neither set alone covers the merged config. Reconciled to the union: warmup copy-detection now tolerates every appended field (the positional optional-copy suffix plus the mid-order compile_max_variants). unslothai#922's report_grad_norm landed as the new last field, so it joins _MLX_CONFIG_OPTIONAL_COPY_FIELDS to keep every post-image_size field copy-optional and preserve the positional-copy suffix invariant. Validation on the merged tree: 242 unit + 16 metal + 25 recorded-validation checks pass (grad-norm and streaming/prefetch loops coexist).
Problem
Long context QLoRA training on Windows and WSL (for example the 64k context Qwen3.5 and Gemma 4 runs in unslothai/unsloth#7203) sits on the first step trying to allocate memory and then OOMs, even for small models on 96GB, although steady state only needs around 22GB. This is CUDA caching allocator fragmentation rather than actual memory exhaustion.
Two issues in the import time allocator block in
unsloth_zoo/__init__.py:Windows and WSL had no fragmentation fallback.
expandable_segmentsis unsupported there, and the Windows and WSL branches deleted every allocator env var, so those platforms ran with no mitigation at all.Off by one on the PyTorch version gate. The unified
PYTORCH_ALLOC_CONFis only read by the caching allocator from torch 2.10. On torch 2.9.x the code wrotePYTORCH_ALLOC_CONF, which 2.9.x ignores, soexpandable_segmentswas silently dropped on every platform, not just Windows and WSL.Verified directly: on torch 2.9.1 a bad key in
PYTORCH_ALLOC_CONFis ignored while the same key inPYTORCH_CUDA_ALLOC_CONFraises; on torch 2.10.0 both raise. So 2.9.x reads only the legacy var and 2.10 reads the unified one.Fix
IS_TORCH_2_9_OR_NEWERtoIS_TORCH_2_10_OR_NEWER(thresholdminor >= 10) so the loader writes the legacyPYTORCH_CUDA_ALLOC_CONFon torch <= 2.9.x and the unifiedPYTORCH_ALLOC_CONFon 2.10 and newer.expandable_segmentsinstead of wiping the whole config, so any other user allocator options survive, then apply aroundup_power2_divisionsfallback on NVIDIA CUDA when nothing else remains.expandable_segmentsagain inside the fallback, since the 2.10 promotion step can re-read the original value (it only cleans it under standby or ROCm).The change is env var only, so it is compatible across transformers, TRL and vLLM versions. It is gated to NVIDIA CUDA and does not run under vLLM standby. Opt out with
UNSLOTH_DISABLE_ALLOC_FALLBACK=1.Behaviour with no user config
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,roundup_power2_divisions:[...]PYTORCH_ALLOC_CONF=expandable_segments:TruePYTORCH_CUDA_ALLOC_CONF=roundup_power2_divisions:[...]PYTORCH_ALLOC_CONF=roundup_power2_divisions:[...]Only torch 2.9.x behaviour changes on Linux, from silently dropped to actually applied. torch 2.6 to 2.8 and torch 2.10+ are unchanged.
Isolation and safety
_SKIP_GPU_INIT.DEVICE_TYPE == "cuda" and not IS_HIP_RUNTIME.vllm_utils.py.Tests
New
tests/test_alloc_conf_platform_matrix.pyrunsimport unsloth_zooin a fresh subprocess per case (the allocator config is import time, process global state, so an in process reload is not used) and asserts the resulting allocator env vars. It covers the torch version boundary (2.6, 2.8, 2.9.1, 2.10, 2.13), the Windows and WSL fallback, user precedence, the promotion strip, the opt out, vLLM standby, and AMD and Intel isolation. All pass.Note on torch 2.9.x
Because torch 2.9.x only reads the legacy
PYTORCH_CUDA_ALLOC_CONF, setting it surfaces torch's own one time deprecation notice for that variable. This is unavoidable and is the cost of the config actually taking effect on 2.9.x, which previously it did not. torch 2.6 to 2.8 and torch 2.10+ do not print it.Addresses unslothai/unsloth#7203