Include closing <|im_end|> in assistant loss mask
The completion mask ended the assistant span just before <|im_end|>, so the stop marker never contributed to the loss and the model wasn't trained to emit it — generations wouldn't terminate cleanly. The docstring and the zero-length-span test already assumed the marker was in the span; only the code disagreed. Extend the span to cover <|im_end|> and add a test that the marker's byte range lands inside the span and stays masked whether it tokenizes as one special token or splits into several (the add_special_tokens=false case a reviewer flagged). Also run rustfmt to fix the import-ordering CI failure. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
keypair34 committed
Jul 11, 2026 at 08:00 UTC
a74062c76c6675062608885119983b70716892a5
1 file changed
+56
-6
src/finetune.rs
+56
-6
@@ -10,7 +10,7 @@ use std::io::{BufRead, BufReader};
10
use std::path::{Path, PathBuf};
11
12
use anyhow::{Context, Result};
13
-use candle_core::{DType, Device, Tensor, Var, D};
13
+use candle_core::{D, DType, Device, Tensor, Var};
14
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder};
15
use tokenizers::Tokenizer;
16
@@ -539,8 +539,9 @@ struct DatasetMessage {
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.
542
+/// spans of every assistant turn (content start through the end of that turn's
543
+/// closing `<|im_end|>`, which is included so the model learns to emit it; the
544
+/// trailing `\n` is not) for loss-mask computation.
545
fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) {
546
let mut output = String::new();
547
let mut assistant_spans = Vec::new();
@@ -554,9 +555,14 @@ fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) {
555
output.push_str(&msg.content);
556
557
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));
559
+ // Span end is *after* `<|im_end|>` so the closing marker is masked
560
+ // in for loss: the model has to learn to emit it, otherwise
561
+ // generations never terminate. Whether the marker tokenizes as one
562
+ // special token or splits into several is irrelevant here — every
563
+ // resulting token's byte offset falls inside this span.
564
+ let span_end = output.len();
565
+ assistant_spans.push((content_start, span_end));
566
} else {
567
output.push_str("<|im_end|>");
568
}
@@ -1230,8 +1236,12 @@ mod tests {
1236
1237
assert_eq!(spans.len(), 1, "should have exactly one assistant span");
1238
let (start, end) = spans[0];
1239
+ // The span covers the assistant content and its closing `<|im_end|>`,
1240
+ // so the model is trained to emit the stop marker.
1241
assert!(rendered[start..end].contains("Hi there"));
1234
- assert!(rendered[end..].starts_with("<|im_end|>"));
1242
+ assert!(rendered[start..end].ends_with("<|im_end|>"));
1243
+ // Only the trailing newline remains after the span.
1244
+ assert!(rendered[end..].starts_with('\n'));
1245
}
1246
1247
#[test]
@@ -1359,4 +1369,44 @@ mod tests {
1369
// [20, 25) overlaps [18, 22) → true
1370
assert_eq!(mask, vec![true, true, true, true, true]);
1371
}
1372
+
1373
+ #[test]
1374
+ fn test_assistant_span_masks_closing_im_end() {
1375
+ // The closing `<|im_end|>` of an assistant turn must be masked in for
1376
+ // loss so the model learns to stop. Its byte range must fall inside the
1377
+ // reported span, and every token it produces must be masked true —
1378
+ // whether the marker tokenizes as one special token or splits into
1379
+ // several literal pieces (which can happen with add_special_tokens=false).
1380
+ let messages = vec![DatasetMessage {
1381
+ role: "assistant".to_string(),
1382
+ content: "Hi".to_string(),
1383
+ }];
1384
+
1385
+ let (rendered, spans) = render_chatml(&messages);
1386
+ assert_eq!(spans.len(), 1);
1387
+ let (span_start, span_end) = spans[0];
1388
+
1389
+ let marker = "<|im_end|>";
1390
+ let marker_start = rendered.find(marker).expect("marker present");
1391
+ let marker_end = marker_start + marker.len();
1392
+
1393
+ // The marker's whole byte range sits within the assistant span.
1394
+ assert!(marker_start >= span_start && marker_end <= span_end);
1395
+
1396
+ // Single-token marker → masked.
1397
+ let single = compute_loss_mask(&[(marker_start, marker_end)], &spans);
1398
+ assert_eq!(single, vec![true]);
1399
+
1400
+ // Split marker (e.g. "<|", "im", "_end", "|>") → every piece masked.
1401
+ let split = compute_loss_mask(
1402
+ &[
1403
+ (marker_start, marker_start + 2),
1404
+ (marker_start + 2, marker_start + 4),
1405
+ (marker_start + 4, marker_start + 8),
1406
+ (marker_start + 8, marker_end),
1407
+ ],
1408
+ &spans,
1409
+ );
1410
+ assert_eq!(split, vec![true, true, true, true]);
1411
+ }
1412
}