--- type: reference title: "Pokémon TCG AI Battle — MLX Training Pipeline (reconciled current)" description: "Reconciled training reference: historical FP16 delivery contract, current strict-FP32 trainer behavior, TBPTT memory, auxiliary-loss reduction and handoff boundary." tags: [pokemon-tcg, training, mlx, fp16, fp32, tbptt, muon, adamw, checkpoint, tensorboard, pipeline] timestamp: "2026-08-15T17:24:00-03:00" --- # 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: ```text 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 ```text 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: ```text optimizer routing: Muon=845,568 hidden-matrix params; AdamW=456,583 embedding/head/vector params ``` ## Gradient accumulation For `accum_steps=K`: ```text 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: ```text 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/_/`. 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: ```bash 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: ```text 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