@ondeinference / cli / commits / 344497f

Mask fine-tune loss to completions and add structured dataset format

Implement loss masking for instruction-tuning to address fact accuracy issues. The trainer previously computed cross-entropy loss over the entire sequence (system prompt + user turn + assistant answer), wasting gradient budget on predicting text that's always provided at inference time. Now the loss is masked to only assistant completions, concentrating gradient signal on what the model needs to learn to generate. Add backward-compatible structured dataset format (`messages: [...]`) alongside existing raw text format. Parse dataset entries as untagged enum: - `Raw { text }` variant: original format, full-sequence loss (existing behavior) - `Messages { messages }` variant: new structured format with turn boundaries Render structured messages as literal Qwen ChatML without Jinja engine (avoids new dependency). Compute per-token loss masks by checking byte-offset overlap with assistant spans. Apply masking via index_select to filter both logits and targets before cross_entropy. Preserve existing performance on Raw-format datasets via fast path that skips index_select when mask is all-true. Add comprehensive unit tests for dataset deserialization, ChatML rendering with UTF-8 handling, and loss mask computation. Edge cases documented: whitespace-normalized tokens straddling span boundaries (included in mask, correct but fuzzy), zero-length assistant content (masked false as it's empty), non-assistant roles (excluded from masking, safe). Expected outcome: better fact recall at same or fewer epochs because gradient budget concentrates on what the model actually needs to generate. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KfXzV7X5JZEenoH6MSdNtg

Claude committed Jul 10, 2026 at 17:08 UTC 344497f0f2a1d800b0d4a5b0c97814cab357241b

1 file changed +368 -17

src/finetune.rs

+368 -17

@@ -10,7 +10,7 @@ use std::io::{BufRead, BufReader};

10 use std::path::{Path, PathBuf};

12 use anyhow::{Context, Result};

13 -use candle_core::{D, DType, Device, Tensor, Var};

13 +use candle_core::{DType, Device, Tensor, Var, D};

14 use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder};

15 use tokenizers::Tokenizer;

@@ -194,7 +194,13 @@ fn run_finetune(

194 .context("reading training data")?;

195

196 let total_lines = raw_lines.len();

197 - let mut token_batches: Vec<Vec<u32>> = Vec::new();

197 +

198 + struct TrainingExample {

199 + tokens: Vec<u32>,

200 + mask: Vec<bool>,

201 + }

202 +

203 + let mut examples: Vec<TrainingExample> = Vec::new();

204

205 for (i, line) in raw_lines.iter().enumerate() {

206 let line = line.trim();

@@ -209,29 +215,75 @@ fn run_finetune(

215 let entry: DataEntry =

216 serde_json::from_str(line).with_context(|| format!("parsing JSONL line {}", i + 1))?;

217

218 + let (rendered_text, assistant_spans) = match &entry {

219 + DataEntry::Raw { text } => (text.clone(), Vec::new()),

220 + DataEntry::Messages { messages } => render_chatml(messages),

221 + };

222 +

223 let encoding = tokenizer

213 - .encode(entry.text.as_str(), false)

224 + .encode(rendered_text.as_str(), false)

225 .map_err(|e| anyhow::anyhow!("tokenizing line {}: {e}", i + 1))?;

226

227 let ids = encoding.get_ids();

217 - let ids: Vec<u32> = if ids.len() > config.max_seq_len {

218 - ids[..config.max_seq_len].to_vec()

228 + let offsets = encoding.get_offsets();

229 +

230 + // Truncate ids and offsets together to max_seq_len

231 + let (ids, offsets) = if ids.len() > config.max_seq_len {

232 + (

233 + ids[..config.max_seq_len].to_vec(),

234 + offsets[..config.max_seq_len].to_vec(),

235 + )

236 } else {

220 - ids.to_vec()

237 + (ids.to_vec(), offsets.to_vec())

238 };

239

240 // Need at least 2 tokens for next-token prediction.

224 - if ids.len() >= 2 {

225 - token_batches.push(ids);

241 + if ids.len() < 2 {

242 + eprintln!(

243 + "[finetune] WARNING: skipping line {} (too short: {} tokens)",

244 + i + 1,

245 + ids.len()

246 + );

247 + let _ = tx.send(FineTuneProgress::Tokenizing {

248 + done: i + 1,

249 + total: total_lines,

250 + });

251 + continue;

252 + }

253 +

254 + let mask = if assistant_spans.is_empty() {

255 + // Raw format: full-sequence loss (backward compatible)

256 + vec![true; ids.len()]

257 + } else {

258 + // Messages format: mask only assistant completions

259 + compute_loss_mask(&offsets, &assistant_spans)

260 + };

261 +

262 + // For Messages entries, check that the shifted mask has at least one true

263 + if !assistant_spans.is_empty() && mask.len() > 1 {

264 + let shifted_mask = &mask[1..];

265 + if !shifted_mask.iter().any(|&m| m) {

266 + eprintln!(

267 + "[finetune] WARNING: skipping line {} (assistant content truncated away)",

268 + i + 1

269 + );

270 + let _ = tx.send(FineTuneProgress::Tokenizing {

271 + done: i + 1,

272 + total: total_lines,

273 + });

274 + continue;

275 + }

276 }

277

278 + examples.push(TrainingExample { tokens: ids, mask });

279 +

280 let _ = tx.send(FineTuneProgress::Tokenizing {

281 done: i + 1,

282 total: total_lines,

283 });

284 }

285

234 - if token_batches.is_empty() {

286 + if examples.is_empty() {

287 anyhow::bail!("no valid training examples found in {:?}", config.data_path);

288 }

289

@@ -246,15 +298,15 @@ fn run_finetune(

298 let mut optimizer =

299 AdamW::new(lora_vars.clone(), adamw_params).context("creating AdamW optimizer")?;

300

249 - let total_steps = token_batches.len();

301 + let total_steps = examples.len();

302 let vocab_size = model_cfg.vocab_size;

303 let max_grad_norm: f64 = 1.0;

304

305 for epoch in 0..config.epochs {

254 - for (step, batch_tokens) in token_batches.iter().enumerate() {

255 - let seq_len = batch_tokens.len();

306 + for (step, example) in examples.iter().enumerate() {

307 + let seq_len = example.tokens.len();

308

257 - let input_ids = Tensor::from_slice(batch_tokens.as_slice(), (1, seq_len), &device)

309 + let input_ids = Tensor::from_slice(example.tokens.as_slice(), (1, seq_len), &device)

310 .context("building input_ids tensor")?;

311

312 let logits = model.forward(&input_ids).context("model forward pass")?;

@@ -265,10 +317,50 @@ fn run_finetune(

317 .and_then(|t| t.reshape((seq_len - 1, vocab_size)))

318 .context("preparing shifted logits")?;

319

268 - let targets = Tensor::from_slice(&batch_tokens[1..], (seq_len - 1,), &device)

320 + let targets = Tensor::from_slice(&example.tokens[1..], (seq_len - 1,), &device)

321 .context("building target tensor")?;

322

271 - let loss = candle_nn::loss::cross_entropy(&logits_shifted, &targets)

323 + let shifted_mask = &example.mask[1..];

324 +

325 + // Fast path: if all tokens are active (typical for Raw entries), skip index_select

326 + let (logits_masked, targets_masked) = if shifted_mask.iter().all(|&m| m) {

327 + (logits_shifted.clone(), targets.clone())

328 + } else {

329 + // Masked path: keep only the positions where mask is true

330 + let keep: Vec<u32> = shifted_mask

331 + .iter()

332 + .enumerate()

333 + .filter_map(|(i, &m)| if m { Some(i as u32) } else { None })

334 + .collect();

335 +

336 + if keep.is_empty() {

337 + eprintln!(

338 + "[finetune] WARNING: no masked tokens at epoch {} step {}, skipping",

339 + epoch + 1,

340 + step + 1

341 + );

342 + let _ = tx.send(FineTuneProgress::Training {

343 + epoch: epoch + 1,

344 + total_epochs: config.epochs,

345 + step: step + 1,

346 + total_steps,

347 + loss: 0.0,

348 + });

349 + continue;

350 + }

351 +

352 + let keep_idx = Tensor::from_slice(keep.as_slice(), (keep.len(),), &device)

353 + .context("building keep indices")?;

354 + let logits_masked = logits_shifted

355 + .index_select(&keep_idx, 0)

356 + .context("indexing logits")?;

357 + let targets_masked = targets

358 + .index_select(&keep_idx, 0)

359 + .context("indexing targets")?;

360 + (logits_masked, targets_masked)

361 + };

362 +

363 + let loss = candle_nn::loss::cross_entropy(&logits_masked, &targets_masked)

364 .context("computing cross-entropy loss")?;

365

366 let loss_val = loss.to_scalar::<f32>().context("reading loss scalar")?;

@@ -433,8 +525,60 @@ struct ModelConfig {

525 }

526

527 #[derive(serde::Deserialize)]

436 -struct DataEntry {

437 - text: String,

528 +#[serde(untagged)]

529 +enum DataEntry {

530 + Messages { messages: Vec<DatasetMessage> },

531 + Raw { text: String },

532 +}

533 +

534 +#[derive(serde::Deserialize)]

535 +struct DatasetMessage {

536 + role: String,

537 + content: String,

538 +}

539 +

540 +/// Renders `messages` as literal Qwen ChatML, matching the format written by

541 +/// hand in existing `Raw` datasets. Returns the rendered string plus byte-offset

542 +/// spans of every assistant turn (content start through that turn's closing

543 +/// `<|im_end|>`, not including the trailing `\n`) for loss-mask computation.

544 +fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) {

545 + let mut output = String::new();

546 + let mut assistant_spans = Vec::new();

547 +

548 + for msg in messages {

549 + output.push_str("<|im_start|>");

550 + output.push_str(&msg.role);

551 + output.push('\n');

552 +

553 + let content_start = output.len();

554 + output.push_str(&msg.content);

555 +

556 + if msg.role == "assistant" {

557 + let end_tag_start = output.len();

558 + output.push_str("<|im_end|>");

559 + assistant_spans.push((content_start, end_tag_start));

560 + } else {

561 + output.push_str("<|im_end|>");

562 + }

563 +

564 + output.push('\n');

565 + }

566 +

567 + (output, assistant_spans)

568 +}

569 +

570 +/// Pure function: true where a token's byte-offset range overlaps any

571 +/// assistant span. Decoupled from tokenizer types so it's unit-testable

572 +/// with synthetic data.

573 +fn compute_loss_mask(offsets: &[(usize, usize)], assistant_spans: &[(usize, usize)]) -> Vec<bool> {

574 + offsets

575 + .iter()

576 + .map(|&(token_start, token_end)| {

577 + assistant_spans

578 + .iter()

579 + .any(|&(span_start, span_end)| token_start < span_end && token_end > span_start)

580 + })

581 + .collect()

582 }

583

584 /// Top-level structure of `model.safetensors.index.json`.

@@ -1003,3 +1147,210 @@ impl LoraQwenModel {

1147 Ok(path)

1148 }

1149 }

1150 +

1151 +// ---------------------------------------------------------------------------

1152 +// Tests

1153 +// ---------------------------------------------------------------------------

1154 +

1155 +#[cfg(test)]

1156 +mod tests {

1157 + use super::*;

1158 +

1159 + #[test]

1160 + fn test_dataentry_raw_deserialization() {

1161 + let json = r#"{"text":"hello world"}"#;

1162 + let entry: DataEntry = serde_json::from_str(json).unwrap();

1163 + match entry {

1164 + DataEntry::Raw { text } => assert_eq!(text, "hello world"),

1165 + _ => panic!("expected Raw variant"),

1166 + }

1167 + }

1168 +

1169 + #[test]

1170 + fn test_dataentry_messages_deserialization() {

1171 + let json = r#"{"messages":[{"role":"user","content":"hi"}]}"#;

1172 + let entry: DataEntry = serde_json::from_str(json).unwrap();

1173 + match entry {

1174 + DataEntry::Messages { messages } => {

1175 + assert_eq!(messages.len(), 1);

1176 + assert_eq!(messages[0].role, "user");

1177 + assert_eq!(messages[0].content, "hi");

1178 + }

1179 + _ => panic!("expected Messages variant"),

1180 + }

1181 + }

1182 +

1183 + #[test]

1184 + fn test_dataentry_invalid() {

1185 + let json = r#"{"invalid":"field"}"#;

1186 + let result: Result<DataEntry, _> = serde_json::from_str(json);

1187 + assert!(result.is_err(), "should reject unrecognized format");

1188 + }

1189 +

1190 + #[test]

1191 + fn test_render_chatml_basic() {

1192 + let messages = vec![

1193 + DatasetMessage {

1194 + role: "system".to_string(),

1195 + content: "You are helpful".to_string(),

1196 + },

1197 + DatasetMessage {

1198 + role: "user".to_string(),

1199 + content: "Hello".to_string(),

1200 + },

1201 + DatasetMessage {

1202 + role: "assistant".to_string(),

1203 + content: "Hi there".to_string(),

1204 + },

1205 + ];

1206 +

1207 + let (rendered, spans) = render_chatml(&messages);

1208 +

1209 + assert!(

1210 + rendered.contains("<|im_start|>system"),

1211 + "should contain system marker"

1212 + );

1213 + assert!(

1214 + rendered.contains("You are helpful"),

1215 + "should contain system content"

1216 + );

1217 + assert!(

1218 + rendered.contains("<|im_start|>user"),

1219 + "should contain user marker"

1220 + );

1221 + assert!(rendered.contains("Hello"), "should contain user content");

1222 + assert!(

1223 + rendered.contains("<|im_start|>assistant"),

1224 + "should contain assistant marker"

1225 + );

1226 + assert!(

1227 + rendered.contains("Hi there"),

1228 + "should contain assistant content"

1229 + );

1230 +

1231 + assert_eq!(spans.len(), 1, "should have exactly one assistant span");

1232 + let (start, end) = spans[0];

1233 + assert!(rendered[start..end].contains("Hi there"));

1234 + assert!(rendered[end..].starts_with("<|im_end|>"));

1235 + }

1236 +

1237 + #[test]

1238 + fn test_render_chatml_multi_assistant() {

1239 + let messages = vec![

1240 + DatasetMessage {

1241 + role: "assistant".to_string(),

1242 + content: "First".to_string(),

1243 + },

1244 + DatasetMessage {

1245 + role: "user".to_string(),

1246 + content: "Second".to_string(),

1247 + },

1248 + DatasetMessage {

1249 + role: "assistant".to_string(),

1250 + content: "Third".to_string(),

1251 + },

1252 + ];

1253 +

1254 + let (_rendered, spans) = render_chatml(&messages);

1255 + assert_eq!(spans.len(), 2, "should have exactly two assistant spans");

1256 + }

1257 +

1258 + #[test]

1259 + fn test_render_chatml_no_assistant() {

1260 + let messages = vec![DatasetMessage {

1261 + role: "user".to_string(),

1262 + content: "Just a question".to_string(),

1263 + }];

1264 +

1265 + let (_rendered, spans) = render_chatml(&messages);

1266 + assert_eq!(spans.len(), 0, "should have no assistant spans");

1267 + }

1268 +

1269 + #[test]

1270 + fn test_render_chatml_utf8_indonesian() {

1271 + let messages = vec![

1272 + DatasetMessage {

1273 + role: "user".to_string(),

1274 + content: "Apa kabar? Ini ada diacritik: é".to_string(),

1275 + },

1276 + DatasetMessage {

1277 + role: "assistant".to_string(),

1278 + content: "Baik-baik saja. Jawaban dengan: ñ dan ü".to_string(),

1279 + },

1280 + ];

1281 +

1282 + let (rendered, spans) = render_chatml(&messages);

1283 + assert_eq!(spans.len(), 1, "should have one assistant span");

1284 +

1285 + let (start, end) = spans[0];

1286 + let assistant_content = &rendered[start..end];

1287 + assert!(assistant_content.contains("Baik-baik saja"));

1288 + assert!(assistant_content.contains("ñ"));

1289 + assert!(assistant_content.contains("ü"));

1290 + }

1291 +

1292 + #[test]

1293 + fn test_compute_loss_mask_all_inside() {

1294 + let offsets = vec![(10, 15), (15, 20), (20, 25)];

1295 + let spans = vec![(0, 30)];

1296 + let mask = compute_loss_mask(&offsets, &spans);

1297 +

1298 + assert_eq!(mask, vec![true, true, true]);

1299 + }

1300 +

1301 + #[test]

1302 + fn test_compute_loss_mask_all_outside() {

1303 + let offsets = vec![(0, 5), (5, 10), (40, 45)];

1304 + let spans = vec![(20, 30)];

1305 + let mask = compute_loss_mask(&offsets, &spans);

1306 +

1307 + assert_eq!(mask, vec![false, false, false]);

1308 + }

1309 +

1310 + #[test]

1311 + fn test_compute_loss_mask_straddling_start() {

1312 + let offsets = vec![(15, 25), (25, 35)];

1313 + let spans = vec![(20, 40)];

1314 + let mask = compute_loss_mask(&offsets, &spans);

1315 +

1316 + // First token [15, 25) overlaps [20, 40) → true

1317 + // Second token [25, 35) overlaps [20, 40) → true

1318 + assert_eq!(mask, vec![true, true]);

1319 + }

1320 +

1321 + #[test]

1322 + fn test_compute_loss_mask_straddling_end() {

1323 + let offsets = vec![(10, 25), (25, 35)];

1324 + let spans = vec![(15, 30)];

1325 + let mask = compute_loss_mask(&offsets, &spans);

1326 +

1327 + // First token [10, 25) overlaps [15, 30) → true

1328 + // Second token [25, 35) overlaps [15, 30) → true

1329 + assert_eq!(mask, vec![true, true]);

1330 + }

1331 +

1332 + #[test]

1333 + fn test_compute_loss_mask_zero_length_span() {

1334 + // A zero-length span (empty assistant content) should not match tokens

1335 + let offsets = vec![(10, 20), (20, 30)];

1336 + let spans = vec![(15, 15)]; // zero-length at byte 15

1337 + let mask = compute_loss_mask(&offsets, &spans);

1338 +

1339 + // Zero-length spans match nothing (start == end, so condition fails)

1340 + assert_eq!(mask, vec![false, false]);

1341 + }

1342 +

1343 + #[test]

1344 + fn test_compute_loss_mask_multiple_spans() {

1345 + let offsets = vec![(0, 5), (5, 10), (10, 15), (15, 20), (20, 25)];

1346 + let spans = vec![(2, 12), (18, 22)];

1347 + let mask = compute_loss_mask(&offsets, &spans);

1348 +

1349 + // [0, 5) overlaps [2, 12) → true

1350 + // [5, 10) overlaps [2, 12) → true

1351 + // [10, 15) overlaps [2, 12) → true

1352 + // [15, 20) overlaps [18, 22) → true

1353 + // [20, 25) overlaps [18, 22) → true

1354 + assert_eq!(mask, vec![true, true, true, true, true]);

1355 + }

1356 +}