three_pool: integrate #540 (async consolidation) + #541 (config fork) (#542)

* three_pool: move checkpoint consolidation off the train loop (async)

The 3-pool save was fully synchronous: rank 0 serially read ~100 GB of
per-rank partials to assemble the canonical ThreePoolTrainingState while
all other ranks waited at a rejoining barrier. At XL that read overran the
NCCL watchdog (killed run 34446) and forced a 30-min PG timeout; over a
200k-step run it stalled the loop for hours.

New flow:
- snapshot() (on the train loop) writes a SELF-CONTAINED partial per rank
  (owned model params via the leader key-partition + optimizer state + PPGD
  sources) plus a rank-0 meta.pth, with one pre-write and one post-write
  barrier. No rank-0 read, no live-NCCL model gather, no model assembly.
- consolidate_step() (new three_pool/consolidate.py, runs in the async
  job, off the loop) reads all partials, assembles model_<step>.pth +
  training_<step>.pth, prunes training_*.pth to the last 3 (all model_*.pth
  kept), deletes the scratch dir. Idempotent.
- checkpoint.py: gather_full_state_dict_to_rank0 (NCCL) replaced by
  assemble_model_state_dict_from_partials (file-only) + leader key-partition
  helpers (owned_model_state_keys / ci_fn_state_keys).
- ThreePoolRunSink.checkpoint -> checkpoint_written(step, final); the 3-pool
  sink now only fires the on_save hook.
- submit_slurm_async_slow_eval -> submit_slurm_async_consolidate_and_eval:
  always submitted for 3-pool (consolidation is mandatory even with no slow
  metrics); async_eval consolidates first (rank 0 + barrier), then evals.
- Resume reads the latest CONSOLIDATED training_<step>.pth; if a run dies
  before consolidation, resume from the prior consolidated step (at most one
  save-interval lag; unconsolidated partials remain on disk).
- _DEFAULT_PG_TIMEOUT 30min -> 10min: the longest on-loop collective gap is
  now the fast-eval pass + a partial-write barrier, not the rank-0 read.

Validated on an 8-GPU 3-block topology (job 34739/34740):
- save snapshot wall time ~450ms/save (was a multi-minute rank-0 read)
- consolidation produces valid model/training files; assembly is
  byte-equivalent to the live ComponentModel state_dict
- resume from a consolidated checkpoint continues training
- prune keeps all model_*.pth, trims training_*.pth to last 3
- run survives a tight 120s PG timeout across 5 saves; a 30s rank-0 sleep
  injected into consolidation runs off-loop without touching the train PG

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: fix async-consolidation glue — don't resolve a checkpoint pre-consolidation

The live train->child chain failed: async_eval.main opened with
SavedLMRun.from_path(run), which eagerly resolves a model_<step>.pth. For a
3-pool save that file doesn't exist yet — this job assembles it — so every
child consolidate job crashed with "No files found with prefix model" before
ever calling consolidate_step.

Read the config straight from run_meta.yaml instead, consolidate, and only
resolve the checkpoint afterwards (for the eval pass). Also add a
PD_ASYNC_EVAL_DP override so the consolidate+eval child's GPU count isn't a
hardcoded 8 (lets a small-topology smoke fit train + child within one node,
and lets a cheap eval not match train width in production).

Found by the live end-to-end smoke (6-GPU train firing real 2-GPU child
jobs) that the prior validation had stubbed out via PD_3POOL_SKIP_ASYNC_ONSAVE.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: consolidate before CUDA/NCCL init (the child job hung otherwise)

The fixed glue got the child past the checkpoint-resolution crash, but it then
hung in assemble_model_state_dict_from_partials: building the full
ComponentModel (incl. the transformer CI fn) AFTER get_device()/init_distributed
makes that construction land on a contended GPU and stall (the same assembly is
2.6s on a clean CPU process).

Move consolidation to a CPU-only prelude that runs BEFORE any CUDA/NCCL init:
rank 0 (read from torchrun's RANK env, since the PG isn't up yet) assembles on
CPU; the other ranks poll the shared FS for training_<step>.pth, then everyone
inits distributed for the eval pass. No process group is needed for the wait.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: Option C config fork — standalone ThreePoolLMExperimentConfig + pd-lm-3pool

Splits the 3-pool path into its own composition root + config types, baking the
3-pool's constraints into the types so misconfigs fail at YAML parse instead of
minutes into a multi-node launch.

- ThreePoolConstrainedPDConfig (subclass of core PDConfig): frozen Literal scalars
  (sampling/n_mask_samples/use_delta_component/identity_decomposition_targets) + a
  typed ThreePoolLosses(faith, imp, stoch, ppgd) struct replacing the generic
  loss_metrics list. Derives loss_metrics from the struct (excluded from dump).
- ThreePoolRuntimeConfig (subclass of core RuntimeConfig) carries topology, so nothing
  enters core's RuntimeConfig.
- ThreePoolLMExperimentConfig: standalone sibling; LMExperimentConfig drops three_pool.
- pd-lm-3pool entry + three_pool_run.py composition root (fresh/resume/--dp submit);
  run.py back to single-pool only. Shared pure builders imported from run.py.
- ThreePoolTrainer reads pd.losses.* directly; deleted _validate_pd_config_for_three_pool
  + the by_type/isinstance dance. Cross-field checks (batch divisibility, rank-0) move to
  a load-time model_validator; site coverage stays in _build_runtime.
- SavedThreePoolLMRun reload sibling; init_pd_run / _build_eval_loop relaxed to structural
  protocols.
- Resume provenance now lives on the config (resume_provenance field) → flows to
  run_meta.yaml + wandb.config; dropped the write-only FS side-file.
- Migrated the two 3-pool YAMLs; parse-rejection + roundtrip + provenance tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: init distributed before consolidation; CPU-pin assembly; file-poll wait

Building target_model requires the distributed state to be initialized
(ensure_cached_and_call asserts via get_distributed_state), so consolidation
can't fully precede init_distributed. Reorder: init_distributed + build_target
first, then consolidate.

To keep the multi-second CPU read+assemble off any GPU collective:
- non-rank-0 ranks WAIT by polling the shared FS for training_<step>.pth (which
  consolidate_step writes last), not on an NCCL barrier;
- rank-0 assembly runs under `torch.device("cpu")` so the full ComponentModel
  buffer (incl. transformer CI fn) is pinned to CPU rather than landing on the
  selected, shared GPU (the original hang).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: fix consolidation hang root cause (unfrozen target) + fail-fast + rerun CLI

Root cause of the hung child consolidate jobs: assemble_model_state_dict_from_partials
builds a ComponentModel from the target returned by build_target, which only
.eval()s it — its params still have requires_grad=True. ComponentModel asserts
the target is frozen, so rank 0 crashed on that assertion. In the multi-rank
child this surfaced as a HANG: rank 0 died mid-consolidation while the other
rank sat in the wait, holding both GPUs until the timeout.

Fixes:
- Root cause: freeze the target (eval + requires_grad_(False)) inside the
  assembly, before constructing the ComponentModel buffer.
- Fail-fast, never hang: rank 0 writes a .consolidate_failed_<step> sentinel and
  re-raises on any consolidation error; the waiting ranks poll for the success
  file OR the sentinel and bail out (raising) on either the sentinel or a bounded
  timeout — so a failed child errors and releases its GPUs instead of wedging.
- Partials persist on failure (already true: consolidate_step rmtrees scratch
  only after a successful write) + a recovery CLI:
    python -m param_decomp_lab.three_pool.consolidate_cli <run> [--step N]
  consolidates every unconsolidated step (partials present, no training_<step>.pth),
  making "it failed, just rerun it" a one-liner. CLI lives in its own module to
  avoid an import cycle with experiments.lm.run.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: 4-GPU single-block 3-pool config for live e2e validation

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: make prune race-tolerant + clean stale scratch on already-consolidated

Two concurrent child consolidation jobs prune the same old training_<step>.pth;
the second crashed with FileNotFoundError (after already writing its checkpoint),
marking the job FAILED and leaving its scratch behind. Use unlink(missing_ok=True)
— a concurrent pruner removing the file first is benign. Also clean any leftover
scratch when consolidate_step early-returns on an already-existing training file,
so a step that died in prune doesn't keep looking unconsolidated.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: 4-GPU 2-save config (peak 8 GPU) for live e2e re-validation

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: resume YAML from live-produced p-livee2e05 checkpoint

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: document consolidation reliability (fail-fast, rerun CLI) + drop throwaway resume YAMLs

CLAUDE.md: add the async-job reliability contract (CPU-only post-init assembly,
target-freeze, rank-0-assembles / others-file-wait-fail-fast, partials-persist +
recovery CLI, concurrency-safe prune) and the consolidate_cli.py table entry.
Remove two resume pointer YAMLs that referenced now-deleted validation runs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* three_pool: migrate #540's validation YAMLs to the new ThreePoolLMExperimentConfig shape

#540's four resumption-validation configs were authored in the pre-fork format
(top-level three_pool, flat loss_metrics). The only 3-pool entry point is now
pd-lm-3pool, so migrate them: loss_metrics -> pd.losses struct, three_pool ->
runtime.topology, drop the frozen-Literal pd fields. All parse + round-trip.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* xl: 200k production config in the new ThreePoolLMExperimentConfig shape

The config the human will launch for the ~5-day run. Recreated from the b48
3-pool production config: steps 200000, PPGD optimizer lr warmup_pct 0.05;
everything else identical. Parses + round-trips under pd-lm-3pool.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(3pool): drop duplicate topology key when validating runtime_config on resume

Integration seam between #540 and #541: the consolidated training_<step>.pth
stores runtime_config as a ThreePoolRuntimeConfig dump (base scalars + topology),
but from_snapshot validated it as base RuntimeConfig -> pydantic extra_forbidden
on `topology`. Topology is already carried by three_pool_config (identical to the
dumped topology), so drop the duplicate key before validating the base scalars.
Caught by the resume gate (#540's consolidated checkpoints are only now read by
#541's from_snapshot path).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: ocg-goodfire

CLAUDE.md

@@ -155,7 +155,8 @@ All declared in `param_decomp_lab/pyproject.toml`.
 |---|---|---|
 | `pd-tms` | `experiments/tms/run.py` | Run TMS experiment from a YAML |
 | `pd-resid-mlp` | `experiments/resid_mlp/run.py` | Run ResidMLP from a YAML |
-| `pd-lm` | `experiments/lm/run.py` | Run LM from a YAML |
+| `pd-lm` | `experiments/lm/run.py` | Run single-pool LM from a YAML |
+| `pd-lm-3pool` | `experiments/lm/three_pool_run.py` | Run 3-pool LM from a `ThreePoolLMExperimentConfig` YAML |
 | `pd-lm-layerwise` | `experiments/lm/layerwise.py` | Split an LM YAML into per-matrix configs, submit as a SLURM array |
 | `pd-pretrain` | `experiments/lm/pretrain/cli.py` | Pretrain target models |
 | `pd-harvest` | `harvest/scripts/run_slurm_cli.py` | Submit harvest SLURM job |

param_decomp/run_sink.py

@@ -17,7 +17,7 @@
 
 from typing import Any, Protocol, runtime_checkable
 
-from param_decomp.training_state import ThreePoolTrainingState, TrainingState
+from param_decomp.training_state import TrainingState
 
 
 @runtime_checkable
@@ -52,14 +52,17 @@ def finish(self) -> None:
 class ThreePoolRunSink(Protocol):
     """Side-effect sink for a 3-pool training run.
 
-    Identical to `OnePoolRunSink` apart from the checkpoint parameter's state
-    type. Two separate protocols rather than a union so each trainer's
-    ``run()`` signature can only accept a sink wired to its own pool's state.
+    Unlike `OnePoolRunSink`, the 3-pool sink does NOT persist a training state
+    from the train loop. The trainer writes self-contained per-rank partials to
+    a shared-FS scratch dir (cheap, no rank-0 read), then calls
+    `checkpoint_written` so the sink can fire the async consolidation+eval job
+    that reads those partials, assembles ``model_<step>.pth`` +
+    ``training_<step>.pth`` off the critical path, and runs the slow eval.
     """
 
     def log(self, metrics: dict[str, Any], step: int) -> None: ...
     def console(self, *lines: str) -> None: ...
-    def checkpoint(self, snapshot: ThreePoolTrainingState, *, final: bool) -> None: ...
+    def checkpoint_written(self, step: int, *, final: bool) -> None: ...
     def finish(self) -> None: ...
 
 

param_decomp_lab/experiments/CLAUDE.md

@@ -12,16 +12,42 @@ the concrete reload class directly.
 ```
 experiments/
 ├── utils.py                 # ExperimentConfig[T,D] generic + EvalConfig + WandbConfig
-│                            # + init_pd_run + RUN_META_FILENAME
+│                            # + init_pd_run + RUN_META_FILENAME + resume_provenance field
 ├── tms/run.py
 ├── resid_mlp/run.py
 └── lm/
-    ├── run.py
+    ├── run.py               # single-pool LM (pd-lm)
+    ├── three_pool_run.py    # 3-pool LM (pd-lm-3pool): ThreePoolLMExperimentConfig
+    ├── three_pool_pd.py     # ThreePoolConstrainedPDConfig + ThreePoolLosses
     ├── layerwise.py         # split LM YAML into per-matrix configs + SLURM-array submit
     ├── data.py
     └── pretrain/            # see lm/pretrain/CLAUDE.md
 ```
 
+## 3-pool config fork (`pd-lm-3pool`)
+
+The 3-pool training path is a standalone sibling of `pd-lm`, not a flag on it. Dispatch
+is by entry point (the repo idiom — no discriminator). `ThreePoolLMExperimentConfig`
+(in `lm/three_pool_run.py`) is a standalone `BaseConfig` with `pd:
+ThreePoolConstrainedPDConfig`, `runtime: ThreePoolRuntimeConfig` (core `RuntimeConfig`
+scalars + `topology: ThreePoolConfig`), and the usual `cadence`/`target`/`data`/`eval`/
+`wandb`. `LMExperimentConfig` is purely single-pool (no `three_pool` field).
+
+`ThreePoolConstrainedPDConfig` (in `lm/three_pool_pd.py`) subclasses core `PDConfig` and
+bakes the 3-pool's invariants into the types — `sampling`/`n_mask_samples`/
+`use_delta_component`/`identity_decomposition_targets` are frozen `Literal` defaults, and
+the generic `loss_metrics` list is replaced by a typed `ThreePoolLosses(faith, imp,
+stoch, ppgd)` struct (it derives the inherited `loss_metrics` list from the struct via a
+before-validator and excludes it from serialization). `ThreePoolTrainer` reads
+`pd.losses.faith` / `.imp` / `.stoch` / `.ppgd` directly — no `by_type` dict, no
+`isinstance` asserts. Cross-field checks coupling `pd` to `runtime.topology` (batch
+divisibility, rank-0 convention) run as a `@model_validator` on
+`ThreePoolLMExperimentConfig`, so 3-pool misconfigs fail at YAML parse, not minutes into
+a multi-node launch. (Site coverage — owned sites ∈ decomposition_targets after pattern
+expansion — stays in `ThreePoolTrainer._build_runtime` since it needs the loaded model.)
+
+Reload: `SavedThreePoolLMRun` (in `three_pool_run.py`) mirrors `SavedLMRun`.
+
 ## YAML schema
 
 One validated pydantic tree (extra keys raise):

param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_async_consol.yaml

@@ -0,0 +1,149 @@
+pd:
+  seed: 0
+  ci_config:
+    mode: global
+    fn_type: global_shared_transformer
+    simple_transformer_ci_cfg:
+      d_model: 768
+      n_blocks: 2
+      mlp_hidden_dim:
+      - 512
+      attn_config:
+        n_heads: 12
+        max_len: 128
+        rope_base: 10000.0
+  sigmoid_type: leaky_hard
+  decomposition_targets:
+  - module_pattern: h.*.attn.q_proj
+    C: 64
+  - module_pattern: h.*.attn.k_proj
+    C: 64
+  batch_size: 2
+  steps: 25
+  components_optimizer:
+    lr_schedule:
+      start_val: 0.0003
+      warmup_pct: 0.03
+      final_val_frac: 0.1
+      fn_type: cosine
+    grad_clip_norm: 0.01
+  ci_fn_optimizer:
+    lr_schedule:
+      start_val: 0.0001
+      warmup_pct: 0.03
+      final_val_frac: 0.1
+      fn_type: cosine
+  faithfulness_warmup_steps: 2
+  faithfulness_warmup_lr: 0.001
+  faithfulness_warmup_weight_decay: 0.0
+  losses:
+    faith:
+      coeff: 100000000.0
+    imp:
+      coeff: 3.2e-05
+      pnorm: 2.0
+      beta: 0.5
+      p_anneal_start_frac: 0.0
+      p_anneal_final_p: 0.3
+      p_anneal_end_frac: 1.0
+      eps: 1.0e-12
+    stoch:
+      coeff: 100.0
+    ppgd:
+      coeff: 1.0
+      optimizer:
+        type: adam
+        beta1: 0.5
+        beta2: 0.99
+        eps: 1.0e-08
+        lr_schedule:
+          start_val: 0.01
+          warmup_pct: 0.025
+          final_val_frac: 1.0
+          fn_type: constant
+      scope:
+        type: per_batch_per_position
+      use_sigmoid_parameterization: false
+      n_warmup_steps: 2
+      n_samples: 1
+cadence:
+  train_log_every: 5
+  save_every: 5
+eval:
+  n_steps: 1
+  batch_size: 2
+  every: 10
+  slow_every: 10
+  slow_on_first_step: true
+  metrics:
+  - type: CI_L0
+    groups:
+      total:
+      - '*'
+  - type: CEandKLLosses
+    rounding_threshold: 0.0
+runtime:
+  autocast_bf16: true
+  device: cuda
+  dp: null
+  topology:
+    use_fused_kl: true
+    ci_ranks:
+    - 6
+    ppgd_ranks:
+    - 7
+    layerwise_block_groups:
+    - ranks:
+      - 0
+      - 1
+      owned_sites:
+      - h.0.attn.q_proj
+      - h.0.attn.k_proj
+      - h.1.attn.q_proj
+      - h.1.attn.k_proj
+      - h.2.attn.q_proj
+      - h.2.attn.k_proj
+      - h.3.attn.q_proj
+      - h.3.attn.k_proj
+    - ranks:
+      - 2
+      - 3
+      owned_sites:
+      - h.4.attn.q_proj
+      - h.4.attn.k_proj
+      - h.5.attn.q_proj
+      - h.5.attn.k_proj
+      - h.6.attn.q_proj
+      - h.6.attn.k_proj
+      - h.7.attn.q_proj
+      - h.7.attn.k_proj
+    - ranks:
+      - 4
+      - 5
+      owned_sites:
+      - h.8.attn.q_proj
+      - h.8.attn.k_proj
+      - h.9.attn.q_proj
+      - h.9.attn.k_proj
+      - h.10.attn.q_proj
+      - h.10.attn.k_proj
+      - h.11.attn.q_proj
+      - h.11.attn.k_proj
+target:
+  spec:
+    kind: hf_weights_in_vendored
+    model_class: param_decomp_lab.experiments.lm.pretrain.models.gpt2_simple.GPT2Simple
+    model_name: openai-community/gpt2
+  output_extract: 0
+  activation_checkpointing: false
+data:
+  tokenizer_name: openai-community/gpt2
+  max_seq_len: 128
+  dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2
+  column_name: input_ids
+  train_split: train
+  eval_split: train
+  is_tokenized: true
+  streaming: true
+  buffer_size: 1000
+  shuffle_each_epoch: true

param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu.yaml

@@ -0,0 +1,141 @@
+pd:
+  seed: 0
+  ci_config:
+    mode: global
+    fn_type: global_shared_transformer
+    simple_transformer_ci_cfg:
+      d_model: 768
+      n_blocks: 2
+      mlp_hidden_dim:
+      - 512
+      attn_config:
+        n_heads: 12
+        max_len: 128
+        rope_base: 10000.0
+  sigmoid_type: leaky_hard
+  decomposition_targets:
+  - module_pattern: h.*.attn.q_proj
+    C: 64
+  - module_pattern: h.*.attn.k_proj
+    C: 64
+  batch_size: 2
+  steps: 25
+  components_optimizer:
+    lr_schedule:
+      start_val: 0.0003
+      warmup_pct: 0.03
+      final_val_frac: 0.1
+      fn_type: cosine
+    grad_clip_norm: 0.01
+  ci_fn_optimizer:
+    lr_schedule:
+      start_val: 0.0001
+      warmup_pct: 0.03
+      final_val_frac: 0.1
+      fn_type: cosine
+  faithfulness_warmup_steps: 2
+  faithfulness_warmup_lr: 0.001
+  faithfulness_warmup_weight_decay: 0.0
+  losses:
+    faith:
+      coeff: 100000000.0
+    imp:
+      coeff: 3.2e-05
+      pnorm: 2.0
+      beta: 0.5
+      p_anneal_start_frac: 0.0
+      p_anneal_final_p: 0.3
+      p_anneal_end_frac: 1.0
+      eps: 1.0e-12
+    stoch:
+      coeff: 100.0
+    ppgd:
+      coeff: 1.0
+      optimizer:
+        type: adam
+        beta1: 0.5
+        beta2: 0.99
+        eps: 1.0e-08
+        lr_schedule:
+          start_val: 0.01
+          warmup_pct: 0.025
+          final_val_frac: 1.0
+          fn_type: constant
+      scope:
+        type: per_batch_per_position
+      use_sigmoid_parameterization: false
+      n_warmup_steps: 2
+      n_samples: 1
+cadence:
+  train_log_every: 5
+  save_every: 5
+eval:
+  n_steps: 1
+  batch_size: 2
+  every: 10
+  slow_every: 10
+  slow_on_first_step: true
+  metrics:
+  - type: CI_L0
+    groups:
+      total:
+      - '*'
+  - type: CEandKLLosses
+    rounding_threshold: 0.0
+runtime:
+  autocast_bf16: true
+  device: cuda
+  dp: null
+  topology:
+    use_fused_kl: true
+    ci_ranks:
+    - 2
+    ppgd_ranks:
+    - 3
+    layerwise_block_groups:
+    - ranks:
+      - 0
+      - 1
+      owned_sites:
+      - h.0.attn.q_proj
+      - h.0.attn.k_proj
+      - h.1.attn.q_proj
+      - h.1.attn.k_proj
+      - h.2.attn.q_proj
+      - h.2.attn.k_proj
+      - h.3.attn.q_proj
+      - h.3.attn.k_proj
+      - h.4.attn.q_proj
+      - h.4.attn.k_proj
+      - h.5.attn.q_proj
+      - h.5.attn.k_proj
+      - h.6.attn.q_proj
+      - h.6.attn.k_proj
+      - h.7.attn.q_proj
+      - h.7.attn.k_proj
+      - h.8.attn.q_proj
+      - h.8.attn.k_proj
+      - h.9.attn.q_proj
+      - h.9.attn.k_proj
+      - h.10.attn.q_proj
+      - h.10.attn.k_proj
+      - h.11.attn.q_proj
+      - h.11.attn.k_proj
+target:
+  spec:
+    kind: hf_weights_in_vendored
+    model_class: param_decomp_lab.experiments.lm.pretrain.models.gpt2_simple.GPT2Simple
+    model_name: openai-community/gpt2
+  output_extract: 0
+  activation_checkpointing: false
+data:
+  tokenizer_name: openai-community/gpt2
+  max_seq_len: 128
+  dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2
+  column_name: input_ids
+  train_split: train
+  eval_split: train
+  is_tokenized: true
+  streaming: true
+  buffer_size: 1000
+  shuffle_each_epoch: true