Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / OURS /src /rl /trainer.py

EmergentTwilight

4 days ago

download

raw

25.7 kB

"""Config-driven fresh-trajectory multimodal GRPO v0 trainer."""

from __future__ import annotations

import faulthandler
import hashlib
import json
import math
import os
import signal
import time
from pathlib import Path

import torch
from torch.utils.tensorboard import SummaryWriter

from shared.artifacts import (
    acquire_run_writer,
    cache_fingerprint,
    claim_run_directory,
    environment_record,
    model_fingerprint,
    resolve_run_directory,
    save_adapter_atomic,
    seed_everything,
    source_fingerprint,
    write_json,
    rng_state,
)
from shared.data import PackedTrajectoryCache, query_indices_for_update
from shared.checkpoint_retention import prune_checkpoints
from shared.modeling import build_optimizer, load_model, load_processor, unwrap_ar_model
from shared.reference import FrozenReference

from .artifacts import save_update_artifacts
from .collection import collect_logical_batch
from .run_records import (
    commit_metric_artifacts,
    logged_metric_ids,
    publish_run_metadata,
    recover_committed_metrics,
    write_attempt_metrics,
)
from .config import RLConfig
from .distributed import (
    DistributedContext,
    launch_configured_workers,
)
from .engine import execute_actor_update
from .evaluation import (
    checkpoint_actor_identity,
    ensure_validation,
    group_metrics,
    make_backend,
    make_self_scorer,
    validation_identity,
)
from .fingerprints import sampling_descriptor, self_scoring_descriptor
from .pipeline import protocol_limits
from .policy import PolicyScorer
from .rollout import GRPORolloutController
from .runtime import (
    latest_complete_checkpoint,
    load_rl_checkpoint,
    save_reference_state,
    save_rl_checkpoint,
    verify_rl_checkpoint,
)
from .verifier import VisualJigsawABCDVerifier


def _scheduler(optimizer, config: RLConfig):
    warmup = config.scheduler.warmup_updates
    if warmup is None:
        warmup = max(1, round(config.trainer.max_updates * config.scheduler.warmup_ratio))

    def scale(step: int) -> float:
        if step < warmup:
            return max(step, 1) / max(warmup, 1)
        progress = (step - warmup) / max(config.trainer.max_updates - warmup, 1)
        return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, scale), warmup


def _evaluation_due(config: RLConfig, step: int) -> bool:
    return step % config.evaluation.every_updates == 0 or (
        config.evaluation.at_end and step == config.trainer.max_updates
    )


def _train_locked(
    config: RLConfig,
    device: torch.device,
    run_dir: Path,
    parallel: DistributedContext,
    stop_after_updates: int | None = None,
) -> Path:
    tensorboard = None
    log_handle = None
    try:
        cache = PackedTrajectoryCache(config.data.cache_root)
        train_indices = cache.indices(config.data.train_split)
        if config.data.fixed_train_indices is not None:
            fixed = tuple(config.data.fixed_train_indices)
            if not set(fixed).issubset(set(train_indices)):
                raise ValueError("fixed RL train index is outside the train split")
            train_indices = fixed
        validation_indices = tuple(
            index
            for split in config.data.validation_splits
            for index in cache.indices(split)
        )
        if config.checkpoint.resume_from:
            verify_rl_checkpoint(config.checkpoint.resume_from)
        processor = load_processor(config.model.path)
        actor = load_model(config, device, for_training=True)
        optimizer, backbone, visual_head = build_optimizer(actor, config)
        parameters = [*backbone, *visual_head]
        scheduler, warmup = _scheduler(optimizer, config)
        reference = FrozenReference(actor)
        verifier = VisualJigsawABCDVerifier()
        sampling = sampling_descriptor(config)
        self_scoring = self_scoring_descriptor(config)
        provenance = {
            "source": source_fingerprint("shared", "rl", "uniar"),
            "cache": cache_fingerprint(config.data.cache_root, cache.rows),
            "model": model_fingerprint(config.model.path),
            "initial_adapter": model_fingerprint(config.model.initial_adapter),
            "verifier": verifier.descriptor()["fingerprint"],
            "sampling": sampling["fingerprint"],
            "self_scoring": self_scoring["fingerprint"],
            "query_schedule": hashlib.sha256(
                (
                    f"seed={config.run.seed}\n"
                    f"shuffle={config.data.shuffle}\n"
                    f"prompts={config.trainer.prompts_per_update}\n"
                    f"group={config.rollout.group_size}\n"
                    + "\n".join(cache.sample_id(index) for index in train_indices)
                ).encode()
            ).hexdigest(),
        }
        state = {
            "attempted_updates": 0,
            "successful_updates": 0,
            "nonfinite_updates": 0,
            "query_update": 0,
        }
        resume_parent = None
        if config.checkpoint.resume_from:
            state = load_rl_checkpoint(
                config.checkpoint.resume_from,
                optimizer=optimizer,
                scheduler=scheduler,
                config=config,
                provenance_fingerprints=provenance,
                reference=reference,
                world_size=parallel.world_size,
                commit_artifacts=parallel.is_main,
                rank=parallel.rank,
            )
            if Path(config.checkpoint.resume_from).resolve().parent.parent != run_dir.resolve():
                resume_parent = Path(config.checkpoint.resume_from).resolve().parent.parent
                if parallel.is_main:
                    save_reference_state(run_dir, reference)
        elif parallel.is_main:
            save_reference_state(run_dir, reference)
        parallel.barrier()
        if state["nonfinite_updates"] > config.optimizer.max_nonfinite_updates:
            raise RuntimeError("restored GRPO state exceeded its non-finite budget")

        recorded_config = config.as_dict()
        recorded_config["checkpoint"]["resume_from"] = None
        if parallel.is_main:
            publish_run_metadata(
                run_dir,
                {
                    "config": recorded_config,
                    "config_fingerprint": config.fingerprint(),
                    "provenance_fingerprints": provenance,
                    "reference": reference.descriptor(),
                    "sampling": sampling,
                    "self_scoring": self_scoring,
                    "warmup_updates": warmup,
                    "train_samples": len(train_indices),
                    "environment": environment_record(device, parallel.world_size),
                    "parameters": {
                        "trainable": sum(parameter.numel() for parameter in parameters),
                        "backbone": sum(parameter.numel() for parameter in backbone),
                        "visual_head": sum(parameter.numel() for parameter in visual_head),
                    },
                    "image_reward": {"backend": "disabled", "loaded": False},
                },
                recorded_config,
            )
        tensorboard = (
            SummaryWriter(run_dir / "tensorboard")
            if parallel.is_main and config.logging.tensorboard
            else None
        )
        log_path = run_dir / "train_log.jsonl"
        logged_ids = logged_metric_ids(log_path) if parallel.is_main else set()
        log_handle = (
            log_path.open("a", encoding="utf-8", buffering=1)
            if parallel.is_main and config.logging.jsonl
            else None
        )
        if parallel.is_main:
            recover_committed_metrics(
                run_dir,
                log_handle,
                tensorboard,
                logged_ids,
                (
                    Path(config.checkpoint.resume_from)
                    if config.checkpoint.resume_from and resume_parent is None
                    else None
                ),
            )
        parallel.barrier()
        def checkpoint(artifact_roots=(), *, permanent=False, reason="rolling"):
            rank_rng = parallel.all_gather_object(rng_state())
            path = None
            error = None
            if parallel.is_main:
                try:
                    path = save_rl_checkpoint(
                        checkpoints_dir=run_dir / "checkpoints", actor=actor,
                        processor=processor, optimizer=optimizer, scheduler=scheduler,
                        trainer_state=state, config=config, provenance_fingerprints=provenance,
                        reference=reference, artifact_roots=tuple(artifact_roots),
                        world_size=parallel.world_size, rank_rng_states=rank_rng,
                        permanent=permanent, reason=reason,
                    )
                    commit_metric_artifacts(run_dir, tuple(artifact_roots), log_handle,
                                            tensorboard, logged_ids)
                    prune_checkpoints(run_dir / "checkpoints", current=path,
                                      keep_latest=config.checkpoint.keep_latest)
                except Exception as exception:
                    error = f"{type(exception).__name__}: {exception}"
            parallel.gather_stage(None, error, stage="checkpoint commit")
            name = parallel.broadcast_object(path.name if path else None)
            print(f"RL_CHECKPOINT_COMPLETE step={state['successful_updates']} rank={parallel.rank}", flush=True)
            return run_dir / "checkpoints" / name

        if config.checkpoint.resume_from is None or resume_parent is not None:
            step = state["successful_updates"]
            validation_due = config.evaluation.at_start or (step > 0 and _evaluation_due(config, step))
            checkpoint(permanent=validation_due or step == config.trainer.max_updates,
                       reason="validation" if validation_due else "initial")
        validation_parent = resume_parent
        if validation_parent is None and config.checkpoint.resume_mode == "execution_change":
            submission = json.loads((run_dir / "_submission/submission.json").read_text())
            declared = (submission.get("parent_run") or {}).get("run_dir")
            validation_parent = Path(declared) if declared else None
        def log_validation(path):
            if tensorboard is not None:
                payload = json.loads(path.read_text())
                step = payload["identity"]["successful_update"]
                for key, value in payload["metrics"].items():
                    if isinstance(value, (int, float)):
                        tensorboard.add_scalar(f"validation/{key}", value, step)
                tensorboard.flush()
        print(f"RL_RESUME successful_updates={state['successful_updates']} "
              f"query_update={state['query_update']} checkpoint={config.checkpoint.resume_from}", flush=True)
        if config.evaluation.at_start:
            identity = validation_identity(
                cache=cache,
                indices=validation_indices,
                config=config,
                provenance=provenance,
                sampling=sampling,
                self_scoring=self_scoring,
                phase="start",
                successful_update=0,
                actor_identity={
                    "kind": "run_start_reference",
                    "snapshot_sha256": reference.descriptor()["snapshot_sha256"],
                },
            )
            validation_kwargs = {
                "run_dir": run_dir,
                "stem": "validation_start",
                "identity": identity,
                "processor": processor,
                "cache": cache,
                "indices": validation_indices,
                "config": config,
                "parallel": parallel,
                "reuse_from": validation_parent,
            }
            if state["successful_updates"]:
                with reference.apply(actor) as run_start_actor:
                    path = ensure_validation(
                        actor=run_start_actor, **validation_kwargs
                    )
            else:
                path = ensure_validation(actor=actor, **validation_kwargs)
            state["validation_start"] = path.name
            log_validation(path)

        restored_step = state["successful_updates"]
        if restored_step and _evaluation_due(config, restored_step):
            if config.checkpoint.resume_from is None:
                raise RuntimeError("restored GRPO update has no actor checkpoint")
            restored_checkpoint = Path(config.checkpoint.resume_from)
            identity = validation_identity(
                cache=cache,
                indices=validation_indices,
                config=config,
                provenance=provenance,
                sampling=sampling,
                self_scoring=self_scoring,
                phase="step",
                successful_update=restored_step,
                actor_identity=checkpoint_actor_identity(
                    restored_checkpoint.parent.parent, restored_checkpoint
                ),
            )
            restored_validation = ensure_validation(
                run_dir=run_dir,
                stem=f"validation_step_{restored_step:06d}",
                identity=identity,
                actor=actor,
                processor=processor,
                cache=cache,
                indices=validation_indices,
                config=config,
                parallel=parallel,
                reuse_from=validation_parent,
            )
            state["last_validation_artifact"] = restored_validation.name
            log_validation(restored_validation)

        controller = GRPORolloutController(
            make_backend(actor, processor, config),
            protocol_limits(config),
        )
        behavior_scorer = make_self_scorer(actor, config)
        policy_scorer = PolicyScorer(
            actor,
            reference,
            text_temperature=config.rollout.text_temperature,
            visual_temperature=config.rollout.visual_temperature,
            text_chunk_size=config.trainer.text_projection_chunk_size,
        )
        last_checkpoint_wall = time.monotonic()
        pending_artifacts: list[Path] = []
        started_step = state["successful_updates"]
        while state["successful_updates"] < config.trainer.max_updates:
            started = time.monotonic()
            attempted_update = state["attempted_updates"] + 1
            prompt_indices = query_indices_for_update(
                train_indices,
                batch_size=config.trainer.prompts_per_update,
                seed=config.run.seed,
                update=state["query_update"],
                shuffle=config.data.shuffle,
            )
            actor.eval()
            groups, collection_metrics = collect_logical_batch(
                controller=controller, cache=cache, prompt_indices=prompt_indices,
                config=config, scorer=behavior_scorer, tokenizer=processor.tokenizer,
                update=state["query_update"], parallel=parallel, run_dir=run_dir,
            )
            policy_started = time.monotonic()
            actor.train()
            unwrap_ar_model(actor).model.visual.eval()
            update = execute_actor_update(
                groups=groups,
                policy_scorer=policy_scorer,
                config=config,
                parameters=parameters,
                actor=actor,
                optimizer=optimizer,
                scheduler=scheduler,
                parallel=parallel,
            )
            artifact_root = None
            if parallel.is_main:
                artifact_root = save_update_artifacts(
                    run_dir=run_dir,
                    attempted_update=attempted_update,
                    groups=groups,
                    policy_records=update.policy_records,
                    config=config,
                    verifier_descriptor=verifier.descriptor(),
                    sampling_descriptor=sampling,
                    self_scoring_descriptor=self_scoring,
                )
            state["attempted_updates"] += 1
            state["query_update"] += 1
            if update.successful:
                state["successful_updates"] += 1
            else:
                state["nonfinite_updates"] += 1
            peak_memory = max(
                parallel.all_gather_object(
                    torch.cuda.max_memory_allocated(device) / 1024**3
                )
            )
            elapsed = max(
                parallel.all_gather_object(time.monotonic() - started)
            )
            metrics = {
                **{
                    key: value
                    for key, value in state.items()
                    if isinstance(value, (int, float, bool))
                },
                **{f"group/{key}": value for key, value in group_metrics(groups).items()},
                **{f"timing/{key}": value for key, value in collection_metrics.items()},
                "timing/policy_seconds": time.monotonic() - policy_started,
                **{f"policy/{key}": value for key, value in update.statistics.items()},
                "grad_norm": update.grad_norm,
                "grad_clipped": update.grad_norm > config.optimizer.max_grad_norm,
                "zero_gradient": update.zero_gradient,
                "skipped_nonfinite": not update.successful,
                "seconds": elapsed,
                "peak_gpu_memory_gib": peak_memory,
                "image_reward": 0.0,
            }
            if parallel.is_main:
                write_attempt_metrics(run_dir, artifact_root, metrics)
                pending_artifacts.append(artifact_root)
                print(
                    "RL_PROGRESS "
                    f"{state['successful_updates']}/{config.trainer.max_updates} "
                    f"reward={metrics['group/reward_mean']:.4f} "
                    f"valid={metrics['group/protocol_valid']}/"
                    f"{metrics['group/rollouts']} "
                    f"grad={update.grad_norm:.4f}",
                    flush=True,
                )
            step = state["successful_updates"]
            validation_due = update.successful and _evaluation_due(config, step)
            stopping = stop_after_updates is not None and step - started_step >= stop_after_updates
            checkpoint_due = parallel.broadcast_object(
                (
                    time.monotonic() - last_checkpoint_wall
                    >= config.checkpoint.every_minutes * 60
                )
                if parallel.is_main
                else None
            )
            checkpoint_path = None
            if not update.successful or (
                step % config.checkpoint.every_updates == 0
                or step == config.trainer.max_updates
                or checkpoint_due
                or validation_due or stopping
            ):
                checkpoint_path = checkpoint(
                    tuple(pending_artifacts),
                    permanent=validation_due or step == config.trainer.max_updates,
                    reason="validation" if validation_due else (
                        "final" if step == config.trainer.max_updates else "rolling"
                    ),
                )
                pending_artifacts.clear()
                last_checkpoint_wall = time.monotonic()
            if validation_due:
                if checkpoint_path is None:
                    raise RuntimeError("validation requires a committed actor checkpoint")
                identity = validation_identity(
                    cache=cache,
                    indices=validation_indices,
                    config=config,
                    provenance=provenance,
                    sampling=sampling,
                    self_scoring=self_scoring,
                    phase="step",
                    successful_update=step,
                    actor_identity=checkpoint_actor_identity(
                        run_dir, checkpoint_path
                    ),
                )
                validation_path = ensure_validation(
                    run_dir=run_dir,
                    stem=f"validation_step_{step:06d}",
                    identity=identity,
                    actor=actor,
                    processor=processor,
                    cache=cache,
                    indices=validation_indices,
                    config=config,
                    parallel=parallel,
                )
                state["last_validation_artifact"] = validation_path.name
                log_validation(validation_path)
            if state["nonfinite_updates"] > config.optimizer.max_nonfinite_updates:
                raise RuntimeError("GRPO exceeded its non-finite update budget")
            parallel.barrier()
            del groups
            if stopping and step < config.trainer.max_updates:
                print(f"RL_PAUSED step={step}", flush=True)
                return run_dir
        peak_memory = max(
            parallel.all_gather_object(
                torch.cuda.max_memory_allocated(device) / 1024**3
            )
        )
        if parallel.is_main:
            save_adapter_atomic(actor, processor, run_dir / "adapter")
        validation = {}
        if parallel.is_main and state.get("last_validation_artifact"):
            validation = json.loads(
                (run_dir / state["last_validation_artifact"]).read_text()
            )
        public_state = {
            key: value for key, value in state.items() if key != "validation_start"
        }
        if parallel.is_main:
            summary = public_state | {
                "adapter": "adapter",
                "reference": "reference",
                "validation_final": validation,
                "peak_gpu_memory_gib": peak_memory,
                "image_reward": {"backend": "disabled", "loaded": False},
            }
            write_json(run_dir / "summary.json", summary)
            print("RL_TRAINING_COMPLETE " + json.dumps(public_state), flush=True)
        parallel.barrier()
        return run_dir
    finally:
        if tensorboard:
            tensorboard.flush()
            tensorboard.close()
        if log_handle:
            log_handle.close()


def train(config: RLConfig, *, stop_after_updates: int | None = None) -> Path:
    if hasattr(signal, "SIGUSR1"):
        faulthandler.register(signal.SIGUSR1, all_threads=True)
    config.validate()
    if not torch.cuda.is_available():
        raise RuntimeError("GRPO training requires scheduler-assigned CUDA")
    parallel = DistributedContext.initialize(config.trainer.world_size)
    device = torch.device("cuda", parallel.local_rank)
    torch.cuda.set_device(device)
    torch.backends.cuda.matmul.allow_tf32 = config.trainer.tf32
    torch.backends.cudnn.allow_tf32 = config.trainer.tf32
    seed_everything(config.run.seed)

    writer_lock = None
    try:
        run_state = None
        def prepare():
            nonlocal writer_lock
            run_dir = resolve_run_directory(config, "rl")
            in_place_resume = (run_dir / ".run-claim.json").is_file()
            claim_run_directory(run_dir, resume=in_place_resume)
            writer_lock = acquire_run_writer(run_dir)
            if in_place_resume:
                if (run_dir / "summary.json").is_file():
                    if not (run_dir / "adapter" / "COMPLETE").is_file():
                        raise RuntimeError("completed GRPO run has no committed adapter")
                    recorded = json.loads((run_dir / "config.json").read_text())
                    if recorded["config_fingerprint"] != config.fingerprint():
                        raise ValueError("completed GRPO run config differs")
                    return str(run_dir), None, True
                try:
                    latest = latest_complete_checkpoint(run_dir / "checkpoints")
                except FileNotFoundError:
                    latest = None
                if latest is not None:
                    # The same frozen argv must prefer this run's newest commit.
                    explicit = Path(config.checkpoint.resume_from) if config.checkpoint.resume_from else None
                    if (explicit is not None and explicit.parent.parent.resolve() == run_dir.resolve()
                            and explicit.resolve() != latest.resolve()):
                        raise ValueError("manual same-run GRPO checkpoint rollback is unsupported")
                    config.checkpoint.resume_from = str(latest)
            return str(run_dir), config.checkpoint.resume_from, False
        run_state = parallel.main_call(prepare)
        if run_state is None:
            raise RuntimeError("main worker did not claim a GRPO run directory")
        run_dir = Path(run_state[0])
        if run_state[2]:
            print(f"RL_ALREADY_COMPLETE {run_dir}", flush=True)
            return run_dir
        config.checkpoint.resume_from = run_state[1]
        return _train_locked(config, device, run_dir, parallel, stop_after_updates)
    finally:
        if writer_lock is not None:
            writer_lock.close()
        parallel.close()


def launch(config: RLConfig, *, stop_after_updates: int | None = None):
    from functools import partial
    return launch_configured_workers(config, partial(train, stop_after_updates=stop_after_updates))


__all__ = ["launch", "train"]

Xet Storage Details

Size:: 25.7 kB
Xet hash:: bf8487a4f7d613929b69b1b8a6d6b4dc977d63cb2df990a2d09e3c5f8723d000

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Config-driven fresh-trajectory multimodal GRPO v0 trainer."""

	from __future__ import annotations

	import faulthandler
	import hashlib
	import json
	import math
	import os
	import signal
	import time
	from pathlib import Path

	import torch
	from torch.utils.tensorboard import SummaryWriter

	from shared.artifacts import (
	acquire_run_writer,
	cache_fingerprint,
	claim_run_directory,
	environment_record,
	model_fingerprint,
	resolve_run_directory,
	save_adapter_atomic,
	seed_everything,
	source_fingerprint,
	write_json,
	rng_state,
	)
	from shared.data import PackedTrajectoryCache, query_indices_for_update
	from shared.checkpoint_retention import prune_checkpoints
	from shared.modeling import build_optimizer, load_model, load_processor, unwrap_ar_model
	from shared.reference import FrozenReference

	from .artifacts import save_update_artifacts
	from .collection import collect_logical_batch
	from .run_records import (
	commit_metric_artifacts,
	logged_metric_ids,
	publish_run_metadata,
	recover_committed_metrics,
	write_attempt_metrics,
	)
	from .config import RLConfig
	from .distributed import (
	DistributedContext,
	launch_configured_workers,
	)
	from .engine import execute_actor_update
	from .evaluation import (
	checkpoint_actor_identity,
	ensure_validation,
	group_metrics,
	make_backend,
	make_self_scorer,
	validation_identity,
	)
	from .fingerprints import sampling_descriptor, self_scoring_descriptor
	from .pipeline import protocol_limits
	from .policy import PolicyScorer
	from .rollout import GRPORolloutController
	from .runtime import (
	latest_complete_checkpoint,
	load_rl_checkpoint,
	save_reference_state,
	save_rl_checkpoint,
	verify_rl_checkpoint,
	)
	from .verifier import VisualJigsawABCDVerifier


	def _scheduler(optimizer, config: RLConfig):
	warmup = config.scheduler.warmup_updates
	if warmup is None:
	warmup = max(1, round(config.trainer.max_updates * config.scheduler.warmup_ratio))

	def scale(step: int) -> float:
	if step < warmup:
	return max(step, 1) / max(warmup, 1)
	progress = (step - warmup) / max(config.trainer.max_updates - warmup, 1)
	return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))

	return torch.optim.lr_scheduler.LambdaLR(optimizer, scale), warmup


	def _evaluation_due(config: RLConfig, step: int) -> bool:
	return step % config.evaluation.every_updates == 0 or (
	config.evaluation.at_end and step == config.trainer.max_updates
	)


	def _train_locked(
	config: RLConfig,
	device: torch.device,
	run_dir: Path,
	parallel: DistributedContext,
	stop_after_updates: int \| None = None,
	) -> Path:
	tensorboard = None
	log_handle = None
	try:
	cache = PackedTrajectoryCache(config.data.cache_root)
	train_indices = cache.indices(config.data.train_split)
	if config.data.fixed_train_indices is not None:
	fixed = tuple(config.data.fixed_train_indices)
	if not set(fixed).issubset(set(train_indices)):
	raise ValueError("fixed RL train index is outside the train split")
	train_indices = fixed
	validation_indices = tuple(
	index
	for split in config.data.validation_splits
	for index in cache.indices(split)
	)
	if config.checkpoint.resume_from:
	verify_rl_checkpoint(config.checkpoint.resume_from)
	processor = load_processor(config.model.path)
	actor = load_model(config, device, for_training=True)
	optimizer, backbone, visual_head = build_optimizer(actor, config)
	parameters = [backbone, visual_head]
	scheduler, warmup = _scheduler(optimizer, config)
	reference = FrozenReference(actor)
	verifier = VisualJigsawABCDVerifier()
	sampling = sampling_descriptor(config)
	self_scoring = self_scoring_descriptor(config)
	provenance = {
	"source": source_fingerprint("shared", "rl", "uniar"),
	"cache": cache_fingerprint(config.data.cache_root, cache.rows),
	"model": model_fingerprint(config.model.path),
	"initial_adapter": model_fingerprint(config.model.initial_adapter),
	"verifier": verifier.descriptor()["fingerprint"],
	"sampling": sampling["fingerprint"],
	"self_scoring": self_scoring["fingerprint"],
	"query_schedule": hashlib.sha256(
	(
	f"seed={config.run.seed}\n"
	f"shuffle={config.data.shuffle}\n"
	f"prompts={config.trainer.prompts_per_update}\n"
	f"group={config.rollout.group_size}\n"
	+ "\n".join(cache.sample_id(index) for index in train_indices)
	).encode()
	).hexdigest(),
	}
	state = {
	"attempted_updates": 0,
	"successful_updates": 0,
	"nonfinite_updates": 0,
	"query_update": 0,
	}
	resume_parent = None
	if config.checkpoint.resume_from:
	state = load_rl_checkpoint(
	config.checkpoint.resume_from,
	optimizer=optimizer,
	scheduler=scheduler,
	config=config,
	provenance_fingerprints=provenance,
	reference=reference,
	world_size=parallel.world_size,
	commit_artifacts=parallel.is_main,
	rank=parallel.rank,
	)
	if Path(config.checkpoint.resume_from).resolve().parent.parent != run_dir.resolve():
	resume_parent = Path(config.checkpoint.resume_from).resolve().parent.parent
	if parallel.is_main:
	save_reference_state(run_dir, reference)
	elif parallel.is_main:
	save_reference_state(run_dir, reference)
	parallel.barrier()
	if state["nonfinite_updates"] > config.optimizer.max_nonfinite_updates:
	raise RuntimeError("restored GRPO state exceeded its non-finite budget")

	recorded_config = config.as_dict()
	recorded_config["checkpoint"]["resume_from"] = None
	if parallel.is_main:
	publish_run_metadata(
	run_dir,
	{
	"config": recorded_config,
	"config_fingerprint": config.fingerprint(),
	"provenance_fingerprints": provenance,
	"reference": reference.descriptor(),
	"sampling": sampling,
	"self_scoring": self_scoring,
	"warmup_updates": warmup,
	"train_samples": len(train_indices),
	"environment": environment_record(device, parallel.world_size),
	"parameters": {
	"trainable": sum(parameter.numel() for parameter in parameters),
	"backbone": sum(parameter.numel() for parameter in backbone),
	"visual_head": sum(parameter.numel() for parameter in visual_head),
	},
	"image_reward": {"backend": "disabled", "loaded": False},
	},
	recorded_config,
	)
	tensorboard = (
	SummaryWriter(run_dir / "tensorboard")
	if parallel.is_main and config.logging.tensorboard
	else None
	)
	log_path = run_dir / "train_log.jsonl"
	logged_ids = logged_metric_ids(log_path) if parallel.is_main else set()
	log_handle = (
	log_path.open("a", encoding="utf-8", buffering=1)
	if parallel.is_main and config.logging.jsonl
	else None
	)
	if parallel.is_main:
	recover_committed_metrics(
	run_dir,
	log_handle,
	tensorboard,
	logged_ids,
	(
	Path(config.checkpoint.resume_from)
	if config.checkpoint.resume_from and resume_parent is None
	else None
	),
	)
	parallel.barrier()
	def checkpoint(artifact_roots=(), *, permanent=False, reason="rolling"):
	rank_rng = parallel.all_gather_object(rng_state())
	path = None
	error = None
	if parallel.is_main:
	try:
	path = save_rl_checkpoint(
	checkpoints_dir=run_dir / "checkpoints", actor=actor,
	processor=processor, optimizer=optimizer, scheduler=scheduler,
	trainer_state=state, config=config, provenance_fingerprints=provenance,
	reference=reference, artifact_roots=tuple(artifact_roots),
	world_size=parallel.world_size, rank_rng_states=rank_rng,
	permanent=permanent, reason=reason,
	)
	commit_metric_artifacts(run_dir, tuple(artifact_roots), log_handle,
	tensorboard, logged_ids)
	prune_checkpoints(run_dir / "checkpoints", current=path,
	keep_latest=config.checkpoint.keep_latest)
	except Exception as exception:
	error = f"{type(exception).__name__}: {exception}"
	parallel.gather_stage(None, error, stage="checkpoint commit")
	name = parallel.broadcast_object(path.name if path else None)
	print(f"RL_CHECKPOINT_COMPLETE step={state['successful_updates']} rank={parallel.rank}", flush=True)
	return run_dir / "checkpoints" / name

	if config.checkpoint.resume_from is None or resume_parent is not None:
	step = state["successful_updates"]
	validation_due = config.evaluation.at_start or (step > 0 and _evaluation_due(config, step))
	checkpoint(permanent=validation_due or step == config.trainer.max_updates,
	reason="validation" if validation_due else "initial")
	validation_parent = resume_parent
	if validation_parent is None and config.checkpoint.resume_mode == "execution_change":
	submission = json.loads((run_dir / "_submission/submission.json").read_text())
	declared = (submission.get("parent_run") or {}).get("run_dir")
	validation_parent = Path(declared) if declared else None
	def log_validation(path):
	if tensorboard is not None:
	payload = json.loads(path.read_text())
	step = payload["identity"]["successful_update"]
	for key, value in payload["metrics"].items():
	if isinstance(value, (int, float)):
	tensorboard.add_scalar(f"validation/{key}", value, step)
	tensorboard.flush()
	print(f"RL_RESUME successful_updates={state['successful_updates']} "
	f"query_update={state['query_update']} checkpoint={config.checkpoint.resume_from}", flush=True)
	if config.evaluation.at_start:
	identity = validation_identity(
	cache=cache,
	indices=validation_indices,
	config=config,
	provenance=provenance,
	sampling=sampling,
	self_scoring=self_scoring,
	phase="start",
	successful_update=0,
	actor_identity={
	"kind": "run_start_reference",
	"snapshot_sha256": reference.descriptor()["snapshot_sha256"],
	},
	)
	validation_kwargs = {
	"run_dir": run_dir,
	"stem": "validation_start",
	"identity": identity,
	"processor": processor,
	"cache": cache,
	"indices": validation_indices,
	"config": config,
	"parallel": parallel,
	"reuse_from": validation_parent,
	}
	if state["successful_updates"]:
	with reference.apply(actor) as run_start_actor:
	path = ensure_validation(
	actor=run_start_actor, **validation_kwargs
	)
	else:
	path = ensure_validation(actor=actor, **validation_kwargs)
	state["validation_start"] = path.name
	log_validation(path)

	restored_step = state["successful_updates"]
	if restored_step and _evaluation_due(config, restored_step):
	if config.checkpoint.resume_from is None:
	raise RuntimeError("restored GRPO update has no actor checkpoint")
	restored_checkpoint = Path(config.checkpoint.resume_from)
	identity = validation_identity(
	cache=cache,
	indices=validation_indices,
	config=config,
	provenance=provenance,
	sampling=sampling,
	self_scoring=self_scoring,
	phase="step",
	successful_update=restored_step,
	actor_identity=checkpoint_actor_identity(
	restored_checkpoint.parent.parent, restored_checkpoint
	),
	)
	restored_validation = ensure_validation(
	run_dir=run_dir,
	stem=f"validation_step_{restored_step:06d}",
	identity=identity,
	actor=actor,
	processor=processor,
	cache=cache,
	indices=validation_indices,
	config=config,
	parallel=parallel,
	reuse_from=validation_parent,
	)
	state["last_validation_artifact"] = restored_validation.name
	log_validation(restored_validation)

	controller = GRPORolloutController(
	make_backend(actor, processor, config),
	protocol_limits(config),
	)
	behavior_scorer = make_self_scorer(actor, config)
	policy_scorer = PolicyScorer(
	actor,
	reference,
	text_temperature=config.rollout.text_temperature,
	visual_temperature=config.rollout.visual_temperature,
	text_chunk_size=config.trainer.text_projection_chunk_size,
	)
	last_checkpoint_wall = time.monotonic()
	pending_artifacts: list[Path] = []
	started_step = state["successful_updates"]
	while state["successful_updates"] < config.trainer.max_updates:
	started = time.monotonic()
	attempted_update = state["attempted_updates"] + 1
	prompt_indices = query_indices_for_update(
	train_indices,
	batch_size=config.trainer.prompts_per_update,
	seed=config.run.seed,
	update=state["query_update"],
	shuffle=config.data.shuffle,
	)
	actor.eval()
	groups, collection_metrics = collect_logical_batch(
	controller=controller, cache=cache, prompt_indices=prompt_indices,
	config=config, scorer=behavior_scorer, tokenizer=processor.tokenizer,
	update=state["query_update"], parallel=parallel, run_dir=run_dir,
	)
	policy_started = time.monotonic()
	actor.train()
	unwrap_ar_model(actor).model.visual.eval()
	update = execute_actor_update(
	groups=groups,
	policy_scorer=policy_scorer,
	config=config,
	parameters=parameters,
	actor=actor,
	optimizer=optimizer,
	scheduler=scheduler,
	parallel=parallel,
	)
	artifact_root = None
	if parallel.is_main:
	artifact_root = save_update_artifacts(
	run_dir=run_dir,
	attempted_update=attempted_update,
	groups=groups,
	policy_records=update.policy_records,
	config=config,
	verifier_descriptor=verifier.descriptor(),
	sampling_descriptor=sampling,
	self_scoring_descriptor=self_scoring,
	)
	state["attempted_updates"] += 1
	state["query_update"] += 1
	if update.successful:
	state["successful_updates"] += 1
	else:
	state["nonfinite_updates"] += 1
	peak_memory = max(
	parallel.all_gather_object(
	torch.cuda.max_memory_allocated(device) / 1024**3
	)
	)
	elapsed = max(
	parallel.all_gather_object(time.monotonic() - started)
	)
	metrics = {
	**{
	key: value
	for key, value in state.items()
	if isinstance(value, (int, float, bool))
	},
	**{f"group/{key}": value for key, value in group_metrics(groups).items()},
	**{f"timing/{key}": value for key, value in collection_metrics.items()},
	"timing/policy_seconds": time.monotonic() - policy_started,
	**{f"policy/{key}": value for key, value in update.statistics.items()},
	"grad_norm": update.grad_norm,
	"grad_clipped": update.grad_norm > config.optimizer.max_grad_norm,
	"zero_gradient": update.zero_gradient,
	"skipped_nonfinite": not update.successful,
	"seconds": elapsed,
	"peak_gpu_memory_gib": peak_memory,
	"image_reward": 0.0,
	}
	if parallel.is_main:
	write_attempt_metrics(run_dir, artifact_root, metrics)
	pending_artifacts.append(artifact_root)
	print(
	"RL_PROGRESS "
	f"{state['successful_updates']}/{config.trainer.max_updates} "
	f"reward={metrics['group/reward_mean']:.4f} "
	f"valid={metrics['group/protocol_valid']}/"
	f"{metrics['group/rollouts']} "
	f"grad={update.grad_norm:.4f}",
	flush=True,
	)
	step = state["successful_updates"]
	validation_due = update.successful and _evaluation_due(config, step)
	stopping = stop_after_updates is not None and step - started_step >= stop_after_updates
	checkpoint_due = parallel.broadcast_object(
	(
	time.monotonic() - last_checkpoint_wall
	>= config.checkpoint.every_minutes * 60
	)
	if parallel.is_main
	else None
	)
	checkpoint_path = None
	if not update.successful or (
	step % config.checkpoint.every_updates == 0
	or step == config.trainer.max_updates
	or checkpoint_due
	or validation_due or stopping
	):
	checkpoint_path = checkpoint(
	tuple(pending_artifacts),
	permanent=validation_due or step == config.trainer.max_updates,
	reason="validation" if validation_due else (
	"final" if step == config.trainer.max_updates else "rolling"
	),
	)
	pending_artifacts.clear()
	last_checkpoint_wall = time.monotonic()
	if validation_due:
	if checkpoint_path is None:
	raise RuntimeError("validation requires a committed actor checkpoint")
	identity = validation_identity(
	cache=cache,
	indices=validation_indices,
	config=config,
	provenance=provenance,
	sampling=sampling,
	self_scoring=self_scoring,
	phase="step",
	successful_update=step,
	actor_identity=checkpoint_actor_identity(
	run_dir, checkpoint_path
	),
	)
	validation_path = ensure_validation(
	run_dir=run_dir,
	stem=f"validation_step_{step:06d}",
	identity=identity,
	actor=actor,
	processor=processor,
	cache=cache,
	indices=validation_indices,
	config=config,
	parallel=parallel,
	)
	state["last_validation_artifact"] = validation_path.name
	log_validation(validation_path)
	if state["nonfinite_updates"] > config.optimizer.max_nonfinite_updates:
	raise RuntimeError("GRPO exceeded its non-finite update budget")
	parallel.barrier()
	del groups
	if stopping and step < config.trainer.max_updates:
	print(f"RL_PAUSED step={step}", flush=True)
	return run_dir
	peak_memory = max(
	parallel.all_gather_object(
	torch.cuda.max_memory_allocated(device) / 1024**3
	)
	)
	if parallel.is_main:
	save_adapter_atomic(actor, processor, run_dir / "adapter")
	validation = {}
	if parallel.is_main and state.get("last_validation_artifact"):
	validation = json.loads(
	(run_dir / state["last_validation_artifact"]).read_text()
	)
	public_state = {
	key: value for key, value in state.items() if key != "validation_start"
	}
	if parallel.is_main:
	summary = public_state \| {
	"adapter": "adapter",
	"reference": "reference",
	"validation_final": validation,
	"peak_gpu_memory_gib": peak_memory,
	"image_reward": {"backend": "disabled", "loaded": False},
	}
	write_json(run_dir / "summary.json", summary)
	print("RL_TRAINING_COMPLETE " + json.dumps(public_state), flush=True)
	parallel.barrier()
	return run_dir
	finally:
	if tensorboard:
	tensorboard.flush()
	tensorboard.close()
	if log_handle:
	log_handle.close()


	def train(config: RLConfig, *, stop_after_updates: int \| None = None) -> Path:
	if hasattr(signal, "SIGUSR1"):
	faulthandler.register(signal.SIGUSR1, all_threads=True)
	config.validate()
	if not torch.cuda.is_available():
	raise RuntimeError("GRPO training requires scheduler-assigned CUDA")
	parallel = DistributedContext.initialize(config.trainer.world_size)
	device = torch.device("cuda", parallel.local_rank)
	torch.cuda.set_device(device)
	torch.backends.cuda.matmul.allow_tf32 = config.trainer.tf32
	torch.backends.cudnn.allow_tf32 = config.trainer.tf32
	seed_everything(config.run.seed)

	writer_lock = None
	try:
	run_state = None
	def prepare():
	nonlocal writer_lock
	run_dir = resolve_run_directory(config, "rl")
	in_place_resume = (run_dir / ".run-claim.json").is_file()
	claim_run_directory(run_dir, resume=in_place_resume)
	writer_lock = acquire_run_writer(run_dir)
	if in_place_resume:
	if (run_dir / "summary.json").is_file():
	if not (run_dir / "adapter" / "COMPLETE").is_file():
	raise RuntimeError("completed GRPO run has no committed adapter")
	recorded = json.loads((run_dir / "config.json").read_text())
	if recorded["config_fingerprint"] != config.fingerprint():
	raise ValueError("completed GRPO run config differs")
	return str(run_dir), None, True
	try:
	latest = latest_complete_checkpoint(run_dir / "checkpoints")
	except FileNotFoundError:
	latest = None
	if latest is not None:
	# The same frozen argv must prefer this run's newest commit.
	explicit = Path(config.checkpoint.resume_from) if config.checkpoint.resume_from else None
	if (explicit is not None and explicit.parent.parent.resolve() == run_dir.resolve()
	and explicit.resolve() != latest.resolve()):
	raise ValueError("manual same-run GRPO checkpoint rollback is unsupported")
	config.checkpoint.resume_from = str(latest)
	return str(run_dir), config.checkpoint.resume_from, False
	run_state = parallel.main_call(prepare)
	if run_state is None:
	raise RuntimeError("main worker did not claim a GRPO run directory")
	run_dir = Path(run_state[0])
	if run_state[2]:
	print(f"RL_ALREADY_COMPLETE {run_dir}", flush=True)
	return run_dir
	config.checkpoint.resume_from = run_state[1]
	return _train_locked(config, device, run_dir, parallel, stop_after_updates)
	finally:
	if writer_lock is not None:
	writer_lock.close()
	parallel.close()


	def launch(config: RLConfig, *, stop_after_updates: int \| None = None):
	from functools import partial
	return launch_configured_workers(config, partial(train, stop_after_updates=stop_after_updates))


	__all__ = ["launch", "train"]