@ondeinference / cli / commits / fc882a6

finetune: train on CPU and fix gradient-norm overflow

candle 0.10.2's Metal backward pass returns garbage gradients for the Qwen forward here: at pristine weights the per-step max element runs 1e6 to 1e29, where the CPU backward gives ~2-7 for the same loss. Two things broke as a result. The inf-norm guard skipped most steps, so training only saw a biased subset and half-learned. And the norm itself was computed by summing squared gradients in f32, which overflows to inf once any element passes ~1.8e19, so even finite-but-large gradients forced a skip. Training now runs on CPU on macOS (Accelerate is also faster per step). ONDE_FINETUNE_METAL=1 re-enables Metal for retesting on newer candle, and ONDE_FINETUNE_GRAD_DEBUG=1 prints per-step loss, max element, and norm. The global norm is computed in a scaled space: divide by the largest element, sum the squares, recover the true norm in f64. Large-but-finite gradients get a finite norm and are clipped instead of skipping the step, and skips are now logged with the offending value. The full-pipeline test gains env knobs (ONDE_TEST_DATA_PATH, EPOCHS, MAX_SEQ_LEN, SYSTEM_PROMPT, and ||-separated EVAL_PROMPT) so it doubles as a headless runner for a custom dataset. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

keypair34 committed Jul 10, 2026 at 14:56 UTC fc882a629d899a7161da77f46f22d6c723232eb3

3 files changed +107 -28

CLAUDE.md

+3 -3

@@ -17,7 +17,7 @@ cargo clippy --all-targets

18 The build **bakes credentials in at compile time** via `build.rs`, which reads `.env` (or the environment in CI). The build fails if any of these are missing: `ONDE_APP_ID`, `ONDE_APP_SECRET`, `GRESIQ_API_KEY`, `GRESIQ_API_SECRET`. `HF_TOKEN` is optional (defaults to empty). Changing `.env` triggers a rebuild. These are exposed in code as `env!(...)` constants in `src/app.rs`.

20 -On macOS the build pulls in candle's `metal` and `accelerate` features, so inference and training run on Metal; elsewhere they fall back to CPU.

20 +On macOS the build pulls in candle's `metal` and `accelerate` features. Inference runs on Metal; **training deliberately runs on CPU everywhere** — candle 0.10.2's Metal backward pass produces garbage gradient magnitudes (1e6–1e29 where CPU gives ~40–65 for identical losses), and Accelerate is ~2x faster per training step anyway. `ONDE_FINETUNE_METAL=1` re-enables Metal training for testing newer candle versions; verify gradients with `ONDE_FINETUNE_GRAD_DEBUG=1` (per-step loss/max_abs/norm) before trusting it.

22 ### Tests that load real models

@@ -28,7 +28,7 @@ cargo test gguf::tests::exported_qwen3_gguf_is_runnable -- --ignored --nocapture

28 cargo test gguf::tests::finetune_merge_export_run -- --ignored --nocapture

29 ```

31 -Env knobs for these: `ONDE_TEST_MODEL_DIR`, `ONDE_TEST_DTYPE` (`f16`/`q8_0`), `ONDE_TEST_LR`.

31 +Env knobs for these: `ONDE_TEST_MODEL_DIR`, `ONDE_TEST_DTYPE` (`f16`/`q8_0`), `ONDE_TEST_LR`, `ONDE_TEST_DATA_PATH`, `ONDE_TEST_EPOCHS`, `ONDE_TEST_MAX_SEQ_LEN`, `ONDE_TEST_EVAL_PROMPT` (`||`-separated prompts), `ONDE_TEST_SYSTEM_PROMPT`. This makes the full-pipeline test double as a headless fine-tune runner for a custom dataset.

33 ### Debugging the TUI

@@ -46,7 +46,7 @@ Background work uses two task kinds deliberately: network/IO uses `tokio::spawn`

47 These modules form the fine-tune-to-deploy chain, each running on a background thread and streaming a `*Progress` enum:

49 -- `finetune.rs` — hand-written LoRA trainer (candle). Builds a Qwen forward pass (RMSNorm, RoPE, GQA, optional Qwen3 QK-norm), trains LoRA A/B on q/v projections in F32, writes `lora_adapter.safetensors`. Gradients are sanitized (non-finite elements zeroed) and globally norm-clipped before each AdamW step; a step is skipped entirely if its global grad norm is non-finite. Skipping this guard corrupts every weight after the first bad step.

49 +- `finetune.rs` — hand-written LoRA trainer (candle). Builds a Qwen forward pass (RMSNorm, RoPE, GQA, optional Qwen3 QK-norm), trains LoRA A/B on q/v projections in F32, writes `lora_adapter.safetensors`. Gradients are sanitized (non-finite elements zeroed) and globally norm-clipped before each AdamW step; the norm is computed in a scaled space so huge-but-finite gradients clip instead of overflowing the f32 sum to `inf`. A step is skipped (loudly) only when the norm is zero or non-finite. Skipping this guard corrupts every weight after the first bad step. Training runs on CPU — see the Metal note above.

50 - `merge.rs` — folds the LoRA adapter back into base weights (`W + scale·(B@A)`), writes a merged `model.safetensors` plus copied config/tokenizer.

51 - `gguf.rs` — **hand-rolled GGUF writer** (no llama.cpp). Conventions that must hold for mistral.rs/candle to load and run the file correctly: tensor dims are written innermost-first (reverse of safetensors shape, because candle reverses on read); `head_dim` comes from config and is emitted as `attention.key_length`/`value_length` (Qwen3 decouples it from `hidden_size/num_heads`); `token_type` is an INT32 array; `general.architecture` is chosen from `model_type` so Qwen3 routes to candle's `quantized_qwen3` loader.

52 - `chat.rs` — loads a local GGUF via `onde::mistralrs::GgufModelBuilder` (the same engine the Onde SDK uses) and streams a multi-turn chat, so a model can be tested before publishing.

src/finetune.rs

+61 -7

@@ -121,8 +121,18 @@ fn run_finetune(

121 // ------------------------------------------------------------------

122 // Step C — Choose compute device

123 // ------------------------------------------------------------------

124 + // Training runs on CPU even on macOS: candle 0.10.2's Metal *backward*

125 + // pass produces garbage gradient magnitudes (1e6–1e29 at pristine weights,

126 + // where the CPU backward yields norms of ~40–65 for identical losses), and

127 + // with Accelerate the CPU backward is also ~2x faster per step. Metal

128 + // remains opt-in via ONDE_FINETUNE_METAL for re-testing on newer candle.

129 + // Forward-only paths (inference, merge, export) are unaffected.

130 #[cfg(target_os = "macos")]

125 - let device = Device::new_metal(0).unwrap_or(Device::Cpu);

131 + let device = if std::env::var("ONDE_FINETUNE_METAL").is_ok() {

132 + Device::new_metal(0).unwrap_or(Device::Cpu)

133 + } else {

134 + Device::Cpu

135 + };

136 #[cfg(not(target_os = "macos"))]

137 let device = Device::Cpu;

138

@@ -297,24 +307,68 @@ fn run_finetune(

307 // AdamW moment estimates, corrupting every weight from that step on.

308 // Replace non-finite elements with zero so the remaining finite

309 // gradient still trains those parameters this step.

300 - let mut total_norm_sq: f64 = 0.0;

310 + // Pass 1 — sanitize each gradient and find the largest |g| element

311 + // across all tensors. sanitize_grad admits elements up to 1e30, but

312 + // squaring anything above ~1.8e19 overflows the f32 sum below, so the

313 + // norm must be computed in a scaled space.

314 + let mut cleaned: Vec<(&Var, Tensor)> = Vec::with_capacity(lora_vars.len());

315 + let mut global_max_abs: f64 = 0.0;

316 for var in &lora_vars {

317 if let Some(grad) = grads.get(var.as_tensor()) {

318 let clean = sanitize_grad(grad).context("sanitizing gradient")?;

304 - let norm_sq = clean

305 - .sqr()

319 + let max_abs = clean

320 + .abs()

321 + .and_then(|t| t.max_all())

322 + .and_then(|t| t.to_scalar::<f32>())

323 + .unwrap_or(0.0) as f64;

324 + global_max_abs = global_max_abs.max(max_abs);

325 + cleaned.push((var, clean));

326 + }

327 + }

328 +

329 + // Pass 2 — L2 norm of the gradients scaled by 1/global_max_abs.

330 + // Scaled elements are ≤ 1 so the f32 accumulation cannot overflow;

331 + // the true norm is recovered in f64. Exploding-but-finite gradients

332 + // thus get a finite norm and are tamed by clipping below instead of

333 + // forcing the whole step to be skipped.

334 + let mut scaled_norm_sq: f64 = 0.0;

335 + if global_max_abs > 0.0 {

336 + let inv = 1.0 / global_max_abs;

337 + for (_, clean) in &cleaned {

338 + let norm_sq = (clean * inv)

339 + .and_then(|t| t.sqr())

340 .and_then(|t| t.sum_all())

341 .and_then(|t| t.to_scalar::<f32>())

342 .unwrap_or(0.0) as f64;

309 - total_norm_sq += norm_sq;

310 - grads.insert(var.as_tensor(), clean);

343 + scaled_norm_sq += norm_sq;

344 }

345 }

313 - let total_norm = total_norm_sq.sqrt();

346 + let total_norm = global_max_abs * scaled_norm_sq.sqrt();

347 + if std::env::var("ONDE_FINETUNE_GRAD_DEBUG").is_ok() {

348 + eprintln!(

349 + "[finetune] grad-debug epoch {} step {} seq_len {} loss {:.4} max_abs {:.3e} norm {:.3e}",

350 + epoch + 1,

351 + step + 1,

352 + seq_len,

353 + loss_val,

354 + global_max_abs,

355 + total_norm

356 + );

357 + }

358 + for (var, clean) in cleaned {

359 + grads.insert(var.as_tensor(), clean);

360 + }

361

362 // Nothing finite to learn from this step — skip the update so AdamW

363 // state stays clean.

364 if total_norm == 0.0 || !total_norm.is_finite() {

365 + eprintln!(

366 + "[finetune] WARNING: global grad norm is {} at epoch {} step {} (seq_len {}), skipping update",

367 + total_norm,

368 + epoch + 1,

369 + step + 1,

370 + seq_len

371 + );

372 continue;

373 }

374

src/gguf.rs

+43 -18

@@ -988,9 +988,11 @@ mod tests {

988 .build()

989 .await

990 .expect("exported GGUF must load in mistral.rs");

991 - let req = RequestBuilder::new()

992 - .set_sampler_max_len(512)

993 - .add_message(TextMessageRole::User, &user);

991 + let mut req = RequestBuilder::new().set_sampler_max_len(512);

992 + if let Ok(system) = std::env::var("ONDE_TEST_SYSTEM_PROMPT") {

993 + req = req.add_message(TextMessageRole::System, &system);

994 + }

995 + let req = req.add_message(TextMessageRole::User, &user);

996 let mut stream = model

997 .stream_chat_request(req)

998 .await

@@ -1029,9 +1031,12 @@ mod tests {

1031 eprintln!("base Qwen3-0.6B not found locally; skipping");

1032 return;

1033 };

1032 - let data_path = dirs::home_dir()

1033 - .unwrap()

1034 - .join(".onde/datasets/Qwen/Qwen3-0.6B/train.jsonl");

1034 + let data_path = match std::env::var("ONDE_TEST_DATA_PATH") {

1035 + Ok(p) => PathBuf::from(p),

1036 + Err(_) => dirs::home_dir()

1037 + .unwrap()

1038 + .join(".onde/datasets/Qwen/Qwen3-0.6B/train.jsonl"),

1039 + };

1040 if !data_path.exists() {

1041 eprintln!("training data {} not found; skipping", data_path.display());

1042 return;

@@ -1055,8 +1060,14 @@ mod tests {

1060 .ok()

1061 .and_then(|v| v.parse().ok())

1062 .unwrap_or(1e-4),

1058 - epochs: 2,

1059 - max_seq_len: 512,

1063 + epochs: std::env::var("ONDE_TEST_EPOCHS")

1064 + .ok()

1065 + .and_then(|v| v.parse().ok())

1066 + .unwrap_or(2),

1067 + max_seq_len: std::env::var("ONDE_TEST_MAX_SEQ_LEN")

1068 + .ok()

1069 + .and_then(|v| v.parse().ok())

1070 + .unwrap_or(512),

1071 },

1072 ft_tx,

1073 );

@@ -1066,7 +1077,16 @@ mod tests {

1077 rt.block_on(async {

1078 while let Some(p) = ft_rx.recv().await {

1079 match p {

1069 - crate::finetune::FineTuneProgress::Training { loss, .. } => last_loss = loss,

1080 + crate::finetune::FineTuneProgress::Training {

1081 + loss,

1082 + step,

1083 + total_steps,

1084 + epoch,

1085 + ..

1086 + } => {

1087 + eprintln!("[e2e] epoch {epoch} step {step}/{total_steps} loss {loss:.4}");

1088 + last_loss = loss;

1089 + }

1090 crate::finetune::FineTuneProgress::Done { adapter_path: a } => {

1091 adapter_path = Some(a);

1092 break;

@@ -1122,15 +1142,20 @@ mod tests {

1142 )

1143 .expect("GGUF export should succeed");

1144

1125 - // 4. Load + run in mistral.rs. Query in the training distribution

1126 - // (the dataset is a Javanese-language assistant) so a lightly

1127 - // fine-tuned model has something to say.

1128 - let reply = run_gguf_inference(&gguf_path, "Apa bahasa Jawa dari \"terima kasih\"?");

1129 - eprintln!("[e2e] fine-tuned GGUF reply: {reply:?}");

1130 - assert!(

1131 - !reply.trim().is_empty(),

1132 - "fine-tuned GGUF produced an empty reply — not runnable"

1133 - );

1145 + // 4. Load + run in mistral.rs. Query in the training distribution so a

1146 + // lightly fine-tuned model has something to say. `ONDE_TEST_EVAL_PROMPT`

1147 + // overrides the default (multiple prompts separated by `||`).

1148 + let eval_prompts = std::env::var("ONDE_TEST_EVAL_PROMPT")

1149 + .unwrap_or_else(|_| "Apa bahasa Jawa dari \"terima kasih\"?".to_string());

1150 + for prompt in eval_prompts.split("||") {

1151 + let reply = run_gguf_inference(&gguf_path, prompt.trim());

1152 + eprintln!("[e2e] eval prompt: {prompt:?}");

1153 + eprintln!("[e2e] fine-tuned GGUF reply: {reply:?}");

1154 + assert!(

1155 + !reply.trim().is_empty(),

1156 + "fine-tuned GGUF produced an empty reply — not runnable"

1157 + );

1158 + }

1159 }

1160

1161 /// End-to-end check: export a Qwen3 safetensors model to GGUF, then load and