#!/usr/bin/env python3 """Materialize an isolated, full-response ChartQA core-OPD recipe. This recipe deliberately removes the recently added reward-gated/ExOPD objectives and disables every auxiliary training route. The only trainable objective is exact full-vocabulary JSD between a student rollout and the privileged frozen teacher distribution over every completion token. """ from __future__ import annotations import argparse import os import sys from pathlib import Path from typing import Any import yaml PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from config.loader import load_config, validate_config EXPECTED_TRAIN_ROWS = 4576 AUXILIARY_ENABLED_PATHS = ( ("teacher_probe",), ("teacher_trajectory",), ("effective_sampling",), ("positive_replay",), ("rollout_replay",), ("signal_utility_routing",), ("effective_group_filter",), ("adaptive_supervision",), ("dynamic_trigger_monitor",), ("global_signal_logging",), ("perception_reward",), ("eval_format_reward",), ("visual_supervision",), ("visual_supervision", "checker"), ("visual_supervision", "refiner"), ("chart_cot_quality_gate",), ) def _require_model(path: Path, *, role: str) -> Path: resolved = path.expanduser().resolve(strict=False) if not (resolved / "config.json").is_file(): raise FileNotFoundError(f"{role} model lacks config.json: {resolved}") weights = tuple(resolved.glob("*.safetensors")) + tuple(resolved.glob("*.bin")) if not weights: raise FileNotFoundError(f"{role} model lacks model weights: {resolved}") return resolved def _disable_auxiliary_routes(opsd: dict[str, Any]) -> None: for path in AUXILIARY_ENABLED_PATHS: node: Any = opsd for key in path: node = node[key] node["enabled"] = False opsd["teacher_correct_repair"]["mode"] = "none" opsd["teacher_probe"]["candidate_log"]["enabled"] = False opsd["teacher_trajectory"].update( enabled=False, context_providers=["visual_facts_deplot"], max_loss_rows_per_rank=0, weight=0.0, ) opsd["teacher_trajectory"]["audit_log"]["enabled"] = False opsd["visual_supervision"]["logging"].update( enabled=False, save_artifacts=False, log_route_binding=False, ) opsd["privileged_debug"]["save_images"] = False def _assert_core_fullresponse_contract(config: dict[str, Any]) -> None: training = config["training"] train_args = training["dyme_args"] opsd = config["opsd"] loss = opsd["loss"] dataset = config["dataset"] checkpoint_eval = config["checkpoint_eval"] failures: list[str] = [] if training.get("stage") != "opd_only" or opsd.get("mode") != "opd_only": failures.append("training must use the isolated opd_only stage") if str(loss.get("loss_type", "")).lower() != "jsd": failures.append("opsd.loss.loss_type must be exact JSD") if "reward_gate" in loss: failures.append("opsd.loss.reward_gate must be absent") if "reference_mode" in loss or "reward_scale" in loss: failures.append("ExOPD reference/reward fields must be absent") if not bool(loss.get("ignore_leading_whitespace_tokens", False)): failures.append("leading whitespace tokens must be excluded from JSD") response_contract = opsd.get("response_contract") or {} if response_contract != { "enabled": True, "profile": "chartqa_five_part_no_gold", }: failures.append("student/core-teacher must share chartqa_five_part_no_gold") if float(train_args.get("temperature", 0.0)) != 1.0: failures.append("student rollout temperature must be 1.0") if float(loss.get("grpo_weight", 1.0)) != 0.0: failures.append("GRPO must be disabled") if float(loss.get("sft_weight", 1.0)) != 0.0: failures.append("SFT must be disabled") if float(loss.get("teacher_traj_fkl_weight", 1.0)) != 0.0: failures.append("teacher trajectory FKL weight must be zero") if dataset.get("teacher_gate"): failures.append("dataset teacher gate must be absent so all 4,576 rows are used") if dataset.get("max_train_samples") is not None: failures.append("dataset.max_train_samples must be null") if int(config["data_validation"].get("expected_samples", -1)) != EXPECTED_TRAIN_ROWS: failures.append(f"expected_samples must be {EXPECTED_TRAIN_ROWS}") if bool(opsd["gate"].get("require_format_for_opsd", True)): failures.append("format filtering must be disabled") if bool(opsd["gate"].get("skip_degenerate_for_opsd", True)): failures.append("degenerate-response filtering must be disabled") if opsd.get("privileged_providers") != ["visual_facts_deplot"]: failures.append("the only privileged provider must be visual_facts_deplot") for path in AUXILIARY_ENABLED_PATHS: node: Any = opsd for key in path: node = node[key] if bool(node.get("enabled", False)): failures.append(f"opsd.{'.'.join(path)} must be disabled") if str(opsd["teacher_correct_repair"].get("mode", "")).lower() != "none": failures.append("teacher_correct_repair must be disabled") if train_args.get("resume_from_checkpoint"): failures.append("official-student initialization must not resume an optimizer checkpoint") if checkpoint_eval.get("enabled", False): if bool(checkpoint_eval.get("initial_eval", False)): failures.append("checkpoint evaluation must not run at step zero") if list(checkpoint_eval.get("gate_steps") or []): failures.append("checkpoint evaluation must not use early gate steps") if int(checkpoint_eval.get("regular_steps", 0)) != 100: failures.append("checkpoint evaluation must run every 100 optimizer steps") if failures: raise ValueError("Invalid full-response OPD recipe:\n- " + "\n- ".join(failures)) def materialize(args: argparse.Namespace) -> dict[str, Any]: config = load_config(str(args.base_config)) student = _require_model(args.student_model, role="student") teacher = _require_model(args.teacher_model, role="teacher") output_dir = args.output_dir.expanduser().resolve(strict=False) model = config["model"] model["pretrained_model_path"] = str(student) model["teacher_model_path"] = str(teacher) model["teacher_device_map"] = "local" model.pop("exopd_reference_model_path", None) training = config["training"] training["stage"] = "opd_only" training["num_gpus"] = int(args.num_gpus) training["num_client"] = int(args.num_gpus) train_args = training["dyme_args"] train_args.update( output_dir=str(output_dir), logging_steps=1, num_generations=4, max_completion_length=int(args.max_completion_length), temperature=1.0, repetition_penalty=1.0, per_device_train_batch_size=2, gradient_accumulation_steps=int(args.gradient_accumulation_steps), num_train_epochs=float(args.num_train_epochs), learning_rate=float(args.learning_rate), bf16=True, gradient_checkpointing=False, ddp_find_unused_parameters=False, ddp_broadcast_buffers=False, max_grad_norm=1.0, save_total_limit=3, weight_decay=0.01, seed=42, ) train_args.pop("resume_from_checkpoint", None) opsd = config["opsd"] opsd.update( enabled=True, mode="opd_only", privileged_profile="hybrid", privileged_providers=["visual_facts_deplot"], strict_teacher_no_gold_context=True, format_only_hint=None, text_include_gold=False, response_contract={ "enabled": True, "profile": "chartqa_five_part_no_gold", }, ) opsd["gate"].update( require_format_for_opsd=False, skip_degenerate_for_opsd=False, online_sft_on_all_wrong=False, signal_aware_routing=False, ) loss = opsd["loss"] loss.update( beta=0.5, opsd_weight=1.0, grpo_weight=0.0, sft_weight=0.0, acc_gate=False, loss_type="jsd", teacher_score_batch_size=4, reuse_teacher_prefix_cache=True, offload_teacher_logits=True, jsd_token_chunk_size=16, # Preserve the generated sequence as conditioning, but remove only # initial whitespace artifacts from the JSD objective. ignore_leading_whitespace_tokens=True, max_leading_whitespace_tokens=4, teacher_traj_fkl_weight=0.0, variance_adaptive=False, assert_teacher_response_prefix_empty=True, assert_qwen_hint_not_consumed=True, assert_privileged_suffix_no_gold=True, ) for legacy_key in ("reward_gate", "reference_mode", "reward_scale"): loss.pop(legacy_key, None) loss["weight_decay"]["enabled"] = False loss["route_cap"]["enabled"] = False _disable_auxiliary_routes(opsd) dataset = config["dataset"] dataset["max_train_samples"] = None # Do not even annotate/filter through the offline teacher gate. Core OPD # consumes every audited source row and remains gold-blind. dataset.pop("teacher_gate", None) config["data_validation"].update( strict=True, require_real_deplot=True, require_qwen_rewrite=False, expected_samples=EXPECTED_TRAIN_ROWS, ) checkpoint_eval = config["checkpoint_eval"] checkpoint_eval.update( enabled=not args.smoke, # A full ChartQA validation is comparatively expensive on two ranks. # Keep the first model-selection point at the normal save cadence; # it establishes the dynamic reference at step 100 without an # extra step-zero pass or short early-step validations. initial_eval=False, gate_steps=[], regular_steps=100, patience=3, patience_start_step=100, tie_policy="count", ) debug = opsd["debug"] debug.update( probe_on_generate=False, probe_first_token_logits=False, probe_log_model_context=False, ) if args.smoke: train_args["max_steps"] = int(args.max_steps) train_args["num_train_epochs"] = 1.0 train_args["save_strategy"] = "no" train_args["warmup_steps"] = 1 debug["detail_every"] = 1 debug["distributed_trace"]["cuda_sync"] = True else: train_args.pop("max_steps", None) train_args["save_strategy"] = "steps" train_args["save_steps"] = 100 train_args["warmup_steps"] = 50 debug["detail_every"] = 10 debug["distributed_trace"]["cuda_sync"] = False validated = validate_config(config, source=str(args.output_config)) _assert_core_fullresponse_contract(validated) target = args.output_config.expanduser().resolve(strict=False) target.parent.mkdir(parents=True, exist_ok=True) temporary = target.with_name(f".{target.name}.tmp-{os.getpid()}") with temporary.open("w", encoding="utf-8") as handle: yaml.safe_dump(validated, handle, allow_unicode=True, sort_keys=False) os.replace(temporary, target) return validated def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--base-config", type=Path, required=True) parser.add_argument("--student-model", type=Path, required=True) parser.add_argument("--teacher-model", type=Path, required=True) parser.add_argument("--output-config", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--num-gpus", type=int, required=True) parser.add_argument("--gradient-accumulation-steps", type=int, default=8) parser.add_argument( "--max-completion-length", type=int, default=96, help="Maximum generated completion tokens for the OPD rollout.", ) parser.add_argument("--num-train-epochs", type=float, default=10.0) parser.add_argument("--learning-rate", type=float, default=2.0e-6) parser.add_argument("--smoke", action="store_true") parser.add_argument("--max-steps", type=int, default=5) return parser.parse_args() def main() -> None: args = parse_args() if args.num_gpus <= 0: raise ValueError("--num-gpus must be positive") if args.gradient_accumulation_steps <= 0: raise ValueError("--gradient-accumulation-steps must be positive") if args.max_completion_length <= 0: raise ValueError("--max-completion-length must be positive") if args.num_train_epochs <= 0: raise ValueError("--num-train-epochs must be positive") if args.learning_rate <= 0: raise ValueError("--learning-rate must be positive") if args.smoke and args.max_steps <= 0: raise ValueError("--max-steps must be positive for smoke runs") config = materialize(args) loss = config["opsd"]["loss"] print( "Materialized ChartQA core full-response OPD recipe: " f"smoke={args.smoke} ranks={args.num_gpus} rows={EXPECTED_TRAIN_ROWS} " f"gradient_accumulation_steps={config['training']['dyme_args']['gradient_accumulation_steps']} " f"max_completion_length={config['training']['dyme_args']['max_completion_length']} " f"epochs={config['training']['dyme_args']['num_train_epochs']} " f"loss={loss['loss_type']} leading_whitespace_masked={loss['ignore_leading_whitespace_tokens']} " f"response_profile={config['opsd']['response_contract']['profile']} " "reward_gate=absent exopd=absent trajectory=disabled " f"student={config['model']['pretrained_model_path']} " f"output={config['training']['dyme_args']['output_dir']}", flush=True, ) if __name__ == "__main__": main()