Blame for code/scripts/dpo.py

#!/usr/bin/env python
"""DPO stage entry: gates -> banner -> train -> arm result manifest.

    torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_argilla.toml
    torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_steppairs.toml \
        --lr=5e-7

One script, two pre-registered arms; the config selects everything. The
objective is section-05's chunked `fused_dpo`, the loop is section-07's
`Trainer`, the reference logprobs are constants read from the cache
scripts/precompute_ref_logprobs.py wrote. What this script owns is the
decision to start and the record of what happened.

Entry gates run in `dpo_stage.GATE_ORDER` and any failure blocks everything
after it, model load included. They are ordered by cost so an invalid stage
dies in milliseconds rather than after putting 14.5GB onto a shared no-swap
box:

  1. section-11's SFT quality-gate receipt exists AND passed;
  2. build-manifest counts match the config, pins were verified, Arm A is not
     marked DESCOPED;
  3. the reference cache's six-component key matches this run exactly;
  4. the path-bias measurement exists and its decision was honoured;
  5. the 20-step throughput probe (wall-clock is DERIVED from it, never
     asserted in advance);
  6. the RSS-aware preflight for the full-FT job class.

On a tripwire fire the arm re-runs ONCE with the pre-registered response from
docs/RESULTS.md, and that configuration IS the arm's result -- written into
the arm manifest, which section-13 reads for arm selection and the paired
evaluation.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("dpo")


def _split_overrides(extra: list[str]) -> list[str]:
    """Leftover argv (what argparse did not claim) -> dotted `key=value`
    overrides for `dpo_stage.load_stage_config`.

    `--lr=5e-7` and `lr=5e-7` are the same override. Anything without an `=`
    is refused: argparse's own positional `nargs='*'` accepted the bare form
    and then nothing ever read it, so an override the operator believed
    applied trained at the TOML's value (the banner printed the TOML value,
    not the intended one). Every override is either applied or rejected.
    """
    overrides = [a.lstrip("-") for a in extra]
    bad = [a for a, o in zip(extra, overrides) if "=" not in o or not o.split("=", 1)[0]]
    if bad:
        raise SystemExit(
            f"unrecognized arguments {bad}: overrides must be dotted key=value pairs "
            "(e.g. --lr=5e-7); space-separated values are not supported"
        )
    return overrides


def _parse_args(argv=None):
    """argparse for the flags, hand-split for the dotted overrides.

    `--lr=5e-7` is not an argparse option, so the declared-positional form
    rejected the documented spelling (exit 2) while accepting the bare
    `lr=5e-7` and then never reading it. Unclaimed arguments are now the
    overrides: applied (struct mode, via `load_stage_config`) or refused.
    """
    ap = argparse.ArgumentParser(allow_abbrev=False, epilog="dotted overrides: --lr=5e-7")
    ap.add_argument("config")
    ap.add_argument(
        "--resume",
        default=os.environ.get("ALIGN_RESUME") == "1",
        action=argparse.BooleanOptionalAction,
        help="resume from the newest valid checkpoint in run_dir instead of rotating "
        "it away as stale (default: fresh arm; the supervisor sets ALIGN_RESUME=1 on retries)",
    )
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
    args, extra = ap.parse_known_args(argv)
    args.overrides = _split_overrides(extra)
    return args


def main() -> int:
    args = _parse_args()

    from align.config import default_full_ft_config

    def _rss_cfg():
        # Project host RSS from the optimizer this stage ACTUALLY runs.
        # default_full_ft_config defaults to offload_adamw (104.7GB
        # projection) while the DPO arms are pinned to adamw8bit (22GB) --
        # the mismatch made the RSS gate refuse a job 5x smaller than its
        # model whenever the host dipped below ~150GB free.
        c = default_full_ft_config()
        c.optimizer.name = "adamw8bit"
        return c

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="dpo",
        strict_host_check=args.strict_host_check,
        rss_cfg=_rss_cfg(),
    )

    # BEFORE anything touches CUDA -- the duplicate-GPU class, third
    # appearance (sft.py, precompute --path-bias, now here): without
    # set_device(local_rank) every rank's bare "cuda" is device 0, and this
    # stage's first real 4-rank launch put four ~5.4GB policy loads on GPU 0.
    from align.utils.dist import init_distributed

    # The run dir (NCCL flight-recorder dumps land under it, never /tmp) comes
    # from the config, which is parsed before the process group so a bad TOML
    # fails in milliseconds.
    from align.train.dpo_stage import load_stage_config as _early_stage_config

    init_distributed(run_dir=_early_stage_config(args.config, args.overrides).run_dir)

    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()

    from align.data.loader import preference_dataset
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.train.dpo_stage import (
        PreferenceBatchSource,
        apply_tripwire_response,
        arm_result_manifest,
        assert_launch_topology,
        build_manifest_gate,
        load_run_ref_caches,
        load_stage_config,
        make_preference_sampler,
        measure_chosen_longer_fraction,
        read_sft_quality_gate,
        reduce_tripwire_alarms,
        resolve_pair_split_dir,
        resolve_sft_checkpoint,
        run_entry_gates,
        write_arm_result,
    )
    from align.train.ref_cache import read_path_bias_report
    from align.utils.guard import assert_paths_off_root

    cfg = load_stage_config(args.config, args.overrides)
    assert_paths_off_root([cfg.run_dir])
    # LAUNCH REFUSAL, before any gate: world_size x micro x grad_accum must be
    # the pre-registered global_batch the banner and horizon are quoted in.
    # Fed by the live process group, so a 2-rank launch of a 4-rank
    # pre-registration dies here in milliseconds, not after the model load.
    logger.info("launch topology ok: %s", assert_launch_topology(cfg))
    with open(cfg.manifest_path) as f:
        manifest = json.load(f)

    state: dict = {"resume": bool(args.resume)}

    # -- the gate sequence ------------------------------------------------
    # Each gate is a closure over `state` so a later gate can use what an
    # earlier one established -- which is the point of the ORDER: reaching
    # gate N means gates 1..N-1 held.
    def gate_sft(_cfg):
        state["sft_gate"] = read_sft_quality_gate(_cfg.sft_quality_gate_path)
        # The trainer writes <run_dir>/checkpoints/micro_NNNNNNNN with no
        # stable `final` alias; the receipt records which one passed. Resolve
        # HERE, at the first gate, so every later consumer (ref-cache key,
        # model load, banner) sees one path -- never a hand-composed one.
        ckpt_path, source = resolve_sft_checkpoint(
            _cfg.sft_checkpoint, state["sft_gate"].get("checkpoint")
        )
        state["sft_checkpoint"] = ckpt_path
        state["sft_gate"]["checkpoint"] = ckpt_path
        return {"checkpoint": ckpt_path, "checkpoint_source": source}

    def gate_manifest(_cfg):
        result = build_manifest_gate(_cfg, manifest)
        state["descoped"] = result["descoped"]
        return result

    def gate_ref_cache(_cfg):
        # Content-hash key (checkpoint WEIGHT BYTES, not the tensor inventory);
        # forward_impl is read off the cache itself, not the path-bias report --
        # see dpo_stage.expected_ref_cache_key for why each.
        state["ref_cache"], state["heldout_cache"] = load_run_ref_caches(
            _cfg, manifest, state["sft_checkpoint"]
        )
        state["ref_cache_key"] = state["ref_cache"].key
        return {"pairs": len(state["ref_cache"]), "heldout_pairs": len(state["heldout_cache"])}

    def gate_path_bias(_cfg):
        state["path_bias"] = read_path_bias_report(_cfg.path_bias_report_path)
        return {"action": state["path_bias"]["action"], "ratio": state["path_bias"]["ratio"]}

    def gate_probe(_cfg):
        # Filled after the model exists; the probe is the one gate that needs
        # the loop, so it records its intent here and the banner reads the
        # measured value below.
        state["probe"] = {
            "step_time_s_median": float("nan"),
            "projected_wall_clock_h": float("nan"),
            "status": "PENDING",
        }
        return {"status": "PENDING until the 20-step probe runs"}

    def gate_rss(_cfg):
        from align.train.probes import projected_job_rss

        state["rss"] = projected_job_rss(_rss_cfg())
        return dict(state["rss"])

    gate_report = run_entry_gates(
        cfg,
        {
            "sft_quality_gate": gate_sft,
            "build_manifest": gate_manifest,
            "ref_cache_key": gate_ref_cache,
            "path_bias": gate_path_bias,
            "throughput_probe": gate_probe,
            "rss_preflight": gate_rss,
        },
    )

    if state["descoped"]:
        # A descoped arm is a RESULT, not an absence: section-13's selection
        # rule reads this manifest and must be able to tell "measured badly"
        # from "never measured".
        out = arm_result_manifest(
            cfg=cfg, descoped=True, summary={}, heldout={}, banner={}, checkpoint=None
        )
        path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out)
        logger.warning("%s is DESCOPED; wrote %s and did not train", cfg.arm, path)
        return 0

    # -- data + model ------------------------------------------------------
    # DISCOVERY, not path composition: the configs point at the real split
    # dirs, and resolve_pair_split_dir re-finds them from the receipts the
    # store wrote -- so a layout change moves this by itself, and a wrong
    # configured path fails naming the splits that ARE present.
    rows = preference_dataset(resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train"))
    heldout_rows = preference_dataset(
        resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held")
    )
    pad_id = int(manifest.get("pinned", {}).get("pad_token_id", 0))

    baselines = manifest.get("length_bias_baseline", {})
    baseline_key = f"{cfg.arm}_chosen_longer_fraction"
    if baseline_key in baselines:
        length_baseline = (baselines[baseline_key], "build_manifest")
    else:
        # Section-03 records only Arm A's baseline. Measure this arm's from
        # its own shards rather than borrow the other's -- a tripwire armed
        # against the wrong distribution fires forever or never.
        length_baseline = (measure_chosen_longer_fraction(rows), "measured_at_launch")
    logger.info("length-drift baseline: %.4f (%s)", *length_baseline)

    model = _load_policy(
        NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation
    )
    # The RESOLVED backend, read back off the model: the constructor routed
    # the request through the gated selector, so this is what actually runs.
    state["attn_backend"] = model.attn_implementation
    run_cfg = _run_config(cfg, length_baseline[0])

    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp

    dist_model = apply_fsdp(apply_activation_checkpointing(model))
    import torch as _torch

    # Post-shard residency: a rank should hold ~3.6GB (its shard). Anything
    # near 18GB means the pre-shard replica leaked past fully_shard.
    logger.info("post-shard cuda allocated: %.2f GiB", _torch.cuda.memory_allocated() / 2**30)

    source = PreferenceBatchSource(
        rows,
        make_preference_sampler(
            rows,
            max_tokens=run_cfg.train.seq_len * cfg.micro_batch_size,
            max_rows=cfg.micro_batch_size,
            seed=run_cfg.train.seed,
        ),
        pad_token_id=pad_id,
    )

    result = _train_arm(
        cfg,
        run_cfg,
        dist_model,
        source,
        state,
        heldout_rows,
        pad_id,
        gate_report,
        manifest,
        length_baseline,
    )

    # -- tripwire response -------------------------------------------------
    # The rerun is collective, so the verdict is drawn from the UNION of every
    # rank's alarms: rank-local alarms could send some ranks into a second
    # model load + FSDP wrap while the rest exit.
    alarms = reduce_tripwire_alarms(result["summary"].get("tripwire_alarms", []))
    rerun_cfg, tripwire = apply_tripwire_response(cfg, alarms)
    if tripwire["rerun_required"]:
        logger.error(
            "chosen-logprob collapse detected; re-running %s with the pre-registered response: %s",
            cfg.arm,
            tripwire["configuration"],
        )
        model = _load_policy(
            NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation
        )
        state["attn_backend"] = model.attn_implementation
        dist_model = apply_fsdp(apply_activation_checkpointing(model))
        result = _train_arm(
            rerun_cfg,
            _run_config(rerun_cfg, length_baseline[0]),
            dist_model,
            source,
            state,
            heldout_rows,
            pad_id,
            gate_report,
            manifest,
            length_baseline,
            suffix="_response",
        )
        cfg = rerun_cfg

    out = arm_result_manifest(
        cfg=cfg,
        descoped=False,
        summary=result["summary"],
        heldout=result["heldout"],
        banner=result["banner"],
        checkpoint=result["checkpoint"],
        tripwire=tripwire,
    )
    out["gate_report"] = gate_report
    path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out)
    logger.info("%s complete: %s", cfg.arm, path)
    return 0


def _load_policy(model_cls, ckpt, path: str, attn_implementation: str):
    """Load the policy with the CONFIG-DRIVEN attention backend.

    The string is never handed raw to a constructor in the sense that
    matters: `NmoMistralForCausalLM.__init__` routes it through
    `align.model.attention.select_attn_implementation`, which raises if the
    request cannot be honoured -- for the pre-registered `triton_fa2` that
    means section-06's equivalence-gate receipt must be valid for this kernel
    source on this GPU. Silent degradation to sdpa is exactly the failure the
    selector exists to prevent.
    """
    import torch

    from align.model.hf_config import read_config

    model_cfg = read_config(path)
    model = model_cls.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)),
        attn_implementation=attn_implementation,
        device="cuda",
    )
    return model.to("cuda", dtype=torch.bfloat16)


def _run_config(cfg, length_baseline: float):
    """Stage config -> section-07 RunConfig. The two are different objects on
    purpose: RunConfig is the trainer's surface and DpoStageConfig is the
    arm's pre-registration, and collapsing them would let a trainer default
    silently become an experimental parameter."""
    from align.config import RunConfig

    run = RunConfig()
    # adamw8bit per the 2026-08-17 pre-registration amendment ("SFT (and the
    # DPO arms) run adamw8bit"). Left unset, RunConfig's offload_adamw
    # default leaked through and the pinned-staging guard refused with a
    # 27GiB request on a no-swap box -- correctly.
    run.optimizer.name = "adamw8bit"
    run.optimizer.lr = cfg.lr
    run.optimizer.grad_accum = cfg.grad_accum
    run.train.epochs = cfg.epochs
    run.train.micro_batch_size = cfg.micro_batch_size
    run.train.save_every_micro_steps = cfg.save_every_micro_steps
    run.train.keep_checkpoints = cfg.keep_checkpoints
    # Forwarded explicitly: RunConfig is built from defaults here, so a stage
    # field that is not copied is a silently disarmed config.
    run.train.per_rank_optimizer_state = cfg.per_rank_optimizer_state
    run.optimizer.max_grad_norm = cfg.max_grad_norm
    run.train.max_wall_seconds = cfg.max_wall_seconds
    run.train.max_consecutive_skipped_steps = cfg.max_consecutive_skipped_steps
    run.train.max_skipped_fraction = cfg.max_skipped_fraction
    run.train.full_ft = True
    run.data.num_workers = cfg.num_workers
    run.tripwires.chosen_longer_baseline = length_baseline
    run.tripwires.length_band = cfg.length_band
    run.tripwires.logps_window = cfg.logps_window
    run.tripwires.armed = tuple(cfg.armed_tripwires)
    # The tripwire may speak exactly when its estimator is full. These are two
    # different quantities -- `length_window_pairs` is the span the adapter
    # averages over, `length_min_pairs` is when the tripwire is allowed to
    # fire -- and they are deliberately the SAME number: a min_pairs below the
    # window size would let the guard speak from a partially-filled window,
    # which is precisely the noisy regime both settings exist to avoid.
    run.tripwires.length_min_pairs = cfg.length_window_pairs
    run.run_name = f"dpo_{cfg.arm}"
    return run


def _train_arm(
    cfg,
    run_cfg,
    dist_model,
    source,
    state,
    heldout_rows,
    pad_id,
    gate_report,
    manifest,
    length_baseline,
    suffix: str = "",
):
    """Probe -> banner -> train -> held-out accuracy. In that order, because
    the banner quotes the PROBE-DERIVED wall-clock and the plan forbids
    asserting a wall-clock in advance."""
    from align.data.loader import collate_preference
    from align.model.attention import banner_value
    from align.train.dpo_stage import (
        DpoBatch,
        dpo_banner,
        make_dpo_stage_adapter,
        render_dpo_banner,
    )
    from align.train.ref_cache import heldout_preference_accuracy
    from align.train.trainer import Trainer

    run_dir = cfg.run_dir + suffix
    # FRESH-ARM HYGIENE: a non-resuming arm run must not share a checkpoint
    # dir with a previous run's leftovers -- stale higher-numbered dirs
    # shadowed the rerun's fresh saves under keep-N pruning (2026-08-20).
    from align.train.dpo_stage import probe_micro_cap, rotate_stale_checkpoints

    ckpt_dir = os.path.join(run_dir, "checkpoints")
    # Rotation is for FRESH arms only. A resuming arm (--resume, or a
    # supervisor retry via ALIGN_RESUME=1) must keep its checkpoints: rotating
    # them would silently restart the arm from step 0 while the supervisor
    # believes it is continuing -- the replay defect, at the stage level.
    resuming = bool(state.get("resume"))
    # Collective: rank 0 renames, every rank waits (a bare per-rank rename
    # let one rank win and killed the other three with FileNotFoundError).
    if not resuming:
        rotate_stale_checkpoints(ckpt_dir)
    adapter = make_dpo_stage_adapter(state["ref_cache"], cfg, require_fused=True)
    trainer = Trainer(run_cfg, dist_model, source, adapter, run_dir=run_dir, stage="dpo")

    # REFUSE a horizon that disagrees with the corpus, before a single step.
    # Arm A ran a 146-step horizon over a corpus needing ~289 per-rank steps:
    # the cosine finished at half an epoch, clamped the learning rate to
    # exactly 0, and the loop kept stepping on frozen weights for ~24
    # GPU-hours. The Trainer now stops AT the horizon, which turns that into a
    # short run rather than a zombie one -- this check is the other half, and
    # catches it before the GPU-hours rather than after.
    from align.train.dpo_stage import assert_dpo_horizon_consistent
    from align.utils.dist import world_size as _world_size

    horizon_report = assert_dpo_horizon_consistent(
        trainer.scheduler.total_steps,
        n_pairs=len(source.rows),
        world_size=_world_size(),
        micro_batch_size=cfg.micro_batch_size,
        grad_accum=cfg.grad_accum,
        epochs=run_cfg.train.epochs,
        global_batch=cfg.global_batch,
    )
    logger.info("horizon ok: %s", horizon_report)

    if resuming:
        resumed_from = trainer.resume_from_latest()
        logger.info(
            "RESUMED from %s at optimizer step %d / micro step %d",
            resumed_from,
            trainer.optimizer_step,
            trainer.micro_step,
        )
        state["resumed_from"] = resumed_from

    tokens_per_step = (
        run_cfg.train.seq_len * cfg.micro_batch_size * 2 * cfg.grad_accum
    )  # 2: both sides of every pair
    # The probe is the FIRST probe_steps real micro-steps of the run itself:
    # a capped train(), report captured, cap removed, training continues from
    # the stream's exact position. The previous construction asked a trainer
    # that had never stepped for its step times, and throughput_report
    # correctly refused an empty probe. Probe steps here are real training
    # (counted in the horizon), so nothing is spent twice; the banner prints
    # after 20 micro-steps instead of before any. The cap is on the ABSOLUTE
    # micro-step counter, so a resumed arm probes its first probe_steps AFTER
    # the resume point (a bare probe_steps cap trained one micro-step and
    # crashed throughput_probe on every supervisor retry).
    run_cfg.train.max_micro_steps = probe_micro_cap(trainer.micro_step, run_cfg.train.probe_steps)
    trainer.train()
    run_cfg.train.max_micro_steps = None
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
        # optimizer-step boundary; the stage is NOT finished. Exit with the
        # documented status instead of evaluating / reporting a half-trained arm.
        raise SystemExit(trainer.stop_exit_code)
    if resuming and len(trainer.step_times) < trainer.accum:
        # Resumed within one optimizer step of the horizon (or after the last
        # step): the probe window cannot fill a group. The run still finishes;
        # the banner says the wall-clock was not measured rather than raising.
        probe = {
            "step_time_s_median": float("nan"),
            "projected_wall_clock_h": float("nan"),
            "status": "UNMEASURED_ON_RESUME",
        }
    else:
        probe = trainer.throughput_probe(tokens_per_step)
    state["probe"] = probe

    banner = dpo_banner(
        cfg=cfg,
        manifest=manifest,
        ref_cache_key=state["ref_cache_key"],
        path_bias=state["path_bias"],
        probe=probe,
        rss=state["rss"],
        backend=banner_value(
            # The RESOLVED backend, never a literal: FSDP2 mutates the module
            # in place so the attribute normally survives wrapping, and the
            # fallback is what the selector resolved at load -- a hardcoded
            # "sdpa" here would let the banner contradict the kernel that ran.
            dist_model.attn_implementation
            if hasattr(dist_model, "attn_implementation")
            else state["attn_backend"]
        ),
        sft_gate=state["sft_gate"],
        length_baseline=length_baseline,
    )
    print(render_dpo_banner(banner))

    summary = trainer.train()
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
        # optimizer-step boundary; the stage is NOT finished. Exit with the
        # documented status instead of evaluating / reporting a half-trained arm.
        raise SystemExit(trainer.stop_exit_code)
    checkpoint = trainer.save_checkpoint()

    import torch as _torch

    _dev = _torch.device("cuda", _torch.cuda.current_device())
    # .to(_dev): the held-out path missed the device move the training
    # adapter got -- CPU labels reached the CUDA gather AFTER the full arm
    # had trained and saved (Arm A rerun, 2026-08-20).
    heldout_batches = [
        collate_preference(heldout_rows[i : i + cfg.micro_batch_size], pad_token_id=pad_id).to(_dev)
        for i in range(0, len(heldout_rows), cfg.micro_batch_size)
    ]
    heldout = heldout_preference_accuracy(
        dist_model, heldout_batches, state["heldout_cache"], beta=cfg.beta
    )
    logger.info(
        "held-out preference accuracy: %.4f [%.4f, %.4f] n=%d",
        heldout["accuracy"],
        heldout["ci_low"],
        heldout["ci_high"],
        heldout["n"],
    )
    return {"summary": summary, "heldout": heldout, "banner": banner, "checkpoint": checkpoint}


if __name__ == "__main__":
    from align.exit_codes import run_with_exit_codes

    raise SystemExit(run_with_exit_codes(main))