Pokémon TCG AI Battle — MLX Training Pipeline (current)

August 14 audit pointer

The runtime contract below is preserved as the August 7 training baseline. For the current source-versus-blueprint boundary, use pokemon_tcg_aug14_architecture_and_handoff_audit. For current Parquet, SQLite and ETL observations, use pokemon_tcg_aug14_data_etl_database_audit. For the later FP32 curriculum and tournament cohorts, use pokemon_tcg_aug14_ablations_tournaments. The future RoPE-ND and MoE phases remain documented and are not silently removed from the project record.

Current source correction at `20d7d0d`

The body below preserves the FP16 delivery snapshot for historical reproducibility. The live trainer now calls model.set_dtype(mx.float32) and rejects any non-FP32 model leaf. The live PyTorch converter likewise validates strict FP32 tensors. Use pokemon_tcg_current_state_reconciliation for the current source contract and this page's dated sections for the historical transition.

The current trainer also carries scratch memory through TBPTT lanes and agent/main.py carries memory between autoregressive action substeps. Statements below that describe discarded scratch output, one-pass topk(count) selection or FP16 as the current runtime are historical and have been superseded by the source-backed pages.

Boundary

This is the canonical reference for the training pipeline running on develop in ~/workdir/pokemon-tcg as of 2026-08-07. It supersedes the phased narrative in pokemon_tcg_mlx_migration, which is now preserved as a migration record. Trainer entry point: scripts/bc/bc_train_mlx.py (invoked as uv run tcg-train).

Runtime split is fixed:

training      ->  MLX on Apple Silicon (M3 Pro, 24 GiB unified memory), current model FP32
inference     ->  PyTorch FP32, arena-side, self-contained submission
recurrence    ->  TBPTT (always on)

No other paths are supported for this phase.

Runtime contract

aspect	commitment
device	`mlx.core.gpu` (Metal)
dtypes	Current source: FP32 parameters, activations, embeddings and QKV; FP32 loss, reductions, accumulation and optimizer moments. Historical FP16 delivery is retained below.
positional info	zone-typed token embeddings + explicit `opt_src_pos`/`opt_tgt_pos` reference gather; no RoPE / sinusoidal / RoPE-ND (RoPE-ND is pokemon_tcg_ladder_and_research backlog, not implemented)
padding	`padding_idx=0` semantics reproduced explicitly (`ids != 0` mask; zero vector, not a learned "absence" embedding)
attention mask	additive (`0` valid / large negative padded), FP16-safe. Not boolean.
MHA bias	on in the attention projections (`bias=True`, matching the PyTorch reference)
static features	`EN_Card_Data.csv` loaded once as a numpy buffer; a learned `Linear` projects it into `d_model` space; the table itself is not a trainable parameter

End-to-end flow

sqlite: days ⋈ datasets
        │
        ▼
   day list resolved
        │  (filters: --top-elo, --max-rows-per-day, --max-rows)
        ▼
pyarrow.dataset over N parquet files
        │
        ▼
_ParquetRowGroupCache (train)   _ParquetRowGroupCache (val)
        │                                │
        ▼                                ▼
TBPTT lane packer (episode_id, side)    val temporal batcher (same)
        │                                │
        ▼                                ▼
     forward     ←──── model (current FP32 params/acts, split heads, aux heads)
        │                                │
        ▼                                ▼
    FP32 loss  (CE + weighted aux)     FP32 loss + metrics
        │                                │
        ▼                                │
   backward → FP32 grad accum            │
        │                                │
        ▼                                │
 in_opt_step():                          │
   clip → optimizer.update → mx.eval     │
        │                                │
        ▼                                ▼
   tensorboard scalar per opt_step     tensorboard scalar per epoch

Data resolution

The trainer never scans directories directly. It reads day-partitioned Parquet paths from the SQLite catalog (datasets ⋈ days; see pokemon_tcg_sqlite_schema_current) and refuses to start if the resolved list is empty or a resolved path is missing on disk.

Day selection is one of three mutually exclusive flags:

--days 2026-07-30,2026-08-01 (explicit list),
--last-n-days 5 (N most recently registered),
--all-days (every day in the catalog).

Row-level filters, applied after day resolution:

--top-elo N — episode-level filter against agent_elo_daily WHERE source='remote'; both player_name and opponent_name must be in the day's top-N agents. See pokemon_tcg_top_elo_curriculum_filter.
--max-rows-per-day N — per-day cap rounded down to the nearest episode boundary.
--max-rows N — global cap across all resolved days.

Parquet layout, column schema, and streaming reads are documented in pokemon_tcg_parquet_dataset.

KV cache

Two _ParquetRowGroupCache instances (train and val) hold row groups resident between microbatches. Three tiers (hot pinned / transient LRU / SSD spill), pressure-driven eviction at 85% host memory, promotion after 2 hits, opt-step protection wrapped around optimizer.update. Full mechanics in pokemon_tcg_kv_cache_hierarchical.

Model

Instantiated via build_token_net_mlx(ct, net_cfg). Architecture and token schema in pokemon_tcg_agent_architecture. Baseline hyperparameters (config default): d_model=128, nhead=4, nlayers=3, ff_dim=512, scratch_registers=16, static=true, split_heads=true, structured=false. Current session config runs nlayers=4, scratch_registers=32.

Parameter counts around this size: ~1.30M total. The historical August 7 baseline set the model to FP16. The current trainer sets it to FP32 via model.set_dtype(mx.float32) and refuses to start unless every model leaf is mlx.core.float32.

Auxiliary heads

Four Linear heads on top of the CLS output (or VALUE_TOK when split_heads=true):

head	activation	target column	loss	flag
`ko_head_aux`	scalar	`aux_ko`	BCE	`--aux-ko-weight`
`prize_head_aux`	scalar	`aux_prize_delta`	MSE	`--aux-prize-weight`
`terminal_head_aux`	scalar	`aux_terminal`	BCE	`--aux-terminal-weight`
`return_head_aux`	scalar	`aux_return`	MSE	`--aux-return-weight`

All four are masked by aux_valid per row. Weights are fixed at whatever the CLI/config provides — there is no automatic uncertainty weighting yet (Kendall & Gal is pokemon_tcg_ladder_and_research backlog).

The current source has an important reduction detail: _aux_loss returns a weighted sum over valid rows despite a stale docstring describing a mean, while validation metrics use masked means per head. The optimizer later divides accumulated gradients by total examples. This scale distinction is part of the Stage 3/4 incident in pokemon_tcg_stage3_training_failure_postmortem and must not be described as a fully normalized multi-task objective.

The dataset builder is required to emit aux_valid, aux_ko, aux_prize_delta, aux_terminal, aux_return. A day without these will not train.

Optimizer: Muon + AdamW routing

--optimizer muon_adamw is the only supported topology. Parameters split by shape:

Muon — hidden 2D matrix weights (attention QKV+O, FFN linears, projections). Currently ~845k params.
AdamW — everything else: embeddings, output heads, type_query/type_bias, scalar/vector parameters, layer norm gains/biases. ~456k params.

Total: ~1.30M trainable. Routing is decided by walking nn.utils.tree_flatten(model.trainable_parameters()) and applying _use_muon_parameter(path, parameter) per leaf. The rule is stable across resumes and stored in the checkpoint's optimizer_contract.

Weight decay: muon_weight_decay for Muon leaves, adamw_weight_decay for AdamW leaves, structured_weight_decay overriding the AdamW value on type_query / type_bias when --structured true.

The optimizer routing is on tensorboard at run start:

optimizer routing: Muon=845,568 hidden-matrix params;
                   AdamW=456,583 embedding/head/vector params

Gradient accumulation

For accum_steps=K:

for _ in K microbatches:
    forward → FP32 CE + aux loss → backward → FP32 grad add
normalize by real example count (including a partial last microbatch)
clip once (max_grad_norm)
optimizer.update
scheduler.step (counts optimizer updates, not forward passes)
mx.eval(model.parameters(), optimizer.state)  # materialize

Current baseline uses accum_steps=1 (one microbatch per opt step). Accumulation is verified to normalize by the real example count even when the last microbatch is short.

Scheduler

The LR scheduler counts optimizer updates, not forward passes:

total_optimizer_steps = epochs × ceil(microbatches / accum_steps)

warmup_steps is clamped to max(total_optimizer_steps / 5, 1) so short runs don't spend the whole schedule warming up. --lr-min-ratio sets the LR floor after decay completes (Orbit-style hold at floor).

On --scheduler-state resume, the scheduler restores scheduler_phase_step from the checkpoint and refuses to start if the resumed horizon does not have enough remaining steps for the current run.

TBPTT

Sequential training over (episode_id, side) lanes in ordered chunks (default 16 decisions). At chunk boundaries: stop_gradient on scratch memory carry, no backprop through the boundary. Reset flag from the row's new_episode column resets scratch memory at match transitions. Full recurrent accounting in pokemon_tcg_tbptt_training_contract.

The trainer's temporal batcher packs lanes into a _TBPTTChunk list until the --batch row budget is hit. Val uses the same lane packer with --val-batch-size as its budget.

Instrumentation

Tensorboard event files at runs/<tag>_<epoch_timestamp>/. Every optimizer step writes:

train/loss — primary CE per opt step
train/aux_loss — combined weighted aux loss per opt step
train/grad_norm — post-clip
train/lr — scheduler-emitted LR
train/scheduler_phase_step
train/step_time_ms
train/examples_per_step
sys/mlx_peak_memory_gib — from mx.metal.get_peak_memory()
sys/host_memory_percent — from psutil.virtual_memory().percent
cache_train/* and cache_val/* — hit_rate_pct, resident_hot, resident_transient, promotions, evictions, ssd_hits, ssd_spills, ssd_resident (see pokemon_tcg_kv_cache_hierarchical)

Every epoch writes:

train/running_loss / train/running_aux_loss / train/gstep
val/loss / val/acc / val/equiv / val/top3 / val/atk / val/ko
val/epoch_time_s
aux/aux_ko_bce / aux/aux_prize_mse / aux/aux_terminal_bce / aux/aux_return_mse (on val)

At run end: summary/best_val_acc, summary/final_gstep.

41 scalar tags total. Sample of one full run (5d_10ep_OFF, 10 epochs, 1410 optimizer steps) is inspectable via:

uv run tensorboard --logdir runs

Val loop

Val split is deterministic given --seed, --val-frac and the resolved day list — it falls on episode boundaries, not row boundaries, so train and val are guaranteed episode-disjoint. Val batches stream through their own KV cache (see pokemon_tcg_kv_cache_hierarchical) — nothing is materialized to RAM. The old ~22 GiB val materialization regime is gone.

Val metrics computed:

CE loss — proper cross-entropy: -(logit[label] - logsumexp(logits)). Not the historical log(raw_logit).
acc — top-1 accuracy on legal-masked logits.
equiv — top-1 among options in the same equivalence class as the label (uses opt_group).
top3 — top-3 accuracy on legal-masked logits.
atk — accuracy conditional on is_attack.
ko — aux_ko_head accuracy vs aux_ko target.

Val runs at the end of every epoch. Best-val checkpoint is saved to --out; the latest is always rolled to <--out>_latest.pkl.

Checkpoints

The checkpoint payload is pickled (not mx.savez, which suffers from flatten/unflatten name mismatches) and carries the full state to resume:

model                    (mlx params, current source FP32; historical checkpoints may be FP16)
optimizer                (mlx optimizer.state)
optimizer_contract       (routing identity)
optimizer_phase_step
arch_config              (model.get_config())
static_card_features     (buffer)
static_feature_contract  (sha256 of the loaded csv)
run_config               (resolved cfg.to_dict + data_days + data_paths + zero_wouldko)
inference_config         (seed, bc_would_ko, wk_nvar, provenance)
dataset_manifest         (per-day sha256, size, aux-target flag)
dataset_build_fingerprint
phase_id                 (from --phase-id)
epoch                    (0-indexed loop var)
gstep                    (global optimizer-step counter)
val_acc / best_val_acc
seed
accum_steps
microbatches_per_epoch / optimizer_steps_per_epoch
scheduler_phase_step / scheduler_total_steps / scheduler_contract / scheduler_state

On --resume PATH --optimizer-state resume --scheduler-state resume: model params load, optimizer.state loads, scheduler_phase_step restores, start_epoch = state.epoch + 1. The resume path validates arch_config matches the current build; a mismatch raises rather than silently loading incompatible weights.

Reproducibility

Seed applies to numpy AND mx.random. A previous bug seeded only numpy, causing untracked variance across runs even at identical CLI. Fixed.

np.random.seed(seed) + mx.random.seed(seed) at trainer entry. --seed 13 is the current session default.

Historical context — what came before

Reading this page in isolation makes today's pipeline look like a single monolithic design. It is not. It is the fifth-generation shape after several substantial pivots the wiki still documents in full:

The sidecar era (2026-07-29) — a separate second model (the "prospective V2" planner with RoPE-ND positional encoding) was designed and shipped to score counterfactual continuations at inference. It ran as an offline .npy sidecar joined into the training set. Full design in pokemon_tcg_prospective_v2; the reasoning for its rejection is captured in the private project memory project-pokemon-tcg-sidecar. Take-away that survived: any prospective-planning signal that returns will be an integrated aux head or an offline GRPO objective, not a second separate model. Take-away that did not survive: the sidecar's dependency footprint, its .npy on-disk format, and its RoPE-ND positional encoding — none is present in the current runtime.
The NPY dataset era (up to 2026-08-03) — the dataset builder wrote .npy shards per episode with an episode_meta.npy sidecar for (episode_id, side, step_id) metadata. Trainer read via memory-mapped slabs (slab_rows was a live knob). Full contract in pokemon_tcg_data_pipeline. Replaced 2026-08-03 by day-partitioned Parquet with metadata as first-class row columns (pokemon_tcg_parquet_dataset). Take-away that survived: the aux-target contract, the episode-boundary val split, would-KO. Take-away that did not survive: the shard layout, the episode_meta.npy sidecar, and the mmapped-slab abstraction (replaced by the Parquet row-group + KV cache combination).
The flat KV cache (2026-08-06 morning) — commit e772fe0 introduced a single-tier row-group cache with FIFO eviction. Rewritten a few hours later (fa38caa) into the current three-tier hierarchical design after we observed under-load thrash. Full mechanics in pokemon_tcg_kv_cache_hierarchical.

The commit-level version of this story is in pokemon_tcg_repository_timeline.

Failure modes not observed in the historical suite

Evidence-based, from the pokemon_tcg_bc_curriculum_ablation runs (10 checkpoints × 10 epochs × ~150–1400 optimizer steps each):

failure mode	evidence it did not happen
encoding bug	training loss decreased monotonically on every run (5d_10ep_OFF: 3.49 → 1.12)
MLX → PyTorch conversion drift	tournaments emitted zero opponent errors in ~5,400 games; agent always chose legal actions
loss misalignment	val_acc improved with data and epochs (up to 0.61)
FP16 overflow/underflow	no NaN observed in any run; `train/grad_norm` bounded
KV cache corruption	val_acc improved monotonically; cache hit rate 99.9% across all runs
optimizer contract broken	aux heads (`ko_bce`, `prize_mse`, `terminal_bce`, `return_mse`) decreased on every run
collapsed model	intra-suite peer round-robin sits at 44–53% across all 10 models — normal Elo dispersion

Within the August 6-7 suite, these observations supported a data- or regime-limited interpretation rather than a broad architecture failure. They do not rule out later objective or validation defects: the Stage 3/4 incident exposed exactly such a training-contract problem. Discussion and cohort scope are recorded in pokemon_tcg_bc_curriculum_ablation and pokemon_tcg_stage3_training_failure_postmortem.

Provenance / anchor commits

9e92d1d (2026-07-25) — feat(C): FP16-native trainer with gradient accumulation and complete checkpoints — landed the FP16-end-to-end contract and the resumable checkpoint payload.
75ae036 (2026-07-26) — feat(F): memory API with persistent registers + TBPTT support — landed TBPTT and the scratch-register memory API. Follow-up fixes d4da903, f5a4bc8, 6f7067b, b216f4f.
2e8fd34 (2026-07-26) — feat: centralized config module + Kaggle data manager with entrypoints — landed TrainConfig, tcg-train, tcg-data.
a942373 (2026-08-03) — sidecar removed, aux heads + meta features + parquet pipeline — the pivot commit: switched the dataset to Parquet, added the four aux heads, integrated meta buckets, removed the prospective sidecar.
504118d (2026-08-03) — smoke pipeline validated: strict semantics, competition_day, streaming TBPTT — end-to-end validation of the post-pivot pipeline on real smoke data.
e772fe0 (2026-08-06) — train loader: KV-style parquet cache, async prefetch, top-elo/per-day filters, tensorboard — added the row-group cache (flat), single-worker prefetch, --top-elo and --max-rows-per-day, and full per-optimizer-step tensorboard.
fa38caa (2026-08-06) — train loader: hierarchical KV cache + opt-step protection + val streaming — replaced the flat cache with the three-tier hierarchical design and moved val to stream through its own cache.
290d6f9 (2026-08-07) — CLAUDE.md: current phase log — snapshot of the current state that this page mirrors.

Full timeline in pokemon_tcg_repository_timeline.

Cross-references

pokemon_tcg_glossary — vocabulary
pokemon_tcg_agent_architecture — model architecture
pokemon_tcg_parquet_dataset — data format
pokemon_tcg_kv_cache_hierarchical — the row-group cache
pokemon_tcg_train_config_reference — flag-level reference
pokemon_tcg_top_elo_curriculum_filter — --top-elo semantics
pokemon_tcg_tbptt_training_contract — TBPTT exact accounting
pokemon_tcg_sqlite_schema_current — schema the trainer reads from
pokemon_tcg_torch_inference — arena-side runtime
pokemon_tcg_bc_curriculum_ablation — 10-config ablation that validated this pipeline
pokemon_tcg_mlx_migration — the migration record (historical)
pokemon_tcg_repository_timeline — full verified timeline, with the older 147-count snapshot preserved