Reduce TP=2 Q6K cold-load time for Qwen3.6-27B (~5 min today) #1

New Issue

Closed

opened 2026-05-27 10:31:42 +00:00 by grenade · 2 comments

grenade commented 2026-05-27 10:31:42 +00:00

Owner

Copy Link

Context

Cold-loading Qwen/Qwen3.6-27B with quant = "q6k" and tensor_parallel = 2 on beast currently takes ~5 minutes. The dominant phase is layer-by-layer Q6K in-situ quantization, single-threaded per rank: ~60–90 s per rank, run serially across 52 transformer layers plus the lm_head. Everything else is small change (subprocess spawn ~20 ms, NCCL init ~1 s, sanity check stubbed, mmap setup, disk read).

Cold-load latency matters because default_models in neuron.toml is what makes a host ready to serve after deploy.sh — and because a poisoned-thread restart pays the full cold-load cost again.

Findings from codebase analysis

• candle_core::quantized::QTensor::quantize is not internally parallel. candle-core-0.10.2/src/quantized/k_quants.rs imports rayon but the Q6K from_float path (lines 1989–2059) is purely sequential per-block. Each 256-element block is independent — embarrassingly parallel — but the implementation processes them serially.
• Layer loop is serial. tp_qwen3_5.rs:875–892 iterates 52 layers one at a time; no rayon::par_iter or tokio::spawn over the per-layer quantization step (MaybeQuantLinear::from_weight at tp_linear.rs:56–68).
• lm_head is materialized and quantized on every rank. Both TpQwen3_5ForCausalLM::load() (tp_qwen3_5.rs:1003–1017) and worker.rs:298–306 call build_lm_head(cfg, vb, &base, quant). Workers then compute logits through it during forward and discard the result (worker.rs:383–386 comment). The Q6K encoding of the ~152K × 12288 ≈ 120M-weight lm_head on every worker is pure waste.
• No per-phase wall-clock timing. log_construction_complete (tp_qwen3_5.rs:1154) brackets the whole load, but there's no per-phase elapsed_ms for NCCL init, mmap setup, each layer load, or lm_head. We're optimizing partly by inference.

Opportunities (ordered by impact-to-effort ratio)

1. Parallelize layer quantization within a rank via rayon::par_iter

Each layer's weight tensors are independent once sliced out of the mmap. Beast has many cores; on a multi-core host this should give 5–10× speedup over the 52 layers. Two implementation shapes:

• Bigger payoff: pre-stage all per-layer byte slices off the mmap onto host buffers, then par_iter the quantization, then assemble the model on the worker thread.
• Smaller payoff: build the model layer-by-layer but par_iter the quantization step within each layer's QKV/MLP set.

Likely the single biggest win.

2. Skip lm_head quantization on worker ranks

Workers don't need a quantized lm_head — they discard its output. Two graded approaches:

• Cheap (~1-line edit): pass quant = None to build_lm_head on rank > 0 in tp_qwen3_5.rs:1052–1076. Keeps the matmul cost but eliminates the Q6K encoding on every worker. Saves ~2–3 s per worker.
• Bigger: skip build_lm_head and the hidden.apply(&self.lm_head) entirely on rank > 0, returning a sentinel. Saves the matmul cost too. Requires touching the forward path.

3. Disk-cache post-ISQ weights

First cold load runs ISQ; subsequent loads memcpy a pre-quantized blob into VRAM. GGUF is exactly this, but the qwen3_5 (Qwen3-Next) arch has no upstream GGUF — hence the runtime ISQ. We could write our own per-rank .helexa-q6k cache keyed on (model_id, quant, rank, world_size, tp_topology). Significantly more code, but warm loads would drop to 30–60 s (just disk → VRAM).

4. Add per-phase timing instrumentation before optimizing more aggressively

Spend the half-hour to add tracing::info!(elapsed_ms = …) around: NCCL init, mmap setup, each layer load+quant, lm_head, post-construct. Without this we can't validate that #1 above captures ~80% vs ~30% of the total. Low-effort, high-information.

5. Upstream candle: rayon-ify Q6K from_float

par_chunks_mut(QK_K) over the block loop would give linear scaling. Benefits every model that uses Q6K ISQ. Local workaround: vendor candle and patch in place.

6. Skip log_vram per layer

tp_qwen3_5.rs:874, 892 queries the driver every layer for diagnostic VRAM. Gate behind cfg!(debug_assertions) or a per-load flag. Marginal.

Not worth pursuing

• Shared mmap across ranks — OS page cache already absorbs the duplication; fd count is 2, not 200.
• Subprocess spawn overhead — ~20 ms vs. minutes.
• Parallel disk read across ranks — already happens.

Suggested sequence

1. #4 first (instrument) so subsequent measurements are real.
2. #2 cheap variant as the first one-line win.
3. #1 (rayon over per-layer quantization) — the main lever.
4. #3 (disk cache) for the long tail once everything else is exhausted.

Related code

• crates/neuron/src/harness/tp/tp_qwen3_5.rs:875–892 — per-layer load loop
• crates/neuron/src/harness/tp/tp_qwen3_5.rs:1052–1076 — build_lm_head
• crates/neuron/src/harness/tp/tp_linear.rs:56–68 — MaybeQuantLinear::from_weight
• crates/neuron/src/harness/tp/worker.rs:298–306, 383–386 — worker load + forward (logits discarded)
• crates/neuron/src/harness/tp/mod.rs:236–296, 316–362, 487–556 — WorkerPool spawn, NCCL init, load_dense_shard
• crates/neuron/src/harness/device_worker/dispatch.rs:594–673 — tp_load_shard_inner

## Context Cold-loading `Qwen/Qwen3.6-27B` with `quant = "q6k"` and `tensor_parallel = 2` on beast currently takes ~5 minutes. The dominant phase is **layer-by-layer Q6K in-situ quantization, single-threaded per rank**: ~60–90 s per rank, run serially across 52 transformer layers plus the lm_head. Everything else is small change (subprocess spawn ~20 ms, NCCL init ~1 s, sanity check stubbed, mmap setup, disk read). Cold-load latency matters because `default_models` in `neuron.toml` is what makes a host ready to serve after deploy.sh — and because a poisoned-thread restart pays the full cold-load cost again. ## Findings from codebase analysis - **`candle_core::quantized::QTensor::quantize` is not internally parallel.** `candle-core-0.10.2/src/quantized/k_quants.rs` imports `rayon` but the Q6K `from_float` path (lines 1989–2059) is purely sequential per-block. Each 256-element block is independent — embarrassingly parallel — but the implementation processes them serially. - **Layer loop is serial.** `tp_qwen3_5.rs:875–892` iterates 52 layers one at a time; no `rayon::par_iter` or `tokio::spawn` over the per-layer quantization step (`MaybeQuantLinear::from_weight` at `tp_linear.rs:56–68`). - **lm_head is materialized and quantized on every rank.** Both `TpQwen3_5ForCausalLM::load()` (tp_qwen3_5.rs:1003–1017) and `worker.rs:298–306` call `build_lm_head(cfg, vb, &base, quant)`. Workers then compute logits through it during forward and **discard the result** (`worker.rs:383–386` comment). The Q6K encoding of the ~152K × 12288 ≈ 120M-weight lm_head on every worker is pure waste. - **No per-phase wall-clock timing.** `log_construction_complete` (`tp_qwen3_5.rs:1154`) brackets the whole load, but there's no per-phase `elapsed_ms` for NCCL init, mmap setup, each layer load, or lm_head. We're optimizing partly by inference. ## Opportunities (ordered by impact-to-effort ratio) ### 1. Parallelize layer quantization within a rank via `rayon::par_iter` Each layer's weight tensors are independent once sliced out of the mmap. Beast has many cores; on a multi-core host this should give 5–10× speedup over the 52 layers. Two implementation shapes: - **Bigger payoff**: pre-stage all per-layer byte slices off the mmap onto host buffers, then `par_iter` the quantization, then assemble the model on the worker thread. - **Smaller payoff**: build the model layer-by-layer but `par_iter` the quantization step within each layer's QKV/MLP set. Likely the single biggest win. ### 2. Skip lm_head quantization on worker ranks Workers don't need a quantized lm_head — they discard its output. Two graded approaches: - **Cheap (~1-line edit)**: pass `quant = None` to `build_lm_head` on rank > 0 in `tp_qwen3_5.rs:1052–1076`. Keeps the matmul cost but eliminates the Q6K encoding on every worker. Saves ~2–3 s per worker. - **Bigger**: skip `build_lm_head` and the `hidden.apply(&self.lm_head)` entirely on rank > 0, returning a sentinel. Saves the matmul cost too. Requires touching the forward path. ### 3. Disk-cache post-ISQ weights First cold load runs ISQ; subsequent loads memcpy a pre-quantized blob into VRAM. GGUF is exactly this, but the qwen3_5 (Qwen3-Next) arch has no upstream GGUF — hence the runtime ISQ. We could write our own per-rank `.helexa-q6k` cache keyed on (model_id, quant, rank, world_size, tp_topology). Significantly more code, but warm loads would drop to 30–60 s (just disk → VRAM). ### 4. Add per-phase timing instrumentation **before** optimizing more aggressively Spend the half-hour to add `tracing::info!(elapsed_ms = …)` around: NCCL init, mmap setup, each layer load+quant, lm_head, post-construct. Without this we can't validate that #1 above captures ~80% vs ~30% of the total. Low-effort, high-information. ### 5. Upstream candle: rayon-ify Q6K `from_float` `par_chunks_mut(QK_K)` over the block loop would give linear scaling. Benefits every model that uses Q6K ISQ. Local workaround: vendor candle and patch in place. ### 6. Skip `log_vram` per layer `tp_qwen3_5.rs:874, 892` queries the driver every layer for diagnostic VRAM. Gate behind `cfg!(debug_assertions)` or a per-load flag. Marginal. ## Not worth pursuing - Shared mmap across ranks — OS page cache already absorbs the duplication; fd count is 2, not 200. - Subprocess spawn overhead — ~20 ms vs. minutes. - Parallel disk read across ranks — already happens. ## Suggested sequence 1. **#4 first** (instrument) so subsequent measurements are real. 2. **#2 cheap variant** as the first one-line win. 3. **#1** (rayon over per-layer quantization) — the main lever. 4. **#3** (disk cache) for the long tail once everything else is exhausted. ## Related code - `crates/neuron/src/harness/tp/tp_qwen3_5.rs:875–892` — per-layer load loop - `crates/neuron/src/harness/tp/tp_qwen3_5.rs:1052–1076` — `build_lm_head` - `crates/neuron/src/harness/tp/tp_linear.rs:56–68` — `MaybeQuantLinear::from_weight` - `crates/neuron/src/harness/tp/worker.rs:298–306, 383–386` — worker load + forward (logits discarded) - `crates/neuron/src/harness/tp/mod.rs:236–296, 316–362, 487–556` — WorkerPool spawn, NCCL init, load_dense_shard - `crates/neuron/src/harness/device_worker/dispatch.rs:594–673` — `tp_load_shard_inner`

grenade referenced this issue 2026-05-27 15:30:16 +00:00

Research: Prototype a quantization kernel in pure Rust via cuda-oxide #2

grenade referenced this issue 2026-06-12 08:57:46 +00:00

Research: Prototype a quantization kernel in pure Rust via cuda-oxide #2

grenade commented 2026-06-12 08:58:13 +00:00

Author

Owner

Copy Link

Still in scope, and elevated by the 2026-06-12 reframing (README positioning + the 7→8 milestone): cold-load latency is one of the operator-felt metrics the benchmark harness (#22) will measure and publish, and it is the recurring cost of #17-style auto-recovery (every poison→reload pays this in full). The parallelization findings in the body (per-block Q6K quantization is embarrassingly parallel; the 52-layer loop is serial) remain the right attack — CPU-side rayon work or a small candle patch, per the closure rationale on #2.

Still in scope, and elevated by the 2026-06-12 reframing (README positioning + the 7→8 milestone): cold-load latency is one of the operator-felt metrics the benchmark harness (#22) will measure and publish, and it is the recurring cost of #17-style auto-recovery (every poison→reload pays this in full). The parallelization findings in the body (per-block Q6K quantization is embarrassingly parallel; the 52-layer loop is serial) remain the right attack — CPU-side rayon work or a small candle patch, per the closure rationale on #2.

grenade added the p2-next label 2026-06-12 09:01:45 +00:00

grenade referenced this issue 2026-06-12 09:02:07 +00:00

tracking: prioritised path to closing every open issue #27

grenade referenced this issue 2026-06-12 10:18:41 +00:00

tracking: prioritised path to closing every open issue #27

grenade referenced this issue 2026-06-12 12:39:54 +00:00

bench: reproducible benchmark harness + published numbers #22

grenade referenced this issue from a commit 2026-06-12 19:48:01 +00:00

perf(neuron): parallel in-situ quantization + cold-load phase timing (#1)

grenade referenced this issue 2026-06-12 19:48:17 +00:00

perf(neuron): parallel in-situ quantization + cold-load phase timing (#1) #40

grenade referenced this issue from a commit 2026-06-12 20:12:47 +00:00

Merge pull request 'perf(neuron): parallel in-situ quantization + cold-load phase timing (#1)' (#40) from perf/1-parallel-isq into main

grenade commented 2026-06-12 21:43:48 +00:00

Author

Owner

Copy Link

Closing numbers. Landed via #40 (parallel in-situ quantization + the per-phase timing this issue asked for first).

Before/after — Qwen3.6-27B Q6K TP=2 cold load on beast (service start → TP model loaded, journal-verified)

	before (serial ISQ)	after (#40)
cold load	221 s (23:07:11 → 23:10:52)	86 s (00:41:11 → 00:42:37)	2.6×

Phase breakdown after (from the new instrumentation): layer loop 79.0 s (uniform ~1.2 s/layer/rank × 64 layers, ranks in parallel), lm_head 5.2 s, everything else ~2 s. Every #17 auto-recovery now pays 86 s instead of 221 s.

What shipped

• harness/tp/isq.rs: quantize_parallel — candle's per-block k-quant math (k_quants::GgmlType::from_float, public API, no fork) fanned across the rayon pool, byte-identical to QTensor::quantize (parity-tested for Q6K/Q5K/Q4K/Q8_0). Device↔host transfers stay on the context-owning thread per the worker discipline.
• Per-layer debug + layer-loop/lm_head info timing on all ranks, so the next person optimizing this sees the phase split instead of inferring it.

Remaining lever (issue body item 3)

With ISQ compute now spread over 32 cores (~2–3 s total), the residual 79 s is dominated by reading the 54 GB bf16 safetensors (~700 MB/s effective with both ranks reading) plus slicing/upload. The post-ISQ disk cache idea would cut that to ~27 GB of pre-quantized blob per fleet-load and skip ISQ entirely (~30–40 s estimate) — worth a fresh issue only if 86 s still hurts in practice. Items 2 (lm_head on workers — now only ~5 s, and the quant=None variant costs ~1.5 GB VRAM per worker) and 6 (per-layer log_vram) are no longer worth their complexity at the new baseline.

Closing numbers. Landed via #40 (parallel in-situ quantization + the per-phase timing this issue asked for first). ## Before/after — Qwen3.6-27B Q6K TP=2 cold load on beast (service start → `TP model loaded`, journal-verified) | | before (serial ISQ) | after (#40) | | |---|---:|---:|---| | cold load | **221 s** (23:07:11 → 23:10:52) | **86 s** (00:41:11 → 00:42:37) | **2.6×** | Phase breakdown after (from the new instrumentation): layer loop **79.0 s** (uniform ~1.2 s/layer/rank × 64 layers, ranks in parallel), lm_head **5.2 s**, everything else ~2 s. Every #17 auto-recovery now pays 86 s instead of 221 s. ## What shipped - `harness/tp/isq.rs`: `quantize_parallel` — candle's per-block k-quant math (`k_quants::GgmlType::from_float`, public API, no fork) fanned across the rayon pool, **byte-identical** to `QTensor::quantize` (parity-tested for Q6K/Q5K/Q4K/Q8_0). Device↔host transfers stay on the context-owning thread per the worker discipline. - Per-layer debug + layer-loop/lm_head info timing on all ranks, so the next person optimizing this sees the phase split instead of inferring it. ## Remaining lever (issue body item 3) With ISQ compute now spread over 32 cores (~2–3 s total), the residual 79 s is dominated by reading the 54 GB bf16 safetensors (~700 MB/s effective with both ranks reading) plus slicing/upload. The post-ISQ disk cache idea would cut that to ~27 GB of pre-quantized blob per fleet-load and skip ISQ entirely (~30–40 s estimate) — worth a fresh issue only if 86 s still hurts in practice. Items 2 (lm_head on workers — now only ~5 s, and the `quant=None` variant costs ~1.5 GB VRAM per worker) and 6 (per-layer `log_vram`) are no longer worth their complexity at the new baseline.

grenade closed this issue 2026-06-12 21:43:55 +00:00

Sign in to join this conversation.

No Branch/Tag Specified

Branches Tags

main

copy/mission-investor-positioning

feat/F6-beta-polish

feat/B7-upstream-packaging

feat/F5-auth-chat

feat/B6-served-usage

feat/F4-account-dashboard

feat/F3-anon-chat

feat/F2-mission

feat/F1-theming-i18n

feat/B5-topup-codes

feat/B3-cortex-upstream-client

feat/B4-account-api

feat/B2-authz-api

feat/F0-helexa-ai-scaffold

feat/B1-helexa-upstream-skeleton

feat/74-outbound-tls-pinning

feat/75-federation-catalogue

feat/73-capacity-aware-dispatch

feat/72-router-topology-poller

fix/71-shared-streaming-proxy

feat/70-router-skeleton

fix/cortex-poll-debounce-retryable

fix/neuron-models-control-plane-starvation

fix/47-anonymous-tolerates-bad-key

feat/47-phase3-prompt-prevalidation

feat/47-phase2d-load-aware-routing

feat/47-phase2c-per-principal-cap

feat/47-phase2b-health-load

feat/47-phase2a-admission-control

feat/47-phase1d-budget-enforcement

feat/47-phase1c-metering

feat/47-phase1b-auth-principal

feat/47-phase1a-entitlement-provider

feat/47-phase0-error-envelope

feat/deploy-bench-on-bob

feat/version-metadata-and-bench

chore/cudarc-bump-63327a2

feat/neuron-17-stage2

feat/neuron-17-fault-recovery

feat/neuron-14-dynamic-resolution

phase-2-preflight

v0.1.16

v0.1.15

v0.1.14

v0.1.12

v0.1.11

v0.1.10

v0.1.9

v0.1.8

v0.1.7

v0.1.6

v0.1.5

v0.1.4

v0.1.3

v0.1.2

v0.1.1

v0.1.0

Labels

Clear labels

out-of-scope

Outside helexa's constraining principles: near-frontier models on consumer hardware, served transparently and predictably

p1-now

Active pain or unblocking work — pick up first

p2-next

Engine levers — start once P1 baselines exist

p3-later

Dependent on earlier work, or demand-driven

No Label p2-next

Milestone

No items

No Milestone

Projects

Clear projects

No project

Assignees

Clear assignees

grenade

No Assignees

1 Participants

Notifications

Subscribe

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: helexa/helexa#1