test+bins: flash grad-check, flash==composed, PyTorch parity, --flash flag
autograd: flash_attention_batched_bwd (dQ/dK/dV finite-diff, seq>tile) + flash_matches_composed_fwd. model/tests/flash.rs: flash==composed on-vs-off (logits/loss/every param grad), fp32 + bf16. parity_dump: XTRAIN_PARITY_FLASH dumps the flash path for the same parity.py oracle (PyTorch SDPA parity at B>1). train + train_ddp get the --flash flag. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -116,6 +116,9 @@ fn main() {
|
||||
// exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
|
||||
// default off stores every activation (unchanged numerics).
|
||||
let recompute = args.iter().any(|a| a == "--recompute");
|
||||
// Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
|
||||
// no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
|
||||
let flash = args.iter().any(|a| a == "--flash");
|
||||
let ckpt: PathBuf = PathBuf::from(
|
||||
args.iter()
|
||||
.position(|a| a == "--ckpt")
|
||||
@@ -183,6 +186,10 @@ fn main() {
|
||||
model = model.with_recompute(true);
|
||||
println!("activation recompute: ON (per-block gradient checkpointing)");
|
||||
}
|
||||
if flash {
|
||||
model = model.with_flash(true);
|
||||
println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
|
||||
}
|
||||
|
||||
// Eval-only mode: load a checkpoint and score it on the held-out val set, then
|
||||
// exit. Used to put an EXISTING model (e.g. v0) and a new one on the same
|
||||
|
||||
Reference in New Issue
Block a user