Stochastic AGPT (v2 cadence and depth controls)

Stochastic AGPT

Status: concluded, negative (reviewed 2026-09-28; this line previously read "active."). The answer is in the front matter above.

Historical result note (updated 2026-09-28): The June training runs below predate a fixed CUDA kernel race. A matched post-fix control reproduced the 100-epoch gap: whole-trie 12.054 versus root-child 5.348 rolling byte PPL at nearly equal training time. Use the post-fix runs for current comparisons; this page preserves the original runs.

Purpose: verify the CUDA v2 attention-based AGPT path before running a larger attention experiment.

Canonical smoke:

d_model=64
n_layers=2
n_heads=4
d_ff=256
seq_len=16
max_depth=16
partition_depth=1
epochs=1
optimizer=Adam
lr=0.0015
chunk_queries=50000

Run:

bin/agpt_experiment \
  --config rnd/stochastic-agpt/d64L2-depth16-pd1-1ep.yml \
  --trainer v2

Notes:

• This uses bin/agpt_train_v2, the attention AGPT trainer.
• The init checkpoint must have seq_len=16; using the older s128 seed is rejected by the trainer header check.
• If running under a sandbox, CUDA may be hidden. In that case run outside the sandbox or on the GPU pod.

Runs

20260611T155913-d64l2-depth16-pd1-1ep

Result: pass.

run dir: rnd/stochastic-agpt/20260611T155913-d64l2-depth16-pd1-1ep
train wall: 7.1s
total wall: 54.0s
lm_eval_rolling_byte_perplexity: 21.3199
lm_eval_rolling_bits_per_byte: 4.4141
agpt_fixed_token_perplexity: 21.3901

This validates the d64/L2 attention f_theta path at depth/seq_len 16 with pd=1, including CUDA training, checkpoint save, HF conversion, lm-eval rolling byte PPL, and AGPT fixed-token PPL.

20260611T160456-d64l2-depth16-pd1-100ep

Result: pass.

run dir: rnd/stochastic-agpt/20260611T160456-d64l2-depth16-pd1-100ep
train wall: 582.1s
total wall: 821.3s
final lm_eval_rolling_byte_perplexity: 5.3359
final lm_eval_rolling_bits_per_byte: 2.4157
final agpt_fixed_token_perplexity: 4.7929

Checkpoint curve:

The curve is still improving at 100 epochs, but the marginal gain from 80 to 100 is small. This is the pd1 baseline for the matching pd0 run.

20260611T162130-d64l2-depth16-pd0-100ep

Result: pass.

run dir: rnd/stochastic-agpt/20260611T162130-d64l2-depth16-pd0-100ep
train wall: 583.9s
total wall: 861.2s
final lm_eval_rolling_byte_perplexity: 12.0540
final lm_eval_rolling_bits_per_byte: 3.5914
final agpt_fixed_token_perplexity: 12.1095

Checkpoint curve:

pd0 vs pd1 at 100 Epochs

The wall time is essentially the same, but the optimizer cadence is not. With pd0, Adam only gets one update per full-tree epoch; with pd1, it gets one update per unigram subtree, about 65 updates per epoch. This run is strong evidence that plain Adam is a poor fit for full-tree pd0 training at this scale unless the optimizer/update rule is changed.

pd6 Descendant Sweep Control

Motivation: compare stochastic full-subtree sampling against a deterministic fine-grained sweep over all depth-6 anchors.

Implementation note: v2 now supports partition_depth > 1 by enumerating radix nodes whose edge covers the requested depth, then training that anchor plus its full descendant subtree. Ancestor-prefix rows are included only to seed/read the K/V cache and are marked context-only, so their query weights are zero.

The first row is invalid as a clean pd6 comparison because it repeatedly trains the shallow ancestor closure across 275k units. The context-only row is the clean control, and it is still far behind pd1 and Lightning random-descendants. A single deterministic pd6 pass produces many optimizer updates but too little effective target mass per update for this v2 attention setup.

Fixed-Depth Loss Control

Motivation: test whether the rolling-vs-fixed eval gap is partly caused by mixing all query depths during training. The control keeps the same depth-16 tree, seq_len=16, pd=1, Adam settings, and checkpoint cadence, but only one query depth contributes direct cross-entropy loss.

Important: the first depth-filter attempts below were run before bin/agpt_train_v2 was rebuilt with experimental.loss_depth_min/max support. Those runs ignored the depth filter and must not be treated as fixed-depth results. After rebuilding, a one-unit check showed the filter is active: loss_depth=1 trained 1 query row in the first unit, while loss_depth=7 trained 9,521 query rows in the same unit.

The first attempted loss_depth=2 run was stopped after epoch 30 so the first single-depth probe could target the more plausible middle depth. That run is also invalid for fixed-depth analysis because it used the stale trainer.

20260611T173031-d64l2-depth16-pd1-100ep-lossdepth7

Result: invalid for fixed-depth analysis; stale trainer ignored experimental.loss_depth_min/max.

run dir: rnd/stochastic-agpt/20260611T173031-d64l2-depth16-pd1-100ep-lossdepth7
loss depth: 7 only
train wall: 610.3s
total wall: 886.9s
final train loss: 1.6587
final lm_eval_rolling_byte_perplexity: 5.3115
final lm_eval_rolling_bits_per_byte: 2.4091
final agpt_fixed_token_perplexity: 4.7700

Checkpoint curve:

Against the normal pd1 baseline, this stale-binary run is slightly better on rolling byte PPL at every checkpoint and slightly better on fixed-token PPL except epoch 30, but it is not evidence about depth-7-only training:

This is best interpreted as another baseline-like run, not as a valid fixed-depth control. The fixed-depth sweep must use the rebuilt trainer.

Corrected Masked Sweep

After rebuilding bin/agpt_train_v2, experimental.loss_depth_min/max is honored. bin/agpt_experiment now also preflights v2 configs with the actual trainer binary before launching training, so stale binaries should fail before spending GPU time.

All rows below use the same d64/L2 depth-16 pd1 setup as the baseline, but only one query depth has nonzero direct CE loss.

Baseline comparison:

Findings:

• Single shallow depths are correctly terrible once the depth mask is actually honored. The stale depth-1 result was a binary/version problem, not a model result.
• The best single-depth rolling-byte score is depth 12, not the 6-8 range.
• The best fixed-token score is depth 16, which matches the fixed eval's 16-token context.
• No single-depth objective beats the normal pd1 baseline. Training all depths is still materially better for both held-out metrics at 100 epochs.