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