diff --git a/param_decomp/run_sink.py b/param_decomp/run_sink.py index 96d92055c..4a58ffc97 100644 --- a/param_decomp/run_sink.py +++ b/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_.pth`` + + ``training_.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: ... diff --git a/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_async_consol.yaml b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_async_consol.yaml new file mode 100644 index 000000000..d8b33bca3 --- /dev/null +++ b/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 diff --git a/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu.yaml b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu.yaml new file mode 100644 index 000000000..fffd0eb36 --- /dev/null +++ b/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 diff --git a/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu_2save.yaml b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu_2save.yaml new file mode 100644 index 000000000..c815bbea3 --- /dev/null +++ b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_4gpu_2save.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: 20 +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 diff --git a/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_e2e.yaml b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_e2e.yaml new file mode 100644 index 000000000..7f68e6ea3 --- /dev/null +++ b/param_decomp_lab/experiments/lm/_resumption_validation/gpt2_small_3pool_live_e2e.yaml @@ -0,0 +1,149 @@ +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: + - 4 + ppgd_ranks: + - 5 + 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 + - ranks: + - 2 + - 3 + owned_sites: + - 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 diff --git a/param_decomp_lab/experiments/lm/async_eval.py b/param_decomp_lab/experiments/lm/async_eval.py index 876efa412..1530c995f 100644 --- a/param_decomp_lab/experiments/lm/async_eval.py +++ b/param_decomp_lab/experiments/lm/async_eval.py @@ -1,8 +1,12 @@ -"""Run eval metrics against a saved LM PD checkpoint; log into the parent's wandb run. +"""Consolidate a 3-pool save + run eval metrics; log into the parent's wandb run. This is an internal sbatch target, not a user-facing CLI. Invoked by -:func:`param_decomp_lab.experiments.lm.run.submit_slurm_async_slow_eval` after a -training save, to compute slow metrics off the critical path of training. +:func:`param_decomp_lab.experiments.lm.run.submit_slurm_async_consolidate_and_eval` +after a training save, off the critical path of training. For 3-pool runs it +first assembles ``model_.pth`` + ``training_.pth`` from the train +loop's per-rank partials (rank 0; +:func:`param_decomp_lab.three_pool.consolidate.consolidate_step`), then runs the +parent's slow metrics against the assembled checkpoint. The training side hands us: * ``--run `` — the parent training run (SavedLMRun-compatible) @@ -19,11 +23,13 @@ """ import gc +import time from pathlib import Path from typing import Any import fire import torch +import torch.nn as nn import wandb from torch.utils.data import DataLoader @@ -43,7 +49,7 @@ ) from param_decomp_lab.eval_metrics import EVAL_METRIC_CLASSES from param_decomp_lab.experiments.lm.run import ( - SavedLMRun, + LMExperimentConfig, _resolve_train_run_id, build_lm_loader, build_target, @@ -51,9 +57,71 @@ ) from param_decomp_lab.experiments.utils import RUN_META_FILENAME, EvalConfig from param_decomp_lab.infra.run_files import resolve_run_files +from param_decomp_lab.infra.settings import PARAM_DECOMP_OUT_DIR from param_decomp_lab.infra.wandb import get_wandb_entity, try_wandb from param_decomp_lab.run_sink import _wandb_value from param_decomp_lab.seed import set_seed +from param_decomp_lab.three_pool.consolidate import ( + DEFAULT_KEEP_LAST_N_TRAINING, + SNAPSHOT_SCRATCH_DIRNAME, + consolidate_step, +) + + +def _consolidate_or_wait( + *, + cfg: "LMExperimentConfig", + out_dir: Path, + target_model: "nn.Module", + run_batch: Any, + step: int, + wait_timeout_s: float = 1800.0, +) -> None: + """Assemble the step's checkpoint (rank 0) or wait for it (other ranks). + + Rank 0 assembles on CPU and writes `training_.pth` last; the other ranks + wait on the shared FS rather than an NCCL barrier (the multi-second CPU + read+assemble shouldn't interleave with a GPU collective). Assembly is pinned + to CPU — building the full ComponentModel buffer on the selected, shared GPU + hangs. + + Fail-fast on BOTH sides — never wedge holding GPUs: + * rank 0 writes a `.consolidate_failed_` sentinel and re-raises if + consolidation throws, so the failure is visible and the GPUs release; + * the other ranks poll for the success file OR the sentinel, and time out; + any of the three exits the wait (the last two by raising). + """ + training_path = out_dir / f"training_{step}.pth" + fail_sentinel = out_dir / f".consolidate_failed_{step}" + fail_sentinel.unlink(missing_ok=True) + if is_main_process(): + try: + with torch.device("cpu"): + consolidate_step( + scratch_dir=out_dir / SNAPSHOT_SCRATCH_DIRNAME, + out_dir=out_dir, + step=step, + target_model=target_model, + run_batch=run_batch, + ci_config=cfg.pd.ci_config, + sigmoid_type=cfg.pd.sigmoid_type, + keep_last_n_training=DEFAULT_KEEP_LAST_N_TRAINING, + ) + except BaseException: + fail_sentinel.touch() + raise + gc.collect() + return + deadline = time.perf_counter() + wait_timeout_s + while not training_path.is_file(): + assert not fail_sentinel.is_file(), ( + f"rank-0 consolidation failed (sentinel {fail_sentinel.name}); aborting wait" + ) + assert time.perf_counter() < deadline, ( + f"timed out after {wait_timeout_s}s waiting for {training_path} " + f"(rank-0 consolidation did not finish)" + ) + time.sleep(2.0) def _resolve_eval_checkpoint_path(run_path: str | Path, step: int | None) -> Path: @@ -171,44 +239,72 @@ def main( to run. If omitted, falls back to the parent run's ``cfg.eval``. group / tags: optional wandb metadata for the resumed run. """ - pd_run = SavedLMRun.from_path(run) + # Read the config from run_meta.yaml directly rather than `SavedLMRun.from_path`: + # for a 3-pool run no `model_.pth` exists yet (this job assembles it + # below), and `from_path` eagerly resolves one, which would crash here. + train_run_id = _resolve_train_run_id(run) + out_dir = PARAM_DECOMP_OUT_DIR / "decompositions" / train_run_id + cfg = LMExperimentConfig.from_file(out_dir / RUN_META_FILENAME) + if eval_config is not None: eval_cfg = EvalConfig.from_file(Path(eval_config)) else: - assert pd_run.cfg.eval is not None, ( + assert cfg.eval is not None, ( f"async_eval requires either --eval-config or the parent config to " f"declare an `eval:` block ({run})" ) - eval_cfg = pd_run.cfg.eval - - checkpoint_path = _resolve_eval_checkpoint_path(run, step) - resolved_step = _step_from_checkpoint_name(checkpoint_path.name) - train_run_id = _resolve_train_run_id(run) + eval_cfg = cfg.eval dist_state = init_distributed() if is_main_process(): logger.info(f"Distributed state: {dist_state}") - logger.info(f"async_eval: {run} @ step {resolved_step} (train run_id={train_run_id})") - set_seed(pd_run.cfg.pd.seed) + set_seed(cfg.pd.seed) device = get_device() - target_model = build_target(pd_run.cfg.target) + target_model = build_target(cfg.target) + run_batch = make_run_batch(cfg.target) + + # Consolidate the 3-pool save. The train loop wrote only per-rank partials; + # this assembles model_.pth + training_.pth off the train-loop + # critical path. Rank 0 assembles on CPU; the other ranks WAIT on the file + # (not an NCCL barrier) so no GPU collective interleaves with the multi-second + # CPU read+assemble. Assembly is forced onto CPU — building the full + # ComponentModel buffer (incl. the transformer CI fn) on the selected, shared + # GPU hangs. + if cfg.three_pool is not None: + assert step is not None, "3-pool async consolidation requires an explicit --step" + _consolidate_or_wait( + cfg=cfg, out_dir=out_dir, target_model=target_model, run_batch=run_batch, step=step + ) + + # Consolidation may be the only job to do — when there are no slow metrics + # the 3-pool save still needs assembling, but there is no eval pass to run. + if not eval_cfg.metrics: + if is_main_process(): + logger.info("async_eval: no metrics; consolidation-only, skipping eval pass") + return + + checkpoint_path = _resolve_eval_checkpoint_path(run, step) + resolved_step = _step_from_checkpoint_name(checkpoint_path.name) + if is_main_process(): + logger.info(f"async_eval: {run} @ step {resolved_step} (train run_id={train_run_id})") + component_model = load_component_model( - pd_config=pd_run.cfg.pd, + pd_config=cfg.pd, checkpoint_path=checkpoint_path, target_model=target_model, - run_batch=make_run_batch(pd_run.cfg.target), + run_batch=run_batch, ) component_model.to(device) eval_loader = build_lm_loader( - pd_run.cfg.target, - pd_run.cfg.data, + cfg.target, + cfg.data, split="eval", device=device, batch_size=eval_cfg.batch_size, dist_state=dist_state, - seed=pd_run.cfg.pd.seed, + seed=cfg.pd.seed, ) eval_metrics = [EVAL_METRIC_CLASSES[m.type](m) for m in eval_cfg.metrics] for m in eval_metrics: @@ -221,13 +317,13 @@ def main( n_steps=eval_cfg.n_steps, device=device, step=resolved_step, - pd_config=pd_run.cfg, + pd_config=cfg, ) if is_main_process(): _log_eval_to_wandb( results, - cfg=pd_run.cfg, + cfg=cfg, train_run_id=train_run_id, step=resolved_step, group=group, diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index 1eb94afb8..069780f25 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/param_decomp_lab/experiments/lm/run.py @@ -12,10 +12,11 @@ directly via `torchrun --standalone --nproc_per_node=N -m param_decomp_lab.experiments.lm.run config.yaml`. -Async slow-eval of saved checkpoints lives in +Async consolidation + slow-eval of 3-pool saves lives in ``param_decomp_lab.experiments.lm.async_eval``; the helper -:func:`submit_slurm_async_slow_eval` below filters the parent's slow metrics into -a temp YAML and submits an sbatch job that invokes that module. +:func:`submit_slurm_async_consolidate_and_eval` below submits an sbatch job that +assembles the canonical checkpoint from the train loop's per-rank partials and +then runs the parent's slow metrics against it. """ import atexit @@ -86,6 +87,7 @@ from param_decomp_lab.run_sink import OnePoolSink, ThreePoolSink from param_decomp_lab.seed import set_seed from param_decomp_lab.three_pool import ThreePoolConfig, ThreePoolTrainer +from param_decomp_lab.three_pool.consolidate import SNAPSHOT_SCRATCH_DIRNAME def _resolve_class(fqn: str) -> type: @@ -510,14 +512,16 @@ def _fresh_main( if cfg.three_pool is not None: run_id = _agree_on_run_id_three_pool(run_id, dist_state) out_dir = PARAM_DECOMP_OUT_DIR / "decompositions" / run_id - scratch_dir = out_dir / ".snapshot_scratch" + scratch_dir = out_dir / SNAPSHOT_SCRATCH_DIRNAME three_sink = init_pd_run( cfg, sink_class=ThreePoolSink, group=group, tags=tags, run_id=run_id, - on_save=lambda step: submit_slurm_async_slow_eval(out_dir, step=step, parent_cfg=cfg), + on_save=lambda step: submit_slurm_async_consolidate_and_eval( + out_dir, step=step, parent_cfg=cfg + ), ) # Multi-pool eval mirrors the train data-handling contract: full eval # batch on every rank, sliced internally by each pool. So we pass @@ -611,14 +615,14 @@ def _resume_main( if effective_cfg.three_pool is not None: run_id = _agree_on_run_id_three_pool(run_id, dist_state) out_dir = PARAM_DECOMP_OUT_DIR / "decompositions" / run_id - scratch_dir = out_dir / ".snapshot_scratch" + scratch_dir = out_dir / SNAPSHOT_SCRATCH_DIRNAME three_sink = init_pd_run( effective_cfg, sink_class=ThreePoolSink, group=group, tags=tags, run_id=run_id, - on_save=lambda step: submit_slurm_async_slow_eval( + on_save=lambda step: submit_slurm_async_consolidate_and_eval( out_dir, step=step, parent_cfg=effective_cfg ), ) @@ -719,9 +723,9 @@ def _build_eval_loop( """Build the `EvalLoop` from `cfg.eval`, or `None` when eval is disabled. Slow metrics (class-attr ``slow=True``) are filtered out — in-train eval is - fast-only. Slow metrics are picked up later by the async eval-only job (which - receives the slow subset via a temp ``EvalConfig`` YAML, see - :func:`_submit_slurm_async_slow_eval`). + fast-only. Slow metrics are picked up later by the async job (which receives + the slow subset via a temp ``EvalConfig`` YAML, see + :func:`submit_slurm_async_consolidate_and_eval`). """ if cfg.eval is None: return None @@ -834,7 +838,7 @@ def _wandb_url_for_config(config_path: str | Path, run_id: str) -> str | None: return f"https://wandb.ai/{entity}/{cfg.wandb.project}/runs/{run_id}" -def submit_slurm_async_slow_eval( +def submit_slurm_async_consolidate_and_eval( run_path: str | Path, *, step: int, @@ -842,32 +846,52 @@ def submit_slurm_async_slow_eval( dp: int = 8, time: str = "2:00:00", partition: str | None = DEFAULT_PARTITION_NAME, - job_name: str = "pd-lm-eval-slow", + job_name: str = "pd-lm-consol-eval", group: str | None = None, tags: str | None = None, no_snapshot: bool = False, ) -> None: - """Filter the parent's slow eval metrics into a temp YAML, submit a SLURM job. - - Called from training right after a save: emits an async SLURM job that runs - *only* the slow metrics against ``model_.pth`` via - ``python -m param_decomp_lab.experiments.lm.async_eval``, logging into the - parent's wandb run at the same step. No-op when the parent has no eval block - or no slow metrics. + """Submit the async job that consolidates a 3-pool save and runs slow eval. + + Called from 3-pool training right after a save. The train loop has written + per-rank partials to the scratch dir; this job (off the critical path) + assembles ``model_.pth`` + ``training_.pth`` from them, prunes old + ``training_*.pth``, then runs the parent's *slow* eval metrics against the + assembled ``model_.pth`` (logging into the parent's wandb run at the + same step). The job is ALWAYS submitted for 3-pool — consolidation is + mandatory even when there are no slow metrics; in that case the eval pass is a + no-op (see ``async_eval``). """ from param_decomp_lab.experiments.utils import EvalConfig - if parent_cfg.eval is None: - return - slow_metrics, _fast = _split_metrics_by_slow(parent_cfg.eval.metrics) - if not slow_metrics: + assert parent_cfg.three_pool is not None, ( + "async consolidate+eval is only for 3-pool runs (consolidation has nothing to do otherwise)" + ) + + # Test hook (never set in production): skip the child-job submission so a + # smoke can drive consolidation/eval out-of-band and stay within a GPU + # budget. The train loop still writes its partials regardless. + if os.environ.get("PD_3POOL_SKIP_ASYNC_ONSAVE", "").strip() in ("1", "true"): + logger.info(f"PD_3POOL_SKIP_ASYNC_ONSAVE set; skipping async job for step {step}") return + # The consolidate+eval job's GPU count. Defaults to `dp`; overridable so a + # small-topology smoke can keep train + child within one node's 8 GPUs (and, + # in production, so a cheap eval doesn't have to match the train width). + dp_override = os.environ.get("PD_ASYNC_EVAL_DP", "").strip() + if dp_override: + dp = int(dp_override) + + slow_metrics = ( + _split_metrics_by_slow(parent_cfg.eval.metrics)[0] if parent_cfg.eval is not None else [] + ) + eval_batch_size = parent_cfg.eval.batch_size if parent_cfg.eval is not None else dp + eval_n_steps = parent_cfg.eval.n_steps if parent_cfg.eval is not None else 1 slow_eval_cfg = EvalConfig( - batch_size=parent_cfg.eval.batch_size, - n_steps=parent_cfg.eval.n_steps, - every=parent_cfg.eval.every, - slow_every=parent_cfg.eval.every, # any value; not consumed by async_eval + batch_size=eval_batch_size, + n_steps=eval_n_steps, + every=1, # any value; not consumed by async_eval + slow_every=1, # any value; not consumed by async_eval slow_on_first_step=True, metrics=slow_metrics, ) diff --git a/param_decomp_lab/run_sink.py b/param_decomp_lab/run_sink.py index dfddb18c7..7b8afd078 100644 --- a/param_decomp_lab/run_sink.py +++ b/param_decomp_lab/run_sink.py @@ -1,10 +1,11 @@ """Concrete `RunSink` classes used by the in-repo experiments and lab tooling. Two pool-specific sinks (`OnePoolSink`, `ThreePoolSink`) share a private base -(`_LabSinkBase`) for the local-files / wandb / console / log plumbing. The -subclasses differ only in `checkpoint`'s typed parameter — 1-pool takes -`TrainingState`, 3-pool takes `ThreePoolTrainingState`. Each delegates to a -shared `_persist` for the actual save. +(`_LabSinkBase`) for the local-files / wandb / console / log plumbing. They +differ in how a checkpoint is persisted: `OnePoolSink.checkpoint` writes the +`TrainingState` synchronously via `_persist`; `ThreePoolSink.checkpoint_written` +persists nothing (the trainer already wrote per-rank partials) and only fires +the async consolidation+eval job via the `on_save` hook. Three constructors on each: @@ -29,7 +30,7 @@ from param_decomp.base_config import BaseConfig from param_decomp.distributed import is_main_process from param_decomp.log import logger -from param_decomp.training_state import ThreePoolTrainingState, TrainingState +from param_decomp.training_state import TrainingState from param_decomp_lab.infra.run_files import save_file from param_decomp_lab.infra.wandb import init_wandb, try_wandb @@ -163,7 +164,7 @@ def finish(self) -> None: if self._wandb_active and wandb.run is not None: wandb.finish() - def _persist(self, snapshot: TrainingState | ThreePoolTrainingState, *, final: bool) -> None: + def _persist(self, snapshot: TrainingState, *, final: bool) -> None: """Save the snapshot as `model_.pth` + `training_.pth`. `model_.pth` is just the component-model state dict — the artifact @@ -202,10 +203,20 @@ def checkpoint(self, snapshot: TrainingState, *, final: bool) -> None: @dataclass(frozen=True) class ThreePoolSink(_LabSinkBase): - """Lab sink for 3-pool runs (satisfies `ThreePoolRunSink`).""" + """Lab sink for 3-pool runs (satisfies `ThreePoolRunSink`). - def checkpoint(self, snapshot: ThreePoolTrainingState, *, final: bool) -> None: - self._persist(snapshot, final=final) + Persists nothing itself: the trainer has already written self-contained + per-rank partials to the scratch dir. This only fires the rank-0 `on_save` + hook, which submits the async job that consolidates those partials into + ``model_.pth`` + ``training_.pth`` and runs the slow eval. + """ + + def checkpoint_written(self, step: int, *, final: bool) -> None: + del final # the async consolidation path is identical for the final save + if self.out_dir is None: + return + if self.on_save is not None: + self.on_save(step) def _wandb_value(v: Any) -> Any: diff --git a/param_decomp_lab/tests/test_three_pool_pg_timeout.py b/param_decomp_lab/tests/test_three_pool_pg_timeout.py index 78638414f..5622a8041 100644 --- a/param_decomp_lab/tests/test_three_pool_pg_timeout.py +++ b/param_decomp_lab/tests/test_three_pool_pg_timeout.py @@ -1,17 +1,27 @@ """Regression tests for the 3-pool checkpoint-save NCCL-watchdog-timeout fix. -Two independent failure surfaces are covered: +Three failure surfaces are covered: * ``_resolve_pg_timeout`` reads the ``PD_3POOL_PG_TIMEOUT_S`` override (the - knob the save-watchdog repro uses to force the bug at small scale) and - otherwise returns the generous default. + knob the watchdog-safe-at-low-timeout test uses to force a tight bound) and + otherwise returns the default. * ``build_world`` threads its ``pg_timeout`` into *every* ``dist.new_group`` call. This is the actual production fix: ``new_group`` does NOT inherit the timeout passed to ``init_process_group`` — with ``timeout=None`` it falls back to the 10-min NCCL default. The 3-pool runs all real collectives on - these subgroups, so a slow checkpoint save trips that 10-min watchdog - unless the timeout is set explicitly here. + these subgroups, so a slow collective trips that watchdog unless the timeout + is set explicitly here. + + * The async-consolidation invariant: the union of every rank's + self-contained partial must exactly cover the full model's V/U + CI-fn + state-dict keys, so ``assemble_model_state_dict_from_partials`` can rebuild + the checkpoint off the train loop with no live ranks. + +The default PG timeout came down from 30 min to 10 min once consolidation moved +off the train loop: the longest on-loop collective gap is now the fast-eval +pass + a partial-write barrier (minutes), not the old ~10-min rank-0 read of +~100 GB of partials. """ import datetime @@ -19,6 +29,10 @@ import pytest import torch.distributed as dist +from param_decomp_lab.three_pool.checkpoint import ( + ci_fn_state_keys, + owned_model_state_keys, +) from param_decomp_lab.three_pool.layout import LayerwiseBlockGroup, build_world from param_decomp_lab.three_pool.optimize import ( _DEFAULT_PG_TIMEOUT, @@ -30,7 +44,7 @@ def test_resolve_pg_timeout_default(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("PD_3POOL_PG_TIMEOUT_S", raising=False) assert _resolve_pg_timeout() == _DEFAULT_PG_TIMEOUT - assert datetime.timedelta(minutes=30) == _DEFAULT_PG_TIMEOUT + assert datetime.timedelta(minutes=10) == _DEFAULT_PG_TIMEOUT def test_resolve_pg_timeout_env_override(monkeypatch: pytest.MonkeyPatch) -> None: @@ -112,3 +126,36 @@ def test_fingerprint_core_catches_topology_mismatch() -> None: } changed_ppgd = {**base, "ppgd_ranks": [5]} assert _rank_invariant_fingerprint_core(base) != _rank_invariant_fingerprint_core(changed_ppgd) + + +def test_leader_partials_exactly_cover_model_state() -> None: + """The per-leader partial slices (LW owned-site V/U + CI fn) must partition + the full model's fillable keys with no gaps or overlaps — the invariant the + async consolidation asserts before assembling the checkpoint.""" + sites = ("h.0.attn.q_proj", "h.1.attn.q_proj", "h.2.mlp.c_fc") + model_keys = { + "_components.h-0-attn-q_proj.V", + "_components.h-0-attn-q_proj.U", + "_components.h-1-attn-q_proj.V", + "_components.h-1-attn-q_proj.U", + "_components.h-2-mlp-c_fc.V", + "_components.h-2-mlp-c_fc.U", + "ci_fn._global_ci_fn.embed.weight", + "ci_fn._global_ci_fn.proj.weight", + # target-model keys (not owned by any leader; come from the fresh buffer) + "target_model.transformer.wte.weight", + "target_model.transformer.h.0.attn.c_attn.weight", + } + target_keys = {k for k in model_keys if k.startswith("target_model.")} + fillable = model_keys - target_keys + + # Two LW blocks: block 0 owns sites 0+1, block 1 owns site 2. CI leader owns ci_fn. + block0 = owned_model_state_keys(model_keys, owned_sites=(sites[0], sites[1])) + block1 = owned_model_state_keys(model_keys, owned_sites=(sites[2],)) + ci = ci_fn_state_keys(model_keys) + + assert block0.isdisjoint(block1) + assert block0.isdisjoint(ci) + assert block1.isdisjoint(ci) + assert block0 | block1 | ci == fillable + assert target_keys.isdisjoint(block0 | block1 | ci) diff --git a/param_decomp_lab/three_pool/CLAUDE.md b/param_decomp_lab/three_pool/CLAUDE.md index d33d2ba67..95169e44a 100644 --- a/param_decomp_lab/three_pool/CLAUDE.md +++ b/param_decomp_lab/three_pool/CLAUDE.md @@ -9,7 +9,9 @@ docstring in `optimize.py` for the data-handling contract. |---|---| | `optimize.py` | `ThreePoolTrainer` + `optimize_three_pool`; the training loop, `snapshot`/`from_snapshot`, config validation | | `layout.py` | `World` topology; `build_world` constructs every process group (threading `pg_timeout` into each); `BatchEdge` — symmetric per-edge batch-slice geometry (CI↔LW, CI↔PPGD) answering routing for both fan directions | -| `checkpoint.py` | `gather_full_state_dict_to_rank0` — rebuilds the full model state on rank 0 | +| `checkpoint.py` | offline state_dict assembly from on-disk partials (`assemble_model_state_dict_from_partials`) + the leader key-partition helpers (`owned_model_state_keys` / `ci_fn_state_keys`) | +| `consolidate.py` | `consolidate_step` — async, off-train-loop assembly of `model_.pth` + `training_.pth` from a step's scratch partials; prunes old `training_*.pth`; deletes the scratch dir. `unconsolidated_steps` lists recoverable steps | +| `consolidate_cli.py` | `python -m …consolidate_cli [--step N]` — manual CPU-only recovery for a failed/preempted async consolidation (separate module to avoid an import cycle with `experiments.lm.run`) | | `config.py` | `ThreePoolConfig` + topology validation | | `role.py` | `PoolRole = CIRole \| LWRole \| PPGDRole` — this rank's pool role; per-pool fields are union variants, not optional attrs | | `context.py` | `PoolContext = CIContext \| LWContext \| PPGDContext` — `world` + `role` + this pool's portals; the trainer matches on it to dispatch step fns | @@ -17,53 +19,108 @@ docstring in `optimize.py` for the data-handling contract. | `step_{ci,layerwise,ppgd}.py` | per-pool step functions | | `eval_step.py` | 3-pool eval pass (PPGD pool runs metrics; others barrier through) | -## Checkpoint-save invariant (do not break) - -`snapshot()` is collective. The ordering is: all ranks gather the model to -rank 0, all ranks write their per-rank partial to the shared-FS scratch dir, -barrier, then **rank 0 alone** serially reads every partial (~100 GB at XL) to -assemble the canonical `ThreePoolTrainingState`. The other ranks have nothing -to do during that read. - -There is a **mandatory rejoining `dist.barrier(group=cross_pool_p2p_group)` -after rank 0 finishes reading.** Without it, the non-rank-0 ranks race ahead -into the next training step while rank 0 is still reading; the next step's -collectives (PPGD V/U all-reduce, cross-pool sends) block on rank 0 and trip -the NCCL collective-timeout watchdog, aborting the whole job. This is exactly -how run 34446 (p-a5b667e9) died at its first checkpoint. The barrier makes all -ranks resume in lock-step. - -## Process-group timeout (do not lower) +## Checkpoint save: partials on the loop, consolidation off it + +The save path is split so the train loop never blocks on a multi-GB read. + +`snapshot()` (on the train loop, all ranks collective): each rank writes a +**self-contained partial** to `scratch_dir/step_/rank_.pth` — its owned +model params (LW block leaders → owned-sites V/U; CI pool leader → CI fn), +its optimizer state (name-keyed), and (PPGD) its sources. Rank 0 also writes +`meta.pth` (configs + fingerprint + `c_per_site` / `all_sites`). There is **one +pre-write barrier** (so rank 0's `mkdir` + `meta` write land before others write +into the dir) and **one post-write rejoin barrier** (so all ranks leave +`snapshot()` together). Both barriers are cheap — no rank does a 100 GB read on +the loop anymore — so neither can overrun the watchdog. NO model NCCL gather, NO +rank-0 assembly here. + +`consolidate_step()` (`consolidate.py`, async SLURM job, off the loop): reads all +of a step's partials → assembles the full `ComponentModel` state_dict + +`ThreePoolTrainingState` → writes `model_.pth` + `training_.pth` → prunes +old `training_*.pth` to the last `DEFAULT_KEEP_LAST_N_TRAINING` (=3; **all +`model_*.pth` are kept**) → deletes `step_/`. It runs inside the async +slow-eval job (`experiments/lm/async_eval.py`), as a CPU-only phase BEFORE the +eval pass, so the assembled `model_.pth` exists before the eval loads it. +Idempotent: a no-op if `training_.pth` already exists (and it cleans any +leftover scratch in that case); the scratch dir is deleted only on success. + +This is the fix for how run 34446 (p-a5b667e9) died: the old synchronous rank-0 +read held the other ranks at a barrier past the NCCL watchdog. There is no +on-loop read to outrun the watchdog now. + +### Consolidation reliability (the async job's contract) + +`async_eval.main` runs consolidation via `_consolidate_or_wait`, which is +engineered to **never hang** (a hung job holds GPUs forever — the one +unacceptable failure mode for an off-loop job): + + * **CPU-only, post-init.** Assembly builds a full `ComponentModel` buffer; it + must run on CPU (under `torch.device("cpu")`) — building it on the job's + selected, shared GPU hangs. It runs AFTER `init_distributed` (so + `build_target`'s `ensure_cached_and_call` has its distributed state) but the + buffer never touches the GPU. `assemble_model_state_dict_from_partials` + freezes the target (`eval()` + `requires_grad_(False)`) before constructing + the buffer — `build_target` only `.eval()`s it, and `ComponentModel` asserts + a frozen target (an unfrozen target was the original "hang": rank 0 hit the + assertion and died while the other rank waited). + * **Rank 0 assembles; others file-wait, fail-fast.** Non-rank-0 ranks poll the + shared FS for `training_.pth` (written last) rather than an NCCL barrier + (so the multi-second CPU read never interleaves with a GPU collective). On + any consolidation error rank 0 writes a `.consolidate_failed_` sentinel + and re-raises; the waiters bail out (raising) on the sentinel OR a bounded + timeout. A failed child **errors and releases its GPUs**, it does not wedge. + * **Partials persist on failure → re-runnable.** `consolidate_step` deletes + `step_/` only after a successful write, so a failed/preempted + consolidation leaves the partials intact. Recover with the CLI: + + python -m param_decomp_lab.three_pool.consolidate_cli [--step N] + + With no `--step` it consolidates every `unconsolidated_steps(out_dir)` (a step + with partials but no `training_.pth`). The prune is concurrency-safe + (`unlink(missing_ok=True)`) since multiple per-step children may prune the + same old file at once. + +## Process-group timeout `build_world` takes a `pg_timeout` and threads it into every `dist.new_group` call. **This is load-bearing:** `new_group` does NOT inherit the timeout passed to `init_process_group` — with `timeout=None` it silently uses the 10-min NCCL library default. The 3-pool runs all of its real collectives on these subgroups (never the default group), so the timeout must be set explicitly or a slow -checkpoint save / eval pass trips the 10-min watchdog. +collective trips the watchdog. -The default is 30 min (`_DEFAULT_PG_TIMEOUT` in `optimize.py`). The invariant is -simply: **the PG timeout must exceed the worst-case rank-0 checkpoint read -time.** Override (seconds) via `PD_3POOL_PG_TIMEOUT_S` — used by the -save-watchdog repro (`scripts/repro_3pool_save_watchdog.sbatch`) to force the -bug at small scale. +The default is **10 min** (`_DEFAULT_PG_TIMEOUT` in `optimize.py`), brought down +from 30 min once consolidation moved off the loop. The invariant is now: **the +PG timeout must exceed the worst-case on-loop collective gap**, which is the +in-train (fast) eval pass plus a checkpoint partial-write barrier — minutes, not +the old ~10-min rank-0 read. Override (seconds) via `PD_3POOL_PG_TIMEOUT_S` — +used by the watchdog-safe-at-low-timeout test to force a tight bound. ## Resume (`from_snapshot`) 3-pool runs persist a `ThreePoolTrainingState` (not the single-pool -`TrainingState`); `read_training_snapshot` returns either and the caller -narrows. `from_snapshot` validates the saved topology against the current one, -but the comparison runs on EVERY rank, so it compares only the -**rank-invariant** core (`world_size` / `ci_ranks` / `ppgd_ranks` / -block count) via `_rank_invariant_fingerprint_core` — never a rank-local view. -That helper also tolerates the pre-fix rank-local fingerprint format baked into -existing production checkpoints (p-a5b667e9). The per-block ranks→sites mapping -is re-derived from the snapshot's `three_pool_config`. +`TrainingState`), written by the **async consolidation job**, not the train loop. +`resolve_step` / `read_training_snapshot` pick the latest consolidated +`training_.pth`. If a run dies after a save but before that step's +consolidation finishes, resume from the previous consolidated step — at most one +save-interval of lost progress (the scratch partials for the unconsolidated step +are left on disk; they can be consolidated manually via `consolidate_step` if +that interval matters). + +`from_snapshot` validates the saved topology against the current one, but the +comparison runs on EVERY rank, so it compares only the **rank-invariant** core +(`world_size` / `ci_ranks` / `ppgd_ranks` / block count) via +`_rank_invariant_fingerprint_core` — never a rank-local view. That helper also +tolerates the pre-fix rank-local fingerprint format baked into existing +production checkpoints (p-a5b667e9). The per-block ranks→sites mapping is +re-derived from the snapshot's `three_pool_config`. Repro/fault-injection env knobs (never set in production): -`PD_3POOL_PG_TIMEOUT_S`, `PD_3POOL_SNAPSHOT_RANK0_SLEEP_S` (inject a sleep into -rank-0's read), `PD_3POOL_DISABLE_REJOIN_BARRIER` (reproduces the pre-fix race). -Regression tests: `param_decomp_lab/tests/test_three_pool_pg_timeout.py`. +`PD_3POOL_PG_TIMEOUT_S`, `PD_3POOL_SNAPSHOT_RANK0_SLEEP_S` (sleeps rank 0 inside +`snapshot()` AFTER the partial write — proves the sleep no longer stalls the +loop now that the read is async), `PD_3POOL_DISABLE_REJOIN_BARRIER` (drops the +post-write rejoin barrier). Regression tests: +`param_decomp_lab/tests/test_three_pool_pg_timeout.py`. ## Cross-pool batch divisibility (bidirectional) diff --git a/param_decomp_lab/three_pool/checkpoint.py b/param_decomp_lab/three_pool/checkpoint.py index 451627341..42e575199 100644 --- a/param_decomp_lab/three_pool/checkpoint.py +++ b/param_decomp_lab/three_pool/checkpoint.py @@ -1,27 +1,22 @@ -"""Distributed checkpoint assembly for 3-pool training. - -Rank 0 only holds its own block's V/U (and a random-init CI fn). To produce a -checkpoint that matches the schema ``load_component_model_from_checkpoint`` -expects (one flat ``state_dict`` with every site's V/U + the trained CI fn), -this module gathers state onto rank 0 from: - - * Every LW block leader → its owned sites' V/U. - * The CI pool leader → all CI fn params. - -Rank 0 assembles the gathered tensors into a temporary "full" ``ComponentModel`` -(with all sites in its decomposition targets) and returns that model's -``state_dict()``. The temp model shares ``target_model`` with rank 0's existing -component model — ``ComponentModel.__init__`` doesn't mutate ``target_model``, -so this is safe. - -Non-leader ranks no-op. All ranks must enter ``gather_full_state_dict_to_rank0`` -in sync so the P2P sends/recvs match up. +"""Offline checkpoint assembly for 3-pool training. + +The train loop never assembles a checkpoint. Instead each rank writes a +self-contained partial to a shared-FS scratch dir (see +``ThreePoolTrainer.snapshot``): its owned model params (LW leaders → owned-sites +V/U; CI leader → CI fn), its optimizer state, and (PPGD) its sources. A separate +async SLURM job then reads every partial for a step and assembles the canonical +artifacts off the training critical path. + +This module owns that offline assembly. It builds a full-decomposition +``ComponentModel`` on CPU as an assembly buffer, copies in every site's V/U + the +CI fn from the partials, and returns its ``state_dict()`` — the same flat schema +``load_component_model_from_checkpoint`` expects. No live ranks, no NCCL: the +assembly is pure file I/O + tensor copies. """ -# pyright: reportArgumentType=false +from typing import Any import torch -import torch.distributed as dist import torch.nn as nn from torch import Tensor @@ -30,111 +25,86 @@ from param_decomp.ci_sigmoids import SigmoidType from param_decomp.component_model import ComponentModel from param_decomp.decomposition_targets import DecompositionTarget -from param_decomp_lab.three_pool.context import CIContext, LWContext, PoolContext -def gather_full_state_dict_to_rank0( - ctx: PoolContext, - component_model: ComponentModel, - target_model: nn.Module, - run_batch: RunBatch, - ci_config: CiConfig, - sigmoid_type: SigmoidType, - c_per_site: dict[str, int], - device: torch.device, -) -> dict[str, Tensor] | None: - """Collect V/U + CI fn onto rank 0 and return the assembled state_dict. +def site_to_component_prefix(site: str) -> str: + """State-dict key prefix for a site's V/U params. - Returns the state_dict on rank 0; ``None`` everywhere else. All ranks must - call this function in sync (the gather uses ordered P2P sends/recvs). + ``ComponentModel`` registers components under an ``nn.ModuleDict`` keyed by + the site with dots replaced by dashes (so dots don't nest module levels), so + ``h.0.attn.q_proj`` → ``_components.h-0-attn-q_proj``. """ - world = ctx.world - match ctx: - case LWContext() if ctx.role.rank == 0: - return _rank0_assemble( - ctx=ctx, - local_component_model=component_model, - target_model=target_model, - run_batch=run_batch, - ci_config=ci_config, - sigmoid_type=sigmoid_type, - c_per_site=c_per_site, - device=device, - ) - case LWContext() if ctx.role.is_block_leader: - for s in ctx.role.owned_sites: - comp = component_model.components[s] - dist.send(comp.V.data.contiguous(), dst=0, group=world.cross_pool_p2p_group) - dist.send(comp.U.data.contiguous(), dst=0, group=world.cross_pool_p2p_group) - return None - case CIContext() if ctx.role.is_pool_leader: - assert component_model.ci_fn is not None, "CI pool must keep its CI fn" - for _, p in component_model.ci_fn.named_parameters(): - dist.send(p.data.contiguous(), dst=0, group=world.cross_pool_p2p_group) - return None - case _: - # Non-leader LW ranks, non-leader CI ranks, all PPGD ranks: no-op. - return None - - -def _rank0_assemble( - ctx: LWContext, - local_component_model: ComponentModel, + return f"_components.{site.replace('.', '-')}" + + +def owned_model_state_keys( + model_state_dict_keys: set[str], *, owned_sites: tuple[str, ...] +) -> set[str]: + """The V/U state-dict keys for ``owned_sites`` (LW block leaders' partial).""" + prefixes = tuple(f"{site_to_component_prefix(s)}." for s in owned_sites) + return {k for k in model_state_dict_keys if k.startswith(prefixes)} + + +def ci_fn_state_keys(model_state_dict_keys: set[str]) -> set[str]: + """The CI-fn state-dict keys (the CI pool leader's partial).""" + return {k for k in model_state_dict_keys if k.startswith("ci_fn.")} + + +def assemble_model_state_dict_from_partials( + *, + partials: list[dict[str, Any]], target_model: nn.Module, run_batch: RunBatch, ci_config: CiConfig, sigmoid_type: SigmoidType, c_per_site: dict[str, int], - device: torch.device, + all_sites: tuple[str, ...], ) -> dict[str, Tensor]: - """Build a full ComponentModel as an assembly buffer; populate from local - rank's V/U + recvs; return its state_dict.""" - world = ctx.world - full_targets = [DecompositionTarget(module_path=s, C=c_per_site[s]) for s in world.all_sites] - # Share target_model with the existing local component_model — ComponentModel - # __init__ doesn't mutate target_model, so two ComponentModels can coexist - # over the same target. + """Assemble the full ComponentModel state_dict from per-rank scratch partials. + + Each partial's ``model_params`` holds the CPU tensors that rank owns (a slice + of its own ``component_model.state_dict()``): LW block leaders contribute their + owned sites' ``_components..*``, the CI pool leader contributes ``ci_fn.*``. + The union of every partial's keys must exactly cover the full-decomposition + model's V/U + CI-fn keys (target-model params come from the fresh buffer). + """ + # ComponentModel asserts the target has no trainable params; build_target only + # `.eval()`s it, so freeze here (the training / load_component_model paths + # freeze at their own construction sites). + target_model.eval() + target_model.requires_grad_(False) + full_targets = [DecompositionTarget(module_path=s, C=c_per_site[s]) for s in all_sites] full_cm = ComponentModel( target_model=target_model, run_batch=run_batch, decomposition_targets=full_targets, ci_config=ci_config, sigmoid_type=sigmoid_type, - ).to(device) - - # Copy rank 0's own trained V/U into the assembly buffer. + ) + + # `ComponentModel` registers `target_model` as a submodule, so its frozen + # weights appear in the state_dict under the `target_model.` prefix. They come + # from the freshly-built buffer (which shares the real target_model), so the + # partials only need to cover the V/U + CI-fn keys — everything NOT under + # `target_model.`. + expected_keys = set(full_cm.state_dict().keys()) + fillable_keys = {k for k in expected_keys if not k.startswith("target_model.")} + + collected: dict[str, Tensor] = {} + for partial in partials: + for k, v in partial["model_params"].items(): + assert k not in collected, f"duplicate model param {k!r} across partials" + collected[k] = v + + assert set(collected.keys()) == fillable_keys, ( + "partials do not cover the full model state_dict:\n" + f" missing: {sorted(fillable_keys - set(collected))}\n" + f" extra: {sorted(set(collected) - fillable_keys)}" + ) + + full_state = full_cm.state_dict() with torch.no_grad(): - for s in ctx.role.owned_sites: - full_cm.components[s].V.data.copy_(local_component_model.components[s].V.data) - full_cm.components[s].U.data.copy_(local_component_model.components[s].U.data) - - # Recv V/U from every other LW block leader. Order on both sides must match - # the iteration order of `bg.owned_sites` for each non-rank-0 block leader. - for bg in world.layerwise_block_groups: - if bg.leader == 0: - continue - for s in bg.owned_sites: - V_template = full_cm.components[s].V.data - U_template = full_cm.components[s].U.data - V_buf = torch.empty_like(V_template) - U_buf = torch.empty_like(U_template) - dist.recv(V_buf, src=bg.leader, group=world.cross_pool_p2p_group) - dist.recv(U_buf, src=bg.leader, group=world.cross_pool_p2p_group) - with torch.no_grad(): - V_template.copy_(V_buf) - U_template.copy_(U_buf) - - # Recv CI fn params from CI pool leader. Same `named_parameters()` iteration - # order on both sides since the CI fn is constructed from the same config. - ci_leader = world.ci_ranks[0] - assert full_cm.ci_fn is not None, "checkpoint reconstruction needs a CI fn" - for _, p in full_cm.ci_fn.named_parameters(): - buf = torch.empty_like(p.data) - dist.recv(buf, src=ci_leader, group=world.cross_pool_p2p_group) - with torch.no_grad(): - p.data.copy_(buf) - - state_dict = {k: v.cpu() for k, v in full_cm.state_dict().items()} - del full_cm - torch.cuda.empty_cache() - return state_dict + for k, v in collected.items(): + full_state[k].copy_(v) + + return {k: v.cpu() for k, v in full_state.items()} diff --git a/param_decomp_lab/three_pool/consolidate.py b/param_decomp_lab/three_pool/consolidate.py new file mode 100644 index 000000000..5f7d28ac3 --- /dev/null +++ b/param_decomp_lab/three_pool/consolidate.py @@ -0,0 +1,193 @@ +"""Offline consolidation of a 3-pool checkpoint from its scratch partials. + +The 3-pool train loop writes per-rank partials to ``scratch_dir/step_/`` (see +``ThreePoolTrainer.snapshot``) and then continues — it never assembles the +canonical checkpoint. This module does that assembly off the critical path, +driven by the async SLURM job that already runs the slow eval. It reads every +partial for a step, builds the full ``ComponentModel`` state_dict + the canonical +``ThreePoolTrainingState``, writes ``model_.pth`` + ``training_.pth``, +prunes old ``training_*.pth`` (keeping all ``model_*.pth``), and removes the +step's scratch dir. + +Idempotent: if ``training_.pth`` already exists, consolidation is a no-op (the +async job may be retried). If the scratch dir is already gone (consolidation ran +before, or the run died), it's also a no-op — the consolidated artifacts are the +source of truth. +""" + +import os +import shutil +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn + +from param_decomp.batch_and_loss_fns import RunBatch +from param_decomp.ci_fns import CiConfig +from param_decomp.ci_sigmoids import SigmoidType +from param_decomp.log import logger +from param_decomp.training_state import ThreePoolTrainingState +from param_decomp_lab.infra.run_files import save_file +from param_decomp_lab.three_pool.checkpoint import assemble_model_state_dict_from_partials + +SNAPSHOT_SCRATCH_DIRNAME = ".snapshot_scratch" +"""Subdir under the run's out_dir that holds the per-step partials. The trainer +writes here; consolidation reads + cleans here.""" + +CONSOLIDATE_META_FILENAME = "meta.pth" + +DEFAULT_KEEP_LAST_N_TRAINING = 3 +"""How many ``training_.pth`` files to keep after consolidation. All +``model_.pth`` files are always kept (they're the downstream artifact).""" + + +def step_scratch_dir(scratch_dir: Path, step: int) -> Path: + return scratch_dir / f"step_{step}" + + +def consolidate_step( + *, + scratch_dir: Path, + out_dir: Path, + step: int, + target_model: nn.Module, + run_batch: RunBatch, + ci_config: CiConfig, + sigmoid_type: SigmoidType, + keep_last_n_training: int, +) -> None: + """Assemble + persist the step-S checkpoint from its scratch partials. + + No-op (returns early) when ``training_.pth`` already exists or the scratch + dir for the step is missing — both mean the work is already done (or was done + by a prior, possibly-retried, invocation). + """ + training_path = out_dir / f"training_{step}.pth" + step_dir = step_scratch_dir(scratch_dir, step) + if training_path.is_file(): + # Already consolidated. Clean up any scratch left behind by a prior run + # that wrote the checkpoint but died before removing it (e.g. crashed in + # prune), so this step stops looking unconsolidated. + logger.info(f"consolidate: {training_path.name} already exists; skipping") + shutil.rmtree(step_dir, ignore_errors=True) + return + + if not step_dir.is_dir(): + logger.warning( + f"consolidate: scratch dir {step_dir} missing and no {training_path.name}; " + f"nothing to consolidate for step {step}" + ) + return + + # PD_3POOL_SNAPSHOT_RANK0_SLEEP_S models the old slow rank-0 read. It now + # lives here, off the train loop, so injecting it proves the train loop's PG + # timeout is unaffected by a slow consolidation (the watchdog-safe test). + sleep_s = os.environ.get("PD_3POOL_SNAPSHOT_RANK0_SLEEP_S", "").strip() + if sleep_s: + logger.info(f"consolidate: sleep {sleep_s}s (fault injection, off train loop)") + time.sleep(float(sleep_s)) + + meta = torch.load(step_dir / CONSOLIDATE_META_FILENAME, map_location="cpu", weights_only=False) + world_size: int = meta["world_size"] + all_sites: tuple[str, ...] = tuple(meta["all_sites"]) + c_per_site: dict[str, int] = meta["c_per_site"] + + partials: list[dict[str, Any]] = [] + for r in range(world_size): + partials.append( + torch.load(step_dir / f"rank_{r}.pth", map_location="cpu", weights_only=False) + ) + logger.info(f"consolidate: read {world_size} partials for step {step}") + + model_state = assemble_model_state_dict_from_partials( + partials=partials, + target_model=target_model, + run_batch=run_batch, + ci_config=ci_config, + sigmoid_type=sigmoid_type, + c_per_site=c_per_site, + all_sites=all_sites, + ) + + components_optimizer: dict[str, dict[str, Any]] = {} + ci_fn_optimizer: dict[str, dict[str, Any]] = {} + ppgd_by_rank: dict[int, dict[str, Any]] = {} + for r, partial in enumerate(partials): + pool: str = partial["pool"] + match pool: + case "layerwise": + components_optimizer.update(partial["optimizer_by_name"]) + case "ci": + ci_fn_optimizer.update(partial["optimizer_by_name"]) + case "ppgd": + if "ppgd" in partial: + ppgd_by_rank[r] = partial["ppgd"] + case _: + raise AssertionError(f"unknown pool {pool!r} in rank-{r} partial") + + state = ThreePoolTrainingState( + step=step, + pd_config=meta["pd_config"], + runtime_config=meta["runtime_config"], + three_pool_config=meta["three_pool_config"], + layout_fingerprint=meta["layout_fingerprint"], + component_model=model_state, + components_optimizer=components_optimizer, + ci_fn_optimizer=ci_fn_optimizer, + ppgd_state_by_rank=ppgd_by_rank, + ) + + model_path = out_dir / f"model_{step}.pth" + save_file(model_state, model_path) + save_file(state, training_path) + logger.info(f"consolidate: wrote {model_path.name} + {training_path.name}") + + _prune_old_training(out_dir, keep_last_n=keep_last_n_training) + + shutil.rmtree(step_dir, ignore_errors=True) + logger.info(f"consolidate: removed scratch {step_dir}") + + +def _prune_old_training(out_dir: Path, *, keep_last_n: int) -> None: + """Delete all but the last ``keep_last_n`` ``training_.pth`` files. + + ``model_.pth`` files are never pruned — they're the downstream artifact + and are cheap relative to the full training state. + """ + steps: list[int] = [] + for p in out_dir.glob("training_*.pth"): + try: + steps.append(int(p.stem.removeprefix("training_"))) + except ValueError: + continue + steps.sort() + if len(steps) <= keep_last_n: + return + for step in steps[: len(steps) - keep_last_n]: + path = out_dir / f"training_{step}.pth" + # missing_ok: a concurrently-running consolidation job for another step + # may have already pruned this same old file. That's benign — the target + # state (this file gone) is reached either way. + path.unlink(missing_ok=True) + logger.info(f"consolidate: pruned old {path.name}") + + +def unconsolidated_steps(out_dir: Path) -> list[int]: + """Steps that have scratch partials on disk but no `training_.pth` yet. + + These are saves the async job never finished consolidating (it crashed, was + preempted, or never ran). They're safe + cheap to consolidate by re-running. + """ + scratch = out_dir / SNAPSHOT_SCRATCH_DIRNAME + if not scratch.is_dir(): + return [] + out: list[int] = [] + for d in scratch.glob("step_*"): + if not d.is_dir(): + continue + step = int(d.name.removeprefix("step_")) + if not (out_dir / f"training_{step}.pth").is_file(): + out.append(step) + return sorted(out) diff --git a/param_decomp_lab/three_pool/consolidate_cli.py b/param_decomp_lab/three_pool/consolidate_cli.py new file mode 100644 index 000000000..2db66177c --- /dev/null +++ b/param_decomp_lab/three_pool/consolidate_cli.py @@ -0,0 +1,65 @@ +"""Manual CPU-only consolidation of a 3-pool run's scratch partials. + +Recovery tool for a failed/preempted async consolidation. The train loop always +leaves a step's partials on disk until that step consolidates successfully, so a +failed consolidation is fixed by re-running this: + + python -m param_decomp_lab.three_pool.consolidate_cli p-xxxxxxxx + python -m param_decomp_lab.three_pool.consolidate_cli p-xxxxxxxx --step 5000 + +With no ``--step`` it consolidates every step that has partials but no +``training_.pth`` yet (see `unconsolidated_steps`). + +Separate module from `consolidate` so the `experiments.lm.run` import here +doesn't form an import cycle (`run` imports `consolidate`). +""" + +from pathlib import Path + +import fire + +from param_decomp.log import logger +from param_decomp_lab.experiments.lm.run import ( + LMExperimentConfig, + build_target, + make_run_batch, +) +from param_decomp_lab.experiments.utils import RUN_META_FILENAME +from param_decomp_lab.infra.settings import PARAM_DECOMP_OUT_DIR +from param_decomp_lab.three_pool.consolidate import ( + DEFAULT_KEEP_LAST_N_TRAINING, + SNAPSHOT_SCRATCH_DIRNAME, + consolidate_step, + unconsolidated_steps, +) + + +def cli( + run: str, + *, + step: int | None = None, + keep_last_n_training: int = DEFAULT_KEEP_LAST_N_TRAINING, +) -> None: + out_dir = Path(run) if Path(run).is_dir() else PARAM_DECOMP_OUT_DIR / "decompositions" / run + cfg = LMExperimentConfig.from_file(out_dir / RUN_META_FILENAME) + target_model = build_target(cfg.target) + run_batch = make_run_batch(cfg.target) + + steps = [step] if step is not None else unconsolidated_steps(out_dir) + assert steps, f"nothing to consolidate under {out_dir} (no unconsolidated scratch steps)" + logger.info(f"consolidate CLI: steps {steps} under {out_dir}") + for s in steps: + consolidate_step( + scratch_dir=out_dir / SNAPSHOT_SCRATCH_DIRNAME, + out_dir=out_dir, + step=s, + target_model=target_model, + run_batch=run_batch, + ci_config=cfg.pd.ci_config, + sigmoid_type=cfg.pd.sigmoid_type, + keep_last_n_training=keep_last_n_training, + ) + + +if __name__ == "__main__": + fire.Fire(cli) diff --git a/param_decomp_lab/three_pool/optimize.py b/param_decomp_lab/three_pool/optimize.py index e00bd2a3f..646cd6489 100644 --- a/param_decomp_lab/three_pool/optimize.py +++ b/param_decomp_lab/three_pool/optimize.py @@ -39,7 +39,6 @@ class boundary. import datetime import itertools import os -import shutil import time from contextlib import nullcontext from pathlib import Path @@ -78,8 +77,15 @@ class boundary. from param_decomp.sdpa_strict import verify_flash_attention_available from param_decomp.torch_helpers import loop_dataloader from param_decomp.training_state import ThreePoolTrainingState -from param_decomp_lab.three_pool.checkpoint import gather_full_state_dict_to_rank0 +from param_decomp_lab.three_pool.checkpoint import ( + ci_fn_state_keys, + owned_model_state_keys, +) from param_decomp_lab.three_pool.config import ThreePoolConfig +from param_decomp_lab.three_pool.consolidate import ( + CONSOLIDATE_META_FILENAME, + step_scratch_dir, +) from param_decomp_lab.three_pool.context import ( CIContext, LWContext, @@ -105,13 +111,13 @@ class boundary. ) from param_decomp_lab.three_pool.step_ppgd import step_ppgd -# Collective timeout for every 3-pool subgroup. Generous (vs the 10-min NCCL -# default) because a checkpoint save at XL serially reads ~100 GB of partials -# on rank 0 while the other 143 ranks wait at the post-save barrier — that read -# can take many minutes, and any rank reaching the next step's collective before -# rank 0 rejoins must not trip the watchdog. Env-overridable (seconds) so the -# save-watchdog repro can force the timeout at small scale. -_DEFAULT_PG_TIMEOUT = datetime.timedelta(minutes=30) +# Collective timeout for every 3-pool subgroup. Now that consolidation is async +# (off the train loop), the longest gap between collectives on the loop is the +# in-train (fast) eval pass plus a checkpoint partial-write barrier — minutes, +# not the old ~10-min rank-0 read of ~100 GB. 10 min covers that worst case with +# generous margin while still being far below "hang forever". Env-overridable +# (seconds) so the watchdog-safe-at-low-timeout test can force a tight bound. +_DEFAULT_PG_TIMEOUT = datetime.timedelta(minutes=10) def _resolve_pg_timeout() -> datetime.timedelta: @@ -161,13 +167,15 @@ class ThreePoolTrainer: the first batch of :meth:`run` because its source tensor shapes depend on the data's sequence dims. - Resume support is rank-local: :meth:`snapshot` produces a self-contained - :class:`~param_decomp.trainer_snapshot.TrainerSnapshot` whose ``resume`` - half carries this rank's slice, and :meth:`from_snapshot` reconstructs - one from it. The lab's resume loader composes per-rank shards across - ranks. The snapshot's ``consumable`` half is the gathered full state - (rank 0 only; ``None`` elsewhere) — the gather happens inside - :meth:`snapshot`, so all ranks must call it in sync. + Checkpointing is split across two phases. :meth:`snapshot` runs on the train + loop: each rank writes a self-contained partial (its owned model params + + optimizer state + PPGD sources) to the shared-FS scratch dir, with a single + barrier — no rank-0 read, no model assembly. The async consolidation job + (:mod:`param_decomp_lab.three_pool.consolidate`) later reads every partial, + assembles the canonical :class:`~param_decomp.training_state.ThreePoolTrainingState` + + ``model_.pth``, and runs the slow eval. :meth:`from_snapshot` + reconstructs a trainer from a consolidated ``ThreePoolTrainingState`` (each + rank loads its own slice). """ pd_config: PDConfig @@ -399,156 +407,107 @@ def _named_params_for_my_optimizer(self) -> list[tuple[str, nn.Parameter]]: case PPGDContext(): return [] - def snapshot(self, scratch_dir: Path) -> ThreePoolTrainingState | None: - """Canonical point-in-time view of the 3-pool trainer. - - On rank 0 (the only rank a sink consumes from), assembles a - topology-free :class:`ThreePoolTrainingState`: - - * ``component_model``: the full gathered model state (every site's V/U - plus the CI fn) from :func:`gather_full_state_dict_to_rank0`. - * ``components_optimizer`` / ``ci_fn_optimizer``: per-parameter - optimizer state keyed by param name, merged across all ranks. - ``components_optimizer`` comes from the layerwise pool (each LW - rank contributes its owned sites); ``ci_fn_optimizer`` comes from - the CI pool (all CI ranks have the same state via in-pool - all-reduce, so any suffices). - * ``ppgd_state_by_rank``: PPGD adversarial sources, keyed by PPGD - rank id. Genuinely rank-coupled for ``PerBatchPerPositionScope`` - sources (sized by the rank-local batch slice). - - Returns ``None`` on non-rank-0 — the lab sink is silent on those - ranks anyway, and the trainer never needs their canonical state. - - Args: - scratch_dir: Shared-filesystem directory all ranks can write - to / read from. Used as the rendezvous for the per-rank - contributions (optimizer state, PPGD sources). Each - snapshot at step ``S`` writes / reads under - ``scratch_dir / f"step_{S}"`` and cleans it up on rank 0. + def snapshot(self, scratch_dir: Path) -> None: + """Write this rank's self-contained checkpoint partial; then return. + + Each rank writes everything needed to reconstruct the checkpoint for + its own slice to ``scratch_dir / f"step_{S}" / f"rank_{rank}.pth"``: + + * ``model_params``: a slice of this rank's ``component_model.state_dict()`` — + LW block leaders contribute their owned sites' ``_components..*``, + the CI pool leader contributes ``ci_fn.*``, everyone else contributes + nothing (their V/U / CI fn are replicas of a leader's). + * ``optimizer_by_name``: this rank's optimizer state, name-keyed. + * ``ppgd``: PPGD adversarial sources (PPGD ranks only). + + Rank 0 also writes ``meta.pth`` (configs + layout fingerprint + the + ``c_per_site`` / ``all_sites`` needed to rebuild the full model). A + single barrier ensures every partial is on the shared FS before the + train loop continues — there is NO rank-0 read and NO model assembly + here. The async consolidation job + (:func:`param_decomp_lab.three_pool.consolidate.consolidate_step`) + reads these partials off the critical path. """ - trace("snapshot: enter gather_full_state_dict_to_rank0") - gathered_model = gather_full_state_dict_to_rank0( - ctx=self.ctx, - component_model=self.component_model, - target_model=self._target_model, - run_batch=self._run_batch, - ci_config=self.pd_config.ci_config, - sigmoid_type=self.pd_config.sigmoid_type, - c_per_site=self.runtime.c_per_site, - device=self._device, - ) - trace("snapshot: gather_full_state_dict_to_rank0 done") + partial = self._build_my_partial() - my_named_params = self._named_params_for_my_optimizer() - my_optimizer_by_name: dict[str, dict[str, Any]] = ( - optimizer_state_by_name(self.optimizer, my_named_params) - if self.optimizer is not None - else {} - ) - my_contribution: dict[str, Any] = { - "pool": self.ctx.kind, - "optimizer_by_name": my_optimizer_by_name, - } - if self.ppgd_state is not None: - my_contribution["ppgd"] = self.ppgd_state.state_dict() - elif self._pending_ppgd_resume_state is not None: - my_contribution["ppgd"] = self._pending_ppgd_resume_state - - # File-based gather rather than dist.gather_object: at XL the aggregate - # pickle payload (LW optimizer state + PPGD sources) is ~8 GB and - # gather_object scales poorly. Each rank writes its contribution to a - # shared-FS path, a barrier ensures all writes are visible, then rank 0 - # reads them all. p2p_group = self.ctx.world.cross_pool_p2p_group - world_size = self.ctx.world.world_size - step_dir = scratch_dir / f"step_{self.step}" + step_dir = step_scratch_dir(scratch_dir, self.step) if self.ctx.role.rank == 0: step_dir.mkdir(parents=True, exist_ok=True) + torch.save(self._build_meta(), step_dir / CONSOLIDATE_META_FILENAME) + # The mkdir + meta write race: rank 0 must finish them before any other + # rank writes its partial into step_dir. trace("snapshot: enter barrier (pre-write)") dist.barrier(group=p2p_group) trace("snapshot: barrier (pre-write) done") partial_path = step_dir / f"rank_{self.ctx.role.rank}.pth" trace(f"snapshot: writing partial {partial_path.name}") - torch.save(my_contribution, partial_path) + torch.save(partial, partial_path) trace("snapshot: partial write done") - trace("snapshot: enter barrier (post-write)") - dist.barrier(group=p2p_group) - trace("snapshot: barrier (post-write) done") - - # Only rank 0 assembles the canonical state — it serially reads every - # rank's partial (~100 GB at XL). Non-rank-0 ranks have nothing left to - # do here, but they MUST NOT advance to the next training step until - # rank 0 finishes: the next step's collectives (PPGD V/U all-reduce, - # cross-pool sends) would block on rank 0 still reading, and at XL that - # read overruns the collective watchdog and aborts the job. The - # rejoining barrier below makes all ranks resume training in lock-step. - gathered_state: ThreePoolTrainingState | None = None - if self.ctx.role.rank == 0: - gathered_state = self._assemble_rank0_state( - step_dir=step_dir, - world_size=world_size, - gathered_model=gathered_model, - ) - - # PD_3POOL_DISABLE_REJOIN_BARRIER reproduces the pre-fix bug (non-rank-0 - # ranks racing ahead while rank 0 reads). For the save-watchdog repro - # only — never set in production. + # One rejoining barrier so all ranks leave snapshot together. This is + # cheap now (just the partial write, no 100 GB read), so it cannot + # overrun the collective watchdog. PD_3POOL_DISABLE_REJOIN_BARRIER skips + # it for the legacy race repro. The old PD_3POOL_SNAPSHOT_RANK0_SLEEP_S + # fault injection now lives in the async consolidate job (off the loop), + # since that is where the slow read moved — the train loop has nothing + # slow left to sleep through. if os.environ.get("PD_3POOL_DISABLE_REJOIN_BARRIER", "").strip() in ("1", "true"): trace("snapshot: REJOIN BARRIER DISABLED (repro mode)") - return gathered_state - trace("snapshot: enter barrier (rejoin after rank-0 read)") + return + trace("snapshot: enter barrier (post-write rejoin)") dist.barrier(group=p2p_group) - trace("snapshot: barrier (rejoin after rank-0 read) done") - return gathered_state + trace("snapshot: barrier (post-write rejoin) done") - def _assemble_rank0_state( - self, - *, - step_dir: Path, - world_size: int, - gathered_model: dict[str, Any] | None, - ) -> ThreePoolTrainingState: - rank0_read_sleep_s = os.environ.get("PD_3POOL_SNAPSHOT_RANK0_SLEEP_S", "").strip() - if rank0_read_sleep_s: - trace(f"snapshot: rank-0 read sleep {rank0_read_sleep_s}s (fault injection)") - time.sleep(float(rank0_read_sleep_s)) - - gathered: list[dict[str, Any]] = [] - for r in range(world_size): - gathered.append(torch.load(step_dir / f"rank_{r}.pth", weights_only=False)) - trace("snapshot: rank-0 read all partials") - shutil.rmtree(step_dir, ignore_errors=True) - components_by_name: dict[str, dict[str, Any]] = {} - ci_fn_by_name: dict[str, dict[str, Any]] = {} - ppgd_by_rank: dict[int, dict[str, Any]] = {} - for r, c in enumerate(gathered): - pool: str = c["pool"] - match pool: - case "layerwise": - components_by_name.update(c["optimizer_by_name"]) - case "ci": - ci_fn_by_name.update(c["optimizer_by_name"]) - case "ppgd": - if "ppgd" in c: - ppgd_by_rank[r] = c["ppgd"] - case _: - raise AssertionError(f"unknown pool {pool!r} in rank-{r} contribution") - assert gathered_model is not None # rank 0 always has the gathered model - return ThreePoolTrainingState( - step=self.step, - pd_config=self.pd_config.model_dump(), - runtime_config=self.runtime_config.model_dump(), - three_pool_config=self.three_pool_config.model_dump(), - layout_fingerprint=_layout_fingerprint(self.ctx), - component_model=gathered_model, - components_optimizer=components_by_name, - ci_fn_optimizer=ci_fn_by_name, - ppgd_state_by_rank=ppgd_by_rank, + def _build_my_partial(self) -> dict[str, Any]: + my_named_params = self._named_params_for_my_optimizer() + my_optimizer_by_name: dict[str, dict[str, Any]] = ( + optimizer_state_by_name(self.optimizer, my_named_params) + if self.optimizer is not None + else {} ) + partial: dict[str, Any] = { + "pool": self.ctx.kind, + "model_params": self._owned_model_params(), + "optimizer_by_name": my_optimizer_by_name, + } + if self.ppgd_state is not None: + partial["ppgd"] = self.ppgd_state.state_dict() + elif self._pending_ppgd_resume_state is not None: + partial["ppgd"] = self._pending_ppgd_resume_state + return partial + + def _owned_model_params(self) -> dict[str, Tensor]: + """The slice of this rank's model state_dict it's responsible for saving. + + Only leaders contribute: LW block leaders own their sites' V/U, the CI + pool leader owns the CI fn. Every other rank holds a replica, so it + contributes nothing — the union across leaders covers the full model. + """ + sd = self.component_model.state_dict() + keys = set(sd.keys()) + match self.ctx: + case LWContext() if self.ctx.role.is_block_leader: + owned = owned_model_state_keys(keys, owned_sites=self.ctx.role.owned_sites) + case CIContext() if self.ctx.role.is_pool_leader: + owned = ci_fn_state_keys(keys) + case _: + owned = set() + return {k: sd[k].cpu() for k in owned} + + def _build_meta(self) -> dict[str, Any]: + return { + "step": self.step, + "world_size": self.ctx.world.world_size, + "all_sites": list(self.ctx.world.all_sites), + "c_per_site": dict(self.runtime.c_per_site), + "pd_config": self.pd_config.model_dump(), + "runtime_config": self.runtime_config.model_dump(), + "three_pool_config": self.three_pool_config.model_dump(), + "layout_fingerprint": _layout_fingerprint(self.ctx), + } @classmethod def from_snapshot( @@ -875,9 +834,8 @@ def _to_device(b: Any) -> Any: ) if cadence.should_save(step): - snap = self.snapshot(scratch_dir) - if snap is not None: - sink.checkpoint(snap, final=False) + self.snapshot(scratch_dir) + sink.checkpoint_written(step, final=False) batch_T = ( batch_T_plus_1 @@ -890,9 +848,8 @@ def _to_device(b: Any) -> Any: profiler.step() self.step = n_steps - snap = self.snapshot(scratch_dir) - if snap is not None: - sink.checkpoint(snap, final=True) + self.snapshot(scratch_dir) + sink.checkpoint_written(self.step, final=True) def optimize_three_pool(