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

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/lm/_resumption_validation/gpt2_small_3pool_async_consol.yaml

@@ -0,0 +1,153 @@
+pd:
+  seed: 0
+  n_mask_samples: 1
+  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
+  sampling: continuous
+  sigmoid_type: leaky_hard
+  decomposition_targets:
+  - module_pattern: h.*.attn.q_proj
+    C: 64
+  - module_pattern: h.*.attn.k_proj
+    C: 64
+  identity_decomposition_targets: null
+  use_delta_component: true
+  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
+  loss_metrics:
+  - type: FaithfulnessLoss
+    coeff: 100000000.0
+  - type: StochasticReconLayerwiseLoss
+    coeff: 100.0
+  - type: ImportanceMinimalityLoss
+    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
+  - type: PersistentPGDReconLoss
+    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
+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
+three_pool:
+  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

param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu.yaml

@@ -0,0 +1,145 @@
+pd:
+  seed: 0
+  n_mask_samples: 1
+  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
+  sampling: continuous
+  sigmoid_type: leaky_hard
+  decomposition_targets:
+  - module_pattern: h.*.attn.q_proj
+    C: 64
+  - module_pattern: h.*.attn.k_proj
+    C: 64
+  identity_decomposition_targets: null
+  use_delta_component: true
+  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
+  loss_metrics:
+  - type: FaithfulnessLoss
+    coeff: 100000000.0
+  - type: StochasticReconLayerwiseLoss
+    coeff: 100.0
+  - type: ImportanceMinimalityLoss
+    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
+  - type: PersistentPGDReconLoss
+    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
+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
+three_pool:
+  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