\

#!/usr/bin/env python3
"""Train or resume the finalized ASTERIZER 1B / 128K-vocab pretraining recipe.

Key properties of this trainer:
  * Sequence PACKING (no padding waste, no document truncation): documents are
    tokenized, joined with an EOS separator, and chunked into dense
    `sequence_length` blocks. Every trained token is a real token.
  * WSD (warmup-stable-decay) learning-rate schedule matching the recipe, which
    keeps the run extensible (the stable phase can be lengthened for more epochs
    and only the final decay window changes).
  * FlashAttention-2 when available (falls back to SDPA), decoupled weight decay
    (norms/embeddings excluded), DDP no_sync during gradient accumulation.
  * Real token-throughput accounting (tokens/sec) surfaced in the run summary so
    the pilot can report how fast a given GPU is before the full run starts.
"""
import argparse
import contextlib
import json
import math
import os
import random
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import pyarrow.parquet as pq
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM


def log(message: str) -> None:
    ts = time.strftime("%H:%M:%S")
    print(f"[{ts}] {message}", flush=True)


def is_distributed() -> bool:
    return dist.is_available() and dist.is_initialized()


def get_rank() -> int:
    return dist.get_rank() if is_distributed() else 0


def get_world_size() -> int:
    return dist.get_world_size() if is_distributed() else 1


def is_main_process() -> bool:
    return get_rank() == 0


def barrier() -> None:
    if is_distributed():
        if torch.cuda.is_available():
            dist.barrier(device_ids=[torch.cuda.current_device()])
        else:
            dist.barrier()


def log_main(message: str) -> None:
    if is_main_process():
        log(message)


def init_runtime() -> torch.device:
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        if torch.cuda.is_available():
            torch.cuda.set_device(local_rank)
            # Bind NCCL to the local device so barrier() runs on the correct stream;
            # avoids the "using the device under current context" ambiguity that can
            # let barriers stall behind pending kernels on other streams.
            device = torch.device("cuda", local_rank)
            dist.init_process_group(
                backend="nccl",
                device_id=device,
                timeout=__import__("datetime").timedelta(minutes=30),
            )
            return device
        dist.init_process_group(backend="gloo")
        return torch.device("cpu")

    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def cleanup_runtime() -> None:
    if is_distributed():
        dist.destroy_process_group()


def set_seed(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_json(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def limit_files(paths: list[Path], max_files: int) -> list[Path]:
    if max_files <= 0:
        return paths
    return paths[:max_files]


def load_texts(split_dir: Path, max_rows: int, text_column: str, max_files: int) -> list[str]:
    texts: list[str] = []
    parquet_files = limit_files(sorted(split_dir.glob("*.parquet")), max_files)
    for parquet_path in parquet_files:
        if len(texts) >= max_rows:
            break
        table = pq.read_table(parquet_path, columns=[text_column])
        for value in table.column(text_column).to_pylist():
            if value is None:
                continue
            text = str(value).strip()
            if not text:
                continue
            texts.append(text)
            if len(texts) >= max_rows:
                break
    return texts


def dir_manifest_sha(d: Path) -> str:
    """SHA256 over the sorted (filename:bytes:rows) list of a train dir. Detects
    altered/replaced data even when the directory name is unchanged (preflight gate:
    'altered data with the same directory name must not pass resume')."""
    import hashlib
    lines = []
    for f in sorted(Path(d).glob("*.parquet")):
        try:
            rows = pq.ParquetFile(f).metadata.num_rows
        except Exception:
            rows = -1
        lines.append(f"{f.name}:{f.stat().st_size}:{rows}")
    return hashlib.sha256("\n".join(lines).encode()).hexdigest()


def _domain_of(src: str) -> str:
    """Bucket a validation row's source into a domain for per-domain val loss
    (preflight gate: aggregate loss hides regressions). Matches prep source names."""
    s = (src or "").lower()
    if "old_pretrain" in s:
        return "old"
    if "megamath" in s or "math" in s:
        return "math"
    if "python_edu" in s or "code" in s:
        return "code"
    if "openthoughts" in s or "reason" in s:
        return "reasoning"
    return "english"


def load_texts_by_domain(split_dir: Path, max_rows: int, text_column: str, max_files: int) -> dict:
    """Like load_texts but groups by domain using the `src` column when present.
    Returns {domain: [texts]}. Falls back to a single 'english' bucket if no src."""
    from collections import defaultdict
    buckets: dict = defaultdict(list)
    n = 0
    for parquet_path in limit_files(sorted(split_dir.glob("*.parquet")), max_files):
        if n >= max_rows:
            break
        table = pq.read_table(parquet_path)
        has_src = "src" in table.column_names
        texts = table.column(text_column).to_pylist()
        srcs = table.column("src").to_pylist() if has_src else [None] * len(texts)
        for text, src in zip(texts, srcs):
            if text is None:
                continue
            text = str(text).strip()
            if not text:
                continue
            buckets[_domain_of(src)].append(text)
            n += 1
            if n >= max_rows:
                break
    return dict(buckets)


def eval_loss_only(model, tokenizer, texts, batch_size, seq_len, device, max_batches) -> float:
    """Mean cross-entropy over up to max_batches of `texts` (no generation). Used for
    the per-domain validation breakdown; runs on the main process only."""
    im = unwrap_model(model)
    im.eval()
    losses, nb = [], 0
    with torch.no_grad():
        for off in range(0, len(texts), batch_size):
            if nb >= max_batches:
                break
            batch = build_eval_batch(texts[off:off + batch_size], tokenizer, seq_len, device)
            losses.append(float(im(**batch).loss.detach().cpu()))
            nb += 1
    return sum(losses) / max(1, len(losses)) if losses else float("nan")


class PackingStreamLoader:
    """Stream documents from parquet shards and emit dense packed token blocks.

    Documents are tokenized (no special tokens), joined with a single EOS id as a
    separator, and cut into fixed-length `seq_len` blocks. There is no padding and
    no truncation: long documents span multiple blocks and short documents are
    packed together. Blocks are the unit consumed by the training loop.

    Resume policy: shuffling is deterministic per (seed, epoch). On resume we fast
    forward to the resumed epoch so the shuffle order matches; within-epoch byte
    position is not restored (blocks already consumed in the current epoch may be
    re-seen). This is standard for streaming pretraining and has no correctness
    impact on the objective.
    """

    def __init__(
        self,
        split_dir: Path,
        text_column: str,
        tokenizer: AutoTokenizer,
        seq_len: int,
        seed: int,
        parquet_batch_rows: int,
        max_files: int,
        rank: int,
        world_size: int,
        start_epoch: int = 0,
    ) -> None:
        files = limit_files(sorted(split_dir.glob("*.parquet")), max_files)
        if not files:
            raise RuntimeError(f"No parquet files found in {split_dir}")
        self.files = files
        self.text_column = text_column
        self.tokenizer = tokenizer
        self.seq_len = int(seq_len)
        self.seed = int(seed)
        self.parquet_batch_rows = parquet_batch_rows
        self.rank = int(rank)
        self.world_size = max(1, int(world_size))
        self.eos_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id
        self.epoch = int(start_epoch)
        self.token_buffer: list[int] = []
        self.active_files: list[Path] = []
        self.file_index = 0
        self.batch_iter = None
        self._start_epoch(self.epoch)

    def _start_epoch(self, epoch: int) -> None:
        self.epoch = epoch
        rng = random.Random(self.seed + epoch)
        order = list(self.files)
        rng.shuffle(order)
        active = order[self.rank :: self.world_size]
        if not active:
            active = [order[self.rank % len(order)]]
        self.active_files = active
        self.file_index = 0
        self.batch_iter = None
        self.token_buffer.clear()

    def _fill_tokens(self) -> None:
        # Keep pulling parquet row batches (across files / epochs) until we have
        # at least one full block of tokens available.
        while len(self.token_buffer) < self.seq_len:
            if self.batch_iter is None:
                if self.file_index >= len(self.active_files):
                    self._start_epoch(self.epoch + 1)
                parquet_path = self.active_files[self.file_index]
                self.file_index += 1
                self.batch_iter = pq.ParquetFile(parquet_path).iter_batches(
                    batch_size=self.parquet_batch_rows,
                    columns=[self.text_column],
                )
            try:
                batch = next(self.batch_iter)
            except StopIteration:
                self.batch_iter = None
                continue

            texts = [
                str(value).strip()
                for value in batch.column(self.text_column).to_pylist()
                if value is not None and str(value).strip()
            ]
            if not texts:
                continue
            encoded = self.tokenizer(texts, add_special_tokens=False)["input_ids"]
            for ids in encoded:
                self.token_buffer.extend(ids)
                self.token_buffer.append(self.eos_id)

    def next_block(self) -> list[int]:
        self._fill_tokens()
        block = self.token_buffer[: self.seq_len]
        del self.token_buffer[: self.seq_len]
        return block

    def next_batch(self, batch_size: int) -> torch.Tensor:
        blocks = [self.next_block() for _ in range(batch_size)]
        return torch.tensor(blocks, dtype=torch.long)


@dataclass
class EvalSnapshot:
    step: int
    train_loss: float
    validation_loss: float
    validation_perplexity: float
    samples: list  # list of {"prompt": str, "generation": str}


DEFAULT_SAMPLE_PROMPTS = [
    "Explain what machine learning is in simple terms.",
    "The chemical symbol for gold is",
    "Write a short paragraph about why the sky appears blue during the day.",
    "Q: What is 47 times 8?\nA:",
    "Once upon a time in a small village near the mountains,",
]


def build_packed_train_batch(
    loader: PackingStreamLoader,
    batch_size: int,
    device: torch.device,
) -> dict[str, torch.Tensor]:
    input_ids = loader.next_batch(batch_size).to(device)
    attention_mask = torch.ones_like(input_ids)
    labels = input_ids.clone()
    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}


def build_eval_batch(
    texts: list[str],
    tokenizer: AutoTokenizer,
    seq_len: int,
    device: torch.device,
) -> dict[str, torch.Tensor]:
    encodings = tokenizer(
        texts,
        truncation=True,
        max_length=seq_len,
        padding="max_length",
        return_tensors="pt",
    )
    labels = encodings["input_ids"].clone()
    labels[encodings["attention_mask"] == 0] = -100
    encodings["labels"] = labels
    return {key: value.to(device) for key, value in encodings.items()}


def reduce_scalar(value: float, device: torch.device) -> float:
    if not is_distributed():
        return float(value)
    tensor = torch.tensor([value], device=device, dtype=torch.float32)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    tensor /= get_world_size()
    return float(tensor.item())


def maybe_apply_liger(mode: str) -> bool:
    """Patch HF Llama with Liger fused kernels (fused linear cross-entropy avoids
    materializing the full [B, S, 131072] logits, the main memory wall for this
    large-vocab model; also fuses RMSNorm/RoPE/SwiGLU). Must run before model init."""
    if mode == "off":
        return False
    try:
        from liger_kernel.transformers import apply_liger_kernel_to_llama
    except Exception as err:  # noqa: BLE001
        if mode == "on":
            raise RuntimeError(f"--use-liger on but liger-kernel is not importable: {err}")
        log_main(f"Liger kernel not available ({err}); using stock HF Llama")
        return False
    apply_liger_kernel_to_llama()
    log_main("Liger kernel applied (fused RMSNorm/RoPE/SwiGLU + fused linear cross-entropy)")
    return True


def build_model(
    model_cfg: dict,
    tokenizer: AutoTokenizer,
    device: torch.device,
    model_dtype: torch.dtype | None = None,
) -> LlamaForCausalLM:
    model_info = model_cfg["model"]
    config = LlamaConfig(
        vocab_size=int(model_cfg["tokenizer"]["vocab_size"]),
        hidden_size=int(model_info["d_model"]),
        intermediate_size=int(model_info["d_ff"]),
        num_hidden_layers=int(model_info["n_layers"]),
        num_attention_heads=int(model_info["n_heads"]),
        num_key_value_heads=int(model_info["n_kv_heads"]),
        max_position_embeddings=int(model_info["context_length"]),
        rms_norm_eps=float(model_info["rms_norm_eps"]),
        rope_theta=float(model_info["rope_theta"]),
        attention_bias=bool(model_info.get("bias", False)),
        attention_dropout=float(model_info.get("dropout", 0.0)),
        hidden_act="silu",
        tie_word_embeddings=bool(model_info["tie_embeddings"]),
        pad_token_id=tokenizer.pad_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=False,
    )

    # Prefer FlashAttention-2 for training throughput; fall back to SDPA, then eager.
    # FlashAttention requires CUDA, so only attempt it on GPU.
    candidate_impls = ("flash_attention_2", "sdpa", "eager") if device.type == "cuda" else ("sdpa", "eager")
    model = None
    for impl in candidate_impls:
        try:
            config._attn_implementation = impl
            model = LlamaForCausalLM(config)
            log_main(f"Attention implementation: {impl}")
            break
        except (ImportError, ValueError) as err:
            log_main(f"Attention impl {impl} unavailable ({err}); trying next")
    if model is None:
        model = LlamaForCausalLM(config)
    if model_dtype is not None:
        model = model.to(dtype=model_dtype)
    return model.to(device)


def build_optimizer(model: torch.nn.Module, training_cfg: dict, learning_rate: float) -> torch.optim.Optimizer:
    """AdamW with weight decay excluded from norms, biases and embeddings."""
    decay_params, no_decay_params = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        lname = name.lower()
        if param.ndim <= 1 or "norm" in lname or "embed" in lname:
            no_decay_params.append(param)
        else:
            decay_params.append(param)
    param_groups = [
        {"params": decay_params, "weight_decay": float(training_cfg["weight_decay"])},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(
        param_groups,
        lr=learning_rate,
        betas=tuple(training_cfg["betas"]),
        eps=float(training_cfg["epsilon"]),
        foreach=False,
    )


def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
    return model.module if isinstance(model, DDP) else model


def build_wsd_scheduler(
    optimizer: torch.optim.Optimizer,
    max_steps: int,
    warmup_steps: int,
    final_decay_steps: int,
    min_lr_ratio: float,
):
    """Warmup-Stable-Decay schedule.

    Linear warmup -> constant peak (stable) -> cosine decay to min_lr over the last
    `final_decay_steps`. The stable phase absorbs any change in total steps (e.g.
    training for more epochs), so only the final decay window is fixed.
    """
    warmup_steps = max(0, int(warmup_steps))
    final_decay_steps = max(0, int(final_decay_steps))
    decay_start = max(warmup_steps, max_steps - final_decay_steps)

    def lr_lambda(step_index: int) -> float:
        step_num = step_index + 1
        if warmup_steps > 0 and step_num <= warmup_steps:
            return max(1e-8, step_num / warmup_steps)
        if step_num <= decay_start:
            return 1.0
        decay_total = max(1, max_steps - decay_start)
        progress = min(1.0, (step_num - decay_start) / decay_total)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)


def choose_precision(name: str, device: torch.device) -> tuple[str, torch.dtype | None]:
    if name == "auto":
        if device.type == "cuda" and torch.cuda.is_bf16_supported():
            return "bf16", torch.bfloat16
        if device.type == "cuda":
            return "fp16", torch.float16
        return "fp32", None
    if name == "bf16":
        return "bf16", torch.bfloat16
    if name == "fp16":
        return "fp16", torch.float16
    return "fp32", None


def save_checkpoint(
    checkpoint_dir: Path,
    model: torch.nn.Module,
    tokenizer: AutoTokenizer,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LambdaLR,
    step: int,
    train_tokens_processed: int,
    data_epoch: int,
    summary_state: dict,
) -> None:
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    unwrapped = unwrap_model(model)
    unwrapped.save_pretrained(checkpoint_dir, safe_serialization=True)
    tokenizer.save_pretrained(checkpoint_dir)
    state = {
        "step": step,
        "train_tokens_processed": train_tokens_processed,
        "data_epoch": data_epoch,
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "python_random_state": random.getstate(),
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state_all": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
        "summary_state": summary_state,
        "config_hash": globals().get("_RUN_CONFIG_HASH", ""),
        "train_dir_name": globals().get("_TRAIN_DIR_NAME", ""),
        "train_data_sha": globals().get("_TRAIN_DATA_SHA", ""),
    }
    torch.save(state, checkpoint_dir / "training_state.pt")


def prune_old_checkpoints(output_dir: Path, keep_last: int) -> None:
    """Keep only the most recent `keep_last` checkpoint-step-* dirs to bound disk use."""
    if keep_last <= 0:
        return
    checkpoints = sorted(
        (p for p in output_dir.glob("checkpoint-step-*") if p.is_dir()),
        key=lambda p: int(p.name.rsplit("-", 1)[-1]),
    )
    import shutil

    for stale in checkpoints[:-keep_last]:
        shutil.rmtree(stale, ignore_errors=True)
        log(f"Pruned old checkpoint {stale.name}")


def _hf_checkpoint_prefix(repo_subfolder: str, run_name: str) -> str:
    return f"{repo_subfolder.rstrip('/')}/{run_name}/"


def upload_checkpoint_to_hf(
    checkpoint_dir: Path,
    repo_id: str,
    repo_subfolder: str,
    write_token: str,
    run_name: str,
    keep_last: int,
) -> None:
    """Upload one checkpoint dir to HF, then prune old remote checkpoints so the
    repo mirrors the local keep-last policy. Blocking; call from a background thread."""
    from huggingface_hub import HfApi

    api = HfApi(token=write_token)
    prefix = _hf_checkpoint_prefix(repo_subfolder, run_name)
    path_in_repo = f"{prefix}{checkpoint_dir.name}"
    try:
        # upload_folder supports path_in_repo (subfolders); upload_large_folder does not.
        # This runs in a background thread so its speed does not block training.
        api.upload_folder(
            folder_path=str(checkpoint_dir),
            repo_id=repo_id,
            repo_type="dataset",
            path_in_repo=path_in_repo,
            commit_message=f"Checkpoint {checkpoint_dir.name} ({run_name})",
        )
        log(f"Uploaded checkpoint to hf://{repo_id}/{path_in_repo}")
    except Exception as err:  # noqa: BLE001 - background upload must never kill training
        log(f"WARNING: checkpoint upload failed for {checkpoint_dir.name}: {err}")
        return

    if keep_last > 0:
        try:
            _, all_steps = _list_hf_checkpoint_steps(api, repo_id, repo_subfolder, run_name, write_token)
            for old_step in sorted(all_steps)[:-keep_last]:
                api.delete_folder(
                    path_in_repo=f"{prefix}checkpoint-step-{old_step:05d}",
                    repo_id=repo_id,
                    repo_type="dataset",
                    commit_message=f"Prune old checkpoint step {old_step}",
                )
                log(f"Pruned remote checkpoint step {old_step}")
        except Exception as err:  # noqa: BLE001
            log(f"WARNING: remote checkpoint prune failed: {err}")


def _list_hf_checkpoint_steps(api, repo_id: str, repo_subfolder: str, run_name: str, token: str):
    """Return (files_by_step, sorted_steps) for uploaded checkpoints in the run folder."""
    files = api.list_repo_files(repo_id=repo_id, repo_type="dataset", token=token or None)
    prefix = _hf_checkpoint_prefix(repo_subfolder, run_name)
    by_step: dict[int, list[str]] = {}
    for path in files:
        if not path.startswith(prefix):
            continue
        head = path[len(prefix):].split("/", 1)[0]
        if head.startswith("checkpoint-step-"):
            try:
                step = int(head.rsplit("-", 1)[-1])
            except ValueError:
                continue
            by_step.setdefault(step, []).append(path)
    return by_step, list(by_step.keys())


def download_latest_hf_checkpoint(
    repo_id: str,
    repo_subfolder: str,
    run_name: str,
    token: str,
    output_dir: Path,
) -> Path | None:
    """Find the newest uploaded checkpoint on HF and download it into output_dir under
    the standard checkpoint-step-XXXXX name so resume treats it like a local one."""
    from huggingface_hub import HfApi, hf_hub_download
    import shutil

    api = HfApi(token=token or None)
    try:
        by_step, steps = _list_hf_checkpoint_steps(api, repo_id, repo_subfolder, run_name, token)
    except Exception as err:  # noqa: BLE001
        log(f"Could not list HF checkpoints: {err}")
        return None
    if not steps:
        return None
    latest = max(steps)
    dest = output_dir / f"checkpoint-step-{latest:05d}"
    dest.mkdir(parents=True, exist_ok=True)
    for remote in by_step[latest]:
        cached = hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=remote, token=token or None)
        shutil.copy(cached, dest / Path(remote).name)
    if (dest / "training_state.pt").exists():
        log(f"Downloaded HF checkpoint step {latest} to {dest}")
        return dest
    log(f"HF checkpoint step {latest} was incomplete (no training_state.pt)")
    return None


def find_latest_local_checkpoint(output_dir: Path) -> Path | None:
    if not output_dir.exists():
        return None
    checkpoints = [p for p in output_dir.glob("checkpoint-step-*") if (p / "training_state.pt").exists()]
    if not checkpoints:
        return None
    return max(checkpoints, key=lambda p: int(p.name.rsplit("-", 1)[-1]))


def move_optimizer_state_to_device(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
    for state in optimizer.state.values():
        for key, value in list(state.items()):
            if torch.is_tensor(value):
                state[key] = value.to(device)


def load_checkpoint_state(
    checkpoint_dir: Path,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LambdaLR,
    device: torch.device,
) -> tuple[int, int, int, dict]:
    unwrapped = unwrap_model(model)
    loaded_model = LlamaForCausalLM.from_pretrained(checkpoint_dir)
    unwrapped.load_state_dict(loaded_model.state_dict())
    del loaded_model

    state = torch.load(checkpoint_dir / "training_state.pt", map_location="cpu")
    # Config-hash guard: refuse to resume a checkpoint written under a different
    # model/recipe configuration (e.g. a pilot checkpoint left in the output dir).
    saved_hash = state.get("config_hash", "")
    run_hash = globals().get("_RUN_CONFIG_HASH", "")
    if saved_hash and run_hash and saved_hash != run_hash:
        if globals().get("_ALLOW_CONFIG_MISMATCH", False):
            print(f"[WARN] config hash mismatch (ckpt {saved_hash[:12]} vs run {run_hash[:12]}) — overridden")
        else:
            raise SystemExit(
                f"REFUSING RESUME: checkpoint config hash {saved_hash[:12]} != current run {run_hash[:12]}. "
                f"This checkpoint was written under a different model/recipe config. "
                f"Move it out of --output-dir or pass --allow-config-mismatch to override.")
    # Data-transition guard (preflight gate): the ONLY sanctioned dataset change on
    # resume is the curriculum swap train_main -> train_anneal at the decay boundary.
    # Any other change (wrong data dir, typo, stale run) is refused so we never
    # silently continue on the wrong corpus.
    saved_dir = state.get("train_dir_name", "")
    cur_dir = globals().get("_TRAIN_DIR_NAME", "")
    saved_sha = state.get("train_data_sha", "")
    cur_sha = globals().get("_TRAIN_DATA_SHA", "")
    allow = globals().get("_ALLOW_CONFIG_MISMATCH", False)
    if saved_dir and cur_dir and saved_dir != cur_dir:
        sanctioned = (saved_dir == "train_main" and cur_dir == "train_anneal")
        if sanctioned:
            print(f"[curriculum] sanctioned data transition {saved_dir} -> {cur_dir} "
                  f"at resume step {int(state['step'])} (bulk -> anneal)")
        elif allow:
            print(f"[WARN] unsanctioned data transition {saved_dir} -> {cur_dir} — overridden")
        else:
            raise SystemExit(
                f"REFUSING RESUME: checkpoint trained on '{saved_dir}' but --train-dir is "
                f"'{cur_dir}'. The only allowed swap is train_main -> train_anneal. "
                f"Fix --train-dir or pass --allow-config-mismatch to override.")
    elif saved_sha and cur_sha and saved_sha != cur_sha:
        # same directory name but different contents => data was altered/replaced
        if allow:
            print(f"[WARN] data manifest SHA changed on '{cur_dir}' "
                  f"({saved_sha[:12]} -> {cur_sha[:12]}) — overridden")
        else:
            raise SystemExit(
                f"REFUSING RESUME: '{cur_dir}' contents changed since the checkpoint "
                f"(manifest sha {saved_sha[:12]} != {cur_sha[:12]}). The dataset was "
                f"altered/replaced under the same name. Restore the exact shards or pass "
                f"--allow-config-mismatch to override.")
    optimizer.load_state_dict(state["optimizer"])
    move_optimizer_state_to_device(optimizer, device)
    scheduler.load_state_dict(state["scheduler"])
    random.setstate(state["python_random_state"])
    torch.set_rng_state(state["torch_rng_state"].cpu())
    if torch.cuda.is_available() and state["cuda_rng_state_all"] is not None:
        torch.cuda.set_rng_state_all([item.cpu() for item in state["cuda_rng_state_all"]])
    return (
        int(state["step"]),
        int(state["train_tokens_processed"]),
        int(state.get("data_epoch", 0)),
        dict(state.get("summary_state", {})),
    )


def evaluate(
    model: torch.nn.Module,
    tokenizer: AutoTokenizer,
    validation_texts: list[str],
    batch_size: int,
    seq_len: int,
    device: torch.device,
    max_batches: int,
    sample_prompts: list,
    sample_max_new_tokens: int,
) -> EvalSnapshot:
    # Use the unwrapped model for eval forward: DDP's forward hooks can leave
    # CUDA streams in states that stall a following barrier() and deadlock rank 1.
    inference_model = unwrap_model(model)
    inference_model.eval()
    losses: list[float] = []
    total_batches = 0
    with torch.no_grad():
        for offset in range(0, len(validation_texts), batch_size):
            if total_batches >= max_batches:
                break
            batch_texts = validation_texts[offset : offset + batch_size]
            batch = build_eval_batch(batch_texts, tokenizer, seq_len, device)
            outputs = inference_model(**batch)
            losses.append(float(outputs.loss.detach().cpu()))
            total_batches += 1

    avg_loss = sum(losses) / max(1, len(losses))
    perplexity = math.exp(avg_loss) if avg_loss < 20 else float("inf")

    samples = []
    generator = inference_model
    for prompt in sample_prompts:
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        with torch.no_grad():
            output_ids = generator.generate(
                **inputs,
                max_new_tokens=sample_max_new_tokens,
                do_sample=True,
                top_k=40,
                top_p=0.95,
                temperature=0.9,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        generation = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        samples.append({"prompt": prompt, "generation": generation})
    return EvalSnapshot(
        step=0,
        train_loss=0.0,
        validation_loss=avg_loss,
        validation_perplexity=perplexity,
        samples=samples,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train or resume the finalized ASTERIZER 1B 128K recipe.")
    parser.add_argument("--model-config", type=Path, required=True)
    parser.add_argument("--recipe-config", type=Path, required=True)
    parser.add_argument("--train-dir", type=Path, required=True)
    parser.add_argument("--validation-dir", type=Path, required=True)
    parser.add_argument("--tokenizer-repo", required=True)
    parser.add_argument("--tokenizer-subfolder", default="")
    parser.add_argument("--hf-token", default="")
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--resume-from-checkpoint", type=Path)
    parser.add_argument("--init-weights-from", type=Path, default=None,
                        help="CPT MODE: load ONLY model weights from this checkpoint dir. "
                             "Optimizer, scheduler, RNG, step counter and token counter all "
                             "start FRESH at step 0 with a new warmup. Mutually exclusive "
                             "with --resume-from-checkpoint. Ignored automatically once the "
                             "CPT run has its own checkpoints and --auto-resume finds one.")
    parser.add_argument("--allow-config-mismatch", action="store_true",
                        help="override the config-hash guard on resume (DANGEROUS)")
    parser.add_argument("--max-old-regression", type=float, default=0.02,
                        help="max allowed fractional rise of old-distribution val loss vs its "
                             "first recorded value before DOMAIN_GATE=FAIL is logged (default 2%)")
    parser.add_argument("--text-column", default="text")
    parser.add_argument("--parquet-batch-rows", type=int, default=2048)
    parser.add_argument("--max-train-files", type=int, default=0)
    parser.add_argument("--max-validation-files", type=int, default=0)
    parser.add_argument("--max-validation-rows", type=int, default=4096)
    parser.add_argument("--max-steps", type=int, default=0)
    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
    parser.add_argument("--per-device-eval-batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
    parser.add_argument("--learning-rate", type=float, default=-1.0)
    parser.add_argument("--min-learning-rate", type=float, default=-1.0)
    parser.add_argument("--precision", choices=["auto", "fp32", "fp16", "bf16"], default="auto")
    parser.add_argument(
        "--use-liger",
        choices=["auto", "on", "off"],
        default="auto",
        help="Use Liger fused kernels if installed. Fused linear cross-entropy removes the "
             "131K-vocab logit-memory wall, allowing much larger micro-batches (higher tok/s). "
             "'auto' uses it when available, 'on' requires it, 'off' disables.",
    )
    parser.add_argument(
        "--activation-checkpointing",
        choices=["config", "on", "off"],
        default="config",
        help="Override activation checkpointing. 'off' is ~30-40%% faster and fits at small "
             "micro-batch under --full-bf16 on 24GB GPUs; 'on' saves memory for larger batches.",
    )
    parser.add_argument(
        "--full-bf16",
        action="store_true",
        help="Store weights + optimizer states in bf16 (no fp32 master). ~2x less memory; "
             "needed to fit a 1.2B/131K-vocab model on 24GB GPUs. On 40GB+ prefer the default "
             "(fp32 master + bf16 autocast) for best training stability/quality.",
    )
    parser.add_argument("--logging-steps", type=int, default=10)
    parser.add_argument("--eval-steps", type=int, default=0)
    parser.add_argument("--save-steps", type=int, default=0)
    parser.add_argument("--max-eval-batches", type=int, default=8)
    parser.add_argument("--sample-prompts", default="",
                        help="Semicolon-separated prompts for generation samples during eval. "
                             "Empty = use built-in 5-prompt suite.")
    parser.add_argument("--sample-max-new-tokens", type=int, default=48)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--run-name", default="production_1b_128k_ready_to_train_v1")
    parser.add_argument("--auto-resume", action="store_true",
                        help="Resume from the latest checkpoint-step-* in --output-dir if present.")
    parser.add_argument("--keep-last-checkpoints", type=int, default=3,
                        help="Prune older local checkpoints, keeping this many. 0 = keep all.")
    parser.add_argument("--hf-checkpoint-repo", default="",
                        help="If set, upload each checkpoint to this HF dataset repo (async).")
    parser.add_argument("--hf-checkpoint-subfolder", default="project_source_phase1_phase2_20260707/runs")
    parser.add_argument("--hf-write-token", default=os.environ.get("HF_WRITE_TOKEN", ""))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    device = init_runtime()
    set_seed(args.seed)
    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    model_cfg = load_json(args.model_config)
    recipe_cfg = load_json(args.recipe_config)
    training_cfg = recipe_cfg["training"]
    # Run-config fingerprint for the resume guard (model config + recipe bytes).
    import hashlib as _hl
    globals()["_RUN_CONFIG_HASH"] = _hl.sha256(
        Path(args.model_config).read_bytes() + Path(args.recipe_config).read_bytes()
    ).hexdigest()
    globals()["_ALLOW_CONFIG_MISMATCH"] = bool(args.allow_config_mismatch)
    globals()["_TRAIN_DIR_NAME"] = args.train_dir.name
    globals()["_TRAIN_DATA_SHA"] = dir_manifest_sha(args.train_dir)
    print(f"DATA_MANIFEST_SHA train_dir={args.train_dir.name} sha={globals()['_TRAIN_DATA_SHA'][:16]}", flush=True)
    # GLOBAL_BATCH_ASSERT (preflight review P0): the run must refuse to start if
    # world x micro x accum x seq != recipe tokens_per_step (e.g. 8 GPUs with the
    # 4-GPU 16x8 settings would silently double tokens/update).
    _ws = int(os.environ.get("WORLD_SIZE", "1"))
    _expected = int(training_cfg.get("tokens_per_step", 1048576))
    _actual = _ws * args.per_device_train_batch_size * args.gradient_accumulation_steps * int(training_cfg["sequence_length"])
    if _actual != _expected:
        raise SystemExit(f"GLOBAL_BATCH_ASSERT failed: expected={_expected} actual={_actual} "
                         f"(world={_ws} micro={args.per_device_train_batch_size} "
                         f"accum={args.gradient_accumulation_steps} seq={training_cfg['sequence_length']})")
    print(f"GLOBAL_BATCH_ASSERT: expected={_expected} actual={_actual} status=OK", flush=True)
    _pb = recipe_cfg.get("phase_boundaries_tokens", {})
    if _pb:
        print(f"PHASE_BOUNDARIES A_end={_pb.get('A_end')} B_end={_pb.get('B_end')} C_end={_pb.get('C_end')}", flush=True)
    globals()["_PHASE_BOUNDARIES"] = _pb
    seq_len = int(training_cfg["sequence_length"])
    max_steps = int(args.max_steps or training_cfg["total_training_steps"])
    eval_steps = int(args.eval_steps or training_cfg["validation_every_steps"])
    save_steps = int(args.save_steps or training_cfg["checkpoint_every_steps"])
    learning_rate = float(args.learning_rate if args.learning_rate > 0 else training_cfg["learning_rate"])
    min_learning_rate = float(args.min_learning_rate if args.min_learning_rate > 0 else training_cfg["min_learning_rate"])

    tokenizer_kwargs = {"token": args.hf_token or None}
    if args.tokenizer_subfolder:
        tokenizer_kwargs["subfolder"] = args.tokenizer_subfolder
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_repo, **tokenizer_kwargs)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    validation_texts = load_texts(
        split_dir=args.validation_dir,
        max_rows=args.max_validation_rows,
        text_column=args.text_column,
        max_files=args.max_validation_files,
    )
    if not validation_texts:
        raise RuntimeError("No validation texts loaded")
    # Per-domain validation (preflight gate #15): break the aggregate loss out by
    # old/english/math/code/reasoning so a regression in one domain can't hide.
    validation_by_domain = load_texts_by_domain(
        split_dir=args.validation_dir,
        max_rows=args.max_validation_rows,
        text_column=args.text_column,
        max_files=args.max_validation_files,
    )
    log_main(f"Per-domain validation buckets: "
             f"{ {k: len(v) for k, v in validation_by_domain.items()} }")

    precision_name, amp_dtype = choose_precision(args.precision, device)
    # Full-bf16 stores weights + optimizer states in bf16 (no fp32 master, no autocast),
    # roughly halving memory so a 1.2B/131K-vocab model fits on 24GB GPUs.
    model_dtype = torch.bfloat16 if args.full_bf16 else None
    if args.full_bf16:
        precision_name = "bf16-full"
        amp_dtype = None
    log_main(f"Device: {device}")
    log_main(f"World size: {get_world_size()}")
    log_main(f"Precision: {precision_name}")
    log_main(f"Sequence packing: enabled (seq_len={seq_len}, no padding, no truncation)")
    log_main(f"Train dir: {args.train_dir}")
    log_main(f"Validation dir: {args.validation_dir}")
    log_main(f"Validation texts loaded: {len(validation_texts)}")

    liger_active = maybe_apply_liger(args.use_liger)
    model = build_model(model_cfg, tokenizer, device, model_dtype=model_dtype)
    if args.activation_checkpointing == "config":
        use_activation_checkpointing = bool(model_cfg["model"]["activation_checkpointing"])
    else:
        use_activation_checkpointing = args.activation_checkpointing == "on"
    if use_activation_checkpointing:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    log_main(f"Activation checkpointing: {'ON' if use_activation_checkpointing else 'OFF'}")

    if is_distributed():
        model = DDP(
            model,
            device_ids=[device.index] if device.type == "cuda" else None,
            find_unused_parameters=False,
        )

    optimizer = build_optimizer(model, training_cfg, learning_rate)
    scheduler = build_wsd_scheduler(
        optimizer=optimizer,
        max_steps=max_steps,
        warmup_steps=int(training_cfg["warmup_steps"]),
        final_decay_steps=int(training_cfg.get("final_decay_steps", 0)),
        min_lr_ratio=float(min_learning_rate / learning_rate),
    )

    use_amp = device.type == "cuda" and amp_dtype is not None
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp and amp_dtype == torch.float16)

    start_step = 0
    train_tokens_processed = 0
    data_epoch = 0
    summary_state = {"eval_history": [], "train_losses": []}

    resume_checkpoint = args.resume_from_checkpoint
    if resume_checkpoint is None and args.auto_resume:
        resume_checkpoint = find_latest_local_checkpoint(args.output_dir)
        if resume_checkpoint is not None:
            log_main(f"Auto-resume: found latest LOCAL checkpoint {resume_checkpoint.name}")
        elif args.hf_checkpoint_repo:
            # Disk wiped / fresh instance: pull the latest checkpoint back from HF.
            log_main("No local checkpoint; checking HF for the latest uploaded checkpoint...")
            resume_token = args.hf_write_token or args.hf_token
            if is_main_process():
                download_latest_hf_checkpoint(
                    args.hf_checkpoint_repo,
                    args.hf_checkpoint_subfolder,
                    args.run_name,
                    resume_token,
                    args.output_dir,
                )
            barrier()
            resume_checkpoint = find_latest_local_checkpoint(args.output_dir)
            if resume_checkpoint is not None:
                log_main(f"Auto-resume: restored HF checkpoint {resume_checkpoint.name}")
            else:
                log_main("No HF checkpoint found either; starting fresh")
        else:
            log_main("Auto-resume requested but no checkpoint found; starting fresh")

    if resume_checkpoint and args.init_weights_from and args.resume_from_checkpoint:
        raise SystemExit("--init-weights-from and --resume-from-checkpoint are mutually exclusive. "
                         "CPT fresh start = --init-weights-from; continue THIS run = --resume-from-checkpoint/--auto-resume.")

    if resume_checkpoint:
        # (CPT note: when --auto-resume finds one of THIS run's own checkpoints,
        # full-state resume is exactly what we want; --init-weights-from is skipped.)
        barrier()
        start_step, train_tokens_processed, data_epoch, summary_state = load_checkpoint_state(
            checkpoint_dir=resume_checkpoint.resolve(),
            model=model,
            optimizer=optimizer,
            scheduler=scheduler,
            device=device,
        )
        log_main(f"Resumed from {resume_checkpoint} at step {start_step} (data_epoch={data_epoch})")
    elif args.init_weights_from:
        # CPT MODE: weights only. Optimizer/scheduler/RNG/step/token counters stay
        # fresh -> run starts at step 0 with the NEW recipe's warmup, and --max-steps
        # means N NEW steps (the old global step 25667 is irrelevant here).
        barrier()
        init_dir = args.init_weights_from.resolve()
        unwrapped = unwrap_model(model)
        loaded = LlamaForCausalLM.from_pretrained(init_dir)
        unwrapped.load_state_dict(loaded.state_dict(), strict=True)
        del loaded
        # exact log contract (preflight review): these lines are grepped by run gates
        log_main(f"CPT INIT: loaded WEIGHTS ONLY from {init_dir}")
        log_main("optimizer_state=NEW")
        log_main("scheduler_state=NEW")
        log_main("global_step=0")
        log_main("cpt_tokens=0")

    train_stream = PackingStreamLoader(
        split_dir=args.train_dir,
        text_column=args.text_column,
        tokenizer=tokenizer,
        seq_len=seq_len,
        seed=args.seed,
        parquet_batch_rows=args.parquet_batch_rows,
        max_files=args.max_train_files,
        rank=get_rank(),
        world_size=get_world_size(),
        start_epoch=data_epoch,
    )

    train_losses: list[float] = list(summary_state.get("train_losses", []))
    eval_history: list[dict] = list(summary_state.get("eval_history", []))
    grad_accum_steps = max(1, int(args.gradient_accumulation_steps))
    tokens_per_optim_step = args.per_device_train_batch_size * grad_accum_steps * seq_len * get_world_size()
    started = time.time()
    window_started = time.time()
    window_tokens = 0
    throughput_samples: list[float] = []

    # Background uploader so pushing multi-GB checkpoints to HF never stalls training.
    upload_executor = None
    if is_main_process() and args.hf_checkpoint_repo and args.hf_write_token:
        from concurrent.futures import ThreadPoolExecutor

        upload_executor = ThreadPoolExecutor(max_workers=1)
        log_main(f"Checkpoint auto-upload enabled -> hf://{args.hf_checkpoint_repo}")

    try:
        for step in range(start_step + 1, max_steps + 1):
            model.train()
            optimizer.zero_grad(set_to_none=True)
            micro_losses: list[float] = []
            for micro_index in range(grad_accum_steps):
                batch = build_packed_train_batch(
                    loader=train_stream,
                    batch_size=args.per_device_train_batch_size,
                    device=device,
                )
                autocast_context = (
                    torch.autocast(device_type=device.type, dtype=amp_dtype) if use_amp else contextlib.nullcontext()
                )
                # Only synchronize DDP gradients on the final micro-step.
                is_last_micro = micro_index == grad_accum_steps - 1
                sync_context = (
                    model.no_sync() if isinstance(model, DDP) and not is_last_micro else contextlib.nullcontext()
                )
                with sync_context:
                    with autocast_context:
                        outputs = model(**batch)
                        loss = outputs.loss
                    micro_losses.append(float(loss.detach().cpu()))
                    scaled_loss = loss / grad_accum_steps
                    if scaler.is_enabled():
                        scaler.scale(scaled_loss).backward()
                    else:
                        scaled_loss.backward()

            if scaler.is_enabled():
                scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), float(training_cfg["grad_clip"]))
            if scaler.is_enabled():
                scaler.step(optimizer)
                scaler.update()
            else:
                optimizer.step()
            scheduler.step()

            mean_loss = sum(micro_losses) / max(1, len(micro_losses))
            reduced_loss = reduce_scalar(mean_loss, device)
            train_losses.append(reduced_loss)
            _prev_tokens = train_tokens_processed
            train_tokens_processed += tokens_per_optim_step
            window_tokens += tokens_per_optim_step
            data_epoch = train_stream.epoch
            # PHASE / PHASE_SWITCH ledger (preflight review P1): phases defined by
            # exact token boundaries in the recipe, logged the step they are crossed.
            _pb = globals().get("_PHASE_BOUNDARIES") or {}
            if _pb:
                def _phase_of(tk):
                    if tk < _pb.get("A_end", 1 << 62):
                        return "A"
                    if tk < _pb.get("B_end", 1 << 62):
                        return "B"
                    return "C"
                _p_prev, _p_now = _phase_of(_prev_tokens), _phase_of(train_tokens_processed)
                if step == 1:
                    log_main(f"PHASE={_p_now} consumed_tokens={train_tokens_processed}")
                if _p_prev != _p_now:
                    log_main(f"PHASE_SWITCH {_p_prev}->{_p_now} consumed_tokens={train_tokens_processed} step={step}")
                    log_main(f"phase_token_ledger A_end={_pb.get('A_end')} B_end={_pb.get('B_end')} "
                             f"C_end={_pb.get('C_end')} consumed={train_tokens_processed}")

            if step == 1 or step % max(1, args.logging_steps) == 0 or step == max_steps:
                if device.type == "cuda":
                    torch.cuda.synchronize()
                elapsed_window = max(1e-6, time.time() - window_started)
                tokens_per_second = window_tokens / elapsed_window
                throughput_samples.append(tokens_per_second)
                window_started = time.time()
                window_tokens = 0
                log_main(
                    f"Step {step}/{max_steps} train_loss={reduced_loss:.4f} "
                    f"lr={scheduler.get_last_lr()[0]:.6e} "
                    f"tokens={train_tokens_processed} "
                    f"tok/s={tokens_per_second:,.0f} epoch={data_epoch}"
                )

            if step % max(1, eval_steps) == 0 or step == max_steps:
                barrier()
                if is_main_process():
                    snapshot = evaluate(
                        model=model,
                        tokenizer=tokenizer,
                        validation_texts=validation_texts,
                        batch_size=args.per_device_eval_batch_size,
                        seq_len=seq_len,
                        device=device,
                        max_batches=args.max_eval_batches,
                        sample_prompts=(
                            [p.strip() for p in args.sample_prompts.split(";") if p.strip()]
                            if args.sample_prompts else DEFAULT_SAMPLE_PROMPTS
                        ),
                        sample_max_new_tokens=args.sample_max_new_tokens,
                    )
                    snapshot.step = step
                    snapshot.train_loss = reduced_loss
                    eval_history.append(asdict(snapshot))
                    log(
                        f"Eval step={step} validation_loss={snapshot.validation_loss:.4f} "
                        f"validation_ppl={snapshot.validation_perplexity:.2f}"
                    )
                    # per-domain val loss (gate #15): old-distribution first so the
                    # forgetting detector is always visible next to the new-mix loss.
                    dom_losses = {}
                    for dom in ("old", "english", "math", "code", "reasoning"):
                        texts_d = validation_by_domain.get(dom)
                        if texts_d:
                            dom_losses[dom] = eval_loss_only(
                                model, tokenizer, texts_d,
                                batch_size=args.per_device_eval_batch_size,
                                seq_len=seq_len, device=device,
                                max_batches=args.max_eval_batches)
                    unwrap_model(model).eval()  # keep eval mode until the shared restore below
                    log("Eval step=%d val_by_domain=%s" % (
                        step, {k: round(v, 4) for k, v in dom_losses.items()}))
                    snapshot_dict = eval_history[-1]
                    snapshot_dict["val_by_domain"] = dom_losses
                    # old-distribution regression gate (preflight item 4): compare old
                    # val loss to its FIRST recorded value; log PASS/FAIL against the
                    # --max-old-regression threshold so LR A/B has a hard criterion.
                    if "old" in dom_losses and not math.isnan(dom_losses["old"]):
                        base_old = globals().get("_OLD_VAL_BASELINE")
                        if base_old is None:
                            globals()["_OLD_VAL_BASELINE"] = dom_losses["old"]
                            log(f"DOMAIN_GATE baseline old_val={dom_losses['old']:.4f} "
                                f"threshold=+{args.max_old_regression:.1%}")
                        else:
                            reg = (dom_losses["old"] - base_old) / max(1e-9, base_old)
                            verdict = "PASS" if reg <= args.max_old_regression else "FAIL"
                            log(f"DOMAIN_GATE step={step} old_val={dom_losses['old']:.4f} "
                                f"regression={reg:+.2%} threshold=+{args.max_old_regression:.1%} "
                                f"status={verdict}")
                    for idx, sample in enumerate(snapshot.samples, 1):
                        log(f"  [{idx}] prompt   : {sample['prompt']!r}")
                        log(f"      generation: {sample['generation']!r}")
                # Both ranks: flush any pending CUDA work, restore train mode,
                # THEN sync. This prevents rank 1 from hitting the barrier while
                # rank 0's stream still has queued eval kernels — the deadlock
                # we hit at step 500 of the first attempt.
                unwrap_model(model).train()
                if device.type == "cuda":
                    torch.cuda.synchronize()
                barrier()

            if step % max(1, save_steps) == 0 or step == max_steps:
                barrier()
                if is_main_process():
                    ckpt_dir = args.output_dir / f"checkpoint-step-{step:05d}"
                    save_checkpoint(
                        checkpoint_dir=ckpt_dir,
                        model=model,
                        tokenizer=tokenizer,
                        optimizer=optimizer,
                        scheduler=scheduler,
                        step=step,
                        train_tokens_processed=train_tokens_processed,
                        data_epoch=data_epoch,
                        summary_state={"eval_history": eval_history, "train_losses": train_losses},
                    )
                    log(f"Checkpoint saved to {ckpt_dir}")
                    if upload_executor is not None:
                        upload_executor.submit(
                            upload_checkpoint_to_hf,
                            ckpt_dir,
                            args.hf_checkpoint_repo,
                            args.hf_checkpoint_subfolder,
                            args.hf_write_token,
                            args.run_name,
                            args.keep_last_checkpoints,
                        )
                    prune_old_checkpoints(args.output_dir, args.keep_last_checkpoints)
                barrier()

        if is_main_process():
            avg_tokens_per_second = sum(throughput_samples) / max(1, len(throughput_samples))
            steps_remaining_full = max(0, int(training_cfg["total_training_steps"]) - max_steps)
            summary = {
                "device": str(device),
                "world_size": get_world_size(),
                "precision": precision_name,
                "sequence_packing": True,
                "train_files_discovered": len(train_stream.files),
                "validation_rows_loaded": len(validation_texts),
                "seq_len": seq_len,
                "per_device_train_batch_size": args.per_device_train_batch_size,
                "per_device_eval_batch_size": args.per_device_eval_batch_size,
                "gradient_accumulation_steps": grad_accum_steps,
                "train_tokens_processed": train_tokens_processed,
                "actual_tokens_per_step": tokens_per_optim_step,
                "max_steps": max_steps,
                "data_epoch_reached": data_epoch,
                "elapsed_seconds": round(time.time() - started, 2),
                "avg_tokens_per_second": round(avg_tokens_per_second, 2),
                "measured_seconds_per_step": round(tokens_per_optim_step / max(1e-6, avg_tokens_per_second), 4),
                "projected_hours_for_full_recipe": round(
                    int(training_cfg["total_training_steps"]) * tokens_per_optim_step
                    / max(1e-6, avg_tokens_per_second) / 3600.0,
                    2,
                ),
                "steps_remaining_for_full_recipe": steps_remaining_full,
                "final_train_loss": train_losses[-1],
                "best_train_loss": min(train_losses),
                "eval_history": eval_history,
                "output_dir": str(args.output_dir),
                "tokenizer_repo": args.tokenizer_repo,
                "tokenizer_subfolder": args.tokenizer_subfolder,
                "model_config_path": str(args.model_config),
                "recipe_config_path": str(args.recipe_config),
                "resume_from_checkpoint": str(args.resume_from_checkpoint) if args.resume_from_checkpoint else None,
            }
            summary_path = args.output_dir / "training_summary.json"
            summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
            log(f"Summary written to {summary_path}")
            print(json.dumps(summary, indent=2))
    finally:
        if upload_executor is not None:
            # Bounded wait: give background uploads up to 30 minutes to drain, then
            # detach so the interpreter can exit even if HF is slow.
            log_main("Waiting up to 30 min for pending checkpoint uploads to finish...")
            import threading
            drained = threading.Event()

            def _drain():
                upload_executor.shutdown(wait=True)
                drained.set()

            threading.Thread(target=_drain, daemon=True).start()
            if drained.wait(timeout=1800):
                log_main("All uploads finished.")
            else:
                log_main("WARNING: uploads still pending after 30 min; detaching (checkpoints remain locally).")
        cleanup_runtime()


if __name__ == "__main__":
    main()