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-distributed/src/bin/train_ddp.rs

@@ -88,6 +88,7 @@ fn main() {
     let val_tokens: usize = flag(&args, "--val-tokens", 0);
     let eval_every: usize = flag(&args, "--eval-every", 0);
     let eval_batches: usize = flag(&args, "--eval-batches", 64);
+    let sft_tsv = args.iter().any(|a| a == "--sft-tsv");
     // Dropout (Phase T18/T21): residual-path dropout prob, active at training time
     // only (inverted scaling), identity at eval/sampling/export. Default 0 = off
     // (forward graph bit-identical to the no-dropout path). Mirrors bin/train; the
@@ -109,6 +110,11 @@ fn main() {
         .position(|a| a == "--ckpt")
         .and_then(|i| args.get(i + 1))
         .map(PathBuf::from);
+    let init_ckpt: Option<PathBuf> = args
+        .iter()
+        .position(|a| a == "--init-ckpt")
+        .and_then(|i| args.get(i + 1))
+        .map(PathBuf::from);
 
     // Use every visible GPU as a rank (CUDA_VISIBLE_DEVICES selects the set;
     // device ordinals are 0..count within it).
@@ -129,12 +135,19 @@ fn main() {
     );
 
     // Reuse the cached token-id stream (v1's u16 cache); never re-tokenize 2GB.
-    let corpus = Corpus::load_cached(&tok_path, &corpus_path);
+    let corpus = if sft_tsv {
+        Corpus::load_sft_tsv_cached(&tok_path, &corpus_path)
+    } else {
+        Corpus::load_cached(&tok_path, &corpus_path)
+    };
     println!(
         "corpus: {} tokens, vocab {}",
         corpus.len(),
         corpus.vocab_size
     );
+    if sft_tsv {
+        println!("SFT TSV: ON (assistant-only loss via ignore-index labels)");
+    }
     let vocab = corpus.vocab_size;
     // Hold out a tail slice for validation (rank 0 evaluates on it).
     let (train_corpus, valid) = if val_tokens > 0 {
@@ -200,6 +213,10 @@ fn main() {
     if dropout > 0.0 {
         println!("dropout: ON (p={dropout}, residual-path, train-only inverted scaling)");
     }
+    if let Some(path) = &init_ckpt {
+        println!("init checkpoint: {}", path.display());
+    }
+    let init_ckpt_for_ranks = init_ckpt.clone();
     let results = launch(
         &devices,
         &train_corpus,
@@ -216,6 +233,10 @@ fn main() {
             if flash {
                 m = m.with_flash(true);
             }
+            if let Some(path) = &init_ckpt_for_ranks {
+                xtrain_train::checkpoint::load_into(path, &m.params())
+                    .expect("load init checkpoint");
+            }
             m
         },
     );

crates/xtrain-distributed/tests/ddp_correctness.rs

@@ -27,6 +27,7 @@ fn synth_corpus(vocab: usize, n_tokens: usize) -> Corpus {
         .collect();
     Corpus {
         tokens,
+        labels: None,
         vocab_size: vocab,
     }
 }

crates/xtrain-distributed/tests/ddp_proc.rs

@@ -37,6 +37,7 @@ fn synth_corpus() -> Corpus {
         .collect();
     Corpus {
         tokens,
+        labels: None,
         vocab_size: VOCAB,
     }
 }