scripts/materialize_chartqa_fullresponse_opd_config.py

#!/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()