sft: assistant-only SFT (ignore-index CE) + chat-prompt greedy eval

Enable assistant-only supervised fine-tuning and a fixed chat-prompt eval path
used by the v12 SFT runs:

- cross_entropy ignores negative targets (-100 ignore-index), normalizing by
  valid rows instead of all rows; CUDA fwd/bwd skip t<0 (ops.rs, nn.cu).
- Corpus gains optional labels + load_sft_tsv_cached: two-column TSV is
  formatted as 'User: .. \nAssistant:' + answer + <|endoftext|>, prompt tokens
  masked to -100 while answer+EOS are supervised; i32 label cache alongside the
  u16 token cache; sample() retries windows that are fully masked; eval uses
  target_window so masking applies to val loss too (data.rs, train_loop.rs).
- train + train_ddp: --sft-tsv selects the TSV loader, --init-ckpt continues
  training from a base checkpoint.
- greedy_sample: --prompts-file/--prompt/--temperature for fixed chat-prompt
  generation eval.

Test fixtures updated for the new Corpus.labels field; dropout.rs carries
incidental rustfmt. Not rebuilt locally (no CUDA toolchain on this checkout);
correctness rests on the documented v12 base+SFT runs on the GPU box.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-train/src/data.rs

@@ -15,6 +15,7 @@ use xserv_tokenizer::Tokenizer;
 /// A tokenized corpus: one flat stream of token ids, plus the vocab size.
 pub struct Corpus {
     pub tokens: Vec<i32>,
+    pub labels: Option<Vec<i32>>,
     pub vocab_size: usize,
 }
 
@@ -33,6 +34,7 @@ impl Corpus {
         let ids: Vec<i32> = tok.encode(text).into_iter().map(|t| t as i32).collect();
         Self {
             tokens: ids,
+            labels: None,
             vocab_size: tok.vocab_size(),
         }
     }
@@ -52,7 +54,11 @@ impl Corpus {
                 tokens.len(),
                 cache.display()
             );
-            return Self { tokens, vocab_size };
+            return Self {
+                tokens,
+                labels: None,
+                vocab_size,
+            };
         }
         let me = Self::load(tokenizer_path, corpus_path);
         write_u16_cache(&cache, &me.tokens);
@@ -64,22 +70,104 @@ impl Corpus {
         me
     }
 
+    /// Load assistant-only SFT data from a two-column TSV:
+    ///
+    /// ```text
+    /// user<TAB>assistant
+    /// ```
+    ///
+    /// Literal `\n` and `\t` escapes are decoded. Each row is formatted as
+    /// `User: ...\nAssistant:` + answer + `<|endoftext|>`. Labels are `-100`
+    /// for prompt tokens and the token id itself for answer/EOS tokens, so the
+    /// cross-entropy op ignores prompt rows while still training the assistant
+    /// answer and stop token.
+    pub fn load_sft_tsv_cached(tokenizer_path: &Path, corpus_path: &Path) -> Self {
+        let token_cache = cache_path(corpus_path);
+        let label_cache = label_cache_path(corpus_path);
+        let vocab_size = Tokenizer::from_file(tokenizer_path).vocab_size();
+        if token_cache.exists() && label_cache.exists() {
+            let tokens = read_u16_cache(&token_cache);
+            let labels = read_i32_cache(&label_cache);
+            assert_eq!(
+                tokens.len(),
+                labels.len(),
+                "SFT cache token/label length mismatch"
+            );
+            println!(
+                "corpus: read {} cached SFT tokens from {} (+ labels {})",
+                tokens.len(),
+                token_cache.display(),
+                label_cache.display()
+            );
+            return Self {
+                tokens,
+                labels: Some(labels),
+                vocab_size,
+            };
+        }
+
+        let tok = Tokenizer::from_file(tokenizer_path);
+        let text = std::fs::read_to_string(corpus_path)
+            .unwrap_or_else(|e| panic!("failed to read SFT corpus {}: {e}", corpus_path.display()));
+        let mut tokens = Vec::new();
+        let mut labels = Vec::new();
+        for (lineno, line) in text.lines().enumerate() {
+            if line.trim().is_empty() {
+                continue;
+            }
+            let (user, assistant) = line
+                .split_once('\t')
+                .unwrap_or_else(|| panic!("SFT TSV line {} missing tab", lineno + 1));
+            let user = decode_tsv_escapes(user);
+            let assistant = decode_tsv_escapes(assistant);
+            let prompt = format!("User: {user}\nAssistant:");
+            let answer = format!(" {assistant}\n<|endoftext|>");
+            let prompt_ids: Vec<i32> = tok.encode(&prompt).into_iter().map(|t| t as i32).collect();
+            let answer_ids: Vec<i32> = tok.encode(&answer).into_iter().map(|t| t as i32).collect();
+            labels.extend(std::iter::repeat(-100).take(prompt_ids.len()));
+            labels.extend(answer_ids.iter().copied());
+            tokens.extend(prompt_ids);
+            tokens.extend(answer_ids);
+        }
+        assert_eq!(tokens.len(), labels.len(), "SFT tokens/labels mismatch");
+        write_u16_cache(&token_cache, &tokens);
+        write_i32_cache(&label_cache, &labels);
+        println!(
+            "corpus: tokenized {} SFT tokens → cached to {} (+ labels {})",
+            tokens.len(),
+            token_cache.display(),
+            label_cache.display()
+        );
+        Self {
+            tokens,
+            labels: Some(labels),
+            vocab_size: tok.vocab_size(),
+        }
+    }
+
     /// Split off the last `n` tokens as a held-out validation corpus, leaving the
     /// rest as the train corpus. Returns `(train, valid)`. Used for periodic val
     /// loss during training without leaking the eval window into training.
     pub fn split_tail(self, n: usize) -> (Self, Self) {
         let n = n.min(self.tokens.len() / 10); // never hand off more than 10%
         let cut = self.tokens.len() - n;
-        let valid = self.tokens[cut..].to_vec();
+        let valid_tokens = self.tokens[cut..].to_vec();
+        let valid_labels = self.labels.as_ref().map(|labels| labels[cut..].to_vec());
         let mut train = self.tokens;
         train.truncate(cut);
+        let train_labels = self.labels.map(|mut labels| {
+            labels.truncate(cut);
+            labels
+        });
         (
             Self {
                 tokens: train,
+                labels: train_labels,
                 vocab_size: self.vocab_size,
             },
             Self {
-                tokens: valid,
+                tokens: valid_tokens,
+                labels: valid_labels,
                 vocab_size: self.vocab_size,
             },
         )
@@ -101,11 +189,27 @@ impl Corpus {
     pub fn sample(&self, seq: usize, rng_state: &mut u64) -> (Vec<i32>, Vec<i32>) {
         assert!(self.tokens.len() > seq + 1, "corpus shorter than a window");
         let max_start = self.tokens.len() - seq - 1;
-        let start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+        let mut start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+        if let Some(labels) = &self.labels {
+            for _ in 0..16 {
+                if labels[start + 1..start + seq + 1].iter().any(|&t| t >= 0) {
+                    break;
+                }
+                start = (next_rand(rng_state) % (max_start as u64 + 1)) as usize;
+            }
+        }
         let input = self.tokens[start..start + seq].to_vec();
-        let target = self.tokens[start + 1..start + seq + 1].to_vec();
+        let target = self.target_window(start, seq);
         (input, target)
     }
+
+    /// Deterministic target labels for an input window starting at `start`.
+    pub fn target_window(&self, start: usize, seq: usize) -> Vec<i32> {
+        match &self.labels {
+            Some(labels) => labels[start + 1..start + seq + 1].to_vec(),
+            None => self.tokens[start + 1..start + seq + 1].to_vec(),
+        }
+    }
 }
 
 /// Drop a leading partial line (before the first newline) and everything after
@@ -127,6 +231,12 @@ fn cache_path(corpus_path: &Path) -> PathBuf {
     PathBuf::from(s)
 }
 
+fn label_cache_path(corpus_path: &Path) -> PathBuf {
+    let mut s = corpus_path.as_os_str().to_os_string();
+    s.push(".labels.i32.bin");
+    PathBuf::from(s)
+}
+
 /// Read a flat little-endian `[u16]` cache into an `i32` id stream.
 fn read_u16_cache(path: &Path) -> Vec<i32> {
     let mut r = BufReader::new(
@@ -140,6 +250,18 @@ fn read_u16_cache(path: &Path) -> Vec<i32> {
         .collect()
 }
 
+fn read_i32_cache(path: &Path) -> Vec<i32> {
+    let mut r = BufReader::new(
+        std::fs::File::open(path).unwrap_or_else(|e| panic!("open cache {}: {e}", path.display())),
+    );
+    let mut buf = Vec::new();
+    r.read_to_end(&mut buf).expect("read cache");
+    assert!(buf.len() % 4 == 0, "corrupt i32 cache (odd byte count)");
+    buf.chunks_exact(4)
+        .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
+        .collect()
+}
+
 /// Write an id stream as a flat little-endian `[u16]` cache. Ids must fit in u16
 /// (GPT-2 vocab = 50257 < 65536); asserts otherwise.
 fn write_u16_cache(path: &Path, tokens: &[i32]) {
@@ -154,6 +276,21 @@ fn write_u16_cache(path: &Path, tokens: &[i32]) {
     w.flush().expect("flush cache");
 }
 
+fn write_i32_cache(path: &Path, labels: &[i32]) {
+    let mut w = BufWriter::new(
+        std::fs::File::create(path)
+            .unwrap_or_else(|e| panic!("create cache {}: {e}", path.display())),
+    );
+    for &t in labels {
+        w.write_all(&t.to_le_bytes()).expect("write cache");
+    }
+    w.flush().expect("flush cache");
+}
+
+fn decode_tsv_escapes(s: &str) -> String {
+    s.replace("\\n", "\n").replace("\\t", "\t")
+}
+
 /// Tiny LCG (same constants as the model tests' deterministic fill) so dataset
 /// sampling is reproducible from a single u64 seed.
 fn next_rand(state: &mut u64) -> u64 {