Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Avoid 8-bit training optimizer on XPU - #10213

Open
ousamabenyounes wants to merge 5 commits into
unslothai:mainfrom
ousamabenyounes:fix/issue-10021
Open

Avoid 8-bit training optimizer on XPU#10213
ousamabenyounes wants to merge 5 commits into
unslothai:mainfrom
ousamabenyounes:fix/issue-10021

Conversation

@ousamabenyounes

@ousamabenyounes ousamabenyounes commented Sep 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • normalize Studio bitsandbytes 8-bit optimizer choices to adamw_torch on Intel XPU
  • apply the normalization before route-level VRAM coordination and backend worker spawn
  • cover default, paged, and BNB 8-bit aliases in XPU backend tests, plus the route coordination path
  • include pre-commit.ci's formatting-only fix in studio/backend/tests/test_anthropic_messages.py

Closes #10021

Test verification (RED -> GREEN)

  • BASE RED on origin/main with only the new tests applied: XPU optimizer tests fail because adamw_8bit, paged-adamw-8bit, and adamw_bnb_8bit stay unnormalized, and the route still passes adamw_8bit into backend/VRAM coordination (20260902_issue_10021_BASE_RED.log).
  • Revert proof: reverting the production patch on this branch makes the same tests fail for the same optimizer mismatch (20260902_issue_10021_REVERT_RED.log).
  • GREEN targeted: 2 passed, 14 warnings, 3 subtests passed.
  • GREEN affected file: 84 passed, 15 warnings, 84 subtests passed.
  • ruff check studio/backend/core/training/training.py studio/backend/routes/training.py studio/backend/tests/test_gpu_selection.py studio/backend/tests/test_anthropic_messages.py: all checks passed.
  • git diff --check origin/main -- studio/backend/core/training/training.py studio/backend/routes/training.py studio/backend/tests/test_gpu_selection.py studio/backend/tests/test_anthropic_messages.py: clean.

Full local validation suite

  • Same ignored repo-local run-ci.sh replayed on origin/main and final commit.
  • Baseline: 82 passed, 15 warnings, 81 subtests passed; ruff passed; diff-check clean (20260902_issue_10021_BASELINE_RUN_CI.log).
  • Final: 84 passed, 15 warnings, 84 subtests passed; ruff passed; diff-check clean (20260902_issue_10021_FINAL_RUN_CI.log).
  • Result: same-or-better than baseline, no new failures.

Remote check notes

  • Current head: 7d9415867ad22c9bb4ac1f0ec44f6fbc96bb0eb8.
  • Terminal rollup: 37 GitHub Actions checks passed, pre-commit.ci passed, 4 checks skipped.
  • workflow-trigger lint (pull_request_target / cache-poisoning) and Repo tests (CPU) fail on tests/security/test_scan_packages.py::test_context_dependent_unsloth_zoo_findings_are_digest_pinned; the exact test also fails on clean origin/main with the same unsloth_zoo/compiler.py digest mismatch.
  • (Python 3.13) failed remotely because test_a_cold_health_call_answers_inside_the_launcher_deadline measured 2.02s against a 2.0s budget; the same test passes locally on Python 3.12 and Python 3.13, and this PR does not change that health path.

Review gates

  • OCR peer review: OCR_PEER_RESULT=PASS.
  • PR readiness check: PRCHECK_RESULT=PASS.
  • Final patch-only review: FINAL_REVIEW_RESULT=PASS.

@ousamabenyounes
fix(studio): avoid 8-bit optimizer on XPU
8a9113d 
Normalize Studio bitsandbytes 8-bit optimizer choices to adamw_torch on XPU before VRAM sizing and worker spawn. This avoids the late optimizer.step crash path while preserving the existing CUDA defaults.

Closes unslothai#10021
@danielhanchen

danielhanchen commented Sep 3, 2026

Copy link
Copy Markdown
Member

Confirmed studio/backend/core/training/training.py still passes adamw_8bit straight through to the worker and to the VRAM coordination path, so this does fix the XPU step crash. Will get this reviewed.

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 7, 2026 •
edited
Loading

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-09-07T15:15:00.728202Z 12b1421 Manual request
ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 7, 2026

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: 1a9da2bd96

ℹ️ 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".

@oobabooga

oobabooga commented Sep 7, 2026

Copy link
Copy Markdown
Member

Thanks for this. The normalization and the route-level placement are right: putting the swap before can_keep_chat_during_training is what makes the VRAM coexistence estimate match what actually runs (OPTIMIZER_BYTES_PER_PARAM is 4 for the bnb entries and 6 for adamw_torch), and _build_training_worker_config is the single chokepoint both start_training entry points go through. I traced resume too: optim is in _RESUME_CHECKPOINT_STRUCTURE_FIELDS, but that loop overwrites request.optim from the stored run config rather than comparing it, so a run started as normalized adamw_torch resumes cleanly.

One gap, which I pushed a follow-up commit for (4ecbfd0): paged_adamw_32bit was left unnormalized, and it is one of the six values in the Studio optimizer dropdown (studio/frontend/src/config/training.ts). Bit width is not the dividing line here. In transformers, OptimizerNames.PAGED_ADAMW = "paged_adamw_32bit" sits in _BITSANDBYTES_OPTIMIZERS and resolves to bitsandbytes.optim.AdamW(optim_bits=32, is_paged=True), whose update_step calls F.optimizer_update_32bit -> torch.ops.bitsandbytes.optimizer_update_32bit. On XPU, bitsandbytes registers that op to its Triton kernel in both registration branches of backends/xpu/ops.py (the native-SYCL branch still does it, under a TODO: Remove the triton register when quantization sycl kernel is ready). So picking "Paged AdamW 32-bit" on an Intel GPU reaches the same Intel Triton find_sycl assertion #10021 reports, at the same first optimizer.step().

The follow-up renames the set to XPU_UNSUPPORTED_BITSANDBYTES_OPTIMIZERS, adds paged_adamw_32bit, and extends your existing subTest loop. Verified against the versions in the Studio venv (transformers 5.5.0, bitsandbytes 0.50.0, torch 2.11.0, triton 3.6.0). The new subtest fails on your commit and passes with the change; studio/backend/tests/test_training*.py plus test_gpu_selection.py, test_mlx_training_worker_config.py, test_studio_train_validation.py and test_chat_load_during_training.py are 3533 passed, 2 skipped locally, and ruff is clean.

Out of scope for this PR, but noting it for a follow-up: _make_lora_optimizer in diffusion_lora_trainer.py and _make_optimizer in diffusion_dit_trainer.py construct bnb.optim.AdamW8bit unconditionally inside a try/except that only catches construction failure, so image and video LoRA training on XPU still hits the same path. That is pre-existing on main and a separate selection path with no user-facing choice.

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 7, 2026

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Keep them coming!

Reviewed commit: 4ecbfd0ca4

ℹ️ 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".

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 7, 2026

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: 12b1421748

ℹ️ 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".

This was referenced Sep 7, 2026
Closed
Closed
oobabooga pushed a commit to oobabooga/unsloth-test that referenced this pull request Sep 7, 2026
@danielhanchen
Add staging CI workflows for unslothai/unsloth#10213
0f87473
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@danielhanchen danielhanchen Awaiting requested review from danielhanchen danielhanchen is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

codex-converged# edge case#

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug] Intel XPU: default adamw_8bit optimizer crashes training at optimizer.step(), default optimizer is compatible with NVIDIA, not Vulkan

3 participants

@ousamabenyounes @danielhanchen @oobabooga