feat(lora): runtime BF16 LoRA adapter support (Holo-3.1 dense + 35B MoE)#263
feat(lora): runtime BF16 LoRA adapter support (Holo-3.1 dense + 35B MoE)#263rsafier wants to merge 15 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.
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`.
…ow-list TransformerLayer trait Fix the two ≤500-LoC violations the LoRA MVP introduced: - spark-runtime/src/buffers.rs (555 -> 393): the LoRA scratch buffers pushed it over. Extract the pointer/size accessor getters (a cohesive cluster, including the lora_xa/lora_delta/lora_hact getters) into a new buffers/accessors.rs as a second impl block (child module reads the private fields). Real code split per Atlas idiom. - spark-model/src/layer/transformer_layer.rs (504 -> 502): the as_any_mut downcast hook (needed by the LoRA install walk) tipped this trait over. It is an atomic interface — ~20 arg-heavy methods whose SSM-prefill defaults are one-line bail/Ok no-ops, so body-extraction is net-neutral (same situation as the already-allow-listed traits/model.rs). Trim the as_any_mut doc and allow-list it with rationale; a supertrait split is tracked as follow-up.
…ntra-doc-links) spark-model has #![deny(warnings)] (implies deny rustdoc::broken_intra_doc_links). Doc comments wrote matrix shapes with no space after the comma — [N,K], [512,1024], [1024,2048], [3584,1024], [1024,3584] — which rustdoc parses as intra-doc links to nonexistent items, failing 'cargo doc --workspace --no-deps'. Wrap them in backticks. Verified: cargo doc --workspace --no-deps now exits 0.
Review — #263 (runtime BF16 LoRA adapter support)Verdict: strong, well-engineered POC — mergeable as the opt-in draft it is (gated behind Reviewed the full path: PEFT parser ( What's genuinely good
Findings1. The "never silent skip" invariant has two runtime holes (the main callout). The PR's stated contract (docs: "Named hard rejections — never silent skips") doesn't hold at runtime in two documented-but-unguarded places:
2. Minor — always-on adapter. 3. Nit — the pool packs all 6 modules for every full-attn layer regardless of targets. A k/v/o-only adapter still allocates gate/up/down pad space. Bounded and intentional (frozen M2 layout), but on the 35B (large Bottom lineExcellent engineering and test discipline for a POC; the docs are honest about the cut lines. It's opt-in and startup-only, so merging it as a draft/POC is low-risk. The gating item before real use is 1b (enforce batch=1 / no prefix-cache when LoRA is active) and ideally 1a (reject FFN targets until applied) — so the "never silent skip" guarantee the PR itself establishes actually holds when someone points traffic at it. |
[BETA] Automated PR review — SeedSource AI pipelineAdvisory only, not a merge gate. AI-generated (internal risk-triage, Atlas serve gate, GLM-5.2 review), attributed per Atlas convention. Triage
Serve gate
GLM-5.2 code reviewDiff truncated — review based on visible portion. Correctness: Parser logic is sound. Risk: Quality: Naming discipline ( Concern: Cannot verify the loader-side enforcement or CUDA merge kernel from truncated diff. The Pipeline: build+serve on GB10, GLM-5.2 via NVIDIA NIM, internal risk-triage classifier. External diff treated as untrusted input. Public repo. |
2 participants
Problem
Atlas serves the Holo / Qwen3.5 family, but there was no way to serve a LoRA fine-tuned
variant: the standard PEFT adapter format wasn't understood, and there was no runtime path to
apply a low-rank delta on top of the NVFP4-quantized base.
Solution
Apply the LoRA update as a runtime BF16 delta, never merged into the (NVFP4) base weights —
y += scale·(x@Aᵀ)@Bᵀat attention k_proj / v_proj / o_proj on the full-attention layersonly, reusing existing kernels (
dense_gemv_bf16/dense_gemm_tc/bf16_scaled_add). Zeronew CUDA kernels; deltas are graph-safe (captured inside the existing decode CUDA graphs).
atlas-core):adapter_config.json→ validated config; every unsupportedfeature is a named
REJECT(...)(DoRA, bias, rank/alpha patterns,modules_to_save,all-linear, absentuse_rslora, non-LORApeft_type). rsLoRA (α/√r) scaling supported.spark-runtime): dedicatedadapter_model.safetensorsloader with host-sideF16→BF16 and F32→BF16 conversion (F32 is PEFT's default save dtype), header-only OOM
preflight, pickle reject.
spark-model): key remap + per-LayerTypeallow-list — onlyfull-attention × k/v/o is accepted;
q_proj(gated Q|gate), GDN, MoE-MLP, and wrong-layertensors are named hard rejections, never silent skips. A/B are rank-padded into one
fixed-address pool with per-module
[max_loras]pointer tables (the frozen layout for futuremulti-adapter routing; v0 fills slot 0). Pool VRAM is budgeted against the KV cache.
spark-server):--lora-adapter NAME=PATH_OR_HF_ID,--max-lora-rank,--max-loras; the adapter name is advertised byGET /v1/modelsand accepted inrequest.model.qwen3_5(Holo-3.1-0.8b) andholo3_1_moe(Holo-3.1-35B-A3B MoE) — ak/v/o-only adapter leaves
ffn_weights = None, so the MoE expert path is never exercised.Validation
Verified end-to-end on a GB10 with exact parity to
peft/transformers:made-up codeword verbatim.
NVFP4 base with the BF16 delta applied to the 10 full-attention layers — persona verbatim.
examples/lora_apply_microtest.rs: a runtime parity oracle that runs the real CUDA apply(shrink → expand → fold), decode and prefill shapes, bisected against a bf16-faithful CPU
reference. Plus a gated
#[ignore]GPU test pinning the F32→BF16 conversion, and the offlineparity oracle for loaded A/B.
Draft: opening for review/CI on upstream. Deferred (documented cut lines): dense-FFN delta,
per-request multi-adapter routing (
lora_bgmv), q_proj/GDN/MLA targets, TP>1.