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

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

18.4 kB

#!/usr/bin/env python
"""Stage-b GRPO entry: gates -> worker -> banner -> train -> stage manifest.
    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
        configs/grpo_stageb.toml
    # tiny end-to-end smoke (2 prompts x G8 x 32 tokens x 10 steps):
    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
        configs/grpo_stageb.toml --smoke
Topology (v1, DESIGN v3 §B.1b): 3 FSDP train ranks + one dedicated vLLM
engine GPU (default 3, override NMO_ROLLOUT_GPU). Rank 0 owns all protocol
I/O; checkpoint saves ARE the weight sync (the worker reloads the engine
whenever a request's weights_dir changes). The loop is section-07's Trainer
verbatim -- GrpoBatchSource yields lazy handles and the loss-adapter wrapper
materialises each optimizer step's rollout at consumption time.
Launch gates, ordered by cost: config validation (denominator arithmetic,
cadence/sync alignment, R3 clipping, R4 canary) -> ceiling receipt
(0/100-perfect-policy trap) -> banding manifest -> topology check ->
A3 conformance + entropy-when-armed (module construction) -> horizon
consistency (scheduler vs configured steps).
"""

from __future__ import annotations

import argparse
import atexit
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("grpo")


#: The rollout worker runs in its OWN session (torchrun's teardown does not reach it) and used to
#: serve forever. If rank 0 is SIGKILLed it is orphaned holding ~85% of its GPU, so it now exits by
#: itself after this long with no request -- far above any step/checkpoint gap, far below "forever".
ROLLOUT_IDLE_EXIT_S = 7200.0


def _spawn_rollout_worker(spool: str, gpu: int, log_path: str):
    """Spawn the vLLM worker and tie its teardown to rank 0's exit.
    Everything between the spawn and the train try/finally (Trainer construction, the horizon
    check, canary data, resume_from_latest) can raise; atexit covers every uncaught-exception path.
    `terminate_worker` is idempotent, so the normal finally and this hook do not conflict.
    """
    from align.rollout import client as rollout_client

    proc = rollout_client.spawn_worker(spool, gpu=gpu, log_path=log_path, idle_exit_s=ROLLOUT_IDLE_EXIT_S)
    atexit.register(rollout_client.terminate_worker, proc)
    return proc


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("config")
    ap.add_argument("--smoke", action="store_true", help="tiny shapes, 10 steps")
    ap.add_argument("--band-file", default=None, help="override the configured band file")
    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="grpo-stageb",
        strict_host_check=args.strict_host_check,
    )

    from align.utils.dist import init_distributed
    from align.rl.stage_config import load_stage_config as _early_stage_config

    # Run dir from the config (parsed before the process group: a bad TOML fails
    # in milliseconds, and the NCCL flight-recorder dump lands under it).
    init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir)

    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()

    import torch
    import torch.distributed as dist

    from align.data.schema import build_prompt, extract_answer, register_answer_suffixes
    from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test
    from align.eval.verifier import OutcomeVerifier, VerifierConfig
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.rl.batch_source import GrpoBatchSource
    from align.rl.canary import TaskAccuracyCanary
    from align.rl.grpo_module import GRPOTrainModule
    from align.rl.rewards import CompositeReward
    from align.rl.stage_config import load_stage_config
    from align.rl.tripwire import CollapseTripwire
    from align.rollout.client import make_rollout_client
    from align.rollout.protocol import RolloutRequest
    from align.train.dpo_stage import resolve_sft_checkpoint
    from align.train.trainer import Trainer
    from align.utils.guard import assert_paths_off_root

    rank = dist.get_rank() if dist.is_initialized() else 0
    world = dist.get_world_size() if dist.is_initialized() else 1

    cfg = load_stage_config(args.config, smoke=args.smoke)
    if args.band_file:
        cfg.data.band_file = args.band_file
        cfg.validate()
    assert_paths_off_root([cfg.run_dir])
    ceiling = cfg.validate_ceiling_receipt()
    cfg.validate_world(world)
    with open(cfg.data.banding_manifest) as f:
        banding_manifest = json.load(f)

    register_answer_suffixes("</s>")  # BEFORE any extraction (golden.py's incident)

    # -- data: the band's prompt rows ---------------------------------------
    prompt_rows = []
    with open(cfg.data.band_file) as f:
        for line in f:
            if line.strip():
                prompt_rows.append(json.loads(line))
    logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows))

    # -- policy + FSDP -------------------------------------------------------
    sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None)
    logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source)
    from align.model.hf_config import read_config
    from align.tokenizer import load_tokenizer

    tokenizer = load_tokenizer(sft_checkpoint)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

    model_cfg = read_config(sft_checkpoint)
    model = NmoMistralForCausalLM.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)),
        attn_implementation=cfg.train.attn_implementation,
        device="cuda",
    ).to("cuda", dtype=torch.bfloat16)

    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp

    dist_model = apply_fsdp(apply_activation_checkpointing(model))
    logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30)

    # -- rollout backend (rank 0 owns it) ------------------------------------
    # 'spool': external vLLM worker process on NMO_ROLLOUT_GPU (own venv).
    # 'native': in-repo torch sampler on a STANDALONE engine copy living on that
    # same dedicated GPU, in rank 0's process. It must NOT be the FSDP2-sharded
    # model: decode forwards are collectives the other ranks would never join
    # (deadlock). Either way the GPU is outside the training ranks' devices.
    rollout_gpu = int(os.environ.get("NMO_ROLLOUT_GPU", "3"))
    spool = os.path.join(cfg.run_dir, "spool")
    worker_proc = None
    client = None
    if cfg.rollout.backend == "spool":
        if rank == 0:
            worker_proc = _spawn_rollout_worker(
                spool, rollout_gpu, os.path.join(cfg.run_dir, "worker.log")
            )
            logger.info("rollout worker spawned (pid %d) on GPU %s", worker_proc.pid, rollout_gpu)
        client = make_rollout_client("spool", spool=spool, timeout_s=cfg.rollout_timeout_s)
    elif rank == 0:
        client = _build_native_client(
            model_cfg, sft_checkpoint, tokenizer, rollout_gpu, world, cfg.rollout.max_response_len
        )
        logger.info("native rollout engine on cuda:%d (standalone copy of %s)", rollout_gpu, sft_checkpoint)

    # -- module + trainer -----------------------------------------------------
    from align.train.offload_adamw import make_optimizer

    run_cfg = _run_config(cfg)
    optimizer = make_optimizer(run_cfg, [p for p in dist_model.parameters() if p.requires_grad])
    collapse = CollapseTripwire()
    module = GRPOTrainModule(
        dist_model,
        optimizer,
        grad_accum=cfg.train.grad_accum,
        grpo_cfg=cfg.grpo,
        temperature=cfg.rollout.temperature,
        collapse_tripwire=collapse,
        max_grad_norm=cfg.train.max_grad_norm,
    )

    source = GrpoBatchSource(
        prompt_rows,
        prompts_per_step=cfg.rollout.prompts_per_step,
        n_samples_per_prompt=cfg.grpo.num_samples_per_prompt_rollout,
        max_response_len=cfg.rollout.max_response_len,
        temperature=cfg.rollout.temperature,
        optimizer_steps=cfg.train.optimizer_steps,
        grad_accum=cfg.train.grad_accum,
        seed=cfg.train.seed,
        prompt_builder=build_prompt,
        dead_group_retries=cfg.dead_group_retries,
    )
    source.bind_pad(pad_id)

    trainer = Trainer(
        run_cfg,
        dist_model,
        source,
        module._grpo_loss_adapter,
        run_dir=cfg.run_dir,
        stage="grpo",
        train_module=module,
    )
    # HORIZON CONSISTENCY (the Arm A zombie-LR class): the scheduler's total
    # must be exactly the configured optimizer_steps -- this loop generates
    # its own data, so any mismatch is arithmetic, not sampling.
    if trainer.scheduler.total_steps != cfg.train.optimizer_steps:
        raise RuntimeError(
            f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs "
            f"configured optimizer_steps={cfg.train.optimizer_steps}"
        )

    # -- canary (R4, rank 0) --------------------------------------------------
    canary = None
    canary_generate = None
    if rank == 0:
        sources = DataSources.from_local_hub()
        gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems]
        canary_problems = [
            {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows
        ]

        def canary_generate(prompts: list[str]) -> list[str]:
            req = RolloutRequest(
                request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}",
                weights_dir=_latest_weights(trainer, sft_checkpoint),
                prompts=prompts,
                n_samples_per_prompt=1,
                max_tokens=cfg.rollout.max_response_len,
                temperature=0.0,
            )
            client.submit(req)
            result = client.collect(req.request_id)
            return [texts[0] for texts in result.texts]

        canary = TaskAccuracyCanary(
            problems=canary_problems,
            generate_fn=canary_generate,
            verifier=OutcomeVerifier(VerifierConfig(workers=4)),
            extract_fn=extract_answer,
            every=cfg.canary.every,
        )

    source.bind(
        client=client,
        reward=CompositeReward(
            format_weight=cfg.reward.format_weight,
            overlong_max_len=cfg.rollout.max_response_len,
            overlong_cache_len=cfg.rollout.overlong_cache_len,
        ),
        verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None,
        old_logprob_fn=module.compute_logprobs_nograd,
        weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint),
        rank=rank,
        world_size=world,
        device="cuda",
        canary=canary,
        collapse_tripwire=collapse,
        trainer=trainer,
        canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"),
    )
    trainer.loss_adapter = source.lazy_adapter(module)

    # -- resume ---------------------------------------------------------------
    # A relaunch continues the newest checkpoint in run_dir, if any. Before
    # this was wired, a relaunch restarted from the SFT substrate. What a
    # resume restores, and where each piece lives:
    #   weights + optimizer + schedule + stream offset   Trainer (resume.pt)
    #   CollapseTripwire windows and latch                module extra_state
    #   prompt order, cursor, dead-group stats, alarms    source state_dict
    #   rollout engine weights                            _latest_weights: the
    #       Trainer registers the resumed checkpoint, so the first post-resume
    #       rollout reloads from it instead of the substrate.
    resumed_from = trainer.resume_from_latest()
    if resumed_from is not None:
        logger.info(
            "RESUMED from %s at optimizer step %d / micro step %d",
            resumed_from,
            trainer.optimizer_step,
            trainer.micro_step,
        )

    banner = {
        "stage": "grpo_stageb" + ("_smoke" if args.smoke else ""),
        "substrate": sft_checkpoint,
        "resumed_from": resumed_from,
        "band_file": cfg.data.band_file,
        "band_prompts": len(prompt_rows),
        "shape": {
            "prompts_per_step": cfg.rollout.prompts_per_step,
            "G": cfg.grpo.num_samples_per_prompt_rollout,
            "max_response_len": cfg.rollout.max_response_len,
            "loss_denominator": cfg.grpo.loss_denominator,
            "optimizer_steps": cfg.train.optimizer_steps,
        },
        "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()},
        "banding_manifest_rows": banding_manifest.get("rows"),
        "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")},
        "rectifications": (
            f"R3 clip={cfg.train.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, "
            "R5 pin-on-fire, reduction-time denominator assert armed"
        ),
    }
    logger.info("BANNER %s", json.dumps(banner, indent=1))

    # -- train ----------------------------------------------------------------
    try:
        summary = trainer.train()
        checkpoint = trainer.save_checkpoint()
    finally:
        if worker_proc is not None:
            from align.rollout.client import terminate_worker

            terminate_worker(worker_proc)
    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)

    result = {
        "stage": banner["stage"],
        "summary": summary,
        "checkpoint": checkpoint,
        "banner": banner,
        "canary_history": canary.history if canary else [],
        "collapse_alarms": source.alarms,
        "collapse_tripwire": collapse.as_log_dict(),
        "dead_groups": source.filter_stats.as_log_dict(),
        "dropped_remainder_rows": source.dropped_remainder_rows,
    }
    if rank == 0:
        path = os.path.join(cfg.run_dir, "stage_result.json")
        with open(path + ".partial", "w") as f:
            json.dump(result, f, indent=1, default=str)
        os.replace(path + ".partial", path)
        logger.info("stage result: %s", path)
    return 0


def _build_native_client(
    model_cfg, sft_checkpoint: str, tokenizer, rollout_gpu: int, world: int, max_response_len: int
):
    """Rank 0's native rollout backend on a dedicated, standalone engine copy.
    The copy is built from the SFT substrate (the policy's step-0 weights) and
    refreshed from the latest saved checkpoint whenever `weights_dir` changes
    (`checkpoint_weight_loader`). Refuses a device a training rank already
    owns: training ranks are bound to local_rank 0..world-1, so the engine GPU
    index must be >= world, and visible.
    """
    import torch

    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.rollout import checkpoint_weight_loader, make_rollout_client
    from align.rollout.native import NativeRolloutConfig

    if rollout_gpu < world or rollout_gpu >= torch.cuda.device_count():
        raise RuntimeError(
            f"NMO_ROLLOUT_GPU={rollout_gpu} is not a free visible device: training ranks own "
            f"cuda:0..{world - 1} and {torch.cuda.device_count()} device(s) are visible. The "
            "engine copy must live on its own GPU (launch with all four GPUs visible)."
        )
    device = torch.device("cuda", rollout_gpu)
    engine_model = (
        NmoMistralForCausalLM.from_state_dict(
            model_cfg,
            ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint, dtype=torch.bfloat16)),
            device=str(device),
        )
        .to(device, dtype=torch.bfloat16)
        .eval()
    )
    return make_rollout_client(
        "native",
        model=engine_model,
        tokenizer=tokenizer,
        native_config=NativeRolloutConfig(
            max_seq_len_cap=2048,
            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
        ),
        load_weights_fn=checkpoint_weight_loader(engine_model),
        device=device,
    )


def _latest_weights(trainer, sft_checkpoint: str) -> str:
    """The engine's weights source: latest saved checkpoint, else the substrate."""
    if trainer.checkpoints_written:
        return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}")
    return sft_checkpoint


def _run_config(cfg):
    from align.config import RunConfig

    run = RunConfig()
    run.optimizer.name = cfg.train.optimizer_name
    run.optimizer.lr = cfg.train.lr
    run.optimizer.grad_accum = cfg.train.grad_accum
    run.train.epochs = 1  # the horizon is optimizer_steps; the source cycles prompts itself
    run.train.micro_batch_size = cfg.train.micro_batch_size
    run.train.save_every_micro_steps = cfg.train.save_every_micro_steps
    run.train.keep_checkpoints = cfg.train.keep_checkpoints
    run.train.per_rank_optimizer_state = cfg.train.per_rank_optimizer_state
    run.train.max_wall_seconds = cfg.train.max_wall_seconds
    run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps
    run.train.max_skipped_fraction = cfg.train.max_skipped_fraction
    run.train.warmup_frac = cfg.train.warmup_frac
    run.train.full_ft = True
    run.train.seed = cfg.train.seed
    run.tripwires.armed = tuple(cfg.armed_tripwires)
    run.run_name = "grpo_stageb"
    return run


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

    raise SystemExit(run_with_exit_codes(main))

	#!/usr/bin/env python
	"""Stage-b GRPO entry: gates -> worker -> banner -> train -> stage manifest.

	CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
	configs/grpo_stageb.toml
	# tiny end-to-end smoke (2 prompts x G8 x 32 tokens x 10 steps):
	CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \
	configs/grpo_stageb.toml --smoke

	Topology (v1, DESIGN v3 §B.1b): 3 FSDP train ranks + one dedicated vLLM
	engine GPU (default 3, override NMO_ROLLOUT_GPU). Rank 0 owns all protocol
	I/O; checkpoint saves ARE the weight sync (the worker reloads the engine
	whenever a request's weights_dir changes). The loop is section-07's Trainer
	verbatim -- GrpoBatchSource yields lazy handles and the loss-adapter wrapper
	materialises each optimizer step's rollout at consumption time.

	Launch gates, ordered by cost: config validation (denominator arithmetic,
	cadence/sync alignment, R3 clipping, R4 canary) -> ceiling receipt
	(0/100-perfect-policy trap) -> banding manifest -> topology check ->
	A3 conformance + entropy-when-armed (module construction) -> horizon
	consistency (scheduler vs configured steps).
	"""

	from __future__ import annotations

	import argparse
	import atexit
	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("grpo")


	#: The rollout worker runs in its OWN session (torchrun's teardown does not reach it) and used to
	#: serve forever. If rank 0 is SIGKILLed it is orphaned holding ~85% of its GPU, so it now exits by
	#: itself after this long with no request -- far above any step/checkpoint gap, far below "forever".
	ROLLOUT_IDLE_EXIT_S = 7200.0


	def _spawn_rollout_worker(spool: str, gpu: int, log_path: str):
	"""Spawn the vLLM worker and tie its teardown to rank 0's exit.

	Everything between the spawn and the train try/finally (Trainer construction, the horizon
	check, canary data, resume_from_latest) can raise; atexit covers every uncaught-exception path.
	`terminate_worker` is idempotent, so the normal finally and this hook do not conflict.
	"""
	from align.rollout import client as rollout_client

	proc = rollout_client.spawn_worker(spool, gpu=gpu, log_path=log_path, idle_exit_s=ROLLOUT_IDLE_EXIT_S)
	atexit.register(rollout_client.terminate_worker, proc)
	return proc


	def main() -> int:
	ap = argparse.ArgumentParser()
	ap.add_argument("config")
	ap.add_argument("--smoke", action="store_true", help="tiny shapes, 10 steps")
	ap.add_argument("--band-file", default=None, help="override the configured band file")
	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="grpo-stageb",
	strict_host_check=args.strict_host_check,
	)

	from align.utils.dist import init_distributed
	from align.rl.stage_config import load_stage_config as _early_stage_config

	# Run dir from the config (parsed before the process group: a bad TOML fails
	# in milliseconds, and the NCCL flight-recorder dump lands under it).
	init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir)

	from align.train.stage_gates import assert_distinct_devices

	assert_distinct_devices()

	import torch
	import torch.distributed as dist

	from align.data.schema import build_prompt, extract_answer, register_answer_suffixes
	from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test
	from align.eval.verifier import OutcomeVerifier, VerifierConfig
	from align.model import checkpoint as ckpt
	from align.model.mistral import NmoMistralForCausalLM
	from align.rl.batch_source import GrpoBatchSource
	from align.rl.canary import TaskAccuracyCanary
	from align.rl.grpo_module import GRPOTrainModule
	from align.rl.rewards import CompositeReward
	from align.rl.stage_config import load_stage_config
	from align.rl.tripwire import CollapseTripwire
	from align.rollout.client import make_rollout_client
	from align.rollout.protocol import RolloutRequest
	from align.train.dpo_stage import resolve_sft_checkpoint
	from align.train.trainer import Trainer
	from align.utils.guard import assert_paths_off_root

	rank = dist.get_rank() if dist.is_initialized() else 0
	world = dist.get_world_size() if dist.is_initialized() else 1

	cfg = load_stage_config(args.config, smoke=args.smoke)
	if args.band_file:
	cfg.data.band_file = args.band_file
	cfg.validate()
	assert_paths_off_root([cfg.run_dir])
	ceiling = cfg.validate_ceiling_receipt()
	cfg.validate_world(world)
	with open(cfg.data.banding_manifest) as f:
	banding_manifest = json.load(f)

	register_answer_suffixes("</s>") # BEFORE any extraction (golden.py's incident)

	# -- data: the band's prompt rows ---------------------------------------
	prompt_rows = []
	with open(cfg.data.band_file) as f:
	for line in f:
	if line.strip():
	prompt_rows.append(json.loads(line))
	logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows))

	# -- policy + FSDP -------------------------------------------------------
	sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None)
	logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source)
	from align.model.hf_config import read_config
	from align.tokenizer import load_tokenizer

	tokenizer = load_tokenizer(sft_checkpoint)
	pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

	model_cfg = read_config(sft_checkpoint)
	model = NmoMistralForCausalLM.from_state_dict(
	model_cfg,
	ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)),
	attn_implementation=cfg.train.attn_implementation,
	device="cuda",
	).to("cuda", dtype=torch.bfloat16)

	from align.train.fsdp import apply_activation_checkpointing, apply_fsdp

	dist_model = apply_fsdp(apply_activation_checkpointing(model))
	logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30)

	# -- rollout backend (rank 0 owns it) ------------------------------------
	# 'spool': external vLLM worker process on NMO_ROLLOUT_GPU (own venv).
	# 'native': in-repo torch sampler on a STANDALONE engine copy living on that
	# same dedicated GPU, in rank 0's process. It must NOT be the FSDP2-sharded
	# model: decode forwards are collectives the other ranks would never join
	# (deadlock). Either way the GPU is outside the training ranks' devices.
	rollout_gpu = int(os.environ.get("NMO_ROLLOUT_GPU", "3"))
	spool = os.path.join(cfg.run_dir, "spool")
	worker_proc = None
	client = None
	if cfg.rollout.backend == "spool":
	if rank == 0:
	worker_proc = _spawn_rollout_worker(
	spool, rollout_gpu, os.path.join(cfg.run_dir, "worker.log")
	)
	logger.info("rollout worker spawned (pid %d) on GPU %s", worker_proc.pid, rollout_gpu)
	client = make_rollout_client("spool", spool=spool, timeout_s=cfg.rollout_timeout_s)
	elif rank == 0:
	client = _build_native_client(
	model_cfg, sft_checkpoint, tokenizer, rollout_gpu, world, cfg.rollout.max_response_len
	)
	logger.info("native rollout engine on cuda:%d (standalone copy of %s)", rollout_gpu, sft_checkpoint)

	# -- module + trainer -----------------------------------------------------
	from align.train.offload_adamw import make_optimizer

	run_cfg = _run_config(cfg)
	optimizer = make_optimizer(run_cfg, [p for p in dist_model.parameters() if p.requires_grad])
	collapse = CollapseTripwire()
	module = GRPOTrainModule(
	dist_model,
	optimizer,
	grad_accum=cfg.train.grad_accum,
	grpo_cfg=cfg.grpo,
	temperature=cfg.rollout.temperature,
	collapse_tripwire=collapse,
	max_grad_norm=cfg.train.max_grad_norm,
	)

	source = GrpoBatchSource(
	prompt_rows,
	prompts_per_step=cfg.rollout.prompts_per_step,
	n_samples_per_prompt=cfg.grpo.num_samples_per_prompt_rollout,
	max_response_len=cfg.rollout.max_response_len,
	temperature=cfg.rollout.temperature,
	optimizer_steps=cfg.train.optimizer_steps,
	grad_accum=cfg.train.grad_accum,
	seed=cfg.train.seed,
	prompt_builder=build_prompt,
	dead_group_retries=cfg.dead_group_retries,
	)
	source.bind_pad(pad_id)

	trainer = Trainer(
	run_cfg,
	dist_model,
	source,
	module._grpo_loss_adapter,
	run_dir=cfg.run_dir,
	stage="grpo",
	train_module=module,
	)
	# HORIZON CONSISTENCY (the Arm A zombie-LR class): the scheduler's total
	# must be exactly the configured optimizer_steps -- this loop generates
	# its own data, so any mismatch is arithmetic, not sampling.
	if trainer.scheduler.total_steps != cfg.train.optimizer_steps:
	raise RuntimeError(
	f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs "
	f"configured optimizer_steps={cfg.train.optimizer_steps}"
	)

	# -- canary (R4, rank 0) --------------------------------------------------
	canary = None
	canary_generate = None
	if rank == 0:
	sources = DataSources.from_local_hub()
	gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems]
	canary_problems = [
	{"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows
	]

	def canary_generate(prompts: list[str]) -> list[str]:
	req = RolloutRequest(
	request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}",
	weights_dir=_latest_weights(trainer, sft_checkpoint),
	prompts=prompts,
	n_samples_per_prompt=1,
	max_tokens=cfg.rollout.max_response_len,
	temperature=0.0,
	)
	client.submit(req)
	result = client.collect(req.request_id)
	return [texts[0] for texts in result.texts]

	canary = TaskAccuracyCanary(
	problems=canary_problems,
	generate_fn=canary_generate,
	verifier=OutcomeVerifier(VerifierConfig(workers=4)),
	extract_fn=extract_answer,
	every=cfg.canary.every,
	)

	source.bind(
	client=client,
	reward=CompositeReward(
	format_weight=cfg.reward.format_weight,
	overlong_max_len=cfg.rollout.max_response_len,
	overlong_cache_len=cfg.rollout.overlong_cache_len,
	),
	verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None,
	old_logprob_fn=module.compute_logprobs_nograd,
	weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint),
	rank=rank,
	world_size=world,
	device="cuda",
	canary=canary,
	collapse_tripwire=collapse,
	trainer=trainer,
	canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"),
	)
	trainer.loss_adapter = source.lazy_adapter(module)

	# -- resume ---------------------------------------------------------------
	# A relaunch continues the newest checkpoint in run_dir, if any. Before
	# this was wired, a relaunch restarted from the SFT substrate. What a
	# resume restores, and where each piece lives:
	# weights + optimizer + schedule + stream offset Trainer (resume.pt)
	# CollapseTripwire windows and latch module extra_state
	# prompt order, cursor, dead-group stats, alarms source state_dict
	# rollout engine weights _latest_weights: the
	# Trainer registers the resumed checkpoint, so the first post-resume
	# rollout reloads from it instead of the substrate.
	resumed_from = trainer.resume_from_latest()
	if resumed_from is not None:
	logger.info(
	"RESUMED from %s at optimizer step %d / micro step %d",
	resumed_from,
	trainer.optimizer_step,
	trainer.micro_step,
	)

	banner = {
	"stage": "grpo_stageb" + ("_smoke" if args.smoke else ""),
	"substrate": sft_checkpoint,
	"resumed_from": resumed_from,
	"band_file": cfg.data.band_file,
	"band_prompts": len(prompt_rows),
	"shape": {
	"prompts_per_step": cfg.rollout.prompts_per_step,
	"G": cfg.grpo.num_samples_per_prompt_rollout,
	"max_response_len": cfg.rollout.max_response_len,
	"loss_denominator": cfg.grpo.loss_denominator,
	"optimizer_steps": cfg.train.optimizer_steps,
	},
	"ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()},
	"banding_manifest_rows": banding_manifest.get("rows"),
	"topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")},
	"rectifications": (
	f"R3 clip={cfg.train.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, "
	"R5 pin-on-fire, reduction-time denominator assert armed"
	),
	}
	logger.info("BANNER %s", json.dumps(banner, indent=1))

	# -- train ----------------------------------------------------------------
	try:
	summary = trainer.train()
	checkpoint = trainer.save_checkpoint()
	finally:
	if worker_proc is not None:
	from align.rollout.client import terminate_worker

	terminate_worker(worker_proc)
	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)

	result = {
	"stage": banner["stage"],
	"summary": summary,
	"checkpoint": checkpoint,
	"banner": banner,
	"canary_history": canary.history if canary else [],
	"collapse_alarms": source.alarms,
	"collapse_tripwire": collapse.as_log_dict(),
	"dead_groups": source.filter_stats.as_log_dict(),
	"dropped_remainder_rows": source.dropped_remainder_rows,
	}
	if rank == 0:
	path = os.path.join(cfg.run_dir, "stage_result.json")
	with open(path + ".partial", "w") as f:
	json.dump(result, f, indent=1, default=str)
	os.replace(path + ".partial", path)
	logger.info("stage result: %s", path)
	return 0


	def _build_native_client(
	model_cfg, sft_checkpoint: str, tokenizer, rollout_gpu: int, world: int, max_response_len: int
	):
	"""Rank 0's native rollout backend on a dedicated, standalone engine copy.

	The copy is built from the SFT substrate (the policy's step-0 weights) and
	refreshed from the latest saved checkpoint whenever `weights_dir` changes
	(`checkpoint_weight_loader`). Refuses a device a training rank already
	owns: training ranks are bound to local_rank 0..world-1, so the engine GPU
	index must be >= world, and visible.
	"""
	import torch

	from align.model import checkpoint as ckpt
	from align.model.mistral import NmoMistralForCausalLM
	from align.rollout import checkpoint_weight_loader, make_rollout_client
	from align.rollout.native import NativeRolloutConfig

	if rollout_gpu < world or rollout_gpu >= torch.cuda.device_count():
	raise RuntimeError(
	f"NMO_ROLLOUT_GPU={rollout_gpu} is not a free visible device: training ranks own "
	f"cuda:0..{world - 1} and {torch.cuda.device_count()} device(s) are visible. The "
	"engine copy must live on its own GPU (launch with all four GPUs visible)."
	)
	device = torch.device("cuda", rollout_gpu)
	engine_model = (
	NmoMistralForCausalLM.from_state_dict(
	model_cfg,
	ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint, dtype=torch.bfloat16)),
	device=str(device),
	)
	.to(device, dtype=torch.bfloat16)
	.eval()
	)
	return make_rollout_client(
	"native",
	model=engine_model,
	tokenizer=tokenizer,
	native_config=NativeRolloutConfig(
	max_seq_len_cap=2048,
	pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0,
	),
	load_weights_fn=checkpoint_weight_loader(engine_model),
	device=device,
	)


	def _latest_weights(trainer, sft_checkpoint: str) -> str:
	"""The engine's weights source: latest saved checkpoint, else the substrate."""
	if trainer.checkpoints_written:
	return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}")
	return sft_checkpoint


	def _run_config(cfg):
	from align.config import RunConfig

	run = RunConfig()
	run.optimizer.name = cfg.train.optimizer_name
	run.optimizer.lr = cfg.train.lr
	run.optimizer.grad_accum = cfg.train.grad_accum
	run.train.epochs = 1 # the horizon is optimizer_steps; the source cycles prompts itself
	run.train.micro_batch_size = cfg.train.micro_batch_size
	run.train.save_every_micro_steps = cfg.train.save_every_micro_steps
	run.train.keep_checkpoints = cfg.train.keep_checkpoints
	run.train.per_rank_optimizer_state = cfg.train.per_rank_optimizer_state
	run.train.max_wall_seconds = cfg.train.max_wall_seconds
	run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps
	run.train.max_skipped_fraction = cfg.train.max_skipped_fraction
	run.train.warmup_frac = cfg.train.warmup_frac
	run.train.full_ft = True
	run.train.seed = cfg.train.seed
	run.tripwires.armed = tuple(cfg.armed_tripwires)
	run.run_name = "grpo_stageb"
	return run


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

	raise SystemExit(run_with_exit_codes(main))