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.
