gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wip: T10 batched forward (validation)

Browse Source
This commit is contained in:
Gahow Wang
2026-06-16 00:19:26 +08:00
parent d2a585c5cb
commit ce9d22ffc2
13 changed files with 421 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
crates/xtrain-model/tests/parity_dump.rs
View File

@@ -53,12 +53,17 @@ fn dump_for_parity() {
);
fs::create_dir_all(&dir).unwrap();
// Fixed config + ids (independent of any text, for reproducibility).
// Fixed config + ids (independent of any text, for reproducibility). B>1 so
// the batched forward is exercised: 2 sequences of length 4, flattened
// sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
// sequence boundary) + per-sequence causal masking (no cross-sequence
// attention) are both checked against PyTorch.
let mut cfg = Config::tiny();
cfg.vocab = 12;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
let batch = 2usize;
let seq = 4usize;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
let seq = ids.len();
// Same deterministic init as the overfit test.
let mut seed = 1u64;
@@ -83,6 +88,7 @@ fn dump_for_parity() {
writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
writeln!(f, "eps {:e}", cfg.eps).unwrap();
writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
writeln!(f, "batch {batch}").unwrap();
writeln!(f, "seq {seq}").unwrap();
}
{
@@ -105,10 +111,11 @@ fn dump_for_parity() {
write_vec(&dir, &format!("w_{name}.txt"), &param_to_host(p), &shape);
}
// Forward logits + loss, then backward → per-param grads.
// Batched forward logits + loss (B sequences as one forward), then backward
// → per-param grads.
let ids_t = ids_tensor(&ids, device);
let targets_t = ids_tensor(&targets, device);
let logits = model.forward(&ids_t);
let logits = model.forward_batched(&ids_t, batch);
write_vec(
&dir,
"logits.txt",
@@ -116,7 +123,7 @@ fn dump_for_parity() {
logits.value().shape(),
);
let loss = model.loss(&ids_t, &targets_t);
let loss = model.loss_batched(&ids_t, &targets_t, batch);
let loss_val = param_to_host(&loss)[0];
{
let mut f = fs::File::create(dir.join("loss.txt")).unwrap();