Commit graph

22 commits

Author SHA1 Message Date
Manan Shah
d65149795b
feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265)
* Add Apple Silicon MLX routing

Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
Extract original GPU init to _gpu_init.py (unchanged)
MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
GPU path unchanged: from ._gpu_init import *

* Add Apple Silicon MLX routing

- Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
- Extract original GPU init to _gpu_init.py (unchanged)
- MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
- GPU path unchanged: from ._gpu_init import *

* mlx with studio

* mlx with studio

* updating temporary install.sh

* updating temporary install.sh

* adding t_v5 path

* adding t_v5 path

* fixing vision training

* fixing vision training

* adding chat

* adding chat

* minor

* minor

* Adding export and fixing training issues, inference with lora adaptors

* Adding export and fixing training issues, inference with lora adaptors

* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM

* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM

* Merge mlx-apple-silicon into main

* update install.sh to point to main branch

* update install.sh to point to main branch

* fix: export returns 3 values (success, message, output_path) matching upstream worker

* fix: export returns 3 values (success, message, output_path) matching upstream worker

* fix(mlx): show training-process peak memory in Studio UI, not system-wide

Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).

Now the trainer's mx.get_peak_memory value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.

* fix(mlx): show training-process peak memory in Studio UI, not system-wide

Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).

Now the trainer's mx.get_peak_memory() value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.

* fix(mlx): make is_bfloat16_supported detect M1/M2 (no native bf16)

M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info.

* fix(mlx): make is_bfloat16_supported() detect M1/M2 (no native bf16)

M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info().

* feat(mlx): wire training_type="Full Finetuning" through MLX worker

Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.

* feat(mlx): wire training_type="Full Finetuning" through MLX worker

Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.

* fix(mlx): pass save_method='merged_16bit' from Studio's export page

Previously the MLX path called save_pretrained_merged with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.

* fix(mlx): pass save_method='merged_16bit' from Studio's export page

Previously the MLX path called save_pretrained_merged() with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.

* fix(studio): pass private to MLX push, return 3-tuples consistently

MLX push_to_hub branch now forwards private=private (matches GPU)
Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
needed') were tripping the route's 3-tuple unpack. Added a None
output_path so the unpack always succeeds.

* fix(studio): pass private to MLX push, return 3-tuples consistently

- MLX push_to_hub branch now forwards private=private (matches GPU)
- Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
  needed') were tripping the route's 3-tuple unpack. Added a None
  output_path so the unpack always succeeds.

* studio wirings

* studio wirings

* Merge pull request #5 from Manan17/feat/quant_config

studio wirings

* fix(mlx): wire train_on_completions for VLM via per-template lookup

Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.

* fix(mlx): wire train_on_completions for VLM via per-template lookup

Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.

* wire in lora rslora, init lora weights, random_state

* wire in lora rslora, init lora weights, random_state

* loftq studio error message fix

* loftq studio error message fix

* handle unknown optim and lr scheduler

* handle unknown optim and lr scheduler

* Merge pull request #6 from Manan17/update/peftkwargs

Update/peftkwargs

* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel

Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel

Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp

UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.

If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp

UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.

If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm

Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.

Text path (_generate_text):
make_sampler now receives top_k in addition to temp/top_p
make_logits_processors built and forwarded when repetition_penalty is
non-trivial (skip when 0.0/1.0 to avoid pointless overhead)

VLM path (_generate_vlm):
Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
forwards them to generate_step's sampler/logits_processor builders)
Rename temp= → temperature= so it's actually consumed

Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm

Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.

Text path (_generate_text):
- make_sampler now receives top_k in addition to temp/top_p
- make_logits_processors built and forwarded when repetition_penalty is
  non-trivial (skip when 0.0/1.0 to avoid pointless overhead)

VLM path (_generate_vlm):
- Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
  forwards them to generate_step's sampler/logits_processor builders)
- Rename temp= → temperature= so it's actually consumed

Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push

export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
(was hardcoded merged_16bit, ignoring the UI choice).
Both export_merged_model and export_base_model now pass save_directory=
to push_to_hub_merged so it reuses the just-written local folder
instead of re-saving under a relative "username/model" directory.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push

- export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
  (was hardcoded merged_16bit, ignoring the UI choice).
- Both export_merged_model and export_base_model now pass save_directory=
  to push_to_hub_merged so it reuses the just-written local folder
  instead of re-saving under a relative "username/model" directory.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* restore install

* restore install

* fix(mlx): restore FastVisionModel as a distinct class

unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.

Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).

Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).

Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.

* fix(mlx): restore FastVisionModel as a distinct class

unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.

Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).

Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train() runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).

Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.

* Studio: harden MLX training and export, restore GPU init guards

Studio export
Restore Tuple[bool, str, Optional[str]] contract on export_merged_model,
export_base_model, export_gguf, and export_lora_adapter, populating
output_path on successful local saves so routes/worker/CLI/frontend
details.output_path is non-empty again.
Lift the GPU save_method assignment out of the local-save branch so
Hub-only merged exports (save_directory='', push_to_hub=True) no longer
hit UnboundLocalError on the push branch.
For MLX merged and base hub-only export, stage to a tempfile.TemporaryDirectory
before push_to_hub_merged instead of passing save_directory=''.
Source _IS_MLX from unsloth instead of recomputing the platform check
(single source of truth, also enforces mlx-package availability).

Studio MLX training/inference
Pass token=hf_token into FastMLXModel.from_pretrained for gated/private
models, matching the inference path.
Strip hf_token and wandb_token from wandb.init(config=...) so secrets
do not leak into the W&B run config.
Replace load_from_disk(local_datasets[0]) with the existing
UnslothTrainer._resolve_local_files / _loader_for_files helpers so
uploaded JSON/JSONL/CSV/Parquet files train through the normal datasets
loader (load_from_disk still used for HF save_to_disk directories).
Make the dataset slice helper inclusive at the end and treat 0 as a real
index instead of "unset", matching the GPU and embedding paths.
Add a status_message -> message alias inside _send so the existing parent
pump (training.py) renders MLX status updates instead of blanks.
Forward min_p through generate_chat_response into _generate_text /
_generate_vlm and into make_sampler / vlm_kwargs so the sampling control
is no longer a no-op on MLX.
Wrap unsloth_zoo.mlx_loader / mlx_trainer imports with a clearer
ImportError pointing users at install.sh for Apple Silicon.
Exit the MLX stop-polling thread on EOFError/OSError instead of
busy-looping when the queue/pipe is permanently closed (one-line
why-safe rationale inline).

Studio frontend
ParamsSection subscribes to platform deviceType via the Zustand hook so
the gradient checkpointing dropdown re-renders after the async device
fetch completes.

Studio hardware
get_gpu_utilization MLX branch now reads _read_apple_gpu_stats once and
derives VRAM totals from psutil, removing the second ioreg subprocess
per utilization poll.

Unsloth core
Restore the os.geteuid == 0 guard around the CUDA ldconfig recovery
that was lost when GPU initialization moved into _gpu_init.py, plus the
non-root manual-fix warning branch. Non-root CUDA users no longer shell
out to ldconfig at import time.
Load dataprep/raw_text via importlib so the MLX import path no longer
pulls torch in through dataprep/__init__.py -> synthetic.py.
FastVisionModel.from_pretrained overrides the inherited delegator only
to inject text_only=False; this is an extension, not a duplication, and
is needed so VLM checkpoint loads keep the vision tower.
Wrap the MLX-branch unsloth_zoo import with a clearer ImportError.

* Studio: regression tests for MLX training/export and GPU init ldconfig guard

tests/python/test_gpu_init_ldconfig_guard.py asserts the geteuid root
check still wraps the ldconfig recovery and the non-root branch warns
bnb users; AST + source-text inspection so the test runs without torch.
tests/studio/test_export_output_path_contract.py covers the
Tuple[bool, str, Optional[str]] return contract on every export method,
the output_path assignment after successful local save, the Hub-only
GPU save_method binding fix, the MLX hub-only TemporaryDirectory
staging, and the single-source `_IS_MLX` import from unsloth.
tests/studio/test_mlx_training_worker_behaviors.py covers token
forwarding to FastMLXModel.from_pretrained, wandb config secret
stripping, file-aware local dataset loading, status_message ->
message aliasing, inclusive slice semantics, EOFError/OSError stop
thread exit, and the friendly mlx_loader / mlx_trainer ImportError.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(mlx): cap inference memory + release wired on unload + tame worker pre-pin

Three memory-hardening fixes for Studio's MLX path:

1. Inference applies the same Metal caps as the trainer.
   load_model previously only called set_wired_limit(100% of recommended)
   with no upper memory_limit, leaving large VLM checkpoints unbounded
   during the loader allocation. Add _configure_memory_limits() that sets
   memory_limit to 85% of recommended and wired_limit to min(recommended,
   memory_limit) — matching MLXTrainer's defaults so behavior is the same
   whether the user trains or just runs inference.

2. unload_model releases pinned memory back to the OS — but only when
   the cache is empty. Without this, pinned wired bytes stayed allocated
   to MLX after the model was gone, starving other apps. The release is
   guarded on `not self.models` so unloading one of several cached
   models doesn't un-pin weights still in use.

3. Worker pre-cap is conservative instead of aggressive.
   The previous pre-pin set_wired_limit(100% of recommended) competed
   with MLXTrainer's later more conservative cap. Replace with the same
   85%-memory / min(rec, memory) pair that the trainer applies later
   (idempotent re-apply). Bounds the model load + LoRA setup window
   without over-pinning.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests/studio: regression tests for the _IS_MLX dispatch gate

Two gates drive every MLX-vs-CUDA dispatch decision in Studio:

  1. unsloth._IS_MLX in unsloth/__init__.py — evaluated once at import
     time, read by Studio worker code to choose the GPU vs MLX trainer
     and inference paths. Defined as
        Darwin AND arm64 AND find_spec("mlx") is not None.

  2. utils.hardware.detect_hardware() — runtime probe with priority
     CUDA > XPU > MLX > CPU. The MLX branch is reached only when both
     CUDA and XPU are unavailable and the host is Apple Silicon and
     mlx is importable.

Neither gate had a direct test. Adds tests/studio/test_is_mlx_dispatch_gate.py
with six tests:

  test_is_mlx_gate_uses_three_required_predicates
      AST-walks unsloth/__init__.py and asserts the _IS_MLX assignment
      is a BoolOp(And) of platform.system()=="Darwin",
      platform.machine()=="arm64", and find_spec("mlx") is not None.
      Catches accidental rewrites that drop a predicate.

  test_is_mlx_gate_true_on_apple_silicon_with_mlx_present
      Spoofs platform to Darwin/arm64, injects a fake mlx module so
      find_spec returns a real ModuleSpec, re-evaluates the gate
      expression. Verifies it flips True under the exact conditions
      Studio expects.

  test_is_mlx_gate_false_when_mlx_missing
      Spoofs Apple Silicon but with mlx absent. Verifies the gate stays
      False (so a Mac without mlx installed does not pretend to have
      MLX support).

  test_is_mlx_gate_false_on_non_apple_silicon
      Canary on the actual Linux+CUDA / AMD / Intel test host: the gate
      must remain False regardless of whether mlx happens to be
      importable. Protects existing GPU users from accidental MLX
      hijack when MLX support evolves.

  test_detect_hardware_picks_mlx_when_only_apple_silicon_available
      Forces torch.cuda and torch.xpu off, spoofs Apple Silicon, injects
      fake mlx and mlx.core. detect_hardware() must return DeviceType.MLX.

  test_detect_hardware_picks_cuda_on_real_host
      Canary: on a real CUDA host detect_hardware() must return
      DeviceType.CUDA. Protects against the MLX branch shadowing CUDA
      dispatch on NVIDIA / AMD ROCm hosts.

Uses the same monkeypatch.setitem(sys.modules, ...) fake-mlx pattern as
the existing test_mlx_inference_backend.py — no new test infrastructure,
no real mlx install required.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add AGPL-3.0 SPDX header to Studio MLX regression tests

Four Studio MLX test files shipped without an SPDX-License-Identifier:

  studio/backend/tests/test_mlx_training_worker_config.py
  tests/studio/test_mlx_training_worker_behaviors.py
  tests/studio/test_export_output_path_contract.py
  tests/studio/test_is_mlx_dispatch_gate.py

They sit in or alongside studio/backend/, which is governed by
studio/LICENSE.AGPL-3.0, and exercise AGPL Studio code. Add the same
"# SPDX-License-Identifier: AGPL-3.0-only" header that's already on
test_mlx_inference_backend.py so the license declaration matches
the code under test rather than defaulting to the repo-root
Apache-2.0.

* Wrap MLX submodule imports with friendly install hint

The _IS_MLX block at the top of unsloth/__init__.py already catches the
missing-package case with a friendly install hint, but the follow-up
"from unsloth_zoo.mlx_trainer import ..." and "from unsloth_zoo.mlx_loader import ..."
lines run unguarded. An Apple Silicon user who has unsloth-zoo installed
but on an older version (e.g. the current PyPI release, before the MLX
modules ship) sees a raw ImportError on the submodule rather than the
hint that points at install.sh.

Wrap the two submodule imports in the same try/except shape so the
friendly install message fires whether the package is missing entirely
or just predates the MLX submodules. No-op once both packages release
together; smooths the transitional window where unsloth/main has merged
but unsloth-zoo on PyPI has not.

---------

Co-authored-by: DoubleMathew <mmathew23@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-05 23:54:58 -07:00
Datta Nimmaturi
09505fcc6e
Update VRAM estimator to cater to broader model configs (#5175)
* Update VRAM estimator to cater to broader model configs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix attn backend check, better support for MoE etc

* Studio: tighten VRAM estimator structured-shape and attention paths

- Conservative attention fallback: when resolve_attention_implementation
  fails, charge the quadratic non-flash activation path instead of
  silently keeping the optimistic flash_attention_2 default.
- Resolve attention on a shallow config copy so _set_attn_impl does not
  mutate the cached config returned by _load_config_for_gpu_estimate.
- Use getattr for AutoModelForCausalLM._model_mapping to avoid raising
  on private-attribute renames in transformers.
- Treat sdpa as O(n) linear attention; PyTorch SDPA dispatches to flash
  or memory-efficient backends, only eager needs the quadratic term.
- Per-layer activation accounting: structured archs (head_dim,
  layer_types, attention_k_eq_v, num_kv_shared_layers, double-wide MLP)
  now flow into compute_activation_bytes via _text_linear_dims, instead
  of using the legacy hidden_size//num_attention_heads KV/MLP shape.
- Exclude MLA configs (q_lora_rank set) from the structured-shape path
  so q_lora low-rank projection formulas keep applying when head_dim is
  also present.
- _build_text_module_elements emits a single MLA self_attn aggregate
  using _compute_attn_elements when q_lora_rank is set, avoiding the
  ~10% overcount that fed into _compute_skipped_quantizable_elements.
- Restrict _module_path_matches to known text-tower prefixes so VLM
  skip names like vision_tower.model.layers.<i>.self_attn.q_proj no
  longer falsely shadow the text alias model.layers.<i>.self_attn.q_proj.
- Pick up enable_moe_block from the config and add the per-layer dense
  MLP alongside the MoE experts in compute_total_params and
  compute_lora_params (Gemma4-style parallel dense + MoE block).
- Single-pass structured layer accounting in _compute_layer_elements,
  removing the duplicate _text_linear_dims walks.
- Drop the now-zero (activations - activations_computed) shard term in
  VramBreakdown.min_gpu_vram and the stale comment that referred to it.
- attention_implementation typed as Optional[str] to match call sites
  that pass None.
- Inline rationale comments on DOUBLE_QUANT_4BIT_FACTOR and
  NON_FLASH_ATTENTION_FACTOR pointing at VRAM_ESTIMATION.md.

* Studio: extend parallel-MoE accounting + non-prefix dense layer support

- Apply enable_moe_block / moe_has_dense_mlp symmetrically: activation
  per-layer MLP size in _layer_qkv_mlp_sizes now adds the parallel dense
  MLP for MoE layers, matching the weight and LoRA accounting added in
  the prior commit. Skip-quantizable mapping in _build_text_module_elements
  now registers both mlp.experts and per-projection mlp.{name} entries
  for MoE layers when the parallel dense block is present, so an
  llm_int8_skip_modules entry like "model.layers.N.mlp" covers both.
- Track dense layer indices as a tuple (dense_layer_indices) extracted
  from first_k_dense_replace or decoder_sparse_step + mlp_only_layers,
  and dispatch dense-vs-MoE accounting through _is_dense_mlp_layer. The
  prior count-based path silently mis-bucketed layers when mlp_only_layers
  was non-prefix (e.g. [3, 5] on an 8-layer model). num_dense_layers is
  derived from len(dense_layer_indices) for backward compatibility.
- Drop the redundant ">0" check in _is_kv_shared_layer so configs with
  num_kv_shared_layers == num_hidden_layers (every layer shared) are
  correctly recognized as shared.
- Refresh VRAM_ESTIMATION.md section 5 to note that sdpa joins
  flash_attention_2 in the linear activation path; refresh the
  VramBreakdown.activations_computed comment now that the activation
  floor is gone.

* Studio: Gemma4 PLE accounting, flex_attention, KV-share guard restore

- Add flex_attention to LINEAR_ATTENTION_IMPLS. Unsloth's
  resolve_attention_implementation returns "flex_attention" when
  HAS_FLASH_ATTENTION is False and the model class supports flex; PyTorch
  FlexAttention is a memory-efficient kernel, not a quadratic eager
  attention path. Without this, activation estimates over-charge ~36x.
- Restore the `> 0` guard in _is_kv_shared_layer. Transformers Gemma4
  (modeling_gemma4.py:1031, modular_gemma4.py:863, :926) uses
  `layer_idx >= first_kv_shared_layer_idx > 0`, so configs that mark
  every layer as KV-shared raise on construction. Reverting the
  unconditional acceptance avoids producing a detailed estimate for a
  shape the actual model code rejects.
- Extend the parallel dense MLP path (`enable_moe_block`) in
  _build_text_module_elements: when the arch is non-structured, use
  arch.intermediate_size for the dense gate/up/down dims instead of
  _text_linear_dims (which returns moe_intermediate_size via
  _get_mlp_size). Prior code under-counted skipped quantizable elements
  for the parallel dense block by up to 8x on GLM-style configs.
- Add Gemma4 per-layer-input (PLE) module accounting:
  per_layer_model_projection (one global Linear) plus per-layer
  per_layer_input_gate and per_layer_projection are added to the
  quantizable text-linear total in _compute_layer_elements;
  post_per_layer_input_norm and per_layer_projection_norm flow into
  the non-quantizable bucket. compute_lora_params adds the same three
  Linear modules to the all-linear total. References:
  transformers_versions/5.7.0/.../gemma4/modular_gemma4.py:1077-1083,
  :1247-1253.
- VRAM_ESTIMATION.md section 5 now lists flex_attention alongside sdpa
  and flash_attention_2 as linear-memory backends.

* Studio: shared-expert variants, mlp_layer_types dispatch, PLE skip, all-linear str, deepcopy resolver

Five targeted estimator corrections:

- _compute_dense_layer_indices now reads `mlp_layer_types` ahead of
  `first_k_dense_replace` / `decoder_sparse_step`. Transformers Exaone-MoE,
  Laguna, Hy_v3, GLM-MoE-DSA, GLM4-MoE-Lite, Ernie4_5_VL_MoE etc. ship the
  per-position list and may omit the prefix-style fields entirely.
- _build_text_module_elements registers per_layer_input_gate /
  per_layer_projection (per layer) and per_layer_model_projection (global)
  in the canonical element map and alias map. The PLE element count was
  added to total_quantizable in a prior commit but skip-module matching
  against names like model.layers.0.per_layer_input_gate produced 0-byte
  delta. Layer aggregate text.layers.<i> now sums all layer modules so
  prefix skip names cover the PLE pieces too.
- _targets_all_linear coerces a bare string `"all-linear"` to `["all-linear"]`
  before set comparison; the previous set comprehension iterated chars.
  PEFT LoraConfig.target_modules accepts the bare-string convention.
- ModelArchConfig gains `shared_expert_intermediate_size`. extract_arch_config
  reads `n_shared_experts` / `num_shared_experts` aliases and infers
  `n_shared_experts=1` when only `shared_expert_intermediate_size` is set.
  _compute_moe_mlp_elements and the structured + non-structured LoRA paths
  size the shared expert with its own intermediate (Qwen3.5-MoE: 512 vs
  routed moe_intermediate_size).
- _determine_attention_impl_for_gpu_estimate uses copy.deepcopy so the
  resolver does not mutate nested text_config on the cached source.
  PreTrainedConfig._attn_implementation setter walks `sub_configs` and the
  prior shallow copy still touched the inner objects.

* Studio: extend MoE/PLE/KV-share accounting to activation and skip-alias paths

Five activation-path corrections plus two LoRA / skip-alias corrections so
that shared-expert, per-layer-input, and KV-shared-layer support is symmetric
across weights, LoRA, skip-quantizable, and activation paths.

- _layer_qkv_mlp_sizes: include shared-expert FFN in mlp_size (live shared
  expert per token alongside routed experts) and keep K/V activation memory
  for KV-shared layers; only the WEIGHT path uses has_k/has_v from
  _layer_attention_dims.
- _per_layer_activation_bytes / compute_activation_bytes: account for
  per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized) per
  layer plus the global per_layer_model_projection [B,S,L,PLI] tensor when
  hidden_size_per_layer_input is set.
- _build_text_module_elements: split mlp.experts into routed and
  mlp.shared_expert canonical entries; register layers.<i>.experts alias for
  Gemma4 enable_moe_block layouts and mlp.shared_experts (plural) alias for
  Exaone-MoE / Laguna / GLM4-MoE-Lite shared-expert variants.
- _compute_moe_mlp_elements: split into _compute_routed_moe_elements and
  _compute_shared_moe_elements; only count shared_expert_gate (hd->1 Linear
  per shared expert) when shared_expert_intermediate_size is set, which is
  the Qwen2-MoE / Qwen3.5-MoE discriminator. Other shared-expert families
  (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna) lack the gate.
- compute_lora_params: when target_modules='all-linear' bare keyword, drop
  routed and shared MoE expert LoRA contributions. PEFT's all-linear targets
  nn.Linear only; Unsloth's get_moe_target_parameters expands MoE expert
  nn.Parameter LoRA only when target_modules contains explicit
  gate_proj/up_proj/down_proj/gate_up_proj names.
- _per_layer_input_lora_params: thread target_modules through and add the
  per-PLE-module contribution when the corresponding name appears, not only
  under all-linear.

* Studio: top-k MoE activations, ERNIE list configs, suffix skips, multimodal full bytes

Six estimator corrections aligning the detailed accounting paths with real
training behavior:

- _layer_qkv_mlp_sizes scales the MoE-layer mlp_size by num_experts_per_tok
  so the active routed-expert intermediate tensors are charged for activations.
  Adds num_experts_per_tok to ModelArchConfig and extracts it from
  num_experts_per_tok / top_k_experts (Gemma4 alias) in extract_arch_config.
- compute_lora_params splits routed and shared MoE LoRA contributions so that
  bare target_modules='all-linear' zeroes routed (nn.Parameter expert tensors,
  which Unsloth's get_moe_target_parameters does NOT enable for the bare
  keyword) but keeps shared-expert LoRA (regular nn.Linear MLPs that
  Unsloth's get_peft_regex DOES match).
- extract_arch_config gains a _first_scalar helper for ERNIE-style
  moe_intermediate_size = [routed, shared] lists, plus moe_num_experts and
  moe_num_shared_experts attribute aliases. When moe_intermediate_size is a
  pair and shared_expert_intermediate_size is unset, the second element is
  treated as the shared-expert intermediate.
- estimate_required_model_memory_gb's detailed branch retains
  max(0, model_size_bytes - compute_total_params(arch) * 2) on top of the
  arch-derived breakdown.model_weights so multimodal models (vision/audio
  towers) and partially-modeled families (Gemma3n AltUp/Laurel etc.) do not
  silently drop bytes that the safetensors total includes.
- _module_path_matches accepts a tail-only match when the skip entry is
  shorter than the alias path. Transformers' BNB quantizer suffix-matches
  short skip entries like ['q_proj'] / ['lm_head'] against full module
  paths; the previous len(skip) < len(alias) early-return missed those.
- _per_layer_input_lora_params drops the all_linear branch and only counts
  PLE LoRA when the user explicitly names per_layer_input_gate /
  per_layer_projection / per_layer_model_projection. Unsloth's
  get_peft_regex requires module names to contain a component tag
  (mlp/attn/...); PLE module names lack any tag, so all-linear training
  does not attach LoRA to them.

* Studio: full-FT extra optimizer/gradient inflation, MoE top-k aliases, ERNIE position dispatch, sibling experts aggregate

When the safetensors total exceeds the text-arch fp16 estimate (multimodal
vision/audio towers, partially-modeled families), only inflate the model
weights line for adapter methods but extend optimizer + gradient bytes
under full fine-tuning, where the extra params are trainable.

DBRX exposes top-k routing as moe_top_k and Hunyuan-V1-MoE as moe_topk;
neither is aliased to num_experts_per_tok via attribute_map, so probe both
when extracting arch config.

ERNIE 4.5 MoE / VL MoE configs declare MoE layers via
moe_layer_start_index / moe_layer_end_index / moe_layer_interval (with -1
meaning the last layer); add the position-style dispatch alongside the
existing mlp_layer_types / first_k_dense_replace / decoder_sparse_step
paths.

When moe_has_dense_mlp is set (Gemma4 enable_moe_block) the routed experts
live as a sibling of self.mlp at layers.<i>.experts in the actual model
layout; keep the layer mlp aggregate to the dense path and add a separate
experts aggregate so a skip module model.layers.<i>.mlp does not collapse
the routed experts as well.

* Studio: extend MoE family extraction (Llama4 / DBRX / Hunyuan / ERNIE) and align dense vs routed MLP widths

- Llama4: pick up `config.moe_layers` (auto-populated from
  interleave_moe_layer_step) so dense layer indices reflect the actual
  is_moe_layer dispatch.
- Llama4: add a separate `dense_intermediate_size` derived from
  `intermediate_size_mlp` (used for the dense feed_forward path) and keep
  `intermediate_size` for the routed/shared expert width. Auto-attach one
  shared expert per MoE layer when the dense-vs-MoE width split is present.
- DBRX: walk the `ffn_config` sub-config when extracting MoE attrs
  (moe_num_experts / moe_top_k / ffn_hidden_size). Without this DBRX is
  misclassified as a dense arch.
- Hunyuan: normalize layer-wise `moe_topk` (and the canonical
  `num_experts_per_tok` lookup it shadows via attribute_map) through a
  worst-case scalar so the int(...) cast cannot crash on list values.
- ERNIE 4.5 MoE: switch the start/end/interval dispatch to the model's
  `(layer_idx + 1) % interval == 0` modulo gate so MoE layers match the
  decoder when interval > 1.
- ERNIE 4.5 VL MoE: drop the heuristic that read
  `moe_intermediate_size[1]` as the shared expert width; in VL configs [1]
  is the vision-routed width and shared experts are sized from [0].
- estimate_fp16_model_size_bytes: prefer the larger of config-derived and
  local-weight bytes so the multimodal extra_bytes correction can fire
  for local VLM directories.

* Add tests for VRAM estimator extensions

* Studio: trim verbose comments in VRAM estimator

Collapse multi-paragraph rationale blocks to 1-3 lines stating the single
load-bearing fact. Fix one inverted "fall through ... last" comment whose
claim disagreed with the surrounding code.

* Consolidate added tests into existing test_vram_estimation.py and test_gpu_selection.py

Move Llama4 / DBRX / ERNIE arch-extraction tests into test_vram_estimation.py
as TestLlama4ArchExtraction / TestDbrxFfnConfigExtraction /
TestErniePhaseModuloDispatch / TestErnieVlSharedExpertWidth classes. Move
estimate_fp16_model_size_bytes prefer-larger-of-config-or-local tests into
test_gpu_selection.py as TestEstimateFp16ModelSizeBytesPrefersLocalWeights.
Drop one redundant Llama4 num_dense_layers assertion already covered by the
moe_layers dispatch test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
2026-05-05 04:12:36 -07:00
Wasim Yousef Said
e35cbfb454
Add native GGUF intake to Studio (#5246)
* feat(studio): add Tauri native GGUF intake

* feat(studio): polish native GGUF intake

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): load backend helpers during local setup

* fix(studio): acquire native load lease before unload

* Studio: harden native path lease verification and Tauri intake

- Wrap path.resolve(strict=True) and Path.stat() in NativePathLeaseError so a deleted or unmounted GGUF returns 400 instead of leaking the full filesystem path through the generic load_model/validate_model handler.
- Re-apply _reject_network_or_device_path to the resolved canonical path for defense in depth after symlink resolution.
- Replace try/except ValueError pattern in the device-path guard with Path.is_relative_to; the previous shape silently swallowed NativePathLeaseError (which subclasses ValueError) so /dev,/proc,/sys were never actually rejected.
- Broaden the lease redaction regex and dict-key check (Python and Rust diagnostics) to cover both native_path_lease and nativePathLease so the camelCase form emitted by Tauri/frontend payloads is also redacted.
- Hoist the redact_native_paths import to module top in loggers/handlers; the recursive filter no longer pays a per-record import lookup.
- Persist activeNativePathToken in the chat runtime store so the rollback branch can mint a fresh lease and reload the previous native GGUF when a new load fails after unload; clear it in clearCheckpoint and overwrite it on each successful load.
- use-native-drop: read options through a ref so the Tauri onDragDropEvent listener is registered once and stays attached across option changes; reject ambiguous multi-file drops up front instead of silently registering only the first GGUF.
- pick_native_model: use an async pick_file with a tokio oneshot channel instead of blocking_pick_file so the Tokio worker is not held for the duration of the OS dialog.
- registerNativeModelPath: drop the duplicate sourceKind argument; the Rust command parameter is source_kind.
- install_python_stack: insert the script directory (studio/) on sys.path; the previous insert pointed at studio/backend/ which does not satisfy `from backend.utils.wheel_utils import ...`.

* install_python_stack: keep _BACKEND_DIR on sys.path

Restore the studio/backend insertion. Although the immediately following `from backend.utils.wheel_utils import (...)` is satisfied by studio/ already being on sys.path[0] when invoked as `python studio/install_python_stack.py`, wheel_utils itself runs `from utils.native_path_leases import ...`, which requires studio/backend/ to be importable. Without the backend insertion, the existing tests/python/test_install_python_stack.py collection fails with ModuleNotFoundError: No module named 'utils'.

* Studio: tighten native path lease lifecycle and Tauri intake IPC

- register_native_model_path now hardcodes NativePathSourceKind::Drop on the Rust side and the frontend stops sending source_kind. The previous JS payload (source_kind only) never reached the Rust deserializer because Tauri's default ArgumentCase::Camel maps the Rust parameter source_kind to the JS key sourceKind, so drag/drop registration silently failed. Hardcoding the source kind also keeps audit metadata trustworthy on this command.
- Add native_path_secret_removed_for_child_start context manager and wrap multiprocessing.Process.start() at the inference, export, training, and data-recipe job spawn sites. The previous wrapper-only scrub left UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET visible to spawn-platform import-time worker code. The wrapper run_without_native_path_secret stays as defense-in-depth inside the child.
- Stop passing exc_info=True from the native-grant load/validate error logs in routes/inference.py. The structlog filter_sensitive_data processor runs before the renderer, so ConsoleRenderer formatted tracebacks bypassed redaction; the redacted str(e) preserves the message text.
- Replace the os.path.normcase string equality on the resolved canonical path with Path.samefile (with a normcase fallback) so Windows leases that differ only in extended-length \\?\ prefix or short-name spelling are accepted.
- Wrap consumeNativePathToken in its own try/catch in the chat runtime rollback. If the previous native-model token has aged out of TOKEN_TTL we now surface a clear modelsError instead of silently swallowing the rollback inside the outer catch.
- Reject non-ASCII lease strings in _split_lease and convert UnicodeEncodeError / binascii.Error / ValueError raised by _b64decode into NativePathLeaseError so verify_native_path_lease never escapes raw exceptions to the route handler.
- Tighten dropStateForPaths to mark multi-file payloads invalid so the overlay matches the post-fix drop handler that rejects the same payload.
- Replace the one-shot fetch in useNativePathLeasesSupported with a delayed-retry loop so the picker/drop becomes available once the backend is up rather than staying disabled for the rest of the session after a transient failure.
- Drop the unused setActiveNativePathToken setter; the value is set via setState directly in use-chat-model-runtime.
- Add a toast on auto-load failure in use-native-drop so a collapsed model selector does not hide the error.
- Burn the lease nonce before _validate_current_stat so a stat-failed lease is single-use even if a later state change happens to match the original size/mtime.

* Studio: cache lease secret, harden native path stat checks, polish intake UX

- Cache the decoded UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET on first verify and validate that it is base64-decodable and at least 32 bytes. Subsequent _decode_secret calls return from the cache and never touch os.environ, so concurrent /api/inference/load and /api/health requests no longer race with native_path_secret_removed_for_child_start scrubbing the env. native_path_leases_supported now wraps _decode_secret so the health flag matches what verify_native_path_lease actually accepts.
- Replace path.is_file()/is_dir() + path.stat() with os.lstat() in _validate_current_stat and explicitly reject S_ISLNK; size and mtime checks now refer to the link itself, closing the same-size+same-mtime symlink-swap window that the prior follow-symlink stat() left open.
- Add an issued_at_ms < expires_at_ms sanity check in _validate_payload to reject internally inconsistent (HMAC-protected) lease payloads.
- Sort _NATIVE_PATH_REDACTIONS by length (descending) before iterating in redact_native_paths so a longer registered path is replaced before a shorter prefix path; otherwise logs containing /foo/X.gguf.bak after only /foo/X.gguf was registered would leak the .bak suffix.
- classify_existing_path now re-checks the canonical path with symlink_metadata after canonicalize, so a regular file that is replaced with a symlink in the small canonicalize window is rejected at registration.
- ModelSelector renders the local file picker as its own block (not in the eject ternary), so a user with an active model can still replace it via the picker rather than only via drag/drop.
- useNativePathLeasesSupported caps the readiness probe at MAX_READINESS_POLLS (60 = ~5 minutes) and aborts the in-flight fetch on unmount via AbortController, so a permanently-disabled backend stops generating sustained traffic and hot-reload no longer leaks open connections.
- useChooseNativeModel returns a stable useCallback closure and guards the OS dialog with a useRef so rapid double-clicks cannot open multiple dialogs and orphan Rust tokens.
- Branch the multi-file drop toast: if no GGUF was present we say "Only .gguf model files can be dropped here." and otherwise "Drop a single .gguf model file." so users dropping non-GGUF attachments get an accurate explanation.

* native_path_leases: lstat the signed canonical path before resolving

The earlier change to lstat inside _validate_current_stat operates on grant.canonical_path, which is the post-resolve target. If the user atomically replaces the originally-signed file with a symlink to a different file of identical size and mtime, path.resolve(strict=True) follows the symlink, samefile returns True (both ends share the new inode), and the lstat in _validate_current_stat sees the regular target file rather than the symlink, so the swap goes undetected.

Add an os.lstat on the signed canonical path before path.resolve(strict=True), and reject S_ISLNK there. The lstat in _validate_current_stat stays as defense-in-depth for swaps that occur strictly between resolve and stat.

* Studio: scrub native lease secret before mp.Queue spawn and tighten lease lifecycle

- Move _CTX.Queue / _CTX.Event / _CTX.Process construction inside native_path_secret_removed_for_child_start at the inference, export, training and data-recipe spawn sites. The first Queue creation lazily spawns Python's multiprocessing.resource_tracker child, so when it ran outside the scrub context the tracker process inherited the lease secret. Reproduced via the proc filesystem environ entry; the wrapped order keeps the tracker clean.
- native_path_secret_removed_for_child_start now refcounts entries: the env var is popped on the first entry and restored only when the last context exits. Concurrent training/inference/export starts no longer serialize on the env lock across the entire proc.start yield, while still guaranteeing the env stays empty for the duration of every overlapping spawn.
- run_without_native_path_secret now also nulls the module-level cached lease secret. With the existing spawn-only multiprocessing context the cache is irrelevant in practice, but a future fork caller would otherwise inherit the in-memory secret even though the env var was scrubbed.
- filter_sensitive_data now applies the native lease key check on the top-level event_dict, not only on nested dicts, so a logger call that includes a lease value as a top-level keyword field actually redacts it (the bare value does not match the prefix-anchored regex).
- chat-page loadNativeModelIntent now passes intent.id to clearModelIntent so a second drag-drop during an in-flight first auto-load is not wiped from the chip area when the first resolves.
- Bump useNativePathLeasesSupported's MAX_READINESS_POLLS from 60 to 720 so first-run installs that compile llama.cpp from source or download large CUDA wheels (well past 5 minutes) don't permanently disable the native picker.

* native_path_leases: serialize first-decode against scrub context

_decode_secret used a separate _SECRET_INIT_LOCK from the env scrub's _NATIVE_PATH_ENV_LOCK, so the very first decode (before the cache is populated) could race a concurrent native_path_secret_removed_for_child_start and read os.environ during the env-empty window, raising "Native path grants require the managed desktop backend." Subsequent calls hit the cache and were already safe.

Acquire _NATIVE_PATH_ENV_LOCK around the env read inside _SECRET_INIT_LOCK and fall back to _SCRUB_SAVED_SECRET when the scrub has temporarily popped the env var. Lock ordering (init then env) is consistent with no other caller, so no deadlock.

* Studio: surface native model load errors and harden native path label cache

- Native model load and validate now bubble up the actual exception (with
  paths redacted) and apply the same friendly-error rewrite the non-native
  path uses, so users see "CUDA OOM", "trust_remote_code required", etc.
  instead of a generic "Failed to load native model: <label>".
- run_without_native_path_secret now also nulls _SCRUB_SAVED_SECRET so a
  forked grandchild that imports native_path_leases cannot recover the
  secret via the scrub-aware fallback in _decode_secret.
- _NATIVE_PATH_LABELS now has its own 10000-entry cap independent of the
  100-entry redaction list, so display_label_for_native_path no longer
  falls back to returning the raw canonical path after 101 native paths
  in one session. Redaction list keeps the 100-entry cap for log-scan
  performance.
- _validate_payload now also rejects null bytes in display_label, which
  is echoed back in HTTP responses and log lines.

* Studio: harden native path lease validation and chained native rollback

- child_env_without_native_path_secret now copies os.environ under
  _NATIVE_PATH_ENV_LOCK so a concurrent scrub-context env pop cannot
  raise RuntimeError: dictionary changed size during iteration in a
  background hardware scan or other env reader.
- _validate_payload and grant construction route every signed numeric
  field (version, issued_at_ms, expires_at_ms, size_bytes, modified_ms)
  through new _required_int / _optional_int helpers that wrap raw int()
  ValueError into NativePathLeaseError. The single upstream catcher
  produces 400 instead of 500 for malformed signed payloads.
- verify_native_path_lease now runs _validate_current_stat before
  _consume_nonce, so a transient stat error on the canonical path no
  longer permanently burns the nonce. Concurrent verifies still
  serialize through _consume_nonce, so single-use is preserved.
- Chained native model rollback now restores activeNativePathToken in
  the chat runtime store after a successful rollback loadModel. Without
  this, a second consecutive failed switch could not re-roll-back
  because the store token had been overwritten by the failed attempt.
- validate_model now applies the same not_supported_hints friendly
  rewrite to native model errors that load_model already does, so a
  native .gguf that fails validation with an upstream "is not supported"
  message gets the same actionable wording as the non-native branch.

* Studio: harden native path log redaction, status disclosure, and chip lifecycle

- structlog processor chain now runs format_exc_info before
  filter_sensitive_data so traceback strings are produced (and then
  redacted) rather than passed through as untouched (type, value, tb)
  tuples that the JSON or console renderer formats after the redaction
  filter has already finished.
- native_path_secret_removed_for_child_start clears _CACHED_LEASE_SECRET
  in addition to popping the env var, so a fork during the scrub window
  cannot inherit the cached bytes via the parent's heap. Parent verify
  calls during the window keep working through the existing scrub-aware
  fallback in _decode_secret.
- load_model's except ValueError handler now redacts native paths and
  uses the native model log label when native_grant_backed is true.
  Previously a ValueError raised after lease verification (e.g. from
  ModelConfig.from_identifier or downstream GGUF parsing) returned the
  raw exception string in the HTTP response body.
- llama_cpp_backend now records the native display label at GGUF load
  time, and /api/inference/status prefers it over the redaction store.
  After a Python backend restart the redaction store is empty; the
  attribute keeps the friendly label, and an absolute model_identifier
  with no other label source falls back to the basename so the canonical
  path no longer appears in active_model.
- reveal_path_token uses native "reveal and select" commands on macOS
  (open -R) and Windows (explorer /select,) so the file is highlighted
  in the file manager. Linux keeps the existing parent-directory open.
- Native model rollback that fails because the previous token cannot be
  consumed now throws a rollback-specific Error, and the outer empty
  catch was replaced with one that re-throws the rollback error. The
  rollback-specific message now reaches the user instead of being
  overwritten by the original load error message.
- NativeModelChip tracks the Rust token's expiresAtMs on a single
  setTimeout, disables the Load button at expiry, and relabels it
  "Select again" with an explanatory tooltip so users do not click into
  a guaranteed-failure path after the 15-minute TTL elapses.

* Studio: tighten native artifact policy, mmproj sibling check, and intake UX

- is_open_safe_artifact no longer grants Open for directories. Reveal
  already handles directory navigation, so the change closes the
  attack surface where a macOS .app artifact could be launched via
  open_path_token + open::that_detached.
- Display labels are sanitized in classify_existing_path. Control
  characters in filenames (newlines, tabs, NUL et al.) are replaced
  with spaces and the label is trimmed and capped, so a file named
  with embedded newlines cannot inject forged log lines or scramble
  the UI status panel.
- validate_entry_path skips the size_bytes/modified_ms equality check
  when the operation is Reveal or Open. Cloud-sync agents (Dropbox,
  iCloud Drive, OneDrive) routinely rewrite extended-attribute
  metadata which bumps mtime, and the user expects Reveal/Open to
  remain available for files in synced folders.
- llama_cpp_backend gains a _native_grant_backed flag at GGUF load
  success. /api/inference/status only applies the absolute-path
  basename fallback when that flag is true, so a non-native absolute
  local GGUF still reports its canonical model_identifier and unload
  by identifier keeps working.
- Native vision GGUFs now run through _validate_native_mmproj_companion
  before llama-server starts: the companion mmproj must be a regular
  file, not a symlink, and must live in the same resolved directory as
  the granted GGUF. This stops a hostile sibling or symlinked mmproj
  from being loaded under a single-file lease.
- Chained native rollback restructured: the rollback loadModel + state
  + refresh runs inside its own try/catch that swallows so the outer
  throw error surfaces the ORIGINAL load failure. The native-token
  consume-failure case still throws the rollback-specific message
  early, before the inner block runs, so its actionable guidance is
  preserved.
- Loading-model state and the duplicate-load guard in the chat runtime
  hook now compare both the model id and the native path token. Two
  drops or picks with the same basename in different folders no longer
  silently dedup; the second token is honored.
- chat-page loadNativeModelIntent awaits selectModel before clearing
  the pending intent. If selectModel returns early via dedup or
  throws, the chip and its token stay so the user can retry instead
  of losing the selection.
- NativeModelChip's Reveal button is disabled when the lease has
  expired (Rust would reject it anyway), and the Load button label
  reads "Expired" instead of "Select again" so the disabled element
  no longer promises an action it cannot perform.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-04 11:46:18 +02:00
Wasim Yousef Said
a5eb2e3d50
Add tauri (#5144)
* add unsloth studio desktop app

* Fix review findings

- studio/src-tauri/tauri.conf.json: retarget updater to staging repo
  (danielhanchen/unsloth-staging-2); switch to unslothai/unsloth on upstream merge.
- studio/src-tauri/linux/postremove.sh: drop the interactive read loop and the
  /home/* iteration. Package maintainer scripts must stay non-interactive and
  must not touch other users' data.
- studio/frontend/src/app/auth-guards.ts: honor tauriAutoAuth() boolean. Failed
  auto-auth now redirects to /login; requireGuest/requirePasswordChangeFlow
  only redirect to /chat when auth succeeds. The new early-return on failed
  auth is intentional so the login / change-password flows remain reachable
  when desktop auth is not yet established.
- studio/frontend/src/config/env.ts: keep fetched=false on health failure so
  later calls retry instead of caching the client-side platform guess.
- studio/src-tauri/src/install.rs: pick the available system package manager
  (apt-get, dnf, zypper, pacman); AppImage bundles run on non-Debian distros.
- studio/frontend/src/lib/open-link.ts + markdown-text/sources callers: return
  boolean from openLink so callers only preventDefault on handled URLs; relative
  hrefs now navigate natively.
- studio/frontend/src/features/settings/tabs/about-tab.tsx: fetch(apiUrl(...))
  so the version request targets the backend port in desktop mode. The bare
  /api/health predates the Tauri webview (blame: the earlier onboarding commit,
  which ran with same-origin frontend/backend); in desktop mode the webview
  origin is tauri://localhost so the bare path fails.
- install.ps1: gate the install_python_stack.py hotfix on a sentinel comment
  instead of a content regex; append the sentinel after applying so reruns
  are unambiguous.
- unsloth_cli/commands/studio.py _write_auth_secret: use the atomic mkstemp +
  os.replace path on Windows too; chmod calls are wrapped in try/except OSError.
- studio/src-tauri/src/preflight.rs probe_existing_backends: fan out the health
  probes concurrently; desktop-auth status still runs sequentially per candidate.
  reqwest::Client is internally Arc-wrapped so the in-loop .clone() is a
  refcount bump, not a deep clone; annotated inline.
- studio/src-tauri/src/preflight.rs run_cli_probe: wait() after kill() to reap
  the child, matching probe_cli_capability.
- studio/src-tauri/src/process.rs + main.rs: add stop_backend_detached and use
  it from the tray quit handler so the 5s graceful-wait does not block the
  Tauri main loop. RunEvent::Exit keeps the synchronous safety-net call.
- studio/backend/main.py: drop the permissive localhost CORS regex in
  api-only mode; the explicit allow_origins list is sufficient.
- .github/workflows/release-desktop.yml: drop max-parallel: 1 so platform
  builds run in parallel, and lift releaseBody to an env var so the three
  tauri-action invocations share one source of truth.

* Fix review findings (loop 2)

- studio/backend/auth/storage.py update_password: clear_desktop_secret()
  alongside clear_bootstrap_password() so rotating the admin password
  also revokes any previously provisioned .desktop_secret. Without this,
  an old local desktop credential keeps minting fresh admin tokens via
  /api/auth/desktop-login after a password rotation.
- studio/src-tauri/src/desktop_auth.rs provision_desktop_auth: wrap
  cmd.output().await in tokio::time::timeout(30s). DESKTOP_AUTH_LOCK is
  held across the whole desktop_auth flow, and previously a hanging
  `unsloth studio provision-desktop-auth` subprocess would pin the lock
  indefinitely and freeze every subsequent desktop_auth call.

* Add review tests

* Consolidate review tests

Merge review-added tests into the existing studio/backend/tests/test_desktop_auth.py
(the PR's authoritative desktop-auth test file). Drops three scaffolding files under
tests/python/ in favor of five focused tests next to the tests they extend:
- test_update_password_clears_desktop_secret (runtime)
- test_update_password_on_unknown_user_leaves_desktop_secret_intact (runtime)
- test_cli_provisioning_delegates_to_storage_create_desktop_secret (source-level)
- test_cli_connect_auth_db_reads_storage_db_path (source-level)
- test_desktop_auth_provision_has_bounded_timeout (Rust source-level)

* Revert auth-guards.ts Tauri branches to unconditional form

The review loop on PR 5144 introduced a regression: the isTauri branch of
requireAuth redirected to /login when tauriAutoAuth() returned false, and
requireGuest / requirePasswordChangeFlow silently fell through on the same
condition. The Tauri desktop app authenticates via a local auto-generated
secret; it must never surface /login or /change-password to the user. A
failed auto-auth should let the startup layer retry, not expose a password
form.

Restore the three Tauri branches to the author's original unconditional
form (requireAuth: return; requireGuest / requirePasswordChangeFlow: throw
redirect({to: '/chat'})). Keep the rest of the review fixes -- the
apiUrl() fetch wrapping, authRedirect helper, and fetchAuthStatus refactor
are all legitimate improvements and are preserved.

* Revert release-desktop.yml to author's version

The review loop's workflow-file tweaks (drop max-parallel: 1, lift releaseBody
to an env var) are cosmetic. OAuth tokens cannot push workflow-file changes,
and fine-grained PATs cannot honor maintainerCanModify on a third-party fork.
Reverting the workflow file to wasimysaid's version lets the push go through
without needing a classic PAT with both repo and workflow scopes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-23 04:50:10 -07:00
Daniel Han
cad8c6ad05
Add AMD ROCm/HIP support across installer and hardware detection (#4720)
* Add ROCm detection to install.sh and expand shell tests

Add AMD ROCm GPU detection to get_torch_index_url() in install.sh.
When nvidia-smi is not found, probe for ROCm via amd-smi, /opt/rocm
version file, hipconfig, dpkg-query, and rpm.

Includes validation guard for malformed _rocm_tag, Debian epoch prefix
stripping, ROCm 7.2+ cap to rocm7.1 index, bitsandbytes AMD install,
and status messaging. Shell tests expanded to 23 cases.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add ROCm torch reinstall support to install_python_stack.py

Add _detect_rocm_version() and _ensure_rocm_torch() to detect when a
Linux host has ROCm but the venv received CPU-only torch, and reinstall
with the correct ROCm wheels. Covers ROCm 6.0 through 7.1 with a
30-second timeout on the torch GPU probe subprocess.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add ROCm support to llama.cpp prebuilt installer

Add has_rocm field to HostInfo, extend detect_host() to probe for ROCm
via hipcc/amd-smi/rocm-smi/ROCM_PATH, and route ROCm hosts to upstream
prebuilts (Linux ROCm 7.2 prebuilt with source fallback, Windows HIP
prebuilt with CPU fallback). Add linux-rocm and windows-hip install
kinds to runtime_patterns_for_choice().

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add IS_ROCM hardware flag and fix AMD error message

Add IS_ROCM flag to hardware.py detect_hardware() (set when
torch.version.hip is present, DeviceType stays CUDA). Export IS_ROCM
from __init__.py. Add "rocm" key to get_package_versions().

Replace "We do not support AMD" error in tokenizer_utils.py with a
helpful message pointing to ROCm installation docs.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add comprehensive ROCm support test suite (68 tests)

Add tests/studio/install/test_rocm_support.py covering all ROCm code
paths across install_llama_prebuilt.py, install_python_stack.py,
hardware.py, tokenizer_utils.py, and install.sh. All tests use mocks
and run without AMD hardware.

Covers: asset selection (11), runtime patterns (5), HostInfo (4),
ROCm version detection (9), torch reinstall (9), index mapping (8),
hardware flag (8), tokenizer message (2), install.sh structure (10),
and live regression (1).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden ROCm support: probe error handling, version cap, validation

Address review findings from 8 independent reviewers:

- Wrap _ensure_rocm_torch() torch probe in try/except for
  TimeoutExpired and OSError so a hung or broken torch import does not
  crash the installer (8/8 reviewers flagged this)
- Add torch>=2.4,<2.11.0 version cap to the ROCm reinstall path to
  prevent installing unsupported torch 2.11.0 from the rocm7.1 index
- Use with-statement for file reads in _detect_rocm_version() to avoid
  resource leaks
- Handle ROCM_PATH="" correctly (use `or "/opt/rocm"` instead of
  default parameter to avoid relative path resolution)
- Strengthen shell validation guard from rocm[0-9] to rocm[1-9] to
  reject rocm0.x tags that would produce nonexistent PyTorch index URLs
- Switch shell version cap from blocklist to allowlist (rocm6.*|rocm7.0*
  |rocm7.1* pass through, everything else caps to rocm7.1) so future
  ROCm 10+ does not fall through to a nonexistent index
- Add sorted() to _ROCM_TORCH_INDEX lookup for defensive ordering
- Fix test_probe_timeout_handled: replace zero-assertion test with
  proper assertions verifying reinstall proceeds after timeout

* Clean up rocm_paths list construction in detect_host()

Filter None from the ROCM_PATH env var lookup at list construction time
instead of relying on the inline `if p` guard in the any() call.

* Require actual AMD GPU presence before selecting ROCm paths

All 8 reviewers across 2 cycles independently flagged that ROCm
detection used toolkit/filesystem hints (hipcc, /opt/rocm, rocm-core)
as a proxy for GPU presence, which would misroute CPU-only or NVIDIA
hosts that happen to have ROCm tools installed.

Now all 3 detection points (install.sh, install_python_stack.py,
install_llama_prebuilt.py) probe for an actual AMD GPU before
entering the ROCm path:

- install.sh: check rocminfo for gfx* GPU names, or amd-smi list
  for device rows, before version detection
- install_python_stack.py: new _has_rocm_gpu() function probes
  rocminfo and amd-smi list before _ensure_rocm_torch() proceeds
- install_llama_prebuilt.py: detect_host() probes rocminfo/amd-smi
  list instead of just checking tool existence or directory paths

Also:
- Shell test mock amd-smi now handles "list" subcommand
- Python tests updated to mock _has_rocm_gpu where needed
- Added test_no_gpu_with_rocm_tools_skips to verify the new guard
- Test index lookups now use sorted() to match production code

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden hipconfig version parsing and torch probe compatibility

- Add parts[1].isdigit() check in hipconfig version parsing to handle
  versions like "6.3-HIP" where the minor component has non-numeric
  suffix (strip "-" prefix before int() conversion)
- Use getattr() in torch probe subprocess to safely handle old or
  custom torch builds that may lack torch.version.hip/cuda attributes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Strengthen AMD GPU detection and add NVIDIA precedence guard

- Change amd-smi list detection from any-non-empty-output to requiring
  "gpu" marker in output, matching the shell-side NR>1 check. Prevents
  false positives from header-only amd-smi list output.
- Add nvidia-smi check at the top of _ensure_rocm_torch() so mixed
  AMD+NVIDIA hosts preserve NVIDIA precedence (matching install.sh and
  install_llama_prebuilt.py behavior).
- Apply the same amd-smi marker fix to install_llama_prebuilt.py
  detect_host() for consistency.

* Add Windows-specific ROCm/HIP detection in detect_host()

The previous detect_host() ROCm check used rocminfo and amd-smi list
which are Linux-only tools. On Windows, has_rocm would always be False,
making the Windows HIP prebuilt path at line 1794 unreachable.

Now detect_host() uses platform-specific detection:
- Linux: rocminfo (check for gfx GPU names) or amd-smi list
- Windows: hipinfo.exe, amd-smi, or amdhip64.dll on PATH

This allows Windows AMD users to get the HIP prebuilt binary instead
of silently falling through to the CPU prebuilt.

* Add AMD ROCm gaps: Mamba/SSM source builds, GPU monitoring, Windows messaging, RDNA expansion

- worker.py: Add HIP detection to causal-conv1d/mamba-ssm probe, check
  for hipcc before ROCm source builds, improve status messages and error
  reporting, add timeout and uv support for the source build fallback
- amd.py: New AMD GPU monitoring module via amd-smi metric --json,
  mirroring nvidia.py structure (utilization, temperature, power, VRAM)
- hardware.py: Branch to amd.py when IS_ROCM is True for GPU utilization,
  visible GPU queries, and physical GPU count
- install_python_stack.py: Detect AMD GPUs on Windows and warn that
  ROCm-enabled PyTorch must be installed manually
- kernels/utils.py: Expand is_rdna() to cover RDNA2 (gfx1030-1032),
  RDNA3 (gfx1102-1103), RDNA3.5 (gfx1150-1152) alongside existing entries
- tests: Add 32 new tests covering all changes (95/95 pass)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden ROCm detection, fix VRAM heuristic, and expand RDNA2 coverage

- Windows ROCm detection: validate actual GPU presence via hipinfo/amd-smi
  output markers instead of just checking tool existence on PATH
- _ensure_rocm_torch: validate nvidia-smi actually reports a GPU before
  giving NVIDIA precedence (fixes AMD-only hosts with stale NVIDIA tools)
- amd.py _parse_numeric: handle dict-shaped metric objects from newer
  amd-smi versions ({"value": 10, "unit": "W"}) and strip MiB/GiB units
- amd.py VRAM heuristic: raise threshold from 100k to 10M to correctly
  handle MI300X (192 GB = 196608 MB) and other high-VRAM GPUs
- amd.py visible GPU: use AMD-reported GPU IDs instead of enumerate index
  so non-dense sets like CUDA_VISIBLE_DEVICES=1,3 report correctly
- install.sh: add ROCm <6.0 minimum version guard (no PyTorch wheels
  exist for older versions); fix rocm7.1* glob to not match rocm7.10+
- is_rdna: add gfx1033-1036 for RDNA2 mobile GPUs (RX 6600M etc.)
- worker.py: increase ROCm source build timeout from 600s to 1800s;
  fix success log message for ROCm source builds
- Tests: update mocks for _has_usable_nvidia_gpu, add RDNA2 target asserts

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add HIP_VISIBLE_DEVICES support, unit-aware VRAM parsing, Windows GPU validation

- hardware.py: check HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm
  before falling back to CUDA_VISIBLE_DEVICES, so multi-GPU AMD setups with
  HIP-specific env vars report the correct visible device set
- amd.py: add _parse_memory_mb() that reads "unit" from dict-shaped amd-smi
  JSON (e.g. {"value": 192, "unit": "GiB"}) and converts to MB correctly;
  fixes MI300X VRAM misreported as 0.19 GB instead of 192 GB
- install_python_stack.py: Windows AMD warning now validates actual GPU
  presence via hipinfo/amd-smi output markers before printing
- install_llama_prebuilt.py: restore amdhip64.dll fallback for Windows HIP
  detection after tool-based checks, so Windows HIP installs without CLI
  tools on PATH are still detected
- hardware.py: fix IS_ROCM comment to accurately describe its role

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix HIP_VISIBLE_DEVICES empty-string handling in GPU visibility spec

Use explicit None checks instead of Python `or` operator when reading
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES, so that an empty string
("") is correctly honored as "no visible GPUs" rather than silently
falling through to CUDA_VISIBLE_DEVICES on mixed ROCm+CUDA systems.

* Fix IS_ROCM test assertion for multi-line formatting

* Cap torchvision/torchaudio versions, remove amdhip64.dll fallback, fix visible GPU count

- Cap torchvision<0.26.0 and torchaudio<2.11.0 alongside torch<2.11.0 in
  both install.sh and install_python_stack.py to prevent resolver from
  selecting incompatible companion packages from ROCm wheel index
- Remove amdhip64.dll fallback in Windows ROCm detection (DLL presence
  without hipinfo/amd-smi is not proof of GPU existence)
- Fix get_visible_gpu_count() to use _get_parent_visible_gpu_spec() which
  respects HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES on ROCm hosts

* Attribute is_rdna() RDNA2/3/3.5/4 expansion to PR #4428

The is_rdna() expansion to cover RDNA2 (gfx1030-1036), RDNA3
(gfx1100-1103), RDNA3.5 (gfx1150-1152), and RDNA4 (gfx1200-1201)
architectures is based on the original work from PR #4428.

Co-authored-by: GoldenGrapeGentleman <yueyuan@amd.com>
Co-authored-by: billishyahao <bill.he@amd.com>

* Support AMD Radeon for studio (#4770)

Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>

* Remove ROCm test files from main PR

Move test_rocm_support.py and shell test additions to a separate PR
to keep the main ROCm support PR focused on implementation changes.

* Fix installer and hardware detection issues for PR #4720

- Fix empty _tri_arg passed to uv pip install in Radeon path (causes
  "Empty field is not allowed for PEP508" error)
- Fix Radeon fallback: use ROCm index instead of CPU-only when
  repo.radeon.com is unreachable (TORCH_INDEX_URL already has ROCm)
- Use $TORCH_CONSTRAINT in fallback paths instead of hardcoded strings
- Fix _pick_radeon_wheel: relax suffix to match manylinux_2_28_x86_64
  wheels (AMD Radeon repo does not use bare linux_x86_64 platform tag)
- Fix IS_ROCM export: use __getattr__ so callers always see the live
  value after detect_hardware() runs
- Fix apply_gpu_ids: set HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES
  on ROCm so _get_parent_visible_gpu_spec picks up narrowed GPU set
- Fix _parse_memory_mb: distinguish GB (1000 MB) from GiB (1024 MiB)
- Add amd-smi version as a fallback in _detect_rocm_version
- Fix trailing whitespace and missing newline at EOF in install.sh

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix GPU detection false positives and add missing health groups

- Fix _has_rocm_gpu() false positive: require "GPU: <number>" data rows
  from amd-smi list, not just header containing "gpu"
- Apply same fix in detect_host() in install_llama_prebuilt.py
- Add runtime_payload_health_groups for linux-rocm and windows-hip so
  partial/corrupt ROCm/HIP prebuilt installs are properly detected
- Add bitsandbytes install to Radeon fallback paths (was only in the
  success path, skipped when repo.radeon.com was unreachable)
- Keep DEVICE/CHAT_ONLY as direct imports in __init__.py (matching main)
  and only use __getattr__ for IS_ROCM

* Fix _ensure_rocm_torch and Windows AMD warning false positives

- _ensure_rocm_torch: only skip when HIP is already present, not for
  CUDA builds (which are unusable on AMD-only hosts). Fixes the case
  where a venv has a stale CUDA wheel and the repair step is skipped.
- Windows AMD warning: use GPU data row check (same as Linux fix) to
  avoid false positives from amd-smi list header-only output.

* Fix amd-smi GPU detection for GPU[N] output format

Older amd-smi versions output "GPU[0] : Card series: ..." instead of
"GPU: 0". The regex now matches both "GPU: <digit>" and "GPU[<digit>"
formats to detect actual GPU data rows.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden AMD GPU detection against false positives

- install.sh: replace weak amd-smi list check (awk 'NR>1 && NF') with
  strict pattern matching GPU data rows (/^GPU[[:space:]]*[:\[]/)
- All files: reject rocminfo gfx000 (CPU HSA agent) by requiring
  gfx[1-9] instead of gfx[0-9] in the rocminfo GPU probe
- Fixes false positives on hosts with ROCm tools but no AMD GPU

* Remove duplicate comment from pre-commit merge

* Refactor: deduplicate AMD detection, consolidate bitsandbytes, clean up imports

- Extract _has_amd_rocm_gpu() shell function to avoid duplicating the
  rocminfo/amd-smi GPU detection logic in get_torch_index_url and
  the Radeon auto-detect block
- Consolidate bitsandbytes install into a single case block after torch
  install (was duplicated 4 times across Radeon success/fallback paths)
- Move math and re imports to top of amd.py (were inline in functions)
- Add _smi_query() helper in hardware.py to centralize IS_ROCM backend
  selection for get_gpu_utilization and get_visible_gpu_utilization

Addresses Gemini code review suggestions.

* Fix VRAM parsing for string values and GB/GiB consistency

- Extract unit from string-valued VRAM fields (e.g. "192 GiB") so
  _parse_memory_mb correctly applies the unit multiplier instead of
  treating the value as bare MB
- Treat GB and GiB identically (both as binary x1024) since GPU tools
  including amd-smi use binary units even when labeling them "GB"
- Fixes incorrect VRAM reporting on MI300-class cards (was showing
  ~0.19 GB instead of 192 GB for string-valued outputs)

* Add --no-cache to uv for ROCm HIP source builds

Avoid stale cache artifacts from partial HIP source builds when
uv is used for causal-conv1d/mamba-ssm compilation on ROCm.
The pip path already uses --no-cache-dir; this adds the uv equivalent
(--no-cache) only when is_hip is True.

* Fix critical: initialize _amd_gpu_radeon before case block

_amd_gpu_radeon was only set inside the */rocm*) case arm, so on
NVIDIA/CPU/macOS paths where TORCH_INDEX_URL does not contain "rocm",
the variable was unbound. With set -u (nounset) enabled, this crashes
the installer for every non-AMD user.

Move initialization to before the case block so it is always defined.

* Fix Windows AMD: route has_rocm hosts to HIP prebuilt path

resolve_release_asset_choice was selecting windows-cpu for all Windows
x86_64 hosts including those with has_rocm=True. Windows AMD users
should fall through to resolve_upstream_asset_choice which tries the
HIP prebuilt first. Add "not host.has_rocm" guard to the published
windows-cpu selection.

* Harden ROCm detection, Radeon wheel fallback, and HIP visibility

Addresses review findings from parallel reviewers on PR #4720:

- install.sh: add _has_usable_nvidia_gpu() helper requiring nvidia-smi -L
  to actually list a GPU before treating the host as NVIDIA. Fixes the
  stale-nvidia-smi-on-PATH regression where AMD-only hosts fell into the
  CUDA branch.
- install.sh: fix hipconfig awk blocks to propagate a non-zero exit code
  when the output is not a recognisable version string, so the ||-chain
  continues to dpkg-query / rpm instead of terminating early.
- install.sh: fail-closed on Radeon wheel fallback. When torch,
  torchvision or torchaudio is missing from the Radeon repo for the
  active Python tag, fall back to the standard ROCm index instead of
  silently mixing Radeon wheels with PyPI defaults. Quote all wheel
  arguments individually so wheel filenames cannot be word-split or
  glob-expanded.
- install_llama_prebuilt.py: detect_host() now requires nvidia-smi -L to
  list a GPU before setting has_physical_nvidia. Routes AMD ROCm hosts
  with a broken leftover nvidia-smi to the ROCm path instead of
  misclassifying them as NVIDIA.
- install_llama_prebuilt.py: scan upstream assets for any rocm-<version>
  prebuilt instead of hard-coding rocm-7.2, so ROCm 6.x / 7.0 / 7.1 / 7.3+
  users pick up a matching upstream prebuilt when one exists.
- install_llama_prebuilt.py: validate_server() adds --n-gpu-layers 1 for
  linux-rocm and windows-hip hosts, so new HIP prebuilts are preflighted
  on the GPU path instead of passing validation on CPU only.
- install_llama_prebuilt.py: restore the published windows-cpu fallback
  for AMD Windows hosts without a HIP prebuilt so hash-approved bundles
  are still preferred over the raw upstream CPU asset.
- install_python_stack.py: drop the /opt/rocm / hipcc gate in
  _ensure_rocm_torch() and rely on _has_rocm_gpu(). Runtime-only ROCm
  installs (package-managed minimal installs, Radeon software) that ship
  amd-smi / rocminfo without hipcc can now repair a CPU-only venv via
  "unsloth studio update". Adds an explicit IS_WINDOWS / IS_MACOS guard.
- studio/backend/utils/hardware/amd.py: honour HIP_VISIBLE_DEVICES /
  ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES in
  get_primary_gpu_utilization(). A process restricted to GPU 2 now
  reports metrics for GPU 2 instead of physical GPU 0. Tighten the plain
  bytes unit detection to an explicit allowlist.
- studio/backend/utils/hardware/hardware.py: route
  get_backend_visible_gpu_info()'s backend_cuda_visible_devices field
  through a helper that reads HIP_VISIBLE_DEVICES on ROCm. Drop the
  unconditional "(rocm=False)" suffix in apply_gpu_ids() logs.

* Fix round 2 regressions: ROCm validate_server and Windows HIP routing

Follow-up to 810b833b addressing review findings on the first round of
hardening commits:

- install_llama_prebuilt.py validate_server: gate --n-gpu-layers on the
  resolved install_kind instead of host.has_rocm. AMD Windows hosts
  without a HIP prebuilt fall back to windows-cpu and must not be
  validated with GPU layers; thread install_kind through from the
  caller.
- install_llama_prebuilt.py resolve_release_asset_choice: reinstate the
  "not has_rocm" guard on the published windows-cpu bundle so AMD
  Windows hosts reach resolve_upstream_asset_choice() where the new
  HIP prebuilt path lives. Prefer a published windows-hip bundle first
  when one exists, fall through to upstream HIP + upstream CPU
  otherwise.
- install_llama_prebuilt.py detect_host: also set has_physical_nvidia
  when the secondary --query-gpu block confirms a working NVIDIA GPU,
  so older nvidia-smi versions without -L support do not silently skip
  the Linux diagnostics that key off has_physical_nvidia.
- install_llama_prebuilt.py: drop redundant "import re as _re" /
  "import re as _re_rocm" local aliases in favour of the existing
  top-level "import re".
- install_python_stack.py _ensure_rocm_torch: run the AMD
  bitsandbytes install unconditionally after the HIP-torch probe so
  "unsloth studio update" on venvs that already have ROCm torch still
  gains the AMD bitsandbytes build.
- install.sh: add a non-x86_64 early-exit to get_torch_index_url() so
  aarch64 / arm64 Linux hosts do not hit the ROCm wheel index
  (PyTorch only publishes ROCm wheels for linux_x86_64).
- install.sh: add bitsandbytes install to the migrated-environment
  branch so upgrades pick it up for ROCm hosts instead of only the
  fresh-install path.
- install.sh: in the Radeon wheel path, pass version constraints +
  --no-index --find-links to uv instead of explicit wheel URLs so a
  version-compatible torch / torchvision / torchaudio triple is
  resolved, rather than picking the highest-version wheel for each
  package independently.
- studio/backend/utils/hardware/amd.py _first_visible_amd_gpu_id: fall
  through to lower-priority visibility env vars when the first entry
  is malformed (leading comma, all-whitespace first token) instead of
  silently returning GPU 0.

* Fix round 3 findings: x86_64 guard, ROCm version clip, Radeon deps

Address issues surfaced by the round 3 reviewers on top of 8636fa63:

- install_python_stack.py _ensure_rocm_torch: add the same `x86_64`
  guard that install.sh already has. Linux aarch64 / arm64 ROCm hosts
  must skip the repair path entirely; PyTorch only publishes ROCm
  wheels for linux_x86_64, and without this guard
  `unsloth studio update` aborts with a missing-wheel error on non
  x86_64 hosts.
- install_llama_prebuilt.py resolve_upstream_asset_choice: add a
  best-effort _detect_host_rocm_version() helper (reading
  /opt/rocm/.info/version, amd-smi version, hipconfig --version) and
  filter rocm_candidates to entries whose major.minor is <= host
  version. Falls back to the newest candidate only when no compatible
  one exists, so a ROCm 6.4 host downloads rocm-6.4 instead of being
  handed the numerically newest rocm-7.2 bundle (which fails preflight
  and forces a source build).
- install.sh: remove the round 2 --no-index switch from the Radeon
  wheel branch. --no-index forced uv to ignore PyPI entirely, which
  broke transitive dependency resolution (filelock, sympy, networkx,
  jinja2, fsspec, setuptools, typing-extensions, ...) on a fresh venv.
  Restore the round 1 explicit wheel URL invocation but add a
  torch / torchvision / torchaudio version-pair sanity check so a
  mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 + torchaudio
  2.9.0) falls back to the standard ROCm index instead of installing a
  broken combination.
- install_python_stack.py _ensure_rocm_torch: restructure the
  "tag is None" path so it no longer short-circuits the bitsandbytes
  install. On a ROCm runtime older than anything in
  _ROCM_TORCH_INDEX, print the "no wheel" warning but still run the
  AMD bitsandbytes install.
- studio/backend/core/training/worker.py: restore the pre-PR
  "no timeout" behaviour for non-HIP causal-conv1d / mamba-ssm source
  builds. The round 2 "timeout = 1800 if is_hip else 300" cap aborts
  slow non-HIP builds (Linux aarch64, unsupported torch/CUDA combos)
  after 5 minutes; omit timeout for the non-HIP branch so the cap
  only applies to ROCm source builds.

* Fix round 4 findings: apply_gpu_ids env inheritance, Radeon X.Y, bitsandbytes gate

Address remaining issues surfaced by the round 4 reviewers:

- studio/backend/utils/hardware/hardware.py apply_gpu_ids: mirror the
  selection into HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES whenever
  the caller already had a ROCm visibility env var set, not only when
  IS_ROCM has already been set by detect_hardware(). Training and
  inference workers call apply_gpu_ids() before detect_hardware()
  runs, so the old guard would leave a forked ROCm worker with a
  stale HIP_VISIBLE_DEVICES mask that no longer matched the
  narrowed CUDA_VISIBLE_DEVICES selection.
- install.sh get_radeon_wheel_url: accept X.Y ROCm versions in
  addition to X.Y.Z. The `/opt/rocm/.info/version` file and some
  hipconfig versions report only two components, and the Radeon
  repository publishes both rocm-rel-X.Y.Z/ and rocm-rel-X.Y/
  directories, so treating X.Y as invalid caused Radeon hosts to fall
  back to the generic ROCm index even when a matching AMD wheel set
  existed.
- install_python_stack.py _ensure_rocm_torch: only install the AMD
  bitsandbytes build when the venv actually has a ROCm-compatible
  torch (either already present or just installed by this function).
  Previously the bitsandbytes install ran unconditionally, which
  could leave an AMD bitsandbytes layered on top of a CPU/CUDA torch
  on hosts where the ROCm runtime is older than any entry in
  _ROCM_TORCH_INDEX. Also add --force-reinstall so an existing
  CPU/CUDA bitsandbytes is replaced by the AMD build during upgrades.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix gemini findings: amd-smi metric envelope validation and dict-wrapped GPU id

Two medium-severity defensive fixes from the gemini-code-assist review on
the AMD monitoring backend:

1. _extract_gpu_metrics may return a dict where every value is None when
   amd-smi succeeds (zero exit) but the JSON envelope contains no usable
   fields (error response, unsupported card). The new _has_real_metrics
   helper lets get_primary_gpu_utilization surface available:False and
   lets get_visible_gpu_utilization skip ghost device rows so the UI
   does not render placeholder cards with empty numbers.

2. Newer amd-smi versions wrap scalar fields as {"value": 0, "unit":
   "none"}, including the per-GPU id. The previous int(raw_id) call
   silently fell back to the enumeration index in that case, losing the
   real GPU id. Routing raw_id through the existing _parse_numeric
   helper handles bare ints, floats, strings, and the dict shape
   uniformly, with a debug log on parse failure.

* Fix gemini round 2 findings: explicit length guard on ROCm version file parser

Both _detect_rocm_version (install_python_stack.py) and
_detect_host_rocm_version (install_llama_prebuilt.py) read /opt/rocm/.info/version
or $ROCM_PATH/lib/rocm_version, split on "." and unconditionally accessed
parts[1]. The surrounding broad `except Exception: pass` already swallowed
the resulting IndexError, so a one-component file like "6\n" did fall
through to the next detection source -- but the control flow relied on
exception handling instead of an explicit check.

Add `if len(parts) >= 2:` guards in both helpers so the loop falls through
on its own without raising. Behaviour is unchanged for the common multi-
component case; the previously-silent IndexError path becomes an explicit
no-op.

* Fix gemini round 3: include has_rocm in validate_server fallback path

When validate_server is called without an explicit install_kind (older
call sites that have not been updated), the fallback was only enabling
--n-gpu-layers for NVIDIA and macOS arm64 hosts. AMD ROCm Linux hosts
fell through to the CPU validation path even though the prebuilt being
exercised was a HIP binary.

Add host.has_rocm to the fallback expression so the GPU offload flag is
applied consistently with the install_kind=='linux-rocm' / 'windows-hip'
branches above.

* Fix gemini round 4: remove risky bytes-vs-MB heuristic in _parse_memory_mb

The previous heuristic divided any bare number above 10_000_000 by
1024*1024 on the assumption that large unit-less values were bytes.
This misclassified small VRAM allocations: 5 MB of used VRAM reported
as 5_242_880 bytes without a unit would be taken at face value and
render as 5_242_880 MB (~5 TB) in the monitoring UI.

Modern amd-smi always provides explicit units (MiB/GiB dict form),
and legacy amd-smi returns bare numbers in MB -- the heuristic never
had a real workload to handle. Drop it and default to MB for bare
numeric input, keeping the existing unit-aware branches for dict /
string inputs unchanged.

The unrelated gemini suggestion to "default minor to 0" in the
amd-smi version awk parser was intentionally NOT applied: rocm7.0
and rocm7.1 ship different wheel sets, so silently substituting 0
for a missing minor could install the wrong wheels. The existing
reject-and-fall-through behaviour is safer.

* Fix gemini round 5: POSIX compliance and leading-comma visibility parsing

Three medium findings from gemini-code-assist addressed in this commit:

1. _pick_radeon_wheel used grep -o and sort -V, both GNU extensions
   that are not in POSIX and break on BSD/BusyBox coreutils. install.sh
   has a #!/bin/sh shebang so the whole pipeline was rewritten as a
   single awk script that extracts all href="..." hits on each line,
   filters to wheels matching the package prefix and python tag, and
   picks the newest version via zero-padded lexical comparison. No
   external sort or grep is needed.

2. _first_visible_amd_gpu_id in the AMD monitoring backend treated a
   leading comma (e.g. HIP_VISIBLE_DEVICES=",1") as "fall through to
   the next env var", which is surprising given the clear intent to
   narrow to device 1. Filter empty tokens after the split and return
   the first real one. An all-commas value ("," / ",,,") still falls
   through because no real tokens exist; the empty-string and "-1"
   explicit-zero cases are unchanged.

The unrelated amd-smi version awk parser suggestion was not applied
(see round 4 commit message for rationale: defaulting a missing minor
to 0 could silently install the wrong ROCm wheel set).

* Fix 20-reviewer.py findings: base drift, Radeon %2B, dpkg/rpm fallback, bnb, backend label

Consolidated fix batch from a 20-parallel reviewer.py run on the current
head. Each fix is drawn from a high-consensus finding and addresses a
real bug or feature gap, not a stylistic preference.

1. install.sh: bump `unsloth>=2026.4.2` -> `unsloth>=2026.4.4` at five
   call sites so this branch no longer regresses main's version floor
   (main bumped to 2026.4.4 in #4876). Without this, merging 4720 would
   silently downgrade the minimum version pin for fresh installs.

2. install.sh: URL-decode Radeon wheel names before extracting the
   torch / torchvision / torchaudio version strings. Real wheel URLs
   from repo.radeon.com are percent-encoded ("torch-2.10.0%2Brocm7.2.0...")
   so the previous `[+-]` terminator in the sed regex never matched,
   `_torch_ver` stayed empty, `_radeon_versions_match` stayed false,
   and every Radeon consumer install silently fell back to the generic
   ROCm index. Now decode %2B -> + first, then extract, then validate.

3. install.sh: the two AMD bitsandbytes install lines were running
   `uv pip install "bitsandbytes>=0.49.1"` without `--force-reinstall`,
   so upgrades where the venv already has a CPU/CUDA bitsandbytes
   satisfying the constraint would keep the stale non-AMD wheel. Add
   `--force-reinstall --no-cache-dir` to both call sites, matching the
   pattern already used in install_python_stack.py::_ensure_rocm_torch.

4. install_python_stack.py and install_llama_prebuilt.py: add
   `dpkg-query -W rocm-core` and `rpm -q rocm-core` fallbacks to the
   Python-side ROCm version detectors so they match the chain in
   install.sh::get_torch_index_url. Package-managed ROCm installs
   (Debian/Ubuntu/RHEL/Fedora distro packages) can expose GPUs via
   rocminfo/amd-smi but still lack /opt/rocm/.info/version, hipconfig,
   or amd-smi `version` output -- without these fallbacks, `unsloth
   studio update` on such hosts returned None and skipped the ROCm
   torch repair. Also strip the dpkg epoch prefix ("1:6.3.0-1") before
   parsing so epoch-annotated packages parse correctly.

5. hardware.py: add a `_backend_label(device)` helper that returns
   "rocm" when IS_ROCM is set and the device is DeviceType.CUDA, and
   use it for every `"backend": ...` emission in JSON responses served
   to the Studio frontend. Internally we still represent ROCm hosts as
   DeviceType.CUDA (ROCm torch reuses the whole torch.cuda.* API
   surface), but the user-facing API now correctly reports "rocm" on
   AMD boxes instead of labeling them as "cuda".

All 250 simulation scenarios pass (was 233 before this batch: added 17
new regression tests covering the version pin, %2B decoding, bnb
force-reinstall flags, dpkg/rpm fallback presence, and the
_backend_label helper's four-way truth table).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix gemini round 6 + URL audit: amd.py defensive checks, rocm6.5+ clip to 6.4

Two rounds of fixes in one commit, plus a full URL audit of every PyPI /
download.pytorch.org / repo.radeon.com reference the PR introduces.

amd.py (4 medium gemini findings on commit b3627bc2):

1. _extract_gpu_metrics used `and vram_total_mb` as part of the vram_util
   gate. The follow-up `vram_total_mb > 0` already handles the division
   guard, but the truthiness check was redundant and slightly surprising
   for a 0.0 valid value. Replace with explicit `is not None and > 0`
   for both vram_util and power_util.

2. get_physical_gpu_count called `data.get("gpu", ...)` without guarding
   for non-dict envelopes. A scalar / string JSON response from amd-smi
   would raise AttributeError. Add an isinstance(data, dict) check and
   return None for unexpected shapes.

3. get_visible_gpu_utilization had the same .get() exposure on the outer
   envelope. Rewrite the gpu_list extraction as an explicit
   list/dict/else cascade so a malformed scalar envelope produces
   gpu_list=[data] and continues without raising.

4. The same function's per-entry loop also called gpu_data.get() on
   whatever was inside gpu_list. If a scalar ever leaks into the list
   (directly or via the previous fix's fallback), _extract_gpu_metrics
   would raise on the first .get() inside the helper. Skip non-dict
   entries in the loop before extracting metrics.

install.sh (URL audit finding, previously flagged by 20-reviewer as #13):

5. get_torch_index_url used `rocm6.*` in the rocm tag case statement,
   which matched rocm6.5 and rocm6.6 and emitted
   download.pytorch.org/whl/rocm6.5 -- which returns HTTP 403 because
   PyTorch only publishes rocm 5.7, 6.0-6.4, 7.0-7.2. Enumerate the
   supported 6.x minors explicitly and add a rocm6.* fallback branch
   that clips to rocm6.4 (the last supported 6.x wheel set).

URL audit results (all URLs PR 4720 references):
- 14/14 download.pytorch.org/whl/{cpu,cu118,cu124,cu126,cu128,cu130,
  rocm6.0..6.4,rocm7.0..7.2} return HTTP 200.
- 9/9 repo.radeon.com/rocm/manylinux/rocm-rel-{5.7,6.0,6.1,6.2,6.3,
  6.4,7.0,7.1,7.2}/ return HTTP 200.
- X.Y.Z patch directories exist for 7.0.2, 7.1.1, 7.2.1 but NOT for
  6.3.0, 6.4.0, 6.2.1 -- install.sh already handles this via the X.Y.Z
  -> X.Y fallback sed in the Radeon wheel install block.
- Docs links (rocm.docs.amd.com, docs.unsloth.ai AMD guide) and the
  llama.cpp GitHub releases API endpoint all return 200.

Test suite: 255 -> 258. New regression coverage:
- U17: get_physical_gpu_count tolerates scalar amd-smi envelope
- U18: get_visible_gpu_utilization tolerates scalar envelope
- U19a-c: vram_util / power_util return None on zero total, but
  vram_total_gb still echoes 0.0 (not None)
- A_rocm{6.5,6.6,6.9}_clips_to_rocm64: install.sh clips unsupported
  6.x minors to rocm6.4 instead of producing a 403 index URL

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reviewer.py round 2: tokenizer AMD multi-GPU, --no-torch bnb, main.py backend label

Three high-confidence findings from a second 20-parallel reviewer.py run
on commit 7effb3ae. Triaged 15 total findings and applied the three that
were confirmed as real bugs; the rest were either false positives (e.g.
"migrated AMD venv not repaired" -- _ensure_rocm_torch runs downstream
via setup.sh regardless), design decisions (e.g. visibility mask env
vars not consulted in installer detection), or edge cases the existing
fallback logic already handles.

1. unsloth/tokenizer_utils.py [6/20]: the multi-GPU guard's shell probe
   runs `nvidia-smi --query-gpu=memory.used`, catches the failure, then
   only raises if `torch.cuda.is_available()` is False. On ROCm torch,
   torch.cuda.is_available() returns True (ROCm reuses the torch.cuda.*
   API), so the guard becomes dead code on AMD hosts and multi-GPU AMD
   setups slip through even though unsloth does not support them yet.
   Add a torch.cuda.device_count() > 1 fallback inside the except so
   AMD multi-visible-device setups are flagged consistently with the
   original CUDA memory check.

2. install.sh [1/20]: the fresh-install bitsandbytes block for AMD ROCm
   ran unconditionally when TORCH_INDEX_URL matched `*/rocm*`, even when
   SKIP_TORCH=true (from --no-torch or Intel Mac auto-detect). A user
   running `install.sh --no-torch` on an AMD host would still pull in
   bitsandbytes despite explicitly asking for GGUF-only mode. Wrap the
   case block in an outer `[ "$SKIP_TORCH" = false ]` guard.

3. studio/backend/main.py [3/20]: the /api/system endpoint returned
   `"device_backend": get_device().value`, which is "cuda" on ROCm
   hosts (because ROCm torch piggybacks on torch.cuda). Other endpoints
   (hardware.py) already use the _backend_label helper which swaps
   "cuda" -> "rocm" when IS_ROCM. Route /api/system through the same
   helper so the Studio UI reports the backend consistently across all
   endpoints.

4. studio/backend/tests/test_utils.py: update test_backend_matches_device
   to call _backend_label(get_device()) instead of raw get_device().value
   so the test matches the new contract and still passes on CUDA hosts.

Tests: 258 -> 261. New regression coverage:
- X08 main.py /api/system uses _backend_label
- X09 tokenizer multi-GPU guard has device_count() fallback
- X10 fresh-install bnb case block gated on SKIP_TORCH=false

* fix: prevent bitsandbytes from overwriting ROCm torch with CUDA wheels

During install, bitsandbytes was installed without --no-deps, causing
uv to resolve torch from PyPI (CUDA build) and silently overwrite the
ROCm wheels that were just installed in the previous step.

This happened in three places:
- install.sh: bitsandbytes install in both migrated and fresh paths
- install_python_stack.py: bitsandbytes install inside _ensure_rocm_torch()

Additionally, multiple install steps in install_python_stack.py (extras,
overrides, studio deps) can pull in CUDA torch via transitive
dependencies. A final _ensure_rocm_torch() call at the end of the
install sequence ensures ROCm torch is always in place at runtime.

All changes are gated behind ROCm-specific conditions and do not affect
NVIDIA, CPU-only, macOS, or Windows install paths.

Tested on AMD Instinct MI300X VF with ROCm 7.2.0 -- confirms
torch==2.10.0+rocm7.1 with HIP 7.1.25424 after install.

* fix: ROCm inference fallback -- skip Unsloth patching and bnb 4-bit on HIP

On AMD ROCm (HIP), two issues prevent the normal Unsloth inference path:

1. Unsloth's global monkey-patching of transformers model classes
   (LlamaRotaryEmbedding, attention modules) triggers
   _assert_async_cuda_kernel crashes on HIP during generation.
   Training uses different code paths and works fine.

2. bitsandbytes 4-bit matmul kernels also trigger HIP assertion
   failures on MI300X (CDNA3 / gfx942), even without Unsloth patching.

This commit adds a ROCm-specific inference fallback that:
- Skips importing Unsloth at module level (prevents global patching)
- Loads models in 16-bit with plain transformers + PEFT instead
- Resolves pre-quantized model names (e.g. "xxx-bnb-4bit" -> "xxx")
  since pre-quantized HF repos still trigger bnb codepaths
- Guards get_chat_template calls (unavailable without Unsloth import)
- Fixes max_seq_length=0 being passed to from_pretrained (GGUF
  semantics don't apply to transformers path)

The NVIDIA path is completely unchanged -- Unsloth import and
for_inference() optimization remain active. GGUF inference (via
llama-server/HIP) is unaffected since it never imports Python model
classes. AMD GPUs typically have large VRAM (e.g. 192GB on MI300X)
so 16-bit loading is practical for inference.

Tested on AMD Instinct MI300X VF (ROCm 7.2, HIP 7.1.25424):
- Simple generation: PASS
- Compare mode (base vs finetuned): PASS
- GGUF inference + tool calling: PASS (unaffected by this change)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: guard audio/vision inference on ROCm, remove unused import

- Add clear RuntimeError for audio/vision model inference on ROCm
  (these paths use Unsloth's FastModel/FastVisionModel which would
  crash on HIP; GGUF inference is the supported path on AMD)
- Remove unused `import os as _os` from the ROCm changes

* fix: amd-smi parsing for newer output format (gpu_data wrapper, mem_usage, temperature)

amd-smi on recent ROCm versions (7.x) wraps metric output in a
{"gpu_data": [...]} envelope instead of returning a raw list. This
caused get_primary_gpu_utilization() and get_visible_gpu_utilization()
to fail silently (returning available=False) because the GPU data
dict was never unwrapped.

Additionally:
- VRAM data moved from "vram" to "mem_usage" with "total_vram" /
  "used_vram" keys. Added fallback key lookup.
- Temperature "edge" sensor returns "N/A" on MI300X VF; the previous
  dict.get() chain returned the "N/A" string instead of falling
  through to "hotspot". Changed to a loop that checks each key until
  a parseable value is found.

Tested on AMD Instinct MI300X VF (ROCm 7.2, amd-smi 24.x):
- GPU utilization: 0% (idle), up to 100% during training
- Temperature: 40-44C (from hotspot sensor)
- VRAM: 0.28/191.69 GB (idle)
- Power: 158-211W draw

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Bug fix detecting radeon (#4940)

* Bug fix detecting radeon

* Expanding GPU target for gfx1100*

* Generalize gfx family-prefix filter to cover gfx10/gfx12 as well

rocminfo on ROCm 6.1+ emits LLVM generic-family ISA lines alongside the
specific GPU (e.g. gfx11-generic next to gfx1100). The outer grep captures
the bare family prefix from the generic line, and passing that to
-DGPU_TARGETS breaks the HIP build because clang only accepts specific
gfxNNN ids.

The previous filter only special-cased gfx11. Generalize it so any bare
2-digit family prefix (gfx10, gfx11, gfx12, ...) is dropped whenever a
specific sibling target is present in the same list. No real AMD GPU has
a 2-digit gfx id, so the filter can only ever drop family prefixes and
never a real target.

Covers the existing gfx11 cases unchanged, and extends the same fix to
gfx10-1-generic / gfx10-3-generic (RDNA1/2) and gfx12-generic (RDNA4),
which would otherwise hit the same build failure on newer rocminfo.

---------

Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>

---------

Co-authored-by: Eda Z <eda.zhou@amd.com>
Co-authored-by: GoldenGrapeGentleman <yueyuan@amd.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: billishyahao <bill.he@amd.com>
Co-authored-by: Iswarya Alex <47045679+iswaryaalex@users.noreply.github.com>
Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
2026-04-10 01:56:12 -07:00
Datta Nimmaturi
3b5a49776b
[studio] multi gpu: revert to balanced for inference. (#4698)
* Revert to balanced for inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused for_inference parameter from get_device_map

Since inference and training both use "balanced" now, the for_inference
flag is dead code. Remove it from the function signature, the call site
in inference.py, and simplify the tests accordingly.

* Remove redundant TestDeviceMapForInference test class

TestGpuAutoSelection already covers the same multi-gpu and single-gpu
device_map assertions. The TestDeviceMapForInference class was left
over from when for_inference had distinct behavior.

* Remove redundant test_get_device_map_multi_gpu_uses_balanced

Its assertions ([0,1] -> balanced, [0] -> sequential) are already
covered by test_get_device_map_uses_explicit_gpu_selection.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-31 01:24:41 -07:00
Datta Nimmaturi
9311df2b29
[Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602)
* [WIP] balanced device map for studio

* gpus as a request parameter

* API for multi GPU stuff

* return multi gpu util in new API

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use balanced_low0 instead of balanced

* Use balanced_low0 instead of balanced

* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests

- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
  gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
  CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
  add required job_id arg to start_training() calls

* Smart GPU determinism using estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disallow gpu selection for gguf for now

* cleanup

* Slightly larger baseline

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Treat empty list as auto

* Verbose logging/debug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup and revert unnecessary deletions

* Cleanup excessive logs and guard against disk/cpu offload

* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* support for non cuda gpus

* Fix multi-GPU auto-selection memory accounting

The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.

Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add sandbox tests for multi-GPU selection logic

24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry

1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
   instead of the FP16 1.3x multiplier. This prevents over-sharding
   quantized models across unnecessary GPUs.

2. When model size estimation fails, auto_select_gpu_ids now falls back to
   all visible GPUs instead of returning None (which could default to
   single-GPU loading for an unknown-size model).

3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
   None) instead of rejecting it as an unsupported explicit request.

4. Training retry path for "could not get source code" now preserves the
   gpu_ids parameter so the retry lands on the same GPUs.

5. Updated sandbox tests to cover the new 4-bit inference estimate branch.

* Remove accidentally added unsloth-zoo submodule

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix UUID/MIG visibility and update test expectations

1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
   visibility APIs now return "unresolved" with empty device lists instead
   of exposing all physical GPUs. This prevents the UI from showing GPUs
   that the backend process cannot actually use.

2. test_gpu_selection.py: Updated test expectations to match the new
   multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
   additional GPUs) and 4-bit inference memory estimation formula.
   All 60 tests now pass.

* Add CPU/disk offload guard to audio inference path

The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.

* Improve VRAM requirement estimates

* Replace balanced_low_0 with balanced

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refine calculations for slightly easier nums

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* adjust estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use nums instead of obj to avoid seralisation error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden nvidia-smi parsing and fix fallback GPU list

1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
   so MIG slices, N/A values, or unexpected nvidia-smi output skip the
   unparseable row instead of aborting the entire GPU list.

2. nvidia.py: Handle GPU names containing commas by using the last
   field as memory instead of a fixed positional index.

3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
   VRAM data) instead of raw devices list, which could include GPUs
   with null VRAM that were excluded from the ranking.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* consolidate raise_if_offload

* Improve MoE support. Guard against nvidia-smi failures

* Improve MoE support. Guard against nvidia-smi failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case

1. vram_estimation.py: compute_lora_params now includes shared experts
   (n_shared_experts) alongside routed experts when computing MoE LoRA
   adapter parameters. Previously only n_experts were counted, causing
   the estimator to undercount adapter, optimizer, and gradient memory
   for DeepSeek/GLM-style models with shared experts.

2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
   reports system-wide VRAM usage) instead of memory_allocated (which
   only reports this process's PyTorch allocations). This prevents
   auto-selection from treating a GPU as mostly free when another
   process is consuming VRAM. Falls back to memory_allocated when
   mem_get_info is unavailable.

3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
   CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
   list inherits the parent visibility, same as None.

4. hardware.py: Upgraded fallback_all GPU selection log from debug to
   warning so operators are notified when the model likely will not fit
   in available VRAM.

* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired

get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.

* Guard get_primary_gpu_utilization and reset GPU caches between tests

1. nvidia.py: get_primary_gpu_utilization now catches OSError and
   TimeoutExpired internally, matching the pattern already used in
   get_visible_gpu_utilization and get_backend_visible_gpu_info. All
   three nvidia-smi callers are now self-contained.

2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
   module-level _physical_gpu_count and _visible_gpu_count caches in
   tearDown. Applied to all test classes that exercise GPU selection,
   device map, or visibility functions. This prevents stale cache
   values from leaking between tests and causing flaky results on
   machines with real GPUs.

* Fix nvidia-smi fallback regression and physical GPU count validation

1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
   get_backend_visible_gpu_info now check result.get("available") before
   returning the nvidia-smi result. When nvidia-smi is unavailable or
   returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
   the functions fall through to the torch-based fallback instead of
   returning an empty result. This fixes a regression where the internal
   exception handling in nvidia.py prevented the caller's except block
   from triggering the fallback.

2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
   validation from physical upper-bound validation. The physical count
   check is only enforced when it is plausibly a true physical count
   (i.e., higher than the largest parent-visible ID), since
   torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
   visible count, not the physical total. The parent-visible-set check
   remains authoritative in all cases. This prevents valid physical IDs
   like [2, 3] from being rejected as "out of range" when nvidia-smi is
   unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
   2 devices.

* Fix UUID/MIG torch fallback to enumerate devices by ordinal

When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.

Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.

Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-30 02:33:15 -07:00
NuoFang
4cedeba8c2
fix(studio): prevent ModuleNotFoundError in dataset.map() on Windows (#4473)
* fix(studio): prevent ModuleNotFoundError in dataset.map() on Windows

On Windows, dataset.map() uses "spawn", which requires workers to
import compiled modules from disk. Previously, clear_unsloth_compiled_cache()
deleted the entire directory, causing workers to crash when looking for
UnslothSFTTrainer.py.

Changes:
1. Added `preserve_patterns` to cache cleanup to keep `Unsloth*Trainer.py`
   on Windows while clearing model-specific files.
2. Added the cache directory to PYTHONPATH for spawn workers.
Linux/macOS behavior is unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix spawn-platform coverage, CWD path mismatch, and race condition for PR #4473

- Extend platform guard from win32-only to include macOS (also uses spawn
  since Python 3.8, same ModuleNotFoundError would occur)
- Replace fragile CWD-based PYTHONPATH registration with centralized
  register_compiled_cache_on_path() that uses the same __file__-relative
  _CACHE_DIRS already used by cache_cleanup -- fixes path mismatch when
  studio is launched from a directory other than the repo root
- Move PYTHONPATH registration to the top of _train_worker(), before any
  dataset.map() call (previously it ran late in config assembly, after
  dataset formatting which also calls dataset.map())
- Update inference.py model-unload to preserve trainer files on spawn
  platforms, preventing a race where unloading a model via inference tab
  would delete UnslothSFTTrainer.py while training workers are importing it

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix cache-dir precedence reversal in register_compiled_cache_on_path()

Iterating _CACHE_DIRS in forward order while calling insert(0) each time
reverses the declared priority: later entries shadow earlier ones. When
multiple compiled-cache directories exist, spawned workers could import a
stale trainer from the wrong cache.

Fix: iterate in reverse so that the highest-priority entry (first in
_CACHE_DIRS) is inserted last and ends up at position 0 in sys.path and
PYTHONPATH.

* fix: harden worker-count helpers against cpu_count=None and desired<=0

- safe_num_proc: guard os.cpu_count() with `or 1`, clamp multi-GPU
  path with max(1, min(4, desired)), clamp return with max(1, desired)
- safe_thread_num_proc: same os.cpu_count() guard and return clamp
- Add regression tests (31 L1 unit + 10 sandbox edge-case tests)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove regression tests from PR

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-22 06:11:24 -07:00
Andrew Barnes
2c5d3c48ec
fix: subprocess crash during map operation on Windows (#4507)
* fix: handle Windows subprocess crash during dataset.map()

Windows uses spawn (not fork) for multiprocessing. Spawned workers
cannot resolve Unsloth's dynamically compiled cache modules from
unsloth_compiled_cache/, causing ModuleNotFoundError and RuntimeError
during dataset.map() tokenization.

Add two platform-guarded patches for sys.platform == "win32":
1. Force HF_DATASETS_MULTITHREADING_MAX_WORKERS=1 and set spawn method
2. Monkey-patch Dataset.map() to force num_proc=None

Fixes #4490

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review: extend spawn fix to macOS, add multiprocess fallback

- Change platform checks from sys.platform == "win32" to
  sys.platform != "linux" so macOS (also spawn-based) is covered
- Wrap multiprocess import in try/except falling back to stdlib
  multiprocessing when the multiprocess package isn't installed
- Rename _win32_safe_map to _spawn_safe_map to reflect broader scope

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: replace global Dataset.map monkey-patch with targeted num_proc routing

The previous approach had issues: Patch 1 set HF_DATASETS_MULTITHREADING_MAX_WORKERS
and forced set_start_method (dead code on platforms already using spawn), and Patch 2
globally monkey-patched Dataset.map() (too broad, missed Dataset.filter()).

Replace with a two-layer fix:

1. Studio layer: Add dataset_map_num_proc() that returns None on spawn platforms
   (Windows, macOS). Unlike num_proc=1 which still creates Pool(1) and spawns a
   worker, num_proc=None runs Dataset.map()/filter() truly in-process.
   Update all dataset.map() callsites to use it. ThreadPoolExecutor callers
   (format_conversion.py) keep using safe_num_proc() since threads are unaffected.

2. Root-cause layer: Propagate UNSLOTH_COMPILE_LOCATION via PYTHONPATH on spawn
   platforms so spawned workers can import compiled modules. Mirrors the .venv_t5
   pattern in worker.py. Does not import unsloth_zoo.compiler (heavy torch/triton
   imports). Completely skipped on Linux.

Also extend safe_num_proc() to return 1 on macOS (was only guarding Windows),
and narrow the transformers 5.x dataloader guard from != "linux" to explicit
("win32", "darwin").

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: add safe_thread_num_proc() for ThreadPoolExecutor callsites

safe_num_proc() correctly caps to 1 on macOS/Windows for process-based
multiprocessing, but format_conversion.py reuses it for ThreadPoolExecutor
workers. Threads share address space and are unaffected by spawn, so
capping to 1 makes image URL downloads sequential -- a real regression.

Add safe_thread_num_proc() that skips the platform guard but keeps the
cpu_count heuristic, and switch both ThreadPoolExecutor callsites in
format_conversion.py to use it.

* fix: remove double-wrap in dataset_num_proc + fix num_proc=1 in datasets route

- trainer.py:3009: Replace safe_num_proc(max(1, os.cpu_count() // 4))
  with max(1, (os.cpu_count() or 1) // 4) to avoid double-wrapping
  inside dataset_map_num_proc which already calls safe_num_proc
- trainer.py:15-20: Clarify comment on PYTHONPATH propagation
- datasets.py:445: Change num_proc=1 to num_proc=None for 10-row
  preview slice (avoids unnecessary multiprocessing overhead)

* fix: guard os.cpu_count() against None in worker-count helpers

os.cpu_count() can return None on some platforms. Use (os.cpu_count() or 1)
to prevent TypeError in safe_num_proc() and safe_thread_num_proc().

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-22 05:21:09 -07:00
Daniel Han
0acd1c7eec
studio: improve onboarding UX, tooltips, and training defaults (#4355)
* studio: improve onboarding UX, tooltips, and training defaults

- Change splash text to "Train and run LLMs locally"
- Add "Chat Only" card with BubbleChatIcon to skip directly to chat
- Add Skip/Skip to Chat buttons in sidebar and footer
- Back button on step 1 returns to splash screen instead of being disabled
- Change "Watch video guide" to "Get started with our guide" with new URL
- Update intro text to mention all model types + chat
- Make all tooltips clickable (in addition to hover) via React context
- Strip surrounding quotes from pasted HF tokens
- Rename "Eval Split" to "Evaluation Split"
- Add SparklesIcon to "Auto Detect" format option
- Change step 4 heading to "Choose your training parameters"
- Default max_steps to 60
- Learning rate displayed in scientific notation with +/- stepper
- Context length options capped by model's max_position_embeddings (via AutoConfig)
- Fix "QLORA"/"LORA" to "QLoRA"/"LoRA" in summary step
- Backend: add max_position_embeddings to model config endpoint

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* compare for 2 diff models

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* resolving gemini comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: disable thinking for Qwen3.5 <9B and always for AI Assist

- Change Qwen3.5 thinking threshold from <=2B to <9B (0.8B, 2B, 4B
  all disable thinking by default; 9B+ enables it)
- Always pass enable_thinking=False in AI Assist helper calls
  (_run_with_helper and _generate_with_backend) regardless of chat
  thinking settings

* studio: address PR review comments

- Extract _get_max_position_embeddings helper to DRY config extraction
- Fix "Skip to Chat" to navigate to /chat on step 1 (was /studio)

* fix: comment out debug print statements

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: skip Shiki highlighting for incomplete SVG code fences

While streaming SVG content, the syntax highlighter (Shiki) re-parses
the entire growing SVG on every token, blocking the main thread and
freezing the code area until the fence closes. Show a plain-text
preview for incomplete SVG fences instead, similar to how Mermaid
diagrams show a placeholder while streaming.

* studio: fix default top_k from 50/40 to 20 for chat inference

Per Qwen3.5 docs (unsloth.ai/docs/models/qwen3.5), top_k should be 20
for both thinking and non-thinking modes. The model-specific config in
inference_defaults.json already had top_k=20 for Qwen3.5, but the
generic fallback defaults were wrong:
- Frontend DEFAULT_INFERENCE_PARAMS.topK: 50 -> 20
- Backend generate_chat_completion top_k: 40 -> 20
- Backend generate_chat_completion_with_tools top_k: 40 -> 20
- Frontend title generation top_k: 40 -> 20

* studio: set universal inference defaults for unknown models

Default params for any model without specific config:
  temperature=0.6, top_p=0.95, top_k=20, min_p=0.01,
  presence_penalty=0.0, repetition_penalty=1.0

Models with entries in inference_defaults.json (Qwen3.5, Gemma-3,
Llama, etc.) override these with their recommended values.

Updated in: frontend DEFAULT_INFERENCE_PARAMS, backend Pydantic
request models, and backend generate_chat_completion defaults.

* studio: only trust_remote_code for unsloth/ models in AutoConfig

Only set trust_remote_code=True when the model name starts with
"unsloth/". All other models default to False for safety.

* studio: move Generating spinner above the composer

The "Generating" spinner was below the send message bar, causing
the bar to jump up and down. Move it above the composer in both
the regular thread view and the welcome/empty view.

* studio: adjust toast close button position away from edge

Move the X close button on toasts (like "Starting model...") from
top-1.5 to top-3 and add right-3, giving more breathing room from
the top-right corner.

* studio: make Think button smaller with tighter icon-text gap

Reduce gap from 1.5 to 0.5, padding from px-2.5/py-1 to px-2/py-0.5,
and icon from size-3.5 to size-3.

* studio: multiple onboarding and chat UX improvements

- Move Generating spinner above composer (fixes jumping send bar)
- Make Think button smaller with tighter icon-text gap
- Chat card now inside grid (same size as Audio/Embeddings cards)
- Rename "Chat Only" to "Chat"
- Chat card requires Continue to proceed (no auto-advance)
- Continue on Chat selection skips onboarding and goes to /chat
- Tooltip (i) click on Chat card doesn't trigger navigation
- Step 1 footer Back button goes back to splash (label is "Back")
- Splash "Skip Onboarding" renamed to "Skip to Chat", navigates to /chat
- Toast close button moved away from edge

* studio: align Skip to Chat button, add Skip to footer

- Sidebar "Skip to Chat" now uses primary (green) Button style with
  arrow icon, full width, aligned like step items. Shows on all steps.
- Footer: added "Skip" outline button next to Continue that goes
  directly to /studio with progress saved (markOnboardingDone)

* studio: change default max steps from 30 to 60 in toggle hook

The DEFAULT_MAX_STEPS in use-max-steps-epochs-toggle.ts was still 30,
used as fallback when toggling from epochs back to max steps.

* studio: extend context length options to 262K

CONTEXT_LENGTHS now includes 65536, 131072, 262144 in addition to
the existing 512-32768 range. The onboarding step filters these by
the model's max_position_embeddings (e.g. Nemotron-3-Nano-4B has
262144), showing powers of 2 up to the model's maximum.

* studio: auto-select LoRA vs QLoRA based on model size and GPU memory

After selecting a model in onboarding, detect the total model weight
file size from HF Hub (safetensors/bin files). Then estimate memory
needed: model_size_gb * 1.5 * context_scale, where context_scale is:
  - <=8192 tokens: 1.0x
  - >8192 tokens: 1.7x
  - >=16384 tokens: 2.0x
  - >=32768 tokens: 4.0x

If the estimate fits in free GPU VRAM, default to LoRA (16-bit).
Otherwise default to QLoRA (4-bit).

Backend changes:
- Add model_size_bytes to ModelDetails (models.py)
- Add _get_model_size_bytes() using HfApi.repo_info (routes/models.py)
- Add vram_free_gb to get_gpu_summary (hardware.py)

Frontend changes:
- Add autoSelectTrainingMethod() in training-config-store.ts
- Called after model defaults are loaded
- Add model_size_bytes to ModelConfigResponse type
- Add vramFreeGb to HardwareInfo hook

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: rename "Importing ML libraries..." to "Importing Unsloth..."

* studio: show model/dataset in training status, fix LoRA/QLoRA casing

- Training status now shows 'Training "model_name"' and 'Dataset = ...'
  instead of generic "Starting training..."
- Fix Studio progress section to show QLoRA/LoRA instead of QLORA/LORA

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: rename 'Skip to Chat' to 'Skip Onboarding' on splash screen

* studio: add presence_penalty support for chat inference

Add presence_penalty as a parameter across the full stack:
- Backend: llama_cpp.py generate_chat_completion/with_tools, Pydantic
  models (inference.py), routes/inference.py pass-through
- Frontend: InferenceParams type, DEFAULT_INFERENCE_PARAMS (0.0),
  chat-adapter.ts payload, chat-settings-sheet.tsx slider (0-2),
  model defaults loading from inference_defaults.json
- Set Qwen3.5 default presence_penalty to 1.5 per official docs
- Default for unknown models is 0.0 (off)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix Chat card deselecting Text and aligning with other cards

* studio: fix presence_penalty not loading from inference defaults

The inference_config.py load_inference_config() was not including
presence_penalty in the returned config dict, so the Qwen3.5
default of 1.5 from inference_defaults.json never reached the
frontend. Added it to the config builder.

* studio: add delete button for cached models in model selector

Add trash icon on each downloaded model row (GGUF and safetensors) with
confirmation dialog. Backend DELETE /api/models/delete-cached endpoint
uses huggingface_hub scan_cache_dir + delete_revisions to cleanly remove
cached repos, refusing if the model is currently loaded.

* studio: restore inference defaults, reasoning, and tools on page refresh

On page refresh with a model already loaded, the frontend was not
re-applying model-specific inference defaults (presence_penalty,
temperature, etc.) or restoring reasoning/tools support flags.

Backend: Add inference config, supports_reasoning, supports_tools,
and context_length to InferenceStatusResponse.

Frontend: In the refresh callback, when an active model is detected,
apply mergeRecommendedInference and restore reasoning/tools flags
with proper Qwen3.5 size-based defaults.

* studio: fix delete dialog closing before async completes

Prevent AlertDialogAction's default close behavior with
e.preventDefault() so the dialog stays open during deletion.
Also block onOpenChange dismiss while deleting is in progress.

* fix: add Dict and Any imports to inference models

* studio: fix Qwen3.5 reasoning threshold in frontend load path

The frontend loadModel handler had the old threshold (<=2) for
disabling reasoning on small Qwen3.5 models. Changed to <9 to
match the backend. This was causing 4B to not properly disable
thinking by default when auto-loaded.

* studio: move GGUF delete to per-variant level

For GGUF repos, the trash icon now appears on each downloaded variant
row inside the quantization expander instead of on the repo-level row.
Backend accepts optional variant param to delete specific GGUF files
(blob + symlink) rather than the entire repo cache.

* studio: restore ggufContextLength on page refresh

The Max Tokens slider was capped at 32768 on page refresh because
ggufContextLength was not restored from the status response.
Now set it from statusRes.context_length on reconnect.

* fix: remove <think> from Qwen3.5 response template marker

The train-on-responses-only feature uses template markers to find
where the assistant response starts. The Qwen3.5 response marker
included '<think>\n' which is only present when thinking mode is
enabled. With thinking disabled (default for <9B), the marker
never matched, causing 100% of samples to be dropped.

Changed response marker from '<|im_start|>assistant\n<think>\n'
to '<|im_start|>assistant\n' which works regardless of thinking mode.

* studio: fix sloth ASCII art alignment in training overlay

* fix: correct sloth ASCII art alignment to match Unsloth banner

* studio: add Python and terminal tool calling to chat

Register python and terminal tools alongside web search. Python
executor validates imports (stdlib only) via unsloth_zoo
rl_environments, runs code in a subprocess sandbox with 5-min
timeout and cancel support. Terminal executor blocks dangerous
commands (rm, sudo, etc.) and runs in a temp directory.

Update llama_cpp tool loop to show tool-specific status messages
and pass cancel_event through to executors. Rename composer
toggle from "Search" to "Tools" and show TerminalIcon for
execution status pills.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix Nemotron/transformers 5.x support, onboarding navigation, port binding

Backend:
- Dynamic transformers 5.x detection via tokenizer_config.json fetch
  (checks for TokenizersBackend class, cached per-model)
- Bump transformers 5.x version from 5.2.0 to 5.3.0 across all workers,
  setup scripts (setup.sh, setup.ps1)
- Auto-enable trust_remote_code for unsloth/* models needing transformers 5.x
  (workaround for NemotronH config parsing bug in transformers)
- Auto-install mamba-ssm/causal-conv1d for SSM models (NemotronH, Falcon-H1)
  with --no-build-isolation --no-deps to avoid torch version conflicts
- Add SO_REUSEADDR to port check in run.py (fixes Colab proxy stale connection
  falsely reporting port as in-use)

Frontend:
- Fix "Skip to Chat" navigation: use window.location.href instead of React
  Router navigate() to bypass useEffect redirect race
- Fix "Skip Onboarding" on splash: navigates to /studio (not /chat)
- Fix onboarding guard: only check isOnboardingDone() on initial mount
- Fix Chat card on step 1: add sr-only spacer for consistent alignment
- Fix Chat+Text both selected: clear RadioGroup value when Chat is selected

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: split tools toggle into Search and Code buttons

Replace the single "Tools" toggle with two independent toggles:
- "Search" (globe icon) enables web search only
- "Code" (terminal icon) enables Python and terminal execution

Add enabled_tools list field to the inference payload so the
backend only registers the tools the user has toggled on. Both
toggles appear in the main composer and the compare composer.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix tool calling import validation and error logging

Replace unsloth_zoo-dependent import checker with a standalone
ast-based validator using sys.stdlib_module_names. This properly
blocks non-stdlib imports (numpy, requests, etc.) and returns a
clear error message to the model so it can rewrite using only
stdlib.

Add full traceback to tool streaming error logs for debugging.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: parse gpt-oss harmony channels for clean safetensors chat output

gpt-oss models emit multi-channel output via harmony protocol tokens
(<|channel|>analysis<|message|>... and <|channel|>final<|message|>...).
TextIteratorStreamer with skip_special_tokens=True strips the special
tokens but leaves channel names concatenated with content, producing
garbled output like "analysisWe need to...assistantfinalHello!".

Add HarmonyTextStreamer that decodes with skip_special_tokens=False,
parses harmony markup via regex, and emits <think>analysis</think>
for the analysis channel and plain text for the final channel --
reusing the existing frontend reasoning UI.

Also expose supports_reasoning=True for non-GGUF gpt-oss models in
the /status endpoint so the frontend enables the Think toggle.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: use unsloth_zoo for Python sandbox validation

Set UNSLOTH_IS_PRESENT=1 and import check_python_modules and
check_signal_escape_patterns directly from unsloth_zoo instead
of a standalone fallback. This gives us the full Unsloth
validation including stdlib-only import checks and signal/timeout
escape pattern detection.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: allow all imports in Python tool sandbox

Remove stdlib-only import restriction. Keep signal escape
pattern detection via unsloth_zoo for safety.

* studio: fix ReadTimeout on tool streaming final pass

The 0.5s read timeout used for cancel-checking during streaming
also fires when waiting for the first response from llama-server
(e.g. reasoning model thinking for 15+ seconds). Add
_stream_with_retry() context manager that retries on ReadTimeout
while checking cancel_event, so the model has unlimited time to
think before producing the first token. Applied to both the
regular streaming path and the tool-calling final pass.

* fix: rewrite HarmonyTextStreamer with stateful incremental parsing

The delta-on-transformed approach had two critical bugs:

1. Before the full <|channel|>X<|message|> pattern was complete, the
   strip-tokens fallback emitted "analysis" as plain text. Then when
   the regex matched, _transform returned a completely different format
   (<think>...</think>) and the delta was computed against the wrong
   base string, producing fragments like "think>", "nk>", ">".

2. Even with full matches, the closing </think> tag shifted position
   as content grew, so text[prev_len:] produced garbled deltas.

Replace with stateful incremental parsing that:
- Buffers until a complete channel+message pair is seen
- Emits <think> once when analysis channel first appears
- Streams analysis content deltas (computed on channel content directly)
- Emits </think> once when final channel first appears
- Streams final content deltas
- Closes open think tags in end()

Also skip the generic all_special_tokens stripping in
_clean_generated_text for gpt-oss since HarmonyTextStreamer already
produces clean output and the generic stripping was mangling <think>
tags.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: strip all <|...|> tokens in gpt-oss cleanup, not just harmony subset

The gpt-oss tokenizer has added tokens like <|return|> (id=200002) that
are not part of the harmony channel protocol but can leak into output.
The previous regex only stripped channel|message|start|end tokens.

Broaden the _clean_generated_text regex for gpt-oss to <\|[a-z_]+\|>
which catches all pipe-delimited tokens (return, constrain, reserved,
etc.) without matching <think>/<\/think> tags.

Verified: gpt-oss all_special_tokens are only <|return|>,
<|reserved_200017|>, <|startoftext|> -- none overlap with <think>.
The harmony tokens (channel, message, start, end) are added_tokens
but not in all_special_tokens.

* fix: hide config-only model repos from cached models list

Repos that only have metadata/config files cached (no .safetensors or
.bin weight files) were showing up in the Downloaded list with tiny
sizes like "1.8 KB" or "24 KB". These are just leftover config
snapshots from architecture checks, not usable models.

Filter the cached-models endpoint to only include repos that contain
actual model weight files (.safetensors or .bin).

* studio: fix toast description text contrast in dark mode

Add explicit !text-muted-foreground to toast description classNames
so secondary text (e.g. "Releases VRAM and resets inference state.")
is readable in dark mode.

* studio: fix Chat card icon alignment with size-4 spacer

Replace sr-only span (takes no space) with a size-4 shrink-0 div
matching the RadioGroupItem dimensions in other cards, so the Chat
icon aligns vertically with Text/Audio/Vision/Embeddings icons.

---------

Co-authored-by: workspace <user@workspace.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Manan17 <shahmanan170602@gmail.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-03-17 07:46:07 -07:00
Daniel Han
eeffa4c065
studio: web search, KV cache dtype, training progress, inference fixes
## Summary
- Add web search tool calling for GGUF models (Search toggle, DuckDuckGo via ddgs)
- Add KV cache dtype dropdown (f16/bf16/q8_0/q5_1/q4_1) in Chat Settings
- Fix Qwen3/3.5 inference defaults per official docs (thinking on/off params)
- Enable reasoning by default for Qwen3.5 4B and 9B
- Replace "Generating" toast with inline spinner
- Fix stop button via asyncio.to_thread (event loop no longer blocked)
- Fix CUDA 12 compat lib paths for llama-server on CUDA 13 systems
- Fix auto-load model name not appearing in selector
- Training progress messages + dataset_num_proc fix

Integrated PRs:
- #4327 (imagineer99): BETA badge alignment (already in tree)
- #4340 (Manan Shah): prioritize training models in model selection
- #4344 (Roland Tannous): setup.sh macOS python version compatibility
- #4345 (Manan Shah): revamp model+dataset checking logic
2026-03-17 00:30:01 -07:00
Roland Tannous
47654cb91c Final cleanup 2026-03-12 18:28:04 +00:00
Roland Tannous
a2baf80511 Update license headers 2026-03-12 17:23:10 +00:00
Roland Tannous
817f2e8dcc feat: integrate structlog, configure workers for prod logging, and migrate print statements 2026-03-11 12:33:16 +00:00
Roland Tannous
d882678fe4 Add AGPL-3.0 SPDX headers to all source files 2026-03-09 20:17:45 +00:00
Roland Tannous
7e021886c8 Force num_proc=1 on Windows to avoid slow spawn overhead 2026-03-01 13:05:10 +00:00
Roland Tannous
d74174f7f5 Cap dataset.map num_proc on multi-GPU machines to prevent fork deadlocks 2026-02-23 14:25:31 +00:00
Manan17
76cd1dc24c fixing the hangup of training after multiple back to back training processes 2026-02-18 08:18:13 +00:00
Roland Tannous
a0ebd9183a feat: add live GPU monitor with nvidia-smi polling during training 2026-02-16 11:47:43 +00:00
Roland Tannous
b20d50e8d0 feat: add GET /api/system/hardware endpoint for GPU info and package versions 2026-02-16 10:29:21 +00:00
Roland Tannous
63c583c54f replace torch MPS with MLX 2026-02-11 16:04:35 +00:00
Roland Tannous
59d5f24eb5 integrate global hardware detection at lifespan entrypoint 2026-02-11 15:34:26 +00:00