--chat_template {auto,grpo,native} matches the flag added to
qwen3_flex_inference so cross-engine comparisons can hold the prompt
template constant per model (auto: GRPO for Qwen3, tokenizer native
otherwise). Threaded through build_prompts() and applied to all three
backends (vllm, tpaged, unsloth_fi_false).
--enforce_eager (vLLM backend only) forwards enforce_eager=True to
FastLanguageModel.from_pretrained so the vLLM engine skips torch.compile
+ cudagraph capture. Needed because vLLM 0.19.1 on torch 2.10 raises
`RuntimeError: Tried to erase Node size_1 but it still had 2 users`
inside compilation.backends.split_graph during the first forward; the
eager path still uses PagedAttention + FlashInfer for decode, so the
measurement stays meaningful (just no graph capture).
The flex_attention + paged KV + CUDA graphs inference engine was
Qwen3-specific in a handful of places, but the underlying engine
(PageTable, PagedKVCache, manual forward walker, decode graph capture,
double-copy LoRA rollout, FA4 capability guard) reads only attributes
that LlamaAttention / LlamaModel also expose. This change makes the
engine run on both Qwen3 and Llama-3.2-3B-Instruct.
Attention forward factory:
- make_flex_qwen3_attention_forward -> make_flex_attention_forward
- Guard the per-head QK RMSNorm call behind hasattr(self, "q_norm").
Qwen3 has it, Llama does not. The Qwen3 path is byte-equivalent to
before: RMSNorm on [B, S, H, D] (per-head) then transpose.
- patch_qwen3_model -> patch_model_attention_forwards.
Chat template selection:
- New --chat_template {auto,grpo,native}. auto picks GRPO for Qwen3
and the tokenizer's shipped template otherwise. grpo forces GRPO
(matches prior Qwen3 baselines). native forces the tokenizer's own
template (Llama-3.2-Instruct only produces coherent completions
with its shipped Instruct template).
Stats JSON:
- backend: "qwen3_flex" -> "flex"; adds "model_name" so multi-arch
runs land in a single schema.
README: one paragraph noting Llama-3.2 support + the --chat_template
native flag.
Measured on B200 (sm_100), n_prompts 64, max_new_tokens 512, 5 rounds,
--capture_cudagraph, double-copy LoRA rank 32:
Qwen3-4B-Base bf16 3975 tok/s 44.2 GB
Qwen3-4B-Base bf16 + LoRA 3656 tok/s 52.0 GB
Qwen3-4B-Base 4bit + LoRA 1734 tok/s 40.6 GB
Llama-3.2-3B-Inst bf16 4216 tok/s 34.7 GB
Llama-3.2-3B-Inst bf16+L 4205 tok/s 40.9 GB
Llama-3.2-3B-Inst 4bit+L 1892 tok/s 31.7 GB
--verify_no_drift passes on both arches (base bit-identical across 10
perturb+refresh cycles, inference hash deterministic).
Llama runs the Triton flex_attention backend instead of FA4:
flash-attn-4 b9's sm_100 kernel raises a NoneType in
handle_block_sparse_empty_tile_correction_sm100 on Llama-3.2's head
shapes. Qwen3 is unaffected. Pass --no-fa4_prefill on Llama; auto-FA4
still enables on Qwen3.
Delete two unreferenced drivers (qwen3_grpo_notebook.py, qwen3_grpo_unified.py)
that duplicated the canonical trio. Port the --compile_mode / --compile_dynamic
flags from unified into qwen3_grpo_naive.py and qwen3_grpo_tpaged.py before
deletion so the torch.compile path is preserved on the training-side backends
(vLLM is excluded because it owns its own inference graph).
Extract the 20-line StepTimer TrainerCallback, the per-step stats JSON writer,
the vLLM GuidedDecodingParams shim, and the optional torch.compile wrapper
into unsloth_grpo_common.py so the three canonical drivers
(qwen3_grpo_{vllm,naive,tpaged}.py) share one implementation. Stats schema is
unchanged: backend, train_wall_s, peak_memory_gb, step_wall_s, losses, rewards,
max_prompt_length, max_completion_length, num_generations, max_steps, plus
backend-specific extras (attn_impl, persistent_cb) passed through write_stats's
extra kwarg.
Add a short paragraph to scripts/benchmarks/README.md describing the new
--compile_mode flag.
Verified:
- python -m py_compile on all four modified files.
- --help on all three drivers shows --compile_mode on naive + tpaged only.
- 2-step tpaged smoke (flash_attention_2, num_generations=2, pdb=2) runs to
completion on B200. Stats JSON schema matches the pre-refactor output exactly.
Net: 7 files changed, +235 / -1093, 21 -> 19 benchmark files.
Remove the raw benchmark output JSONs from the PR diff. The writeup
markdowns under scripts/benchmarks/results/*.md keep their filename
anchors as a record of what each table was measured from; anyone who
wants the raw numbers can re-run the benchmark scripts.
41 files removed, 6068 lines deleted.
Removed JSON files under scripts/benchmarks/results/stats/ that no
writeup markdown in scripts/benchmarks/results/*.md referenced.
Intermediate debugging dumps (cb_sync_smoke, cb_tpaged_64_lora_4bit*,
flex_256x512, flex_32x512_eager, flex_32x512_mauto_nocg,
flex_64_lora_4bit*, flex_64_lora_bf16_{fusedmerge,mergeadapter,nomerge},
flex_verify_fusedmerge, unsloth_fi_true_64_lora_4bit). No writeups
edited -- every `results/*.md` reference still resolves.
--fa4_prefill now accepts three states: True (force on, warn + fall
back on sub-Hopper), False (force off), None / default (auto-enable
where supported). Argparse switches to BooleanOptionalAction so both
--fa4_prefill and --no-fa4_prefill work, with the default being
auto-detect from torch.cuda.get_device_capability.
Adds a cu13 / cu12 install section and a per-GPU support matrix to
scripts/benchmarks/README.md.
Adds tests/test_fa4_capability_guard.py covering the nine
combinations of (explicit-on / auto / explicit-off) x (sm_80 / sm_90
/ sm_100 / sm_120). Monkey-patches get_device_capability and stubs
PageTable / patch_qwen3_model so it runs without CUDA.
Inductor was specialising create_block_mask on the first prefill shape
it saw (warmup with 16 prompts -> small total L). When round-0 prefill
ran at bs=64 with a much larger packed L, the cached triton block-mask
kernels launched with the wrong shape constants and hit CUDA illegal
memory access inside the document_causal mask construction, even
though the fused GEMMs and graph-captured decode path were fine.
torch.compile(create_block_mask, dynamic=True) keeps L as a runtime
arg so the same kernels work across the warmup and full-batch prefill
shapes.
Add the drift-verification and end-to-end rollout stats for the fused
addmm path: base bit-identical across 10 cycles, and bs=64 + LoRA +
capture_cudagraph + decode_kernel_options reaches 5057 tok/s median,
5224 tok/s best on B200 at 52 GB peak -- within noise of the prior
5785 tok/s baseline. Variance is round-0/1 warmup (3298, 3607 tok/s)
rather than steady-state (4940, 5082, 5057 tok/s).
Replace the copy+merge_adapter pair in refresh_lora_merge_from_pristine
with one torch.addmm(pristine, B, A, alpha=scaling, out=W_inf) per
LoraLayer, then set merged_adapters directly so PEFT's forward
short-circuits to base_layer(x).
Previously each refresh did two passes per weight: a bf16 copy from
pristine, then PEFT merge_adapter which materialises a full [out, in]
fp32 delta via get_delta_weight and in-place adds it back. The fused
path skips the transient delta allocation and runs one cuBLAS GEMM
instead. cuBLAS accumulates the bf16 matmul in fp32 internally, so the
numerical result stays within 1 bf16 ULP of PEFT's path (verified on
the rank-32 Qwen3-4B adapter: max abs diff 1.22e-04).
DoRA, fan_in_fan_out, and lora_bias=True layers fall back to PEFT's
get_delta_weight/merge path via a single trailing merge_adapter call
after restoring their base_layer.weight from pristine. rslora is not a
fallback -- PEFT folds alpha/sqrt(r) into module.scaling[adapter], so
the fused addmm picks it up transparently via alpha=.
Drift verification still passes: base bit-identical across 10
perturb+refresh cycles, inference state deterministic after LoRA
restore.
PEFT's merge/unmerge pair is asymmetric at bf16 and leaks ~1 ULP per
cycle onto base_layer.weight. Across hundreds of GRPO refreshes the
base drifts, so the adapter trains against a moving target.
Keep a pristine base_model on GPU and a deep-copied inference_model
wrapped by PEFT. Before each rollout, restore the inference copy's
LoRA-target base_layer weights in-place from pristine and call
merge_adapter fresh. Never call unmerge_adapter.
Adds --verify_no_drift which hashes base params before/after N
perturb+refresh cycles and asserts bit-identical, and checks that the
merged inference state is deterministic after restoring the LoRA.
Update flex_vs_vllm.md with the double-copy row and memory cost.
Prior default was `peft_model.merge_and_unload()` which bakes LoRA into
the base and destroys the adapter. Same inference speed, but the adapter
is unrecoverable so you can't train on it for the next rollout -- which
GRPO explicitly needs.
Switch default to `peft_model.merge_adapter()`, which:
- Folds LoRA into `base_layer.weight` non-destructively.
- Keeps `lora_A` / `lora_B` parameters intact.
- Flips a `merged` flag inside each `LoraLayer` so its forward
short-circuits to just `base_layer(x)`, giving identical inference speed
to the destructive merge.
- Is fully reversible via `unmerge_adapter()` (bf16 round-trip error ~6e-5).
Measured end-to-end at batch 64 + LoRA rank 32:
- merge_adapter: 5785 tok/s best (was 5744 with merge_and_unload)
- merge+unmerge cycle: ~48 ms total for the 36-layer 7-target adapter,
which is <1 % of a ~5-7 s rollout -- fully amortizable per iteration.
This is *the* rollout path GRPO should use. vLLM's LoRARequest achieves
the same outcome via double-copy (pristine base + materialized base+LoRA
copy) or Punica-style fused kernels, but from a throughput standpoint
both get you to "near-merged speed with adapter separable for training".
Reframes the writeup: removes the previous panic correction that claimed
flex was 35 % of vLLM. The 35 % row is what you'd get with a naive PEFT
wrapper (3 matmuls per projection) -- a path nobody should actually
use. The real headline is still flex reaches 74 % of vLLM at 3.5 x less
memory under proper LoRA semantics.
`--no_merge_lora` flag preserved for the unmerged-PEFT path; documented as
reference only. 4-bit still uses the unmerged path (bnb merging is
unsupported).
Previous bf16 + LoRA rank 32 runs called `peft_model.merge_and_unload()`,
which bakes LoRA into the base weights and destroys the adapter. Every
subsequent forward is then plain bf16 with no LoRA-active cost -- one
matmul per projection. vLLM's LoRARequest path keeps LoRA dynamic (base
matmul + rank-r adapter matmuls + add), which is what GRPO actually needs
because the adapter has to be updateable between rollouts and training
steps.
Adds `--no_merge_lora` flag and runs the honest comparison:
| Backend | tok/s best | vs vLLM |
|--------------------------------|-----------:|--------:|
| vLLM (LoRARequest) | 7775 | 100 % |
| flex -- LoRA merged (prior) | 5744 | 74 % |
| flex -- LoRA active (no merge) | 2683 | 35 % |
The 74 % number in the earlier writeup was only meaningful if you can eat
the merge/unmerge cost between rollouts and training steps (which is not
free). The real flex-vs-vLLM gap under GRPO semantics is ~35 %, not 72 %.
vLLM wins its dynamic-LoRA number via Punica-style fused kernels that
avoid the extra matmul roundtrip. flex has no equivalent and runs
base + LoRA_A + LoRA_B as three separate matmuls per projection.
Writeup updated.
Adds --load_in_4bit (+ --model_name_4bit override) to both the flex
benchmark script and the vllm/tpaged benchmark script. When set, loads the
pre-quantized Unsloth bnb-4bit shard (e.g.
unsloth/Qwen3-4B-Base-unsloth-bnb-4bit) and keeps the LoRA adapter as a
PEFT wrapper instead of merging, because merging into 4-bit weights is not
supported.
Ties lm_head.weight to model.embed_tokens.weight post-load in both scripts,
because the bnb-4bit shards ship without an lm_head parameter even though
tie_word_embeddings is True in the config, so transformers leaves it
randomly initialised otherwise (garbage generations).
Results at batch 64 + LoRA rank 32:
| Backend | tok/s | peak mem | output |
|-------------------------------|------:|---------:|-----------|
| Unsloth fast_inference (vLLM) | 4515 | 159 GB | coherent |
| flex (this PR) | 1738 | 40.6 GB | coherent |
| transformers CB (sdpa) | 504 | 124 GB | gibberish |
4-bit costs ~40 % throughput on the vLLM path vs bf16 and ~70 % on flex.
flex regresses worse because PEFT-without-merge doubles the matmuls per
projection (base + LoRA add) on top of bnb dequant, whereas bf16 flex
merges LoRA into the base. Peak memory barely moves for vLLM because KV
cache at gpu_memory_utilization=0.8 dominates regardless of base size.
transformers CB (generate_batch) at 4-bit + LoRA produces garbage even
with lm_head tied. Likely PEFT-over-bnb + batched CB interaction; not
debugged further -- it was always the 10 % reference path.
Writeup updated in scripts/benchmarks/results/flex_vs_vllm.md with a new
"Same workload at load_in_4bit=True" section.
Wired up two suggestions from the FlashAttention-4 blog + attention-gym:
1. `--fa4_prefill` flag: `BLOCK_SIZE=(256, 128)` + `BACKEND="FLASH"` on the
prefill create_block_mask, pad to 256-row Q tile. Confirmed FA4 kernel
fires on Blackwell (torch 2.11 + flash-attn CuTeDSL). Output is coherent
but 4617 tok/s vs 5744 baseline at batch 64 + LoRA.
Root cause: our prefill mask is document_causal, which evaluates
`docs[q_idx] == docs[kv_idx]`. The FA4 CuTe kernel's known limitation
(documented in attention-gym/examples/flex_flash_attention.py) is that
"Indexing by kv_idx is a large perf hit". The doc mask hits that
slow path directly. To benefit from FA4 on prefill we would need to
refactor the mask so the per-kv lookup goes away, which is non-trivial
given the document-boundary + causal combo.
2. flex_autotune_replay.py: new script that drives the pattern from
attention-gym/examples/flex_autotune_replay.py -- sets
`TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE` + runs with
`mode="max-autotune-no-cudagraphs"`, parses the JSON log (handling
symbolic dims like `s40`), picks the decode-shape entry (Q_LEN=1),
and writes best fwd_* kernel options as JSON.
Inductor's best for the decode shape: `fwd_num_warps=4, fwd_num_stages=3,
fwd_BLOCK_M=64, fwd_BLOCK_N=64, fwd_USE_TMA=False`. Applied end-to-end:
4827 tok/s vs 5744 manual baseline. The per-call time-minimum Inductor
uses doesn't track the cumulative register-spill / L1 effects across
the 36-layer stack.
Kept `--fa4_prefill` and flex_autotune_replay.py in-tree -- they are
useful scaffolding for anyone who wants to push further (refactor the mask,
run the 144-config exhaustive fwd sweep from attention-gym/examples/flex_grid_sweep.py,
etc.). Default config is unchanged.
Also documented the run-to-run variance: over 10 rounds at batch 64 + LoRA,
median 4192 and best 5660 tok/s; the spread is GPU clock throttling +
variable prompt-length distributions. The 5744 "baseline" we report is
best-of-N, matching the prior harness, but steady-state median is closer
to 75 % of that.
Writeup update in scripts/benchmarks/results/flex_vs_vllm.md.
Added the final two entries to results/flex_vs_vllm.md:
| Batch | LoRA | vLLM tok/s | flex tok/s | flex/vLLM |
|-------|------|-----------:|-----------:|----------:|
| 256 | no | 21170 | 7074 | 33 %|
flex scales sub-linearly past batch ~128 (7074 @ 256 vs 6501 @ 128 is only a
9 % jump for 2x batch), while vLLM keeps climbing (14996 -> 21170). That's
expected: vLLM's chunked prefill + per-step kernel packing is more
efficient at huge batches. For GRPO's realistic batch range (4-64
concurrent seqs) the flex path sits at 30-55 % of vLLM.
Also tried `FLEX_COMPILE_MODE=max-autotune-no-cudagraphs` on the
flex_attention compile (gated via env var). Same throughput as default
compile (2186 vs 2189 at batch 32). max-autotune with cudagraphs crashes
because it nests its own cudagraph_trees inside our CUDA graph capture
and hits `Cannot prepare for replay during capturing stage`.
Expanded benchmark sweep with the flex_attention + paged-KV path:
| Batch | LoRA | vLLM tok/s | flex tok/s | flex / vLLM |
|-------|------|-----------:|-----------:|------------:|
| 32 | no | 7224 | 2189 | 30 % |
| 32 | yes | 4581 | 2334 | 51 % |
| 64 | yes | 7775 | 4279 | 55 % |
| 128 | no | 14996 | 6501 | 43 % |
Before this PR, transformers CB topped out at 9.2 % of vLLM on the
reference (batch 32 + LoRA) workload. The flex path reaches 51 % on the
same config and 55 % at batch 64.
Details in scripts/benchmarks/results/flex_vs_vllm.md plus raw stats for
each run. Output coherence verified by sampling the first three
completions; see `sample_completions` in the stats JSONs.
qwen3_flex_inference.py: added sample_completions + decode_tps_best to
the output JSON so the PR writeup can cite both median and steady-state
numbers without rerunning.
Memory: flex uses 44-81 GB depending on batch, vs vLLM's 156 GB at every
configuration. That's half to a fifth of vLLM's footprint.
Remaining gap is kernel-level (vLLM uses FlashInfer / TRTLLM kernels
tuned for sm_100, flex uses Inductor-generated Triton) plus chunked
prefill (flex still does separate prefill passes per new batch). Closing
those is out of scope for this PR.
Instead of fighting transformers CB's Python-heavy dispatch, build a minimal
paged-attention decode loop on top of `torch.nn.attention.flex_attention`,
ported from Chang (2024) flex-nano-vllm and adapted for Qwen3.
### Numbers (B200, Qwen3-4B-Base, bf16, 32 prompts x 512 new tokens, no LoRA)
| Backend | Decode tok/s | % of vLLM |
|----------------------------------|--------------|-----------|
| vLLM (fast_inference) | 4581 | 100 % |
| **qwen3_flex + CUDA graphs** | **1618-2192**| **35-48%**|
| qwen3_flex eager | 372-532 | 8-12 % |
| unsloth_fi_false | 641 | 14 % |
| CB paged+FA4 persistent | 422 | 9.2 % |
| CB sdpa_paged persistent | 434 | 9.5 % |
The plan's 30% target is met. Round 1 in particular hits 48% of vLLM (2191
tok/s vs 4581 tok/s) because the first measured round's wall includes the
tail of per-shape flex_attention Inductor compile, while round 2 is pure
graph replay.
### Why this works
Three architectural choices from flex-nano-vllm:
1. **Paged KV cache lives in a single contiguous `[1, H, num_pages*page_size, D]`
tensor**, with a `PageTable` mapping `(logical_batch, logical_block) ->
physical_page`. `flex_paged_attention.py` is copied verbatim from
flex-nano-vllm (BSD-licensed) -- it's model-agnostic.
2. **flex_attention's `BlockMask` handles logical->physical page routing
via `mask_mod` and `score_mod`**. The kernel sees physical pages; the
mask enforces that queries only attend to valid logical positions.
Crucially, flex_attention is designed for `torch.compile` so the whole
attention forward traces cleanly.
3. **One CUDA graph per batch-size bucket** (1, 2, 4, 8, 16, 32 ...) captured
during warmup. Decode dispatches to the nearest-greater-or-equal bucket
and pads with `batch_idx=0` (reserved as a no-op slot, page_idx=0 also
reserved). Graph replay is the lever that closes the gap to vLLM.
### Files
- `scripts/benchmarks/flex_paged_attention.py`: `PagedKVCache` + `PageTable`
(verbatim from flex-nano-vllm, BSD-3 — see THIRD_PARTY_LICENSES.md of the
source repo).
- `scripts/benchmarks/qwen3_flex_inference.py`: adapts to Qwen3-4B.
Monkey-patches `Qwen3Attention.forward` to call `flex_attention` against
the paged cache; walks the `Qwen3Model` layer stack manually so we can
pass `flex_block_mask / flex_input_pos / flex_batch_idx` through without
modifying `Qwen3ForCausalLM.forward`. `FlexInference.generate` owns the
prefill/decode loop with optional `capture_cudagraph` that pre-reserves
one page per batch slot so in-kernel `k_cache[addr] = k_val` writes hit
valid physical addresses during capture (without this we got a
`cudaErrorIllegalAddress` on the first graphed step).
### CB sync driver side-note
`scripts/benchmarks/cb_sync_driver.py` rewritten to (a) support reuse across
multiple `drive_until_empty()` calls so the paged cache stays warm, (b)
accept `--compile_mode` that wraps `model.forward` with torch.compile. Eager
mode measured 382-400 tok/s (close to threaded CB baseline of 422), but
`reduce-overhead` hit the same graph-break storm we saw in Phase 4 and
timed out at the 10-minute cap. The flex_attention path sidesteps that
entirely.
### Next steps
- Try LoRA rank 32 through the flex_attention path (PR's canonical workload).
- Scale to max_batch_size=64 to see if throughput keeps climbing.
- Integrate into TRL GRPO's rollout path for a full end-to-end speedup.
30-step results:
| Backend | Train wall | Median step | Peak mem | % of vLLM |
|----------------|-----------|-------------|----------|-----------|
| vLLM | 215.9 s | 5.14 s | 159 GB | 100 % |
| fi_false | 1165.4 s | 41.30 s | 10.7 GB | 12.4 % |
| cb_paged | 1564.5 s | 39.82 s | 61.9 GB | 12.9 % |
Pairwise diff vs vLLM (30 steps, compare_grpo_runs.py):
| Pair | max |loss| | max |kl| | max |reward| |
|---------------------------------|------------|------------|---------------|
| vLLM vs unsloth_fi_false | 0.39 | **0.015** | 9.25 (noisy) |
| vLLM vs cb_paged | 0.83 | (missing) | 6.25 (noisy) |
KL trajectory match between vLLM and unsloth_fi_false is the load-bearing
equivalence signal: both stay in [0, 0.015] across all 30 steps, so the
policy drift guardrail behaves the same. Reward diffs of ~3-9 are expected
because the rollout backends produce different completions even at
temperature=0.1 (kernel-level non-determinism).
Two caveats documented in results/grpo_equivalence.md:
- cb_paged's StatisticsCallback doesn't capture TRL's kl log entry because
TRL emits kl on a separate log call that doesn't include loss.
- cb_paged's grad_norm (~200-900) is unclipped pre-optimizer, while vLLM's
goes through Unsloth's internal max_grad_norm=1.0. Not a correctness bug,
just not apples-to-apples until cb_paged sets max_grad_norm in GRPOConfig.
Also updates the report with Phase 3 + Phase 4 status (CB sync driver eager
works; CUDA graph capture hangs on output_ids slice pending fix; torch.compile
on training step incompatible with both unsloth_fi_false and cb_paged).
Phase 4 (torch.compile on training step) -- negative result documented:
- unsloth_fi_false + compile_mode=default: crashes immediately with
`PeftModel_fast_forward() got multiple values for argument 'input_ids'`.
Unsloth's monkey-patched forward and Dynamo's argument rebinding don't
compose.
- cb_paged + compile_mode=default: Dynamo emits 700+ graph breaks /
recompiles during the first optimizer step and never makes progress
(killed after 10 minutes at step 0/10). Root cause trail:
`Tensor.requires_grad_()` inside `modeling_utils.make_inputs_require_grads`
triggers GB0125 (no Dynamo support), which propagates up through TRL's
`_compute_loss`. Fixing this would require restructuring GRPO's LoRA
gradient enablement path -- out of scope for this PR.
- vllm is excluded from Phase 4 because vLLM owns its own compile pipeline.
Net: torch.compile on the training step is not a quick win for the non-vLLM
paths in this stack. Phase 3 (CUDA graph on the rollout decode step) remains
the right lever for closing the CB <-> vLLM gap, and Phase 1 already
demonstrates that Unsloth's fast_inference=False path narrows the gap to
~14% of vLLM at 1/7th the peak memory without any compile.
Helper script `scripts/benchmarks/compare_grpo_runs.py`:
- Wraps `torch_debugging_utils.compare_training_runs` (loss/grad-norm diff)
- Adds reward/KL/time pairwise diffs with `max_abs` and `mean_abs`
- Reads the StatisticsCallback-emitted JSON written by
`qwen3_grpo_unified.py`
Sample output on Phase 2 vibe (10 steps):
vllm vs unsloth_fi_false:
max_loss_diff = 0.40, max_kl_diff = 0.009 (both tiny), reward_diff
mean 1.85 (different rollouts expected across backends at temp=0.1)
vllm vs cb_paged:
max_loss_diff = 0.29, max_grad_norm_diff = 715 (cb_paged has no
gradient clipping on the vanilla-HF path; vLLM path is clipped to
1.0 by Unsloth internally -- apples-to-oranges without matching
clipping, tracked for the 30-step run).
Phase 2 results (`scripts/benchmarks/results/grpo_equivalence.md`):
- vLLM: 74.4s train, 4.14s median step, 158 GB peak (100%)
- unsloth_fi_false: 355s train, 23.95s median, 10.7 GB peak (17%)
- cb_paged (via sdpa_paged load): 466s train, 36s median, 55.6 GB peak (11.5%)
Coherence gate passes on all three backends: losses finite, rewards in the
expected early-GRPO range, KL trajectories qualitatively matched between
vLLM and unsloth_fi_false in [0, 0.015]. Memory story is striking:
unsloth_fi_false uses 15x less memory than vLLM.
qwen3_grpo_unified.py fixes:
- Auto-adjust per_device_train_batch_size -> num_generations for vanilla-HF
backends (Unsloth's loader does this automatically; TRL on the HF path
doesn't and crashes on the divisibility check).
- cb_paged now loads with sdpa_paged (not paged_attention). The FA4
paged_attention kernel requires cu_seq_lens_q on every forward, but the
GRPO training forward feeds a dense batch without them. sdpa_paged
gracefully falls back to plain SDPA in that case and still exercises the
paged path during the CB rollout.
cb_sync_driver.py fixes:
- FIFOScheduler no longer accepts manual_eviction in its signature; dropped.
- drive_until_empty used to check has_pending_requests() before calling
prepare_next_batch(), which returned False at startup because nothing had
yet been pulled from the input_queue. Now the loop drains the input queue
first and exits only when both queues + scheduler are empty.
Smoke test on GPU 1 (8 prompts, 64 tokens): eager path produces 512 correct
tokens; CUDA-graph path hangs during first-step capture (PagedAttentionCache
probably allocates on first use). Tracked for the next commit.
Phase 1 results (`scripts/benchmarks/results/lora_rollout_baselines.md`):
- vLLM+LoRA: 4581 decode tok/s, 156 GB peak (gold, 100%)
- unsloth_fi_false+LoRA: 641 tok/s, 15.8 GB peak (14%, 7x lower mem)
- CB paged+FA4 persistent+LoRA: 422 tok/s (9.2%)
- CB sdpa_paged persistent+LoRA: 434 tok/s (9.5%)
Headline finding: Unsloth's `fast_inference=False` path (custom HF inference
kernels with cached fp16 LoRA in `fast_linear_forward`) is 1.5x faster than
CB at 1/7th the peak memory. Phase 2 will include it as a first-class backend.
cb_vs_vllm_generation.py:
- New --lora_adapter flag. vLLM uses LoRARequest; tpaged uses
PeftModel.from_pretrained (no merge_adapter so we measure LoRA-active
inference); unsloth_fi_false copies the adapter weights into Unsloth's
get_peft_model wrapper (with key normalization so PEFT's base_model.model.
prefix and Unsloth's .default. wrapper both match).
- New unsloth_fi_false backend: batched generate across all 32 prompts in
a single call after FastLanguageModel.for_inference(model).
- Exposed sampling knobs (temp/top_p/min_p/top_k); defaults are the
equivalence params.
cb_sync_driver.py (Phase 3 scaffold):
- SyncCBDriver owns PagedAttentionCache + ContinuousBatchProcessor +
FIFOScheduler on the main thread. Never calls manager.start() so there is
no background thread.
- slice_inputs=False => fixed-shape buffer views each step => CUDA graph
replay is safe.
- use_cuda_graph=True path: 2-step eager warmup, then capture one decode
step, then replay. `_is_pure_decode()` keeps prefill out of the graphed
path since those have varying shapes.
- Greedy sampling only (CUDA-graph-safe); stochastic sanity checks stay in
the non-graphed path.
- Standalone benchmark harness at the bottom.
qwen3_grpo_unified.py (Phase 4 scaffold):
- Single entrypoint for vllm / unsloth_fi_false / cb_paged / cb_sdpa /
naive_trl backends sharing dataset, reward funcs, sampling, and the
torch_debugging_utils StatisticsCallback.
- New --compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}
that compiles `trainer.model.forward` and `trainer.ref_model.forward`
after the trainer is built. CompileDebugger tracks graph breaks and
recompiles. Skipped for vLLM since vLLM owns its own compile pipeline.
- Post-warmup median (skip first 3 steps) is computed and saved alongside
the full per-step logs.
make_lora_adapter.py: writes a canonical PEFT adapter to
outputs/lora_rank32_fresh. Re-initializes lora_B with a tiny gaussian so
the adapter isn't a no-op (PEFT's default zero-init would let LoRA kernels
short-circuit).
qwen3_grpo_notebook.py (Phase 0): notebook-to-script port with
StatisticsCallback and equivalence sampling. 10-step reference reported in
scripts/benchmarks/results/notebook_ref_10.md (median step 5.80s,
peak 158.9 GB).
Phase 0 (canonical reference):
- scripts/benchmarks/qwen3_grpo_notebook.py: notebook-to-script port of
Qwen3_(4B)-GRPO.ipynb with StatisticsCallback from torch_debugging_utils
and equivalence-friendly sampling (temp=0.1, top_p=0.97, min_p=0.5, top_k=5).
- scripts/benchmarks/results/notebook_ref_10.md: 10-step reference run table
(median step post-warmup = 5.80s, peak 158.9 GB).
- scripts/benchmarks/results/stats/notebook_ref_10.json: full per-step logs
for downstream compare_training_runs checks.
Phase 1 (rollout-only LoRA comparison scaffold):
- scripts/benchmarks/make_lora_adapter.py: one-shot that materializes a
rank-32 LoRA at outputs/lora_rank32_fresh. Re-initializes lora_B with a
tiny gaussian so the adapter isn't a no-op (otherwise LoRA kernels can
short-circuit and we'd be measuring the base model).
- scripts/benchmarks/cb_vs_vllm_generation.py: extended with --lora_adapter
for vLLM (LoRARequest) and tpaged (peft.PeftModel.from_pretrained,
no merge_adapter), plus a new unsloth_fi_false backend that exercises the
custom HF inference path (cached fp16 LoRA via fast_linear_forward).
Sampling knobs are exposed and default to equivalence params.
Phase 2 scaffold:
- scripts/benchmarks/qwen3_grpo_unified.py: single entry point for all 5
backends (vllm, unsloth_fi_false, cb_paged, cb_sdpa, naive_trl) sharing
dataset, reward funcs, sampling, and StatisticsCallback. Skips the first
3 steps when reporting median step wall.
No unsloth internals touched.
New scripts under scripts/benchmarks/:
- flash_attn_fa4_shim.py: monkey-patches that let transformers CB dispatch to
Flash Attention 4 on Blackwell (sm_100). CB's ContinuousBatchProcessor
otherwise emits a 4D paged attention mask for flash_attention_2 (which then
breaks _flash_attention_forward's _upad_input branch), and passes
max_seqlen_q instead of max_length_q. The shim skips the mask for FA and
accepts both names.
- persistent_cb.py: replaces model.generate_batch with a version that reuses
one ContinuousBatchingManager across calls, avoiding the per-step
PagedAttentionCache realloc. Wired up behind --persistent_cb on the tpaged
and standalone scripts.
- qwen3_grpo_naive.py: vanilla HF model.generate + TRL GRPOTrainer, no vLLM
and no CB. Mirrors the TRL docs example. Useful as a third column in the
comparison and also as a "will this at least converge" sanity check.
Adds --attn_impl and --persistent_cb flags to the existing generation and
training scripts. No changes to Unsloth internals.
Updated README.md with the FA install recipe (flash-attn-4==4.0.0b9, plus
a small site-packages shim that re-exports FA4's cute.* symbols under the
FA2 flash_attn namespace so transformers' is_flash_attn_2_available() and
_lazy_imports("flash_attention_2") succeed on B200).
Benchmark numbers on a single B200, Qwen3-4B-Base LoRA rank 32, bf16:
Generation microbenchmark (32 prompts, 512 new tokens):
vLLM 7224 decode tok/s (100%)
CB paged|sdpa 527 decode tok/s ( 7.3%)
CB paged|flash_attention_2 (FA4) 709 decode tok/s ( 9.8%)
CB paged|flash_attention_2 persistent 529 decode tok/s ( 7.3%)
GRPO training (max_steps=20, num_generations=2, per_device_batch=2):
vLLM colocated 136.6 s peak 157 GB
naive TRL (HF generate) 910.0 s peak 15 GB
CB SDPA 1521.5 s peak 98 GB (prior run)
CB FA4 1470.7 s peak 82 GB
CB FA4 + persistent 1562.1 s peak 87 GB
CB FA4 + ng=4 persistent 1597.0 s peak 94 GB
FA4 is a real ~1.4x improvement over SDPA for CB decode throughput but the
50% of vLLM target is still not reached. The remaining gap is driven by
CUDA graph capture (which ContinuousBatchingManager still NotImplementedErrors
on) and vLLM's scheduler being more efficient for decode-heavy GRPO rollouts.
Naive TRL generate is the honest small-rig baseline: 6.7x slower than vLLM
at 10% of the VRAM footprint, and ~1.7x faster than CB here.
Adds reproducible scripts under scripts/benchmarks/ that compare vLLM
colocated rollouts against the transformers continuous batching API
(model.generate_batch, paged attention) for GRPO training on Qwen3-4B.
Contents:
- unsloth_grpo_common.py: shared dataset, reward functions, and GRPO
hyperparameters so the two backends differ only in the rollout engine.
- qwen3_grpo_vllm.py: baseline training entry using fast_inference=True
and TRL use_vllm=True, vllm_mode=colocate.
- qwen3_grpo_tpaged.py: candidate using a vanilla HF Qwen3 + PEFT LoRA
with TRL use_transformers_paged=True.
- cb_vs_vllm_generation.py: standalone generation microbenchmark.
- README.md: integration notes, reproduction steps, and observed numbers.
On a single B200 with Unsloth Qwen3-4B-Base at LoRA rank 32 bf16,
transformers continuous batching reaches 7-9 percent of vLLM throughput
on this workload. The README documents the integration sharp edges
(top_k=-1, PagedAttentionCache default upper bounds, Unsloth's
Qwen3Attention_fast_forward bypassing the functional attention
interface, TRL importing GuidedDecodingParams from a newer vLLM that no
longer exports it, and UnslothGRPOTrainer expecting for_training /
for_inference hooks on the model).
The scripts are intentionally self-contained so they are easy to rerun
after either upstream change that could close the throughput gap
(flash-attn availability, CUDA graphs in ContinuousBatchingManager,
persistent paged caches across generate_batch calls, or a
paged-compatible Unsloth attention forward).
* fix(studio): forward OpenAI tools/tool_choice to llama-server (#4999)
Studio's /v1/chat/completions silently stripped standard OpenAI `tools`
and `tool_choice` fields, so clients using standard function calling
(opencode, Claude Code, Cursor, Continue, ...) never got structured
tool_calls back. Adds a client-side pass-through path mirroring the
existing Anthropic /v1/messages flow: when `tools` is present without
Studio's `enable_tools` shorthand, the request is forwarded to
llama-server verbatim so the client sees native id, finish_reason
("tool_calls"), delta.tool_calls, and accurate usage tokens.
Also wires Anthropic tool_choice forwarding: /v1/messages previously
accepted tool_choice on the request model but silently dropped it with
a warning. Translate the four Anthropic shapes to OpenAI format and
forward them so agentic clients can actually enforce tool use.
- ChatCompletionRequest: add tools, tool_choice, stop; extra="allow"
- ChatMessage: accept role="tool", optional tool_call_id / tool_calls /
name; content is now optional (assistant with only tool_calls)
- routes/inference.py: _openai_passthrough_stream /
_openai_passthrough_non_streaming helpers, routing branch in
openai_chat_completions, vision+tools via content-parts injection
- _build_passthrough_payload: tool_choice parameter (default "auto")
- anthropic_compat: anthropic_tool_choice_to_openai() translator
- tests/test_openai_tool_passthrough.py: Pydantic + translator unit tests
- tests/test_studio_api.py: 5 new E2E tests (non-stream, stream,
multi-turn, OpenAI SDK, Anthropic tool_choice=any regression)
* fix(studio): surface httpx transport errors from OpenAI passthrough
When the managed llama-server subprocess crashes mid-request, the
async pass-through helpers in routes/inference.py used to return a
bare 500 (non-streaming) or an "An internal error occurred" SSE chunk
(streaming) because _friendly_error only recognized the sync path's
"Lost connection to llama-server" substring -- httpx transport
failures (ConnectError / ReadError / RemoteProtocolError /
ReadTimeout) stringify differently and fell through to the generic
case.
- _friendly_error: map any httpx.RequestError subclass to the same
"Lost connection to the model server" message the sync chat path
emits. Placed before the substring heuristics so the streaming path
automatically picks it up via its existing except Exception catch.
- _openai_passthrough_non_streaming: wrap the httpx.AsyncClient.post
in a try/except httpx.RequestError and re-raise as HTTPException
502 with the friendly detail.
- tests/test_openai_tool_passthrough.py: new TestFriendlyErrorHttpx
class pinning the mapping for ConnectError, ReadError,
RemoteProtocolError, ReadTimeout, and confirming non-httpx paths
(context-size heuristic, generic fallback) are unchanged.
* fix(studio): close aiter_bytes/aiter_lines explicitly in passthroughs
The httpcore asyncgen cleanup fix in 5cedd9a5 is incomplete on Python
3.13 + httpcore 1.0.x: it switched to manual client/response lifecycle
but still used anonymous `async for raw_line in resp.aiter_lines():`
patterns in all three streaming paths. Python's async for does NOT
auto-close the iterator on break/return, so the aiter_lines /
aiter_bytes async generator remains alive, reachable only from the
surrounding coroutine frame. Once `_stream()` returns the frame is
GC'd and the orphaned asyncgen is finalized on a LATER GC pass in a
DIFFERENT asyncio task, where httpcore's
HTTP11ConnectionByteStream.aclose() enters anyio.CancelScope.__exit__
with a mismatched task and prints "Exception ignored in: <async
generator>" / "async generator ignored GeneratorExit" / "Attempted
to exit cancel scope in a different task" to the server log.
User observed this on /v1/messages after successful (status 200)
requests, with the traceback pointing at HTTP11ConnectionByteStream
.__aiter__ / .aclose inside httpcore.
Fix: save resp.aiter_lines() / resp.aiter_bytes() as a variable and
explicitly `await iter.aclose()` in the finally block BEFORE
resp.aclose() / client.aclose(). This closes the asyncgen inside the
current task's event loop, so the internal httpcore byte stream is
cleaned up before Python's asyncgen GC hook has anything orphaned to
finalize. Each aclose is wrapped in try/except Exception so nested
anyio cleanup noise can't bubble out.
Applied to all three streaming passthrough paths:
- _anthropic_passthrough_stream (/v1/messages client-side tool path)
- _openai_passthrough_stream (/v1/chat/completions client-side tool
path, new in this PR)
- openai_completions (/v1/completions bytes proxy from PR #4956)
* fix(studio): default ChatCompletionRequest.stream to false per OpenAI spec
OpenAI's /v1/chat/completions spec defaults `stream` to false, so
clients that omit the field (naive curl, minimal integrations) expect
a single JSON response back. Studio was defaulting to true, silently
switching those clients into SSE and breaking any parser that didn't
also handle streaming. ResponsesRequest and AnthropicMessagesRequest
already default to false correctly; only ChatCompletionRequest was
wrong.
Studio's own frontend always sets `stream` explicitly on every
chat-adapter / chat-api / runtime-provider call site, so the flip has
no UI impact. SDK users (OpenAI Python/JS SDK, opencode, Claude Code,
Cursor, Continue) also always pass `stream` explicitly, so they're
unaffected. The only clients feeling the change are raw-curl users
who were relying on the wrong default -- those get the correct OpenAI
behavior now.
Added a regression test pinning the default so it can't silently
flip back.
* fix(studio): reject images in OpenAI tool passthrough for text-only GGUFs
The new tool passthrough branch runs before _extract_content_parts,
skipping the existing not is_vision guard. Requests combining tools
with an image on a text-only tool-capable GGUF were forwarded to
llama-server, producing opaque upstream errors instead of the
pre-existing clear 400. Restore the guard inline at the dispatch
point, checking both legacy image_base64 and inline image_url parts.
* fix(studio): require tool_call_id on role=tool chat messages
Enforce the OpenAI spec rule that role="tool" messages must carry a
tool_call_id. Without it, upstream backends cannot associate a tool
result with the assistant's prior tool_calls entry and the request
fails in non-obvious ways through the passthrough path. Reject at the
request boundary with a 422 instead.
* fix(studio): harden OpenAI tool passthrough validation and error surfacing
Three related fixes called out by the PR review:
1. Preserve upstream status codes in the streaming passthrough. The
httpx request is now dispatched before the StreamingResponse is
constructed. Non-200 upstream responses and httpx RequestError
transport failures raise HTTPException with the real status
instead of being buried inside a 200 SSE error frame, so OpenAI
SDK clients see APIError/BadRequestError/... as expected.
2. Require non-empty content on user/system/tool messages. Per the
OpenAI spec, content may only be omitted on assistant messages
that carry tool_calls; enforce that at the request boundary so
malformed messages never reach the passthrough path.
3. Role-constrain tool-call metadata. tool_calls is only valid on
role=assistant, tool_call_id and name only on role=tool. Without
this, a user/system message with tool_calls would flip the
passthrough branch on and be forwarded to llama-server, surfacing
as an opaque upstream error.
* fix(studio): normalize image mode and passthrough JSON verbatim
Two Gemini-code-assist review findings on PR #5099:
1. Unconditionally convert decoded images to RGB before PNG encoding.
The prior code only handled RGBA, letting CMYK/I/F images crash
at img.save(format="PNG") and surface as opaque 400s. Applied to
both the passthrough helper and the non-passthrough GGUF path
that originally carried this pattern, keeping the two sites in
sync.
2. Return the upstream JSON body as raw bytes via Response rather
than parse-then-re-serialize with JSONResponse. Matches the
passthrough helper's "verbatim" contract and drops a redundant
round-trip.
---------
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* unsloth gemma4 support files
* some fixes
* Fixing cache.empty() calls (#4813)
* Fixing cache.empty() calls
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: Manan Shah <mananshah@Manans-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Fix/gemma4 mlx (#4816)
* Fixing cache.empty() calls
* fixing for mlx versions
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: Manan Shah <mananshah@Manans-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* removed bidirectional check for 31b (#4839)
Co-authored-by: Manan17 <shahmanan170602@gmail.coml>
* Add Gemma 4 26B MoE support (MLX) (#4844)
* removed bidirectional check for 31b
* Change gemma4_text for moe
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: Manan Shah <mananshah@Manans-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* fix(gemma4): cast RoPE offset to int before mx.arange() (#4901)
* fix(gemma4): cast RoPE offset to int before mx.arange()
* fix(gemma4): use zero-based arange + offset to avoid CPU-GPU sync
* qwen3.6 patches for multi-turn chat
* qwen3.6 script
* removing unnecessary scripts
* displaying errors for not installed packages
---------
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Manan Shah <mananshah@Manans-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Manan17 <shahmanan170602@gmail.coml>
Co-authored-by: Théophile Lafargue <138336683+eauchs@users.noreply.github.com>
* Add Qwen3.6 inference defaults for Studio
Add qwen3.6 family entry to inference_defaults.json with the
recommended sampling parameters from Qwen's documentation:
temperature=0.7, top_p=0.8, top_k=20, min_p=0.0,
presence_penalty=1.5, repetition_penalty=1.0.
Without this, Qwen3.6 models fall through to the generic qwen3
pattern which uses different defaults (temperature=0.6,
top_p=0.95, no presence_penalty).
* Add Qwen3.6-35B-A3B-GGUF to default model lists
* Add Qwen3.5/3.6 presence_penalty to thinking toggle and small-model disable logic
- Thinking toggle (on-load + button click) now sets presencePenalty: 1.5 for
Qwen3.5 and Qwen3.6 models (both thinking-ON and thinking-OFF states)
- Small-model thinking-disable check (<9B defaults to no-thinking) extended
from Qwen3.5-only to also cover Qwen3.6, in all 3 locations:
frontend on-load, frontend refresh, backend llama_cpp.py
* fix: multi-GPU inference crash for bnb 4-bit/8-bit models
When load_in_4bit or load_in_8bit is used with device_map="sequential"
and max_memory constraints that place weights across multiple GPUs (or
entirely on a non-default GPU like cuda:1), the bitsandbytes loading
path in transformers never calls dispatch_model. No AlignDevicesHook is
installed, and the first forward/generate call crashes with:
RuntimeError: Expected all tensors to be on the same device
This adds _attach_bnb_multidevice_hooks() which is called after
from_pretrained returns. It infers a device map from actual parameter
placements and calls dispatch_model(force_hooks=True) to install the
missing hooks. The function is a complete no-op for the common
single-GPU cuda:0 case.
Call sites: FastBaseModel.from_pretrained (vision.py) and
FastLlamaModel.from_pretrained (llama.py).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: align with PR #5053 final review improvements
- Add hook call to the bnb quantized loading branch in llama.py (the
primary load_in_4bit path), not just the non-fast-inference fallback
- Expand bnb detection: also check model.is_loaded_in_4bit,
model.is_loaded_in_8bit, model.quantization_method
- Pass explicit main_device and skip_keys to dispatch_model
- Use logger.info instead of print for the success message
- Use kwargs.get("load_in_8bit", False) at llama.py call sites
* [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>
* auth: default to chat
* settings: relaunch onboarding
* onboarding: return to launch page
* studio: stop auto guided tour
* ui: soften global radius
* cleanup: rename onboarding exit prop
* fix onboarding redirect safety
* Show real Unsloth version in settings
* [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>
* feat(studio): replace navbar navigation with collapsible sidebar
Add an app-wide sidebar with hover-expand and pin-to-dock behavior.
Navigation items (Studio, Recipes, Export, Chat) move from the center
pill navbar to the sidebar. Chat threads and recipes render as
collapsible sub-lists. Navbar simplified to logo + update + close.
- Extend SidebarProvider with pinned/hovered state model
- New AppSidebar with animated active indicator, sloth profile menu,
theme toggle, guided tour, back/forward navigation
- Chat page refactored to URL-driven view state via search params
- Extract reusable hooks for chat thread and recipe sidebar data
- Guard startViewTransition for browser compatibility
- Wrap chat deletions in Dexie transaction for data integrity
* feat(studio): move logo to sidebar and make navbar overlay
- Sidebar is now full-height with logo in SidebarHeader
- Collapsed sidebar shows sticker.png, expanded shows full logo
- Navbar is absolute-positioned overlay (no layout space)
- Main content extends to top, aligning with navbar controls
* feat(studio): full-height sidebar with recents, edge-to-edge nav buttons
- Sidebar outside max-w-7xl, pinned to left edge
- Remove sidebar rounding, menu buttons rounded-md
- Nav buttons flush to sidebar edges with no left rounding
- Replace collapsible recipes/chat with flat nav items
- Add Recents section with chat history (1 item when not on chat, full on chat)
- New Chat as first nav item with PencilEdit02Icon
- Cursor pointer on all sidebar buttons
- Navbar temporarily hidden for screenshots
* fix(studio): fix chat scroll, action bar hover, collapsible recents
- Fix sticky composer by removing `relative` override on viewport footer
- Action bar buttons only show on hover (autohide=always)
- Remove floating border/shadow from action bar
- Add scroll space above composer for last message actions
- Back/forward buttons use router history (stay in-app)
- Recents section collapsible with chevron on chat route
- Set html/body/#root height for proper h-full chain
* fix(studio): address review feedback, clean up unused code
- Unhide navbar (was left hidden from screenshot)
- Remove unused imports: SidebarMenuSub*, BubbleChatIcon, ColumnInsertIcon
- Remove unused vars: recipeItems, activeRecipeId, canCompare, recipesOpen
- Include compare query id in active sidebar selection
- Use store type for contextUsage instead of inline type
- Simplify noop in sidebar.tsx
- Remove empty className prop
* feat(studio): add mobile sidebar, recent runs section, and misc UX fixes
* feat(studio): scaffold settings feature module with dialog store
* feat(studio): add tri-state theme store for settings
* feat(chat): add clear-all-chats and export-chat-history utils
* feat(studio): add settings dialog shell with tab rail
* feat(studio): add appearance tab with theme and sidebar pin
* feat(studio): add settings general tab with hf token, auto-title, reset prefs
* feat(studio): add settings chat tab with export and clear
* feat(studio): add api keys tab with list and revoke flow
* feat(studio): add create-key form and reveal dialog
* feat(studio): add usage examples panel to api keys tab
* feat(studio): add settings about tab with update and shutdown
* feat(studio): add settings dropdown item and cmd-comma shortcut
* feat(studio): remove legacy api-keys route and chat-sheet preference rows
* fix(studio): settings dialog a11y + polish pass
* feat(studio): inline api key reveal card replacing nested dialog
* fix(studio): hide revoked keys from settings list
* refactor(studio): strip navbar and hoist training unload guard
* feat(studio): explicit sidebar toggle, remove hover-open and pin icons
* fix(studio): use SidebarRight01Icon for collapsed sidebar open toggle
* fix(studio): address code review findings for settings dialog
* feat(studio): collapsible navigate group with standalone new-chat and compare
* fix(studio): chat-only standalone actions, use ColumnInsertIcon for compare
* fix(studio): sidebar new-chat/compare state reset and icon-mode collapsible
* feat(studio): add compact logo assets for sidebar header
* Fixed sidebar design
* fix(studio): sidebar delete icon hover contrast and sizing
* feat(studio): route-gate sidebar recents (chats off /studio, runs on /studio)
* feat(studio): add chat search store
* feat(studio): add chat search index hook with snapshot-on-open
* feat(studio): add chat search command dialog with global shortcut
* feat(studio): wire chat search into sidebar
* fix(studio): trim hf token on save, add show/hide toggle, commit on close
* revert(studio): restore original sidebar/border colors, brighten sidebar
* feat(studio): forward overlayClassName through CommandDialog
* fix(studio): wrap search dialog in Command context, redesign as flat 635px card
* fix(studio): reserve right padding on recent items so delete icon stops overlapping title
* fix(studio): skip hf token unmount-commit during reset-prefs reload
* chore(studio): drop unused icon import and unreachable runs navigate branch
* fix(studio): chat search index filters archived before limit, batches message query, picks up reasoning text
* fix(studio): keep CommandEmpty in tree so empty state renders correctly
* fix(studio): cap system prompt and chat template textareas so they scroll instead of growing
* fix(studio): attach chat-compare tour anchor to sidebar compare button
* fix(studio): persist system theme explicitly so next-themes does not clobber on reload
* fix(studio): auto-switch to history tab when selecting a recent run from sidebar
* UI overhaul: chatbox, scrollbar, sidebar, and compare view
UI Changes:
- Redesigned the Compare UI with general cleanup
- Redesigned the Chatbox UI
- Reduced the width of the user chat bubble for improved readability
- Narrowed the user chat box across the content page
- Adjusted thinking-box text color to be slightly darker
- Removed faded text effect from chat messages
- Removed faded text effect from the thinking box
- Added a small LLM chat safety note at the bottom of the chatbox
- Restyled the scrollbar
Layout & Behavior:
- Reworked the scrollbar to span the full height of the page (no top/bottom padding) and remain persistently visible when content is scrollable, rather than only on hover
- Reworked the Configuration sidebar to span full height — removed rounded corners and borders, with the scrollbar adjusted to match the full top-to-bottom layout
- Adjusted the top menu and bottom chatbox content areas to work correctly with the new full-page scroll behavior
- Made chat content match the chatbox width, with content sliding slightly behind the chatbox when scrolling
- Aligned chat text width with the chatbox for visual consistency, including how far the text extends behind the chatbox
Fixes:
- Fixed the chatbox not auto-expanding when typing multi-line input while bottom-positioned during an active chat (previously only worked before a chat had started)
- Fixed positioning and design of the user chat hover menu buttons to match the assistant chat box — now displayed below the chat bubble instead of on the left side
* Fix user message layout in thread component
* swap code icon
* fix compare layout
* fix compare pane flex
* Sidebar improvements and fixes
- Added scrolling support to the sidebar so menus and recent chats no longer get hidden
- Recent chats are now always visible in the sidebar, not hidden when in Studio, Recipes, or Export
- Recent chat is now deselected when selecting other navigations
- Fixed sidebar glitch where browser resize could make the sidebar and expand button disappear completely
- Fixed glitch where the open-sidebar hover tooltip appeared above the logo when clicking expand sidebar
- Reduced sidebar width on mobile to around 2/3 of the screen (was too wide)
- Made the close-sidebar hover tooltip consistent with the rest of the design
- Removed sidebar collapse/expand animation
- Small adjustment to chat width
* Fix route scrolling, polling, and theme sync issues
* Fix Studio page scrolling
---------
Co-authored-by: sneakr <hauzin@hotmail.com>
* Studio: Ollama support, recommended folders, Custom Folders UX polish
Backend:
- Add _scan_ollama_dir that reads manifests/registry.ollama.ai/library/*
and creates .gguf symlinks under <ollama_dir>/.studio_links/ pointing
at the content-addressable blobs, so detect_gguf_model and llama-server
-m work unchanged for Ollama models
- Filter entries under .studio_links from the generic models/hf/lmstudio
scanners to avoid duplicate rows and leaked internal paths in the UI
- New GET /api/models/recommended-folders endpoint returning LM Studio
and Ollama model directories that currently exist on the machine
(OLLAMA_MODELS env var + standard paths, ~/.lmstudio/models, legacy
LM Studio cache), used by the Custom Folders quick-add chips
- detect_gguf_model now uses os.path.abspath instead of Path.resolve so
the readable symlink name is preserved as display_name (e.g.
qwen2.5-0.5b-Q4_K_M.gguf instead of sha256-abc...)
- llama-server failure with a path under .studio_links or .cache/ollama
surfaces a friendlier message ("Some Ollama models do not work with
llama.cpp. Try a different model, or use this model directly through
Ollama instead.") instead of the generic validation error
Frontend:
- ListLabel supports an optional leading icon and collapse toggle; used
for Downloaded (download icon), Custom Folders (folder icon), and
Recommended (star icon)
- Custom Folders header gets folder icon on the left, and +, search,
and chevron buttons on the right; chevron uses ml-auto so it aligns
with the Downloaded and Recommended chevrons
- New recommended folder chips render below the registered scan folders
when there are unregistered well-known paths; one click adds them as
a scan folder
- Custom folder rows that are direct .gguf files (Ollama symlinks) load
immediately via onSelect instead of opening the GGUF variant expander
(which is for repos containing multiple quants, not single files)
- When loading a direct .gguf file path, send max_seq_length = 0 so the
backend uses the model's native context instead of the 4096 chat
default (qwen2.5:0.5b now loads at 32768 instead of 4096)
- New listRecommendedFolders() helper on the chat API
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: log silent exceptions and support read-only Ollama dirs
Replace silent except blocks in _scan_ollama_dir and the
recommended-folders endpoint with narrower exception types plus debug
or warning logs, so failures are diagnosable without hiding signal.
Add _ollama_links_dir helper that falls back to a per-ollama-dir hashed
namespace under Studio's own cache (~/.unsloth/studio/cache/ollama_links)
when the Ollama models directory is read-only. Common for system installs
at /usr/share/ollama/.ollama/models and /var/lib/ollama/.ollama/models
where the Studio process has read but not write access. Previously the
scanner returned an empty list in that case and Ollama models would
silently not appear.
The fallback preserves the .gguf suffix on symlink names so
detect_gguf_model keeps recognising them. The prior "raw sha256 blob
path" fallback would have missed the suffix check and failed to load.
* Address review: detect mmproj next to symlink target for vision GGUFs
Codex P1 on model_config.py:1012: when detect_gguf_model returns the
symlink path (to preserve readable display names), detect_mmproj_file
searched the symlink's parent directory instead of the target's. For
vision GGUFs surfaced via Ollama's .studio_links/ -- where the weight
file is symlinked but any mmproj sidecar lives next to the real blob
-- mmproj was no longer detected, so the model was misclassified as
text-only and llama-server would start without --mmproj.
detect_mmproj_file now adds the resolved target's parent to the scan
order when path is a symlink. Direct (non-symlink) .gguf paths are
unchanged, so LM Studio and HF cache layouts keep working exactly as
before. Verified with a fake layout reproducing the bug plus a
regression check on a non-symlink LM Studio model.
* Address review: support all Ollama namespaces and vision projector layers
- Iterate over all directories under registry.ollama.ai/ instead of
hardcoding the "library" namespace. Custom namespaces like
"mradermacher/llama3" now get scanned and include the namespace
prefix in display names, model IDs, and symlink names to avoid
collisions.
- Create companion -mmproj.gguf symlinks for Ollama vision models
that have an "application/vnd.ollama.image.projector" layer, so
detect_mmproj_file can find the projector alongside the model.
- Extract symlink creation into _make_symlink helper to reduce
duplication between model and projector paths.
* Address review: move imports to top level and add scan limit
- Move hashlib and json imports to the top of the file (PEP 8).
- Remove inline `import json as _json` and `import hashlib` from
function bodies, use the top-level imports directly.
- Add `limit` parameter to `_scan_ollama_dir()` with early exit
when the threshold is reached.
- Pass `_MAX_MODELS_PER_FOLDER` into the scanner so it stops
traversing once enough models are found.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: Windows fallback, all registry hosts, collision safety
_make_link (formerly _make_symlink):
- Falls back to os.link() hardlink when symlink_to() fails (Windows
without Developer Mode), then to shutil.copy2 as last resort
- Uses atomic os.replace via tmp file to avoid race window where the
.gguf path is missing during rescan
Scanner now handles all Ollama registry layouts:
- Uses rglob over manifests/ instead of hardcoding registry.ollama.ai
- Discovers hf.co/org/repo:tag and any other host, not just library/
- Filenames include a stable sha1 hash of the manifest path to prevent
collisions between models that normalize to the same stem
Per-model subdirectories under .studio_links/:
- Each model's links live in their own hash-keyed subdirectory
- detect_mmproj_file only sees the projector for that specific model,
not siblings from other Ollama models
Friendly Ollama error detection:
- Now also matches ollama_links/ (the read-only fallback cache path)
and model_identifier starting with "ollama/"
Recommended folders:
- Added os.access(R_OK | X_OK) check so unreadable system directories
like /var/lib/ollama/.ollama/models are not advertised as chips
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: filter ollama_links from generic scanners
The generic scanners (models_dir, hf_cache, lmstudio) already filter
out .studio_links to avoid duplicate Ollama entries, but missed the
ollama_links fallback cache directory used for read-only Ollama
installs. Add it to the filter.
* Address review: idempotent link creation and path-component filter
_make_link:
- Skip recreation when a valid link/copy already exists (samefile or
matching size check). Prevents blocking the model-list API with
multi-GB copies on repeated scans.
- Use uuid4 instead of os.getpid() for tmp file names to avoid race
conditions from concurrent scans.
- Log cleanup errors instead of silently swallowing them.
Path filter:
- Use os.sep-bounded checks instead of bare substring match to avoid
false positives on paths like "my.studio_links.backup/model.gguf".
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: drop copy fallback, targeted glob, robust path filter
_make_link:
- Drop shutil.copy2 fallback -- copying multi-GB GGUFs inside a sync
API request would block the backend. Log a warning and skip the
model when both symlink and hardlink fail.
Scanner:
- Replace rglob("*") with targeted glob patterns (*/*/* and */*/*/*)
to avoid traversing unrelated subdirectories in large custom folders.
Path filter:
- Use Path.parts membership check instead of os.sep substring matching
for robustness across platforms.
Scan limit:
- Skip _scan_ollama_dir when _generic already fills the per-folder cap.
* Address review: sha256, top-level uuid import, Path.absolute()
- Switch hashlib.sha1 to hashlib.sha256 for path hashing consistency.
- Move uuid import to the top of the file instead of inside _make_link.
- Replace os.path.abspath with Path.absolute() in detect_gguf_model
to match the pathlib style used throughout the codebase.
* Address review: fix stale comments (sha1, rglob, copy fallback)
Update three docstrings/comments that still referenced the old
implementation after recent changes:
- sha1 comment now says "not a security boundary" (no hash name)
- "rglob" -> "targeted glob patterns"
- "file copies as a last resort" -> removed (copy fallback was dropped)
* Address review: fix stale links, support all manifest depths, scope error
_make_link:
- Drop size-based idempotency shortcut that kept stale links after
ollama pull updates a tag to a same-sized blob. Only samefile()
is used now -- if the link doesn't point at the exact same inode,
it gets replaced.
Scanner:
- Revert targeted glob back to rglob so deeper OCI-style repo names
(5+ path segments) are not silently skipped.
Ollama error:
- Only show "Some Ollama models do not work with llama.cpp" when the
server output contains GGUF compatibility hints (key not found,
unknown architecture, failed to load). Unrelated failures like
OOM or missing binaries now show the generic error instead of
being misdiagnosed.
---------
Co-authored-by: Daniel Han <info@unsloth.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
* Fix review findings for PR #49
1. Sandbox fallback Jinja env in _VariantTokenizerProxy.apply_chat_template
(use SandboxedEnvironment, matching _derive_assistant_prefix_by_render)
2. Unwrap benign outer-If guards in _template_ends_with_toplevel_for so
templates like {% if messages %}{% for ... %}{% endfor %}{% endif %}
are still repairable (preserves Qwen3-Guard rejection via else-branch
and add_generation_prompt-name checks)
3. Preserve raw name_or_path in _VariantTokenizerProxy._source_path so
local-path detection works for dict/list variant tokenizers
4. Context-aware strict-mode messages: omit "will still load" and
"Set UNSLOTH_STRICT_CHAT_TEMPLATE=1" when already raising
* [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>
