feat(lora): runtime BF16 LoRA adapter MVP (M0 + M1-attention)#8
Closed
rsafier wants to merge 13 commits into
Closed
feat(lora): runtime BF16 LoRA adapter MVP (M0 + M1-attention)#8rsafier wants to merge 13 commits into
rsafier wants to merge 13 commits into
Conversation
Tiny generated PEFT adapter (r=8, alpha=16) targeting k/v/o_proj + dense-FFN gate/up/down on the 6 full-attention layers (3,7,11,15,19,23) of Holo-3.1-0.8B. gen_lora_adapter.py builds the adapter (nonzero B so the delta is observable); reference_deltas.py emits scale*(B@A) as the offline-parity oracle. Verified: 36/36 modules, scaling=2.0, 0.0 rel-err. Large fp32 reference deltas gitignored.
parsers/lora.rs: PeftAdapterConfig {r, lora_alpha, target_modules, use_rslora,
layers_to_transform} + scaling() (alpha/r | alpha/sqrt(r)). Hard-fail Result
(unlike the Option-lenient quantization parser) with named REJECT(...) reasons
for every unsupported PEFT feature: peft_type!=LORA, use_dora, bias, rank/alpha
_pattern, modules_to_save, all-linear, q_proj (attn_output_gate), GDN modules,
embeddings, absent use_rslora, r=0. ModelConfig.adapter_max_rank (scratch sizing
seam). 11 unit tests green.
…-runtime) weights/adapter.rs: load_adapter_safetensors (adapter_model.safetensors -> WeightStore, host F16->BF16, header-only OOM preflight, named pickle reject) — dedicated loader because SafetensorsLoader only probes model.safetensors* and WeightDtype rejects F16. BufferSizes/BufferArena: lora_xa/lora_delta/lora_hact persistent scratch (NULL when adapter_max_rank==0), sized max(hidden,inter).
…rk-model)
lora/mod.rs: classify_key (PEFT key -> layer/module/A|B with named hard rejects:
gated-q-proj, gdn-target, non-full-attention-layer, shape/pair/bidirectional
audits) + load_lora_adapters_generic (rank-padded fixed-address pool, memset-0
padded-K correctness, slot-0 pack, frozen per-module u64 ptr tables). ops/
lora_delta.rs: canonical LoraPair{a,b:DenseWeight,rank,k_in,n_out,scale} +
LoraAttnWeights/LoraFfnWeights/LoraKernels + apply_lora_delta (M1 workhorse,
no callers yet). Trait hook load_lora_adapters (working generic default). Layer
+ model set_lora_weights install walk (downcast by global index, logs 'installed
on N layers'); TransformerLayer::as_any_mut hook. Factory LoraBuildArgs + build_
model wiring: pool alloc after loader / before BufferArena+free_memory so it's
KV-budgeted; install after construction. cargo check + clippy green.
…park-server) --lora-adapter NAME=PATH (Vec + parse_lora_adapter_spec; duplicate rejected), --max-lora-rank (64), --max-loras (8). resolve_adapter_dir (adapter_config.json marker; adapter_model.bin pickle reject). load_lora_adapter serve phase (dedicated adapter loader, r>max-rank + len>1 named bails, r/alpha/scaling log). world_size>1 guard (v0 TP=1). Threads LoraBuildArgs through serve_phases::build_model. AppState .adapter_name; /v1/models lists adapter first + get_model matches either; chat warns on base-name requests. cargo check --workspace green.
apply_lora_delta wired at prefill k/v (paged_qkv prefill_one_proj) + prefill o (paged_oproj), decode k/v (attention_forward, before norm/rope/kv-cache so the cache stores adapted k/v) + decode o (attention_forward_oproj, post-gate input). Guarded on layer self.lora (zero-cost on base + GDN layers). ATLAS_LORA_EAGER hatch disables decode-graph capture when set (deltas otherwise graph-safe: pool/scratch/scale load-time-fixed). FIX: contract expand/shrink at padded max_rank (B row stride = max_rank), not real rank — a k=rank expand misreads every B row past the first when r<max_rank. Deferred (documented cut): dense-FFN compute, cache_skip warm-prefix, multi-seq decode.
Delta max|.| now ~0.30 (comparable to base weight max ~0.32) so the served adapter unambiguously shifts greedy output for the M1 differ-from-base test. Offline parity unchanged: 36/36 modules, 0.0 rel-err, scaling=2.0.
…ora base) cuda_backend copy_d2d_2d_async uses the runtime symbol cudaMemcpy2DAsync; the normal CUDA link branch linked only cuda+cublasLt, so the spark bin failed to link on GB10/CUDA-13 (libcudart.so lives in targets/sbsa-linux/lib, not lib64). Newer branches use the cu* driver API; this makes the research/lora base build.
What ships, the end-to-end verification (offline parity 36/36 @ 0.0 rel-err; served base 'Paris.' vs adapter differ; graph==eager), and the documented cut lines (dense-FFN compute, cache-skip, multi-seq, per-request routing, TP).
Run rustfmt over the LoRA files; add the missing lora_args (None) arg to the build_model call in spark-server/tests/integration.rs (the new trailing param). cargo fmt --all --check clean on all changed files; cargo clippy --workspace --all-targets -- -D warnings exit 0; release build links; 11 parser tests green.
rsafier
force-pushed
the
research/lora-mvp
branch
from
July 4, 2026 23:43
3991145 to
98edbb7
Compare
PEFT saves LoRA adapters in F32 by default (trainable params stay fp32), but the adapter loader only converted F16->BF16; F32 fell through as WeightDtype::FP32. The pool pack (lora/mod.rs, BF16_BYTES) and dense_gemv_bf16 assume BF16 (2 B/elem) unconditionally, so F32 data was read at half stride -> garbage delta. Load succeeded, so it corrupted silently; only the BF16 test fixture had ever been served, so it hid. Add an F32->BF16 host conversion branch mirroring the F16 one. Atlas now reproduces the peft/transformers output verbatim (STARFALL-4728 demo). Also add examples/lora_apply_microtest.rs: a runtime parity oracle that runs the real CUDA apply_lora_delta (shrink/expand/fold) at holo k/v/o shapes for both decode (m=1 dense_gemv) and prefill (m>1 dense_gemm), bisecting each stage vs a bf16-faithful CPU reference. This is the runtime check the offline reference_deltas.py never provided. Plus docs/lora-adapter-mode-settings.md: the full train->upload->serve recipe + the resolved-bug writeup.
…ader test The LoRA family allow-list hard-rejected everything except qwen3_5 dense (holo-3.1-0.8b), so serving a LoRA on Holo-3.1-35B-A3B (model_type holo3_1_moe, 256 experts) failed at model build before any adapter tensor was examined. Relax the gate to also admit holo3_1_moe: it routes to Qwen35WeightLoader so its full-attention layers are Qwen3AttentionLayer (what the install walk downcasts to), and a k/v/o-only adapter leaves ffn_weights=None so the MoE FFN path is never exercised. The per-tensor classify_key + shape audit + full-attention gating still hard-reject q_proj / GDN / MoE-MLP / wrong-layer tensors, so a mistargeted adapter still fails loudly. Also add a gated (#[ignore], GPU-only) unit test in adapter.rs that writes an F32 PEFT adapter, loads it via load_adapter_safetensors, and asserts the tensors come back BF16 and round-trip bit-exact — pinning the F32->BF16 conversion invariant so the dtype bug can't silently regress.
The metal / no-CUDA CI build (`cargo test --no-default-features --features
metal`) tried to compile two CUDA-only additions and failed:
- examples/lora_apply_microtest.rs uses AtlasCudaBackend + ptx_modules()
- the adapter.rs #[cfg(test)] mod uses crate::cuda_backend::AtlasCudaBackend
(a #[cfg(feature="cuda")] module)
Gate both like every other CUDA microtest in the tree: add a [[example]]
block with required-features = ["cuda", "gpu-examples"], and change the
test module cfg to #[cfg(all(test, feature = "cuda"))]. cuda is a default
feature, so the gated test still runs on a normal `cargo test`.
Author
|
Superseded by the upstream draft PR Avarok-Cybersecurity#263 — the LoRA MVP now goes directly to upstream. Closing this fork PR. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Atlas has no LoRA-adapter support: it can only serve its own kernel-packed base
models, and there's no way to run a fine-tune as a runtime delta. The design work
landed on
research/lora(docs/lora-mvp-proposal.md+docs/lora-codebase-brief.md);this PR implements the MVP.
Solution
Serve one PEFT LoRA adapter, loaded at startup, applied to every request as a runtime
BF16 delta (
y += scale·(x@Aᵀ)@Bᵀ) — never merged (the NVFP4/FP8 base forbids it).Zero new CUDA kernels: deltas reuse
dense_gemv_bf16/dense_gemm_tc/bf16_scaled_addand are captured inside the existing CUDA decode graphs.M0 — load / validate / account.
--lora-adapter/--max-lora-rank/--max-loras;a hard-fail PEFT
adapter_config.jsonparser; a dedicated adapter safetensors loader(host F16→BF16, OOM preflight); per-
LayerTypekey remap with named hard rejections(
gated-q-proj,gdn-target,non-full-attention-layer, bidirectional tensor↔target +shape audits — never a silent skip); A/B packed rank-padded into one fixed-address pool
with per-module device pointer tables (the frozen M2 layout); pool VRAM budgeted against
the KV cache.
M1 — runtime delta at the attention insertion points.
apply_lora_deltawired atprefill + decode k/v/o (decode k/v before norm/RoPE/kv-cache-write, so the cache stores
adapted values).
ATLAS_LORA_EAGER=1debugging hatch; deltas are otherwise graph-safe./v1/modelsadvertises the adapter; base-name requests warn.Only the full-attention layers × {k,v,o,gate,up,down} are eligible; holo-3.1-0.8b's
GDN layers and its gated
q_projare rejected by design (v0 scope).Verification (holo-3.1-0.8b on GB10 / CUDA 13)
36/36 modules, scaling=2.0, 0.0 rel-err.
atlas-coreparser: 11 unit tests (accept + every named reject) green.installed on 6 layers (pool=117.0 MiB); base"Paris."→ adapter outputdiffers; graph == eager byte-identical (deltas capture correctly in CUDA graphs).
The
test_data/lora-holo-tiny/fixture is a generated PEFT adapter (no communityadapter exists for Atlas's custom NVFP4 bases; a controllable one also exercises the
reject/parity paths). It's deliberately strong, so the served effect is unambiguous — a
trained adapter would give coherent adapted text rather than the demo's degenerate output.
Deferred (documented in
docs/lora-implementation-status.md)Dense-FFN delta compute (types/install already present), prefix-cache warm-hit + multi-seq
decode (run the POC at batch size 1), per-request routing + multi-adapter +
lora_bgmv(M2), TP>1 (startup-guarded to
world_size=1).Note
crates/spark-runtime/build.rsnow linkslibcudart(+ CUDA-13 SBSA lib path) — theresearch/lorabase still calls the runtimecudaMemcpy2DAsync; harmless no-op on thedriver-API
main.🤖 Generated with Claude Code