code/scripts/precompute_ref_logprobs.py

align / code /scripts /precompute_ref_logprobs.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

16.7 kB

#!/usr/bin/env python
"""Precompute DPO reference log-probabilities from the frozen SFT checkpoint.
    # one no-grad pass per (arm, split), single GPU
    uv run python scripts/precompute_ref_logprobs.py configs/dpo_argilla.toml
    # the MANDATORY path-bias dual pass (needs the sharded forward, so torchrun)
    torchrun --nproc_per_node=4 scripts/precompute_ref_logprobs.py \
        configs/dpo_argilla.toml --path-bias
Its own heavy job: it takes the exclusive flock and is SERIALIZED before any
DPO run. Precompute and training must never be concurrent -- a data job
running against a live training job is exactly what killed the predecessor's
first SFT run.
What it produces, per arm, on /apps:
  * `<arm>-train.safetensors` and `<arm>-held.safetensors` -- pair_id ->
    (ref_logp_chosen, ref_logp_rejected, len_chosen, len_rejected), with the
    full six-component cache key in the header. The held-out split is computed
    in the SAME pass because held-out preference accuracy is a primary metric
    and a cache missing it produces an accuracy over a silently smaller
    denominator.
  * `<arm>-path_bias.json` -- the single-GPU-vs-sharded delta distribution and
    the decision derived from it. scripts/dpo.py refuses to launch without it,
    and refuses to launch if it fired and was not honoured.
Why the reference is cached at all: keeping a second 7B model resident would
cost 14.5GB of a 23GB card for no gradient. With dropout disabled a no-grad
pass over a fixed dataset is exactly reproducible, so the numbers can be
computed once and consumed as constants -- see nmo/train/ref_cache.py for the
six things the cache key has to cover for that to stay true.
"""

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("precompute_ref_logprobs")

#: Pairs pushed through both forwards for the path-bias measurement. 100 is
#: the plan's number: enough to see a systematic offset against the
#: early-margin scale, small enough to run inside the pre-launch window.
PATH_BIAS_EXAMPLES = 100


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("config", help="configs/dpo_argilla.toml | dpo_steppairs.toml")
    ap.add_argument(
        "--out-dir", default=None, help="defaults to the config's ref_cache_path directory"
    )
    ap.add_argument(
        "--batch-size",
        type=int,
        default=8,
        help="forward batch size. Regrouping perturbs the fp32 "
        "reduction order by ~1e-7 relative (measured); it is "
        "recorded in the receipt, not in the cache key.",
    )
    ap.add_argument(
        "--path-bias",
        action="store_true",
        help="run the dual single-GPU/sharded pass and write the "
        "decision report. Requires torchrun.",
    )
    ap.add_argument(
        "--sharded-cache",
        action="store_true",
        help="compute the caches under the FSDP2-SHARDED forward (requires "
        "torchrun). This is how a fired path-bias decision "
        "(rerun_precompute_sharded) is honoured: training consumes the "
        "sharded forward, so the reference must too.",
    )
    ap.add_argument(
        "--batch-invariance",
        action="store_true",
        help="measure batched-vs-single-row precompute deltas over "
        "100 pairs and write the receipt that decides whether "
        "batch_size belongs in the cache key. Single process.",
    )
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
    args = ap.parse_args()

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="precompute_ref_logprobs",
        strict_host_check=args.strict_host_check,
    )

    import torch

    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 (
        is_hf_checkpoint,
        load_stage_config,
        read_sft_quality_gate,
        resolve_pair_split_dir,
        resolve_sft_checkpoint,
    )
    from align.train.ref_cache import (
        build_ref_cache_key,
        path_bias_report_cache_forward_impl,
        precompute_ref_logprobs,
        write_path_bias_report,
        write_ref_cache,
    )
    from align.utils.guard import assert_paths_off_root

    cfg = load_stage_config(args.config)
    out_dir = args.out_dir or os.path.dirname(cfg.ref_cache_path)
    assert_paths_off_root([out_dir])

    with open(cfg.manifest_path) as f:
        manifest = json.load(f)

    # Resolve the frozen SFT checkpoint the same way scripts/dpo.py does.
    # The trainer writes <run_dir>/checkpoints/micro_NNNNNNNN with no stable
    # `final` alias; the quality-gate receipt records which one passed. A
    # configured path that IS a checkpoint is an explicit pin; otherwise the
    # receipt is read, and a missing/failed receipt or a missing checkpoint
    # refuses LOUDLY here -- a cache silently computed from the wrong (or no)
    # reference would poison both arms' margins.
    if is_hf_checkpoint(cfg.sft_checkpoint):
        sft_checkpoint, ckpt_source = cfg.sft_checkpoint, "config"
    else:
        gate = read_sft_quality_gate(cfg.sft_quality_gate_path)
        sft_checkpoint, ckpt_source = resolve_sft_checkpoint(
            cfg.sft_checkpoint, gate.get("checkpoint")
        )
    logger.info("frozen SFT checkpoint: %s (source=%s)", sft_checkpoint, ckpt_source)

    # Shard split dirs by DISCOVERY (the receipts the store wrote), never by
    # hand-composed `<root>-<split>` paths -- the attempt-1 defect class.
    train_dir = resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train")
    held_dir = resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held")

    # The key, assembled from facts rather than from intent, by the SAME
    # function scripts/dpo.py's reader uses (build_ref_cache_key). `checkpoint_hash`
    # is the WEIGHT-BYTES hash -- one streamed read of the checkpoint -- not the
    # tensor-inventory hash this script used to write, which could not tell two
    # SFT checkpoints of one architecture apart (2026-10-12 review). Caches
    # written under the old key are refused by dpo.py: re-run this script.
    # `manifest_hash` covers the shard rebuild that nothing else would notice.
    key = build_ref_cache_key(
        sft_checkpoint,
        manifest,
        forward_impl="fsdp2_sharded" if (args.path_bias or args.sharded_cache) else "single_gpu",
    )
    logger.info("cache key: %s", key.to_json())

    # Rank-pinned device: a bare "cuda" is device 0 for EVERY torchrun rank,
    # and the --path-bias 4-rank pass put four ~5.5GB reference shards on
    # GPU 0 (the duplicate-GPU class scripts/sft.py already fixed). This
    # mode had never been run before 2026-08-18.
    if torch.cuda.is_available():
        from align.utils.dist import init_distributed, local_rank

        init_distributed()
        torch.cuda.set_device(local_rank())
        device = f"cuda:{local_rank()}"  # string: downstream does device.startswith(...)
    else:
        device = "cpu"
    model = _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device)

    if args.sharded_cache:
        import torch.distributed as dist

        from align.train.fsdp import apply_fsdp

        if not (dist.is_available() and dist.is_initialized()):
            raise RuntimeError(
                "--sharded-cache computes the reference under the SHARDED "
                "forward and must run under torchrun; single-process would "
                "silently produce the single_gpu numbers under a sharded key."
            )
        model = apply_fsdp(model)
        model.eval()

    written = []
    # In --path-bias mode the caches are NOT recomputed: the dedicated
    # single-GPU run owns them. Before this guard, all four torchrun ranks
    # each recomputed and rewrote the same cache files concurrently -- and
    # stamped forward_impl="fsdp2_sharded" into caches computed single-GPU.
    cache_splits = (
        ()
        if args.path_bias
        else (
            ("train", train_dir, cfg.ref_cache_path),
            ("held", held_dir, cfg.heldout_ref_cache_path),
        )
    )
    for split, shard_dir, cache_path in cache_splits:
        rows = preference_dataset(shard_dir)
        cache = precompute_ref_logprobs(
            model,
            rows,
            key=key,
            pad_token_id=_pad_id(manifest),
            batch_size=args.batch_size,
            device=torch.device(device),
        )
        from align.utils.dist import is_rank0 as _is_rank0

        if _is_rank0():
            # Sharded forwards are collectives (every rank computes); the
            # file is written once.
            write_ref_cache(cache_path, cache)
        written.append({"split": split, "path": cache_path, "pairs": len(cache)})
        logger.info("%s/%s: %d reference pairs", cfg.arm, split, len(cache))

    if args.batch_invariance:
        from align.train.ref_cache import batch_invariance_report

        rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
        report = batch_invariance_report(
            model,
            rows,
            key=key,
            pad_token_id=_pad_id(manifest),
            beta=cfg.beta,
            early_margin_scale=cfg.path_bias_early_margin_scale,
            threshold=cfg.path_bias_threshold,
        )
        out_path = os.path.join(out_dir, f"{cfg.arm}-batch_invariance.json")
        with open(out_path, "w") as f:
            json.dump(report, f, indent=2)
        logger.info(
            "batch invariance -> %s (belongs_in_cache_key=%s)",
            out_path,
            report["belongs_in_cache_key"],
        )

    if args.path_bias:
        # OWNERSHIP TRANSFER: pass the model in a box and drop main's name.
        # A plain argument keeps main's reference alive, so the between-pass
        # free inside _path_bias freed nothing and the reload OOM'd with the
        # old 14.5GB replica still resident (measured 21.66GB, batch-size
        # independent).
        model_box = [model]
        del model
        report = _path_bias(
            model_box,
            cfg,
            manifest,
            key,
            args.batch_size,
            device,
            train_dir,
            reload_fn=lambda: _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device),
        )
        from align.utils.dist import is_rank0

        if is_rank0():
            # Rank-0-only: four ranks renaming one report is the ledger race.
            # `cache_forward_impl` must describe the CACHE, not this run. The
            # old code wrote key.forward_impl, which is "fsdp2_sharded" for
            # every --path-bias run -- but this mode recomputes no cache, so the
            # field named the measuring run and contradicted the single-GPU
            # cache sitting on disk (2026-10-12 review). Read it off the cache.
            write_path_bias_report(
                cfg.path_bias_report_path,
                report,
                n_examples=PATH_BIAS_EXAMPLES,
                forward_impl=path_bias_report_cache_forward_impl(cfg.ref_cache_path),
            )
        logger.info("path bias: %s (ratio %s)", report["action"], report["ratio"])

    logger.info("precompute complete: %s", written)
    return 0


def _manifest_hash(manifest: dict) -> str:
    """Hash of the build's own content hashes. Moves on any shard rebuild --
    which is the point: the section-03 marker fix retokenizes every pair while
    leaving the checkpoint, tokenizer and max_length untouched."""
    from align.data.receipts import content_hash

    return content_hash([manifest["hashes"][k] for k in sorted(manifest.get("hashes", {}))])


def _max_length(manifest: dict) -> int:
    return int(manifest.get("pinned", {}).get("max_len", 2048))


def _pad_id(manifest: dict) -> int:
    """Pad id for collation. Fully masked out of every logprob (prompt and pad
    positions carry IGNORE_INDEX), so its value cannot reach a number -- but
    it is read from the build rather than assumed, because 'cannot reach a
    number' is the kind of claim that stops being true after a refactor."""
    return int(manifest.get("pinned", {}).get("pad_token_id", 0))


def _load_reference(model_cls, ckpt, path: str, device: str):
    """The FROZEN SFT checkpoint, bf16 on GPU, eval mode.
    eval() is asserted downstream rather than merely called here: Mistral-v0.3
    has attention_dropout 0.0, so a model left in train mode would produce the
    same numbers today and silently different ones after any config change.
    """
    import torch
    from align.eval.golden import resolve_model_config

    model_cfg = resolve_model_config(path)
    model = model_cls.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)),
        attn_implementation="sdpa",
        device=device,
    )
    model = model.to(
        device, dtype=torch.bfloat16 if str(device).startswith("cuda") else torch.float32
    )
    model.eval()
    for module in model.modules():
        if isinstance(module, torch.nn.Dropout) and module.p > 0:
            raise RuntimeError(
                f"an active Dropout(p={module.p}) is present in the reference "
                "model. Cached reference logprobs must be exactly "
                "reproducible; dropout makes them a sample."
            )
    return model


def _path_bias(
    model_box, cfg, manifest, key, batch_size: int, device: str, train_dir: str, reload_fn=None
) -> dict:
    """The mandatory dual pass: single-GPU forward vs FSDP2-sharded forward.
    Requires torchrun. Without a process group there IS no sharded forward to
    compare against, and the honest report is then no report -- `path_bias_
    decision` treats an empty delta set as firing, because not measuring is
    not passing and a stage that skipped this must not silently proceed on a
    single-GPU cache.
    """
    import torch
    import torch.distributed as dist

    from align.data.loader import preference_dataset
    from align.train.fsdp import apply_fsdp
    from align.train.ref_cache import path_bias_decision, precompute_ref_logprobs

    if not (dist.is_available() and dist.is_initialized()):
        raise RuntimeError(
            "--path-bias needs the SHARDED forward, so it must run under "
            "torchrun (--nproc_per_node=4). Running it single-process would "
            "compare the single-GPU forward with itself and record a "
            "reassuring zero, which is worse than not measuring: it would "
            "look like evidence."
        )

    rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
    model = model_box.pop()  # sole reference from here on
    single = precompute_ref_logprobs(
        model,
        rows,
        key=key,
        pad_token_id=_pad_id(manifest),
        batch_size=batch_size,
        device=torch.device(device),
    )
    # FREE the single-pass replica before sharding: wrapping a resident
    # 14.5GB model in-place stacks the wrap transients ON TOP of it and
    # OOMs every rank at 21.8GB (measured, 2026-08-18). Rebuilding alone
    # and then sharding is the stage launcher's proven peak.
    if reload_fn is None:
        raise RuntimeError("--path-bias needs reload_fn to rebuild the model between passes")
    del model
    import gc

    gc.collect()
    torch.cuda.empty_cache()
    model = reload_fn()
    sharded_model = apply_fsdp(model)
    sharded_model.eval()
    sharded = precompute_ref_logprobs(
        sharded_model,
        rows,
        key=key,
        pad_token_id=_pad_id(manifest),
        batch_size=batch_size,
        device=torch.device(device),
    )

    # The delta of the reference MARGIN, not of either logprob: the margin is
    # what enters the loss, and a bias common to both sides cancels there.
    deltas = [
        (sharded.entries[p][0] - sharded.entries[p][1])
        - (single.entries[p][0] - single.entries[p][1])
        for p in single.entries
    ]
    return path_bias_decision(
        deltas,
        beta=cfg.beta,
        early_margin_scale=cfg.path_bias_early_margin_scale,
        threshold=cfg.path_bias_threshold,
    )


if __name__ == "__main__":
    raise SystemExit(main())

	#!/usr/bin/env python
	"""Precompute DPO reference log-probabilities from the frozen SFT checkpoint.

	# one no-grad pass per (arm, split), single GPU
	uv run python scripts/precompute_ref_logprobs.py configs/dpo_argilla.toml

	# the MANDATORY path-bias dual pass (needs the sharded forward, so torchrun)
	torchrun --nproc_per_node=4 scripts/precompute_ref_logprobs.py \
	configs/dpo_argilla.toml --path-bias

	Its own heavy job: it takes the exclusive flock and is SERIALIZED before any
	DPO run. Precompute and training must never be concurrent -- a data job
	running against a live training job is exactly what killed the predecessor's
	first SFT run.

	What it produces, per arm, on /apps:
	* `<arm>-train.safetensors` and `<arm>-held.safetensors` -- pair_id ->
	(ref_logp_chosen, ref_logp_rejected, len_chosen, len_rejected), with the
	full six-component cache key in the header. The held-out split is computed
	in the SAME pass because held-out preference accuracy is a primary metric
	and a cache missing it produces an accuracy over a silently smaller
	denominator.
	* `<arm>-path_bias.json` -- the single-GPU-vs-sharded delta distribution and
	the decision derived from it. scripts/dpo.py refuses to launch without it,
	and refuses to launch if it fired and was not honoured.

	Why the reference is cached at all: keeping a second 7B model resident would
	cost 14.5GB of a 23GB card for no gradient. With dropout disabled a no-grad
	pass over a fixed dataset is exactly reproducible, so the numbers can be
	computed once and consumed as constants -- see nmo/train/ref_cache.py for the
	six things the cache key has to cover for that to stay true.
	"""

	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("precompute_ref_logprobs")

	#: Pairs pushed through both forwards for the path-bias measurement. 100 is
	#: the plan's number: enough to see a systematic offset against the
	#: early-margin scale, small enough to run inside the pre-launch window.
	PATH_BIAS_EXAMPLES = 100


	def main() -> int:
	ap = argparse.ArgumentParser()
	ap.add_argument("config", help="configs/dpo_argilla.toml \| dpo_steppairs.toml")
	ap.add_argument(
	"--out-dir", default=None, help="defaults to the config's ref_cache_path directory"
	)
	ap.add_argument(
	"--batch-size",
	type=int,
	default=8,
	help="forward batch size. Regrouping perturbs the fp32 "
	"reduction order by ~1e-7 relative (measured); it is "
	"recorded in the receipt, not in the cache key.",
	)
	ap.add_argument(
	"--path-bias",
	action="store_true",
	help="run the dual single-GPU/sharded pass and write the "
	"decision report. Requires torchrun.",
	)
	ap.add_argument(
	"--sharded-cache",
	action="store_true",
	help="compute the caches under the FSDP2-SHARDED forward (requires "
	"torchrun). This is how a fired path-bias decision "
	"(rerun_precompute_sharded) is honoured: training consumes the "
	"sharded forward, so the reference must too.",
	)
	ap.add_argument(
	"--batch-invariance",
	action="store_true",
	help="measure batched-vs-single-row precompute deltas over "
	"100 pairs and write the receipt that decides whether "
	"batch_size belongs in the cache key. Single process.",
	)
	ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
	ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
	args = ap.parse_args()

	apply_resource_guards(
	cpu_threads=8,
	heavy=args.heavy_lock,
	label="precompute_ref_logprobs",
	strict_host_check=args.strict_host_check,
	)

	import torch

	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 (
	is_hf_checkpoint,
	load_stage_config,
	read_sft_quality_gate,
	resolve_pair_split_dir,
	resolve_sft_checkpoint,
	)
	from align.train.ref_cache import (
	build_ref_cache_key,
	path_bias_report_cache_forward_impl,
	precompute_ref_logprobs,
	write_path_bias_report,
	write_ref_cache,
	)
	from align.utils.guard import assert_paths_off_root

	cfg = load_stage_config(args.config)
	out_dir = args.out_dir or os.path.dirname(cfg.ref_cache_path)
	assert_paths_off_root([out_dir])

	with open(cfg.manifest_path) as f:
	manifest = json.load(f)

	# Resolve the frozen SFT checkpoint the same way scripts/dpo.py does.
	# The trainer writes <run_dir>/checkpoints/micro_NNNNNNNN with no stable
	# `final` alias; the quality-gate receipt records which one passed. A
	# configured path that IS a checkpoint is an explicit pin; otherwise the
	# receipt is read, and a missing/failed receipt or a missing checkpoint
	# refuses LOUDLY here -- a cache silently computed from the wrong (or no)
	# reference would poison both arms' margins.
	if is_hf_checkpoint(cfg.sft_checkpoint):
	sft_checkpoint, ckpt_source = cfg.sft_checkpoint, "config"
	else:
	gate = read_sft_quality_gate(cfg.sft_quality_gate_path)
	sft_checkpoint, ckpt_source = resolve_sft_checkpoint(
	cfg.sft_checkpoint, gate.get("checkpoint")
	)
	logger.info("frozen SFT checkpoint: %s (source=%s)", sft_checkpoint, ckpt_source)

	# Shard split dirs by DISCOVERY (the receipts the store wrote), never by
	# hand-composed `<root>-<split>` paths -- the attempt-1 defect class.
	train_dir = resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train")
	held_dir = resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held")

	# The key, assembled from facts rather than from intent, by the SAME
	# function scripts/dpo.py's reader uses (build_ref_cache_key). `checkpoint_hash`
	# is the WEIGHT-BYTES hash -- one streamed read of the checkpoint -- not the
	# tensor-inventory hash this script used to write, which could not tell two
	# SFT checkpoints of one architecture apart (2026-10-12 review). Caches
	# written under the old key are refused by dpo.py: re-run this script.
	# `manifest_hash` covers the shard rebuild that nothing else would notice.
	key = build_ref_cache_key(
	sft_checkpoint,
	manifest,
	forward_impl="fsdp2_sharded" if (args.path_bias or args.sharded_cache) else "single_gpu",
	)
	logger.info("cache key: %s", key.to_json())

	# Rank-pinned device: a bare "cuda" is device 0 for EVERY torchrun rank,
	# and the --path-bias 4-rank pass put four ~5.5GB reference shards on
	# GPU 0 (the duplicate-GPU class scripts/sft.py already fixed). This
	# mode had never been run before 2026-08-18.
	if torch.cuda.is_available():
	from align.utils.dist import init_distributed, local_rank

	init_distributed()
	torch.cuda.set_device(local_rank())
	device = f"cuda:{local_rank()}" # string: downstream does device.startswith(...)
	else:
	device = "cpu"
	model = _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device)

	if args.sharded_cache:
	import torch.distributed as dist

	from align.train.fsdp import apply_fsdp

	if not (dist.is_available() and dist.is_initialized()):
	raise RuntimeError(
	"--sharded-cache computes the reference under the SHARDED "
	"forward and must run under torchrun; single-process would "
	"silently produce the single_gpu numbers under a sharded key."
	)
	model = apply_fsdp(model)
	model.eval()

	written = []
	# In --path-bias mode the caches are NOT recomputed: the dedicated
	# single-GPU run owns them. Before this guard, all four torchrun ranks
	# each recomputed and rewrote the same cache files concurrently -- and
	# stamped forward_impl="fsdp2_sharded" into caches computed single-GPU.
	cache_splits = (
	()
	if args.path_bias
	else (
	("train", train_dir, cfg.ref_cache_path),
	("held", held_dir, cfg.heldout_ref_cache_path),
	)
	)
	for split, shard_dir, cache_path in cache_splits:
	rows = preference_dataset(shard_dir)
	cache = precompute_ref_logprobs(
	model,
	rows,
	key=key,
	pad_token_id=_pad_id(manifest),
	batch_size=args.batch_size,
	device=torch.device(device),
	)
	from align.utils.dist import is_rank0 as _is_rank0

	if _is_rank0():
	# Sharded forwards are collectives (every rank computes); the
	# file is written once.
	write_ref_cache(cache_path, cache)
	written.append({"split": split, "path": cache_path, "pairs": len(cache)})
	logger.info("%s/%s: %d reference pairs", cfg.arm, split, len(cache))

	if args.batch_invariance:
	from align.train.ref_cache import batch_invariance_report

	rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
	report = batch_invariance_report(
	model,
	rows,
	key=key,
	pad_token_id=_pad_id(manifest),
	beta=cfg.beta,
	early_margin_scale=cfg.path_bias_early_margin_scale,
	threshold=cfg.path_bias_threshold,
	)
	out_path = os.path.join(out_dir, f"{cfg.arm}-batch_invariance.json")
	with open(out_path, "w") as f:
	json.dump(report, f, indent=2)
	logger.info(
	"batch invariance -> %s (belongs_in_cache_key=%s)",
	out_path,
	report["belongs_in_cache_key"],
	)

	if args.path_bias:
	# OWNERSHIP TRANSFER: pass the model in a box and drop main's name.
	# A plain argument keeps main's reference alive, so the between-pass
	# free inside _path_bias freed nothing and the reload OOM'd with the
	# old 14.5GB replica still resident (measured 21.66GB, batch-size
	# independent).
	model_box = [model]
	del model
	report = _path_bias(
	model_box,
	cfg,
	manifest,
	key,
	args.batch_size,
	device,
	train_dir,
	reload_fn=lambda: _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device),
	)
	from align.utils.dist import is_rank0

	if is_rank0():
	# Rank-0-only: four ranks renaming one report is the ledger race.
	# `cache_forward_impl` must describe the CACHE, not this run. The
	# old code wrote key.forward_impl, which is "fsdp2_sharded" for
	# every --path-bias run -- but this mode recomputes no cache, so the
	# field named the measuring run and contradicted the single-GPU
	# cache sitting on disk (2026-10-12 review). Read it off the cache.
	write_path_bias_report(
	cfg.path_bias_report_path,
	report,
	n_examples=PATH_BIAS_EXAMPLES,
	forward_impl=path_bias_report_cache_forward_impl(cfg.ref_cache_path),
	)
	logger.info("path bias: %s (ratio %s)", report["action"], report["ratio"])

	logger.info("precompute complete: %s", written)
	return 0


	def _manifest_hash(manifest: dict) -> str:
	"""Hash of the build's own content hashes. Moves on any shard rebuild --
	which is the point: the section-03 marker fix retokenizes every pair while
	leaving the checkpoint, tokenizer and max_length untouched."""
	from align.data.receipts import content_hash

	return content_hash([manifest["hashes"][k] for k in sorted(manifest.get("hashes", {}))])


	def _max_length(manifest: dict) -> int:
	return int(manifest.get("pinned", {}).get("max_len", 2048))


	def _pad_id(manifest: dict) -> int:
	"""Pad id for collation. Fully masked out of every logprob (prompt and pad
	positions carry IGNORE_INDEX), so its value cannot reach a number -- but
	it is read from the build rather than assumed, because 'cannot reach a
	number' is the kind of claim that stops being true after a refactor."""
	return int(manifest.get("pinned", {}).get("pad_token_id", 0))


	def _load_reference(model_cls, ckpt, path: str, device: str):
	"""The FROZEN SFT checkpoint, bf16 on GPU, eval mode.

	eval() is asserted downstream rather than merely called here: Mistral-v0.3
	has attention_dropout 0.0, so a model left in train mode would produce the
	same numbers today and silently different ones after any config change.
	"""
	import torch
	from align.eval.golden import resolve_model_config

	model_cfg = resolve_model_config(path)
	model = model_cls.from_state_dict(
	model_cfg,
	ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)),
	attn_implementation="sdpa",
	device=device,
	)
	model = model.to(
	device, dtype=torch.bfloat16 if str(device).startswith("cuda") else torch.float32
	)
	model.eval()
	for module in model.modules():
	if isinstance(module, torch.nn.Dropout) and module.p > 0:
	raise RuntimeError(
	f"an active Dropout(p={module.p}) is present in the reference "
	"model. Cached reference logprobs must be exactly "
	"reproducible; dropout makes them a sample."
	)
	return model


	def _path_bias(
	model_box, cfg, manifest, key, batch_size: int, device: str, train_dir: str, reload_fn=None
	) -> dict:
	"""The mandatory dual pass: single-GPU forward vs FSDP2-sharded forward.

	Requires torchrun. Without a process group there IS no sharded forward to
	compare against, and the honest report is then no report -- `path_bias_
	decision` treats an empty delta set as firing, because not measuring is
	not passing and a stage that skipped this must not silently proceed on a
	single-GPU cache.
	"""
	import torch
	import torch.distributed as dist

	from align.data.loader import preference_dataset
	from align.train.fsdp import apply_fsdp
	from align.train.ref_cache import path_bias_decision, precompute_ref_logprobs

	if not (dist.is_available() and dist.is_initialized()):
	raise RuntimeError(
	"--path-bias needs the SHARDED forward, so it must run under "
	"torchrun (--nproc_per_node=4). Running it single-process would "
	"compare the single-GPU forward with itself and record a "
	"reassuring zero, which is worse than not measuring: it would "
	"look like evidence."
	)

	rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
	model = model_box.pop() # sole reference from here on
	single = precompute_ref_logprobs(
	model,
	rows,
	key=key,
	pad_token_id=_pad_id(manifest),
	batch_size=batch_size,
	device=torch.device(device),
	)
	# FREE the single-pass replica before sharding: wrapping a resident
	# 14.5GB model in-place stacks the wrap transients ON TOP of it and
	# OOMs every rank at 21.8GB (measured, 2026-08-18). Rebuilding alone
	# and then sharding is the stage launcher's proven peak.
	if reload_fn is None:
	raise RuntimeError("--path-bias needs reload_fn to rebuild the model between passes")
	del model
	import gc

	gc.collect()
	torch.cuda.empty_cache()
	model = reload_fn()
	sharded_model = apply_fsdp(model)
	sharded_model.eval()
	sharded = precompute_ref_logprobs(
	sharded_model,
	rows,
	key=key,
	pad_token_id=_pad_id(manifest),
	batch_size=batch_size,
	device=torch.device(device),
	)

	# The delta of the reference MARGIN, not of either logprob: the margin is
	# what enters the loss, and a bias common to both sides cancels there.
	deltas = [
	(sharded.entries[p][0] - sharded.entries[p][1])
	- (single.entries[p][0] - single.entries[p][1])
	for p in single.entries
	]
	return path_bias_decision(
	deltas,
	beta=cfg.beta,
	early_margin_scale=cfg.path_bias_early_margin_scale,
	threshold=cfg.path_bias_threshold,
	)


	if __name__ == "__main__":
	raise SystemExit(main())