code/scripts/sft.py · viai957/align at main

Add files using upload-large-folder tool

bcc8d35 verified 15 days ago

11.4 kB

#!/usr/bin/env python
"""SFT stage entry (section-11): gates -> wrap -> train -> eval -> quality gate.
    torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml --optimizer.lr=1e-5
Thin by design. Everything this drives is already built and tested elsewhere;
what lives here is ORDER, and the order is the product:
    golden_regression -> smoke_7b -> probe_20step -> rss_preflight
                      -> train -> eval -> quality_gate
IMPORT ORDER IS LOAD-BEARING. `apply_resource_guards()` sets
PYTORCH_CUDA_ALLOC_CONF, the /apps cache paths and the thread caps, and torch
reads all of them AT IMPORT. So this module imports only stdlib, the guard,
the config, and `stage_gates` (which is asserted torch-free); torch and every
align module that pulls it in are imported INSIDE main(), after the guards run.
A test enforces this by AST rather than by grep -- a literal search passes on
a file that mentions the guard in a comment and imports torch on the next line.
Resume semantics: rerunning after a kill picks up the gate ledger, skips gates
that already passed, and resumes the Trainer from the latest micro-step
checkpoint. Partial checkpoints are deleted at startup.
"""

from __future__ import annotations

import logging
import os
import sys
from pathlib import Path

# The package is not installed in the venv. Without this a bare `torchrun scripts/sft.py ...` (the launch
# line in docs/P2.md section 11 and in the config header) died with ModuleNotFoundError on all four ranks
# unless the operator remembered PYTHONPATH -- found by the P2 integration smoke loop.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.config import load_config
from align.train.stage_gates import (
    GateLedger,
    assemble_sft_banner,
    assert_build_is_current,
    assert_probe_resolved,
    eval_source_hash,
    evaluate_and_gate,
    parse_args,
    record_golden_from_report,
    require_golden_receipt,
    run_entry_gates,
)
from align.utils.guard import apply_resource_guards

logger = logging.getLogger(__name__)


def build_model(cfg):
    """Imported here, not at module scope: this pulls torch in.
    ARCH DISPATCH. `model.arch` selects the family; an unrecognised value RAISES (it must not
    fall through to the Mistral converters, which would either fail late on a Qwen3 snapshot's
    keys or load a different family into the wrong module).
    """
    if cfg.model.arch == "mistral":
        if cfg.lora.enabled:
            raise ValueError(
                "lora.enabled with model.arch='mistral': the dense-Mistral SFT path is full-FT only "
                "(LoRA Mistral runs live in scripts/ppo.py). Refusing rather than training full-FT "
                "while the config says LoRA."
            )
        from align.train.stage_smoke import build_stage_model

        return build_stage_model(cfg)
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import build_qwen3_stage_model

        return build_qwen3_stage_model(cfg)
    raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")


def run_gates(cfg, ledger, banner_only: bool = False):
    """The M5 entry sequence. Refuses to proceed on a stale data build or an
    unvalidated eval stack before it touches a GPU.
    Non-Mistral arches have their own entry (see `qwen3_entry_banner`: the golden/MATH/GSM8K
    machinery certifies a Mistral checkpoint and does not apply)."""
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import qwen3_entry_banner

        return qwen3_entry_banner(cfg, ledger)
    if cfg.model.arch != "mistral":
        raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")
    assert_build_is_current(cfg.data.manifest_path, cfg)
    # Bridge the golden RUN into the ledger by re-deciding from its persisted
    # report, immediately before the gate that reads the ledger. Without this
    # a PASSING golden run still leaves golden_regression PENDING.
    record_golden_from_report(ledger, cfg=cfg)
    require_golden_receipt(ledger, cfg)
    banner = run_entry_gates(cfg, ledger)
    assert_probe_resolved(banner)
    return banner


def train(cfg, ledger, model, banner):
    """One SFT epoch through the section-07 Trainer."""
    from align.data.loader import collate_sft
    from align.train.sft_loss import make_sft_loss_adapter
    from align.train.stage_smoke import build_training_sampler
    from align.train.trainer import Trainer

    # The SAME sampler the entry gates measured the horizon from. Building a
    # second one here is how a stage ends up timing, projecting and training
    # on three different batch streams.
    sampler, dataset = build_training_sampler(cfg)

    def adapter(m, batch):
        rows = [dataset[i] for i in batch]
        collated = collate_sft(rows, pad_token_id=cfg.data.pad_token_id).to("cuda")
        return make_sft_loss_adapter(require_fused=True)(m, (collated.input_ids, collated.labels))

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    extras = {}
    if cfg.model.arch == "qwen3":
        from align.train.qwen3_stage import qwen3_trainer_extras

        extras = qwen3_trainer_extras(cfg, model, adapter, run_dir)
    trainer = Trainer(
        cfg=cfg,
        # LoRA: the trunk-movement probe must watch the adapter's B matrices (the frozen base can
        # never move and A provably does not on step one), so it gets a view, not the raw model.
        model=extras.pop("probe_model", model),
        sampler=sampler,
        loss_adapter=adapter,
        run_dir=run_dir,
        stage="sft",
        **extras,
    )
    trainer.resume_from_latest()
    summary = trainer.train()
    if cfg.model.arch == "mistral":
        record_train_gate(ledger, trainer, summary)
    else:
        # The ledger's fixed sequence starts at the Mistral golden regression; a non-Mistral
        # stage has not passed those gates, so recording "train" would raise (or be a lie).
        logger.info("train finished (arch=%s): optimizer_step=%s", cfg.model.arch, summary.get("optimizer_step"))
    return trainer


def record_train_gate(ledger, trainer, summary) -> dict:
    """Record the 'train' gate for what `trainer.train()` ACTUALLY did.
    THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13). `train()`
    returns normally after a COORDINATED STOP (SIGTERM, or train.max_wall_seconds)
    -- it checkpoints and returns -- and the gate used to be recorded
    `passed=True` unconditionally. `main` then (correctly) skipped the eval and
    exited 128+n / 82, but the ledger on disk now certified a half-trained
    checkpoint as a finished 'train' stage: anything that reads the ledger
    instead of the exit status (a resumed `evaluate_and_gate`, an operator
    reading `stage_gates.json`) was told the stage had trained.
    A stopped run is neither finished nor failed: it is RESUMABLE. The ledger
    has no third state, so it is recorded `passed=False` with a `STOPPED
    (resumable)` detail; a rerun resumes from the checkpoint and records the
    real verdict at the end of the real run.
    """
    step = summary.get("optimizer_step")
    if getattr(trainer, "stop_exit_code", 0):
        return ledger.record(
            "train",
            passed=False,
            detail=f"STOPPED (resumable): {trainer.stop_reason} stop at optimizer_step={step}; "
            "checkpointed, not finished",
        )
    return ledger.record("train", passed=True, detail=str(step))


def main(argv=None) -> int:
    args = parse_args(argv)

    # FIRST: every env var torch reads at import, plus the heavy-job flock and
    # the /apps path assertions. `heavy=True` also runs the host preflight;
    # `rss_cfg` adds the full-FT RSS projection on top of the generic floor.
    cfg = load_config(args.config, args.overrides)
    apply_resource_guards(cpu_threads=8, heavy=True, label="sft-stage-11", rss_cfg=cfg)

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )

    # BEFORE anything touches CUDA. THE DUPLICATE-GPU DEFECT: torchrun starts
    # four processes but assigns no devices, so without this every rank's
    # `.cuda()` lands on device 0 and NCCL dies with "Duplicate GPU detected:
    # rank 2 and rank 0 both on CUDA device 1b0" -> ncclInvalidUsage. Section-01
    # wrote `init_distributed` for exactly this (init_process_group +
    # set_device(local_rank)); this script simply never called it.
    from align.utils.dist import init_distributed

    run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
    # run_dir: the NCCL flight-recorder dump lands under it (never /tmp).
    init_distributed(run_dir=run_dir)

    # Microseconds, and it converts the whole duplicate-GPU class into a named
    # launch failure instead of an ncclInvalidUsage 400 lines into a forward.
    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()
    ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())

    banner = run_gates(cfg, ledger)
    from align.registry import render_banner

    print(render_banner(banner, title=f"align SFT stage :: {cfg.run_name}"))

    model = build_model(cfg)
    trainer = train(cfg, ledger, model, banner)
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved and
        # returned: the stage is NOT finished, so do not fall into its
        # post-train eval or announce completion. 82 = wall budget (the
        # supervisor will not relaunch it), 128+n = the signal.
        print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr)
        return trainer.stop_exit_code

    # BARRIER, then RELEASE, then eval -- both are load-bearing:
    # 1. train() returns minutes earlier on ranks 1..N-1 than on rank 0 (the
    #    sole checkpoint writer). Without this barrier they raced into the
    #    eval, OOM'd building a second 7B, and torchrun SIGTERM'd rank 0
    #    MID-SAVE -- whose cleanup handler then deleted the .partial. The
    #    epoch's tail died twice this way.
    # 2. run_stage_eval builds its OWN model from the checkpoint; the
    #    training replica + optimizer must be freed first or the engine's
    #    KV-cache preflight refuses (1.2GB free vs 4.8GB required).
    from align.utils.dist import barrier

    barrier()
    checkpoint = trainer.latest_checkpoint()
    if cfg.model.arch != "mistral":
        # The post-train MATH-500/GSM8K gate certifies the Mistral stack; a qwen3 checkpoint is
        # judged by the Lean loop. Say so instead of silently skipping.
        print(f"SFT stage complete (arch={cfg.model.arch}; no MATH/GSM8K gate applies). Checkpoint: {checkpoint}")
        return 0
    trainer.release()  # model, optimizer AND the probe's references (U01-11)
    del model
    import gc

    import torch

    gc.collect()
    torch.cuda.empty_cache()
    decision = evaluate_and_gate(cfg, ledger, checkpoint=checkpoint)
    if not decision.passed:
        print(decision.halt_message, file=sys.stderr)
        return 1

    print(f"SFT stage complete. DPO base checkpoint: {checkpoint}")
    return 0


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

    raise SystemExit(run_with_exit_codes(main))

	#!/usr/bin/env python
	"""SFT stage entry (section-11): gates -> wrap -> train -> eval -> quality gate.

	torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml --optimizer.lr=1e-5

	Thin by design. Everything this drives is already built and tested elsewhere;
	what lives here is ORDER, and the order is the product:

	golden_regression -> smoke_7b -> probe_20step -> rss_preflight
	-> train -> eval -> quality_gate

	IMPORT ORDER IS LOAD-BEARING. `apply_resource_guards()` sets
	PYTORCH_CUDA_ALLOC_CONF, the /apps cache paths and the thread caps, and torch
	reads all of them AT IMPORT. So this module imports only stdlib, the guard,
	the config, and `stage_gates` (which is asserted torch-free); torch and every
	align module that pulls it in are imported INSIDE main(), after the guards run.
	A test enforces this by AST rather than by grep -- a literal search passes on
	a file that mentions the guard in a comment and imports torch on the next line.

	Resume semantics: rerunning after a kill picks up the gate ledger, skips gates
	that already passed, and resumes the Trainer from the latest micro-step
	checkpoint. Partial checkpoints are deleted at startup.
	"""

	from __future__ import annotations

	import logging
	import os
	import sys
	from pathlib import Path

	# The package is not installed in the venv. Without this a bare `torchrun scripts/sft.py ...` (the launch
	# line in docs/P2.md section 11 and in the config header) died with ModuleNotFoundError on all four ranks
	# unless the operator remembered PYTHONPATH -- found by the P2 integration smoke loop.
	sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

	from align.config import load_config
	from align.train.stage_gates import (
	GateLedger,
	assemble_sft_banner,
	assert_build_is_current,
	assert_probe_resolved,
	eval_source_hash,
	evaluate_and_gate,
	parse_args,
	record_golden_from_report,
	require_golden_receipt,
	run_entry_gates,
	)
	from align.utils.guard import apply_resource_guards

	logger = logging.getLogger(__name__)


	def build_model(cfg):
	"""Imported here, not at module scope: this pulls torch in.

	ARCH DISPATCH. `model.arch` selects the family; an unrecognised value RAISES (it must not
	fall through to the Mistral converters, which would either fail late on a Qwen3 snapshot's
	keys or load a different family into the wrong module).
	"""
	if cfg.model.arch == "mistral":
	if cfg.lora.enabled:
	raise ValueError(
	"lora.enabled with model.arch='mistral': the dense-Mistral SFT path is full-FT only "
	"(LoRA Mistral runs live in scripts/ppo.py). Refusing rather than training full-FT "
	"while the config says LoRA."
	)
	from align.train.stage_smoke import build_stage_model

	return build_stage_model(cfg)
	if cfg.model.arch == "qwen3":
	from align.train.qwen3_stage import build_qwen3_stage_model

	return build_qwen3_stage_model(cfg)
	raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")


	def run_gates(cfg, ledger, banner_only: bool = False):
	"""The M5 entry sequence. Refuses to proceed on a stale data build or an
	unvalidated eval stack before it touches a GPU.

	Non-Mistral arches have their own entry (see `qwen3_entry_banner`: the golden/MATH/GSM8K
	machinery certifies a Mistral checkpoint and does not apply)."""
	if cfg.model.arch == "qwen3":
	from align.train.qwen3_stage import qwen3_entry_banner

	return qwen3_entry_banner(cfg, ledger)
	if cfg.model.arch != "mistral":
	raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')")
	assert_build_is_current(cfg.data.manifest_path, cfg)
	# Bridge the golden RUN into the ledger by re-deciding from its persisted
	# report, immediately before the gate that reads the ledger. Without this
	# a PASSING golden run still leaves golden_regression PENDING.
	record_golden_from_report(ledger, cfg=cfg)
	require_golden_receipt(ledger, cfg)
	banner = run_entry_gates(cfg, ledger)
	assert_probe_resolved(banner)
	return banner


	def train(cfg, ledger, model, banner):
	"""One SFT epoch through the section-07 Trainer."""
	from align.data.loader import collate_sft
	from align.train.sft_loss import make_sft_loss_adapter
	from align.train.stage_smoke import build_training_sampler
	from align.train.trainer import Trainer

	# The SAME sampler the entry gates measured the horizon from. Building a
	# second one here is how a stage ends up timing, projecting and training
	# on three different batch streams.
	sampler, dataset = build_training_sampler(cfg)

	def adapter(m, batch):
	rows = [dataset[i] for i in batch]
	collated = collate_sft(rows, pad_token_id=cfg.data.pad_token_id).to("cuda")
	return make_sft_loss_adapter(require_fused=True)(m, (collated.input_ids, collated.labels))

	run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
	extras = {}
	if cfg.model.arch == "qwen3":
	from align.train.qwen3_stage import qwen3_trainer_extras

	extras = qwen3_trainer_extras(cfg, model, adapter, run_dir)
	trainer = Trainer(
	cfg=cfg,
	# LoRA: the trunk-movement probe must watch the adapter's B matrices (the frozen base can
	# never move and A provably does not on step one), so it gets a view, not the raw model.
	model=extras.pop("probe_model", model),
	sampler=sampler,
	loss_adapter=adapter,
	run_dir=run_dir,
	stage="sft",
	**extras,
	)
	trainer.resume_from_latest()
	summary = trainer.train()
	if cfg.model.arch == "mistral":
	record_train_gate(ledger, trainer, summary)
	else:
	# The ledger's fixed sequence starts at the Mistral golden regression; a non-Mistral
	# stage has not passed those gates, so recording "train" would raise (or be a lie).
	logger.info("train finished (arch=%s): optimizer_step=%s", cfg.model.arch, summary.get("optimizer_step"))
	return trainer


	def record_train_gate(ledger, trainer, summary) -> dict:
	"""Record the 'train' gate for what `trainer.train()` ACTUALLY did.

	THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13). `train()`
	returns normally after a COORDINATED STOP (SIGTERM, or train.max_wall_seconds)
	-- it checkpoints and returns -- and the gate used to be recorded
	`passed=True` unconditionally. `main` then (correctly) skipped the eval and
	exited 128+n / 82, but the ledger on disk now certified a half-trained
	checkpoint as a finished 'train' stage: anything that reads the ledger
	instead of the exit status (a resumed `evaluate_and_gate`, an operator
	reading `stage_gates.json`) was told the stage had trained.

	A stopped run is neither finished nor failed: it is RESUMABLE. The ledger
	has no third state, so it is recorded `passed=False` with a `STOPPED
	(resumable)` detail; a rerun resumes from the checkpoint and records the
	real verdict at the end of the real run.
	"""
	step = summary.get("optimizer_step")
	if getattr(trainer, "stop_exit_code", 0):
	return ledger.record(
	"train",
	passed=False,
	detail=f"STOPPED (resumable): {trainer.stop_reason} stop at optimizer_step={step}; "
	"checkpointed, not finished",
	)
	return ledger.record("train", passed=True, detail=str(step))


	def main(argv=None) -> int:
	args = parse_args(argv)

	# FIRST: every env var torch reads at import, plus the heavy-job flock and
	# the /apps path assertions. `heavy=True` also runs the host preflight;
	# `rss_cfg` adds the full-FT RSS projection on top of the generic floor.
	cfg = load_config(args.config, args.overrides)
	apply_resource_guards(cpu_threads=8, heavy=True, label="sft-stage-11", rss_cfg=cfg)

	logging.basicConfig(
	level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
	)

	# BEFORE anything touches CUDA. THE DUPLICATE-GPU DEFECT: torchrun starts
	# four processes but assigns no devices, so without this every rank's
	# `.cuda()` lands on device 0 and NCCL dies with "Duplicate GPU detected:
	# rank 2 and rank 0 both on CUDA device 1b0" -> ncclInvalidUsage. Section-01
	# wrote `init_distributed` for exactly this (init_process_group +
	# set_device(local_rank)); this script simply never called it.
	from align.utils.dist import init_distributed

	run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft")
	# run_dir: the NCCL flight-recorder dump lands under it (never /tmp).
	init_distributed(run_dir=run_dir)

	# Microseconds, and it converts the whole duplicate-GPU class into a named
	# launch failure instead of an ncclInvalidUsage 400 lines into a forward.
	from align.train.stage_gates import assert_distinct_devices

	assert_distinct_devices()
	ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash())

	banner = run_gates(cfg, ledger)
	from align.registry import render_banner

	print(render_banner(banner, title=f"align SFT stage :: {cfg.run_name}"))

	model = build_model(cfg)
	trainer = train(cfg, ledger, model, banner)
	if trainer.stop_exit_code:
	# A coordinated stop (SIGTERM / train.max_wall_seconds) saved and
	# returned: the stage is NOT finished, so do not fall into its
	# post-train eval or announce completion. 82 = wall budget (the
	# supervisor will not relaunch it), 128+n = the signal.
	print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr)
	return trainer.stop_exit_code

	# BARRIER, then RELEASE, then eval -- both are load-bearing:
	# 1. train() returns minutes earlier on ranks 1..N-1 than on rank 0 (the
	# sole checkpoint writer). Without this barrier they raced into the
	# eval, OOM'd building a second 7B, and torchrun SIGTERM'd rank 0
	# MID-SAVE -- whose cleanup handler then deleted the .partial. The
	# epoch's tail died twice this way.
	# 2. run_stage_eval builds its OWN model from the checkpoint; the
	# training replica + optimizer must be freed first or the engine's
	# KV-cache preflight refuses (1.2GB free vs 4.8GB required).
	from align.utils.dist import barrier

	barrier()
	checkpoint = trainer.latest_checkpoint()
	if cfg.model.arch != "mistral":
	# The post-train MATH-500/GSM8K gate certifies the Mistral stack; a qwen3 checkpoint is
	# judged by the Lean loop. Say so instead of silently skipping.
	print(f"SFT stage complete (arch={cfg.model.arch}; no MATH/GSM8K gate applies). Checkpoint: {checkpoint}")
	return 0
	trainer.release() # model, optimizer AND the probe's references (U01-11)
	del model
	import gc

	import torch

	gc.collect()
	torch.cuda.empty_cache()
	decision = evaluate_and_gate(cfg, ledger, checkpoint=checkpoint)
	if not decision.passed:
	print(decision.halt_message, file=sys.stderr)
	return 1

	print(f"SFT stage complete. DPO base checkpoint: {checkpoint}")
	return 0


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

	raise SystemExit(run_with_exit_codes(main))