scripts/train_math_ink_06_federated_online.py

"""승인된 P-track source를 rehearsal해 0.6 shared online encoder의 source 일반화를 개선한다."""

from __future__ import annotations

import argparse
from collections import Counter
from hashlib import sha256
import json
import random
import sys
from pathlib import Path

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))
if str(PROJECT_ROOT / "scripts") not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT / "scripts"))

from math_grid_drawer.research.ink06_federation import (
    CrossSourceLabelBatchSampler06, FederatedPairedInk06Dataset, augment_online_features06,
    cross_source_supervised_contrastive_loss06, federation_provenance06, interpolate_state_dict06, load_product_federation06,
    resolve_training_device06, source_label_balanced_sampler06,
)
from math_grid_drawer.research.math_ink_06 import MathInk06Engine, fuse_hypothesis_class_logits06
from train_math_ink_06_candidate import _balanced_subset
from train_math_ink_06_virtual_adapter import _fused_logits


def _partition(record: dict, buckets: int = 5) -> int:
    """필요 변수: writer 또는 origin. 작동 원리: writer가 있으면 writer 단위, 없으면 origin 단위의 안정 partition을 만든다."""

    key = str(record.get("writer_key") or record.get("origin_id") or record["sample_id"])
    return int.from_bytes(sha256(key.encode("utf-8")).digest()[:8], "big") % buckets


def _source_subset(records: list[dict], maximum: int, seed: int) -> list[dict]:
    """필요 변수: 한 source record. 작동 원리: label round-robin으로 희귀 class를 보존한다."""

    return _balanced_subset(records, maximum, seed)


def _parse_source_train_caps06(value: str) -> dict[str, int]:
    """필요 변수: `source_id=count` CSV. 작동 원리: 공통 상한으로 재현할 수 없는 source별 training cap을 명시적으로 검증해 반환한다."""

    if not value.strip():
        return {}
    output: dict[str, int] = {}
    for item in value.split(","):
        source_id, separator, raw_count = item.strip().partition("=")
        if not separator or not source_id or not raw_count.isdigit() or int(raw_count) <= 0:
            raise ValueError("source train cap은 `source_id=양의정수` CSV여야 합니다.")
        if source_id in output:
            raise ValueError(f"source train cap이 중복되었습니다: {source_id}")
        output[source_id] = int(raw_count)
    return output


def _training_manifest06(records: list[dict]) -> dict[str, object]:
    """필요 변수: 최종 선택된 training record. 작동 원리: source별 실제 수와 source/sample/origin 식별자의 SHA-256을 고정해 동일 count의 다른 표본 재학습을 막는다."""

    rows = sorted(
        f"{row['source']}\x1f{row['sample_id']}\x1f{row.get('origin_id', '')}"
        for row in records
    )
    digest = sha256("\n".join(rows).encode("utf-8")).hexdigest()
    return {
        "samples": len(records),
        "source_counts": dict(sorted(Counter(str(row["source"]) for row in records).items())),
        "sample_origin_sha256": digest,
    }


def _evaluate_source(
    engine: MathInk06Engine, records: list[dict], exact_to_index: dict[str, int],
    family_to_index: dict[str, int], batch_size: int,
) -> dict[str, float]:
    """필요 변수: source별 고정 record. 작동 원리: online과 raster top-k를 동시에 측정한다."""

    loader = DataLoader(
        FederatedPairedInk06Dataset(records, exact_to_index, family_to_index),
        batch_size=batch_size, shuffle=False, num_workers=0,
    )
    totals = {"samples": 0, "online_top1": 0, "online_top5": 0, "raster_top1": 0, "raster_top5": 0}
    engine.model.eval()
    with torch.inference_mode():
        for online, raster, _coordinates, _states, target, _family, _source in loader:
            online, raster, target = online.to(engine.device), raster.to(engine.device), target.to(engine.device)
            online_logits, _online_family = engine.model.forward_online(online)
            raster_output = engine.model.forward_raster(raster)
            raster_logits = _fused_logits(engine, raster_output)
            totals["samples"] += len(target)
            totals["online_top1"] += int((online_logits.argmax(dim=1) == target).sum())
            totals["online_top5"] += int((online_logits.topk(5, dim=1).indices == target[:, None]).any(dim=1).sum())
            totals["raster_top1"] += int((raster_logits.argmax(dim=1) == target).sum())
            totals["raster_top5"] += int((raster_logits.topk(5, dim=1).indices == target[:, None]).any(dim=1).sum())
    count = max(int(totals.pop("samples")), 1)
    return {"samples": count, **{key: value / count for key, value in totals.items()}}


def _evaluate_sources(
    engine: MathInk06Engine, groups: dict[str, list[dict]], exact_to_index: dict[str, int],
    family_to_index: dict[str, int], batch_size: int,
) -> dict[str, dict[str, float]]:
    """필요 변수: source→record. 작동 원리: source별 분리 지표와 macro online top-1을 반환한다."""

    return {
        source_id: _evaluate_source(engine, records, exact_to_index, family_to_index, batch_size)
        for source_id, records in groups.items()
    }


def _macro_online(metrics: dict[str, dict[str, float]]) -> float:
    """필요 변수: source별 지표. 작동 원리: 표본 수와 무관한 source macro online top-1을 계산한다."""

    return float(np.mean([row["online_top1"] for row in metrics.values()]))


def _file_sha25606(path: Path) -> str:
    """필요 변수: UTF-8과 무관한 checkpoint byte 경로. 작동 원리: resume이 다른 base 모델을 섞지 않도록 SHA-256을 계산한다."""

    digest = sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _resume_contract06(
    checkpoint: Path, *, seed: int, source_ids: list[str], train_samples: int,
    max_train_per_source: int, max_validation_per_source: int, max_test_per_source: int,
    cross_source_pair_fraction: float = 0.0, cross_source_contrastive_weight: float = 0.0,
    cross_source_temperature: float = 0.20, source_train_caps: dict[str, int] | None = None,
    training_sample_origin_sha256: str = "",
) -> dict[str, object]:
    """필요 변수: base checkpoint·고정 source split·sampling/loss 설정. 작동 원리: epoch resume이 다른 data 또는 cross-source 정렬 계약을 이어붙이는 일을 fail-closed로 막는다."""

    return {
        "schema_version": "aiflow-math-ink-0.6-federated-online-resume1",
        "base_checkpoint_sha256": _file_sha25606(checkpoint),
        "seed": seed,
        "source_ids": sorted(source_ids),
        "train_samples": train_samples,
        "max_train_per_source": max_train_per_source,
        "max_validation_per_source": max_validation_per_source,
        "max_test_per_source": max_test_per_source,
        "cross_source_pair_fraction": cross_source_pair_fraction,
        "cross_source_contrastive_weight": cross_source_contrastive_weight,
        "cross_source_temperature": cross_source_temperature,
        "source_train_caps": dict(sorted((source_train_caps or {}).items())),
        "training_sample_origin_sha256": training_sample_origin_sha256,
    }


def _restore_optimizer_device06(optimizer: torch.optim.Optimizer, device: torch.device) -> None:
    """필요 변수: CPU로 저장된 optimizer state·실행 device. 작동 원리: resume 뒤 Adam moment를 parameter와 같은 device로 되돌린다."""

    for state in optimizer.state.values():
        for key, value in state.items():
            if isinstance(value, torch.Tensor):
                state[key] = value.to(device)


def _save_epoch_state06(
    path: Path, *, contract: dict[str, object], epoch: int, student_state: dict[str, torch.Tensor],
    best_state: dict[str, torch.Tensor], anchor_state: dict[str, torch.Tensor], optimizer: torch.optim.Optimizer,
    best_metrics: dict[str, dict[str, float]], best_macro: float, best_epoch: int, best_alpha: float,
    history: list[dict], sampler: object,
) -> None:
    """필요 변수: epoch 모델·선택 state·sampler/RNG. 작동 원리: 다음 실행이 같은 sampler 순서와 augmentation 난수에서 정확히 이어지게 저장한다."""

    generator = getattr(sampler, "generator", None)
    sampler_state = generator.get_state() if isinstance(generator, torch.Generator) else None
    payload = {
        "schema_version": "aiflow-math-ink-0.6-federated-online-epoch-state1",
        "contract": contract,
        "last_epoch": epoch,
        "student_state": {key: value.detach().cpu().clone() for key, value in student_state.items()},
        "best_state": {key: value.detach().cpu().clone() for key, value in best_state.items()},
        "anchor_state": {key: value.detach().cpu().clone() for key, value in anchor_state.items()},
        "optimizer_state": optimizer.state_dict(),
        "best_metrics": best_metrics,
        "best_macro": best_macro,
        "best_epoch": best_epoch,
        "best_alpha": best_alpha,
        "history": history,
        "sampler_state": sampler_state,
        "python_rng_state": random.getstate(),
        "numpy_rng_state": np.random.get_state(),
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(payload, path)


def _restore_epoch_state06(
    path: Path, *, contract: dict[str, object], student: MathInk06Engine,
    optimizer: torch.optim.Optimizer, sampler: object,
) -> dict:
    """필요 변수: epoch-state·현재 data contract·student/optimizer. 작동 원리: contract 일치 때만 파라미터·optimizer·난수를 원자적으로 복원한다."""

    payload = torch.load(path, map_location="cpu", weights_only=False)
    if payload.get("schema_version") != "aiflow-math-ink-0.6-federated-online-epoch-state1":
        raise ValueError("지원하지 않는 federated online epoch-state입니다.")
    if payload.get("contract") != contract:
        raise ValueError("resume epoch-state의 checkpoint/source/split 계약이 현재 실행과 다릅니다.")
    student.model.load_state_dict(payload["student_state"])
    optimizer.load_state_dict(payload["optimizer_state"])
    _restore_optimizer_device06(optimizer, student.device)
    generator = getattr(sampler, "generator", None)
    sampler_state = payload.get("sampler_state")
    if isinstance(generator, torch.Generator) and isinstance(sampler_state, torch.Tensor):
        generator.set_state(sampler_state)
    random.setstate(payload["python_rng_state"])
    np.random.set_state(payload["numpy_rng_state"])
    torch.set_rng_state(payload["torch_rng_state"])
    if torch.cuda.is_available() and payload.get("cuda_rng_state") is not None:
        torch.cuda.set_rng_state_all(payload["cuda_rng_state"])
    return payload


def main() -> None:
    """필요 변수: base 0.6·승인 source. 작동 원리: 명시적 writer split을 우선한 distillation rehearsal과 holdout 평가를 수행한다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 federated online rehearsal")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--registry", type=Path, default=PROJECT_ROOT / "research/dataset_registry.json")
    parser.add_argument("--source-registry", type=Path, default=PROJECT_ROOT / "research/math_ink_06_source_registry.json")
    parser.add_argument("--commercial", type=Path, default=PROJECT_ROOT / "research/data/external_trajectory_v1/commercial_ccby4.jsonl.gz")
    parser.add_argument("--hwrt", type=Path, default=PROJECT_ROOT / "research/data/open_pretrain/hwrt_expanded_v2/hwrt_expanded.jsonl.gz")
    parser.add_argument("--approval", type=Path, default=PROJECT_ROOT / "research/approvals/HWRT-ODBL-USE-APPROVAL-v1.json")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--max-train-per-source", type=int, default=2000)
    parser.add_argument("--max-validation-per-source", type=int, default=300)
    parser.add_argument("--max-test-per-source", type=int, default=500)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--learning-rate", type=float, default=2e-5)
    parser.add_argument("--distillation-weight", type=float, default=0.75)
    parser.add_argument("--raster-distillation-weight", type=float, default=1.0)
    parser.add_argument("--raster-supervised-weight", type=float, default=0.0)
    parser.add_argument("--raster-family-supervised-weight", type=float, default=0.0)
    parser.add_argument("--hwrt-tolerance", type=float, default=0.01)
    parser.add_argument("--interpolation-alphas", default="0.25,0.5,0.75,1.0")
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument("--device", default="auto", help="auto|cpu|cuda[:index]")
    parser.add_argument("--online-rotation-degrees", type=float, default=0.0)
    parser.add_argument("--online-scale-jitter", type=float, default=0.0)
    parser.add_argument("--cross-source-pair-fraction", type=float, default=0.0)
    parser.add_argument("--cross-source-contrastive-weight", type=float, default=0.0)
    parser.add_argument("--cross-source-temperature", type=float, default=0.20)
    parser.add_argument("--source-train-caps", default="", help="선택 source별 training cap: source_id=count,...")
    parser.add_argument("--resume-state", type=Path, help="이전 epoch-state에서 안전하게 재개한다.")
    parser.add_argument(
        "--save-epoch-state", action=argparse.BooleanOptionalAction, default=True,
        help="각 epoch의 재개 가능한 state를 저장한다.",
    )
    args = parser.parse_args()
    if args.raster_supervised_weight < 0 or args.raster_family_supervised_weight < 0:
        raise ValueError("raster supervised weight는 0 이상이어야 합니다.")
    if args.online_rotation_degrees < 0 or args.online_scale_jitter < 0:
        raise ValueError("online augmentation 범위는 0 이상이어야 합니다.")
    if not 0.0 <= args.cross_source_pair_fraction <= 1.0:
        raise ValueError("cross-source pair 비율은 0과 1 사이여야 합니다.")
    if args.cross_source_contrastive_weight < 0.0 or args.cross_source_temperature <= 0.0:
        raise ValueError("cross-source contrastive weight는 0 이상, temperature는 양수여야 합니다.")
    if args.cross_source_contrastive_weight > 0.0 and args.cross_source_pair_fraction <= 0.0:
        raise ValueError("contrastive loss에는 0보다 큰 cross-source pair 비율이 필요합니다.")
    source_train_caps = _parse_source_train_caps06(args.source_train_caps)

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = resolve_training_device06(args.device)
    student = MathInk06Engine(args.checkpoint, device=device)
    teacher = MathInk06Engine(args.checkpoint, device=str(student.device))
    print(json.dumps({"device": str(student.device), "cuda": torch.cuda.is_available()}), flush=True)
    for parameter in student.model.raster_encoder.parameters():
        parameter.requires_grad_(False)
    for parameter in student.model.virtual_decoder.parameters():
        parameter.requires_grad_(False)
    if student.model.auxiliary_virtual_decoder is not None:
        for parameter in student.model.auxiliary_virtual_decoder.parameters():
            parameter.requires_grad_(False)
    for parameter in student.model.virtual_adapter.parameters():
        parameter.requires_grad_(False)
    trainable = [
        *student.model.trajectory_encoder.parameters(), *student.model.exact_head.parameters(),
        *student.model.family_head.parameters(),
    ]
    exact_to_index = {label: index for index, label in enumerate(student.labels)}
    family_to_index = {label: index for index, label in enumerate(student.family_labels)}
    sources = load_product_federation06(
        registry_path=args.registry, commercial_path=args.commercial, hwrt_path=args.hwrt,
        approval_path=args.approval, allowed_labels=student.labels, source_registry_path=args.source_registry,
    )
    training_records = []
    validation_groups = {}
    test_groups = {}
    for source_index, source in enumerate(sources):
        eligible = [row for row in source.records if row.get("eligible_for_training")]
        explicit_validation = [row for row in source.records if str(row.get("split")) in {"validation", "valid", "val"}]
        if explicit_validation:
            train_candidates = eligible
            validation_candidates = explicit_validation
        else:
            train_candidates = [row for row in eligible if _partition(row) >= 2]
            validation_candidates = [row for row in eligible if _partition(row) == 0]
        training_records.extend(_source_subset(
            train_candidates, source_train_caps.get(source.source_id, args.max_train_per_source), args.seed + source_index,
        ))
        validation_groups[source.source_id] = _source_subset(
            validation_candidates, args.max_validation_per_source, args.seed + 20 + source_index,
        )
        test_candidates = [row for row in source.records if row.get("split") == "test"]
        test_groups[source.source_id] = _source_subset(
            test_candidates, args.max_test_per_source, args.seed + 40 + source_index,
        )
    if any(not rows for rows in validation_groups.values()) or any(not rows for rows in test_groups.values()):
        raise ValueError("source validation/test partition이 비었습니다.")
    training_manifest = _training_manifest06(training_records)
    dataset = FederatedPairedInk06Dataset(training_records, exact_to_index, family_to_index)
    if args.cross_source_pair_fraction > 0.0:
        epoch_sampler = CrossSourceLabelBatchSampler06(
            training_records, batch_size=args.batch_size, samples=len(training_records), seed=args.seed,
            pair_fraction=args.cross_source_pair_fraction,
        )
        loader = DataLoader(
            dataset, batch_sampler=epoch_sampler, num_workers=0,
        )
    else:
        sampler = source_label_balanced_sampler06(
            training_records, seed=args.seed, samples=len(training_records),
        )
        epoch_sampler = sampler
        loader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler, num_workers=0)
    source_to_index = {source_id: index for index, source_id in enumerate(sorted({str(row["source"]) for row in training_records}))}
    optimizer = torch.optim.AdamW(trainable, lr=args.learning_rate, weight_decay=2e-3)
    resume_contract = _resume_contract06(
        args.checkpoint, seed=args.seed, source_ids=list(validation_groups), train_samples=len(training_records),
        max_train_per_source=args.max_train_per_source, max_validation_per_source=args.max_validation_per_source,
        max_test_per_source=args.max_test_per_source,
        cross_source_pair_fraction=args.cross_source_pair_fraction,
        cross_source_contrastive_weight=args.cross_source_contrastive_weight,
        cross_source_temperature=args.cross_source_temperature,
        source_train_caps=source_train_caps,
        training_sample_origin_sha256=str(training_manifest["sample_origin_sha256"]),
    )
    args.output.mkdir(parents=True, exist_ok=True)
    baseline_validation = _evaluate_sources(
        student, validation_groups, exact_to_index, family_to_index, args.batch_size,
    )
    # 내부 test는 선택에 사용하지 않고 최종 비교 기준만 미리 고정한다.
    baseline_test = _evaluate_sources(
        student, test_groups, exact_to_index, family_to_index, args.batch_size,
    )
    best_metrics = baseline_validation
    best_macro = _macro_online(baseline_validation)
    best_epoch = 0
    best_alpha = 0.0
    best_state = {key: value.detach().cpu().clone() for key, value in student.model.state_dict().items()}
    anchor_state = {key: value.clone() for key, value in best_state.items()}
    interpolation_alphas = tuple(float(value) for value in args.interpolation_alphas.split(",") if value.strip())
    if not interpolation_alphas or any(not 0.0 < value <= 1.0 for value in interpolation_alphas):
        raise ValueError("interpolation alpha는 0보다 크고 1 이하여야 합니다.")
    history = [{"epoch": 0, "macro_online_top1": best_macro, "validation": baseline_validation}]
    start_epoch = 1
    if args.resume_state is not None:
        resumed = _restore_epoch_state06(
            args.resume_state, contract=resume_contract, student=student, optimizer=optimizer, sampler=epoch_sampler,
        )
        best_state = resumed["best_state"]
        anchor_state = resumed["anchor_state"]
        best_metrics = resumed["best_metrics"]
        best_macro = float(resumed["best_macro"])
        best_epoch = int(resumed["best_epoch"])
        best_alpha = float(resumed["best_alpha"])
        history = list(resumed["history"])
        start_epoch = int(resumed["last_epoch"]) + 1
        if start_epoch > args.epochs:
            raise ValueError("resume epoch-state가 요청 epochs보다 이미 앞서 있습니다.")
        print(json.dumps({"resumed_from": str(args.resume_state), "start_epoch": start_epoch}, ensure_ascii=False), flush=True)
    temperature = 2.0
    for epoch in range(start_epoch, args.epochs + 1):
        student.model.train()
        total_loss = seen = contrastive_anchors = 0
        for online, raster, _coordinates, _states, exact_target, family_target, source_ids in loader:
            online, raster = online.to(student.device), raster.to(student.device)
            exact_target, family_target = exact_target.to(student.device), family_target.to(student.device)
            optimizer.zero_grad(set_to_none=True)
            augmented_online = augment_online_features06(
                online,
                rotation_degrees=args.online_rotation_degrees,
                scale_jitter=args.online_scale_jitter,
            )
            embedding = student.model.encode_trajectory(augmented_online)
            exact, family = student.model.exact_head(embedding), student.model.family_head(embedding)
            with torch.inference_mode():
                teacher_exact, _teacher_family = teacher.model.forward_online(online)
                teacher_raster = teacher.model.forward_raster(raster)["exact_logits"]
            supervised = nn.functional.cross_entropy(exact, exact_target, label_smoothing=0.02)
            supervised = supervised + 0.10 * nn.functional.cross_entropy(family, family_target)
            distillation = nn.functional.kl_div(
                nn.functional.log_softmax(exact / temperature, dim=-1),
                nn.functional.softmax(teacher_exact / temperature, dim=-1), reduction="batchmean",
            ) * (temperature * temperature)
            student_raster_output = student.model.forward_raster(raster)
            student_raster = student_raster_output["exact_logits"]
            student_raster_fused = _fused_logits(student, student_raster_output)
            student_raster_family_fused = fuse_hypothesis_class_logits06(
                student_raster_output["family_logits"], student_raster_output["hypothesis_scores"],
                mode=str(student.raster_fusion["mode"]),
                score_weight=float(student.raster_fusion["score_weight"]),
            )
            raster_distillation = nn.functional.kl_div(
                nn.functional.log_softmax(student_raster / temperature, dim=-1),
                nn.functional.softmax(teacher_raster / temperature, dim=-1), reduction="batchmean",
            ) * (temperature * temperature) / student_raster.shape[1]
            loss = (
                supervised + args.distillation_weight * distillation
                + args.raster_distillation_weight * raster_distillation
                + args.raster_supervised_weight * nn.functional.cross_entropy(
                    student_raster_fused, exact_target, label_smoothing=0.02,
                )
                + args.raster_family_supervised_weight * nn.functional.cross_entropy(
                    student_raster_family_fused, family_target, label_smoothing=0.02,
                )
            )
            if args.cross_source_contrastive_weight > 0.0:
                source_tensor = torch.tensor(
                    [source_to_index[str(value)] for value in source_ids], dtype=torch.long, device=student.device,
                )
                contrastive, anchors = cross_source_supervised_contrastive_loss06(
                    embedding, exact_target, source_tensor, temperature=args.cross_source_temperature,
                )
                loss = loss + args.cross_source_contrastive_weight * contrastive
                contrastive_anchors += anchors
            loss.backward()
            nn.utils.clip_grad_norm_(trainable, 1.0)
            optimizer.step()
            seen += len(exact_target)
            total_loss += float(loss.detach()) * len(exact_target)
        trained_state = {key: value.detach().cpu().clone() for key, value in student.model.state_dict().items()}
        interpolation = []
        for alpha in interpolation_alphas:
            mixed_state = interpolate_state_dict06(anchor_state, trained_state, alpha=alpha)
            student.model.load_state_dict(mixed_state)
            metrics = _evaluate_sources(
                student, validation_groups, exact_to_index, family_to_index, args.batch_size,
            )
            macro = _macro_online(metrics)
            hwrt_ok = all(
                metrics["hwrt"][metric] >= baseline_validation["hwrt"][metric] - args.hwrt_tolerance
                for metric in ("online_top1", "online_top5", "raster_top1", "raster_top5")
            )
            interpolation.append({"alpha": alpha, "macro_online_top1": macro, "hwrt_gate": hwrt_ok, "validation": metrics})
            if hwrt_ok and macro > best_macro:
                best_macro, best_metrics, best_epoch, best_alpha = macro, metrics, epoch, alpha
                best_state = {key: value.clone() for key, value in mixed_state.items()}
        student.model.load_state_dict(trained_state)
        final_row = interpolation[-1]
        row = {
            "epoch": epoch, "loss": total_loss / max(seen, 1), "cross_source_contrastive_anchors": contrastive_anchors,
            "macro_online_top1": final_row["macro_online_top1"], "hwrt_gate": final_row["hwrt_gate"],
            "validation": final_row["validation"], "interpolation": interpolation,
        }
        history.append(row)
        if args.save_epoch_state:
            state_path = args.output / f"federated_online_epoch_{epoch:03d}.pt"
            _save_epoch_state06(
                state_path, contract=resume_contract, epoch=epoch, student_state=trained_state,
                best_state=best_state, anchor_state=anchor_state, optimizer=optimizer,
                best_metrics=best_metrics, best_macro=best_macro, best_epoch=best_epoch,
                best_alpha=best_alpha, history=history, sampler=epoch_sampler,
            )
            row["resume_state"] = state_path.name
        print(json.dumps(row, ensure_ascii=False), flush=True)
    student.model.load_state_dict(best_state)
    final_test = _evaluate_sources(
        student, test_groups, exact_to_index, family_to_index, args.batch_size,
    )
    output_payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    output_payload["state_dict"] = best_state
    output_payload["model_version"] = "aiflow-math-ink-0.6-federated-online1"
    provenance = federation_provenance06(sources, args.source_registry)
    output_payload.update(provenance)
    output_payload["federated_online"] = {
        "sources": sorted(validation_groups), "selected_epoch": best_epoch,
        "seed": args.seed,
        "selected_interpolation_alpha": best_alpha,
        "distillation_weight": args.distillation_weight,
        "raster_distillation_weight": args.raster_distillation_weight,
        "raster_supervised_weight": args.raster_supervised_weight,
        "raster_family_supervised_weight": args.raster_family_supervised_weight,
        "online_rotation_degrees": args.online_rotation_degrees,
        "online_scale_jitter": args.online_scale_jitter,
        "cross_source_pair_fraction": args.cross_source_pair_fraction,
        "cross_source_contrastive_weight": args.cross_source_contrastive_weight,
        "cross_source_temperature": args.cross_source_temperature,
        "source_train_caps": source_train_caps,
        "training_manifest": training_manifest,
        "resumed_from": str(args.resume_state) if args.resume_state is not None else None,
    }
    checkpoint = args.output / "math_ink_06_candidate.pt"
    torch.save(output_payload, checkpoint)
    report = {
        "checkpoint": checkpoint.name, "bytes": checkpoint.stat().st_size,
        "seed": args.seed,
        "device": str(student.device),
        **provenance,
        "train_samples": len(training_records), "source_count": len(sources),
        "training_manifest": training_manifest,
        "baseline_validation": baseline_validation, "selected_validation": best_metrics,
        "selected_epoch": best_epoch, "selected_interpolation_alpha": best_alpha,
        "baseline_test": baseline_test, "test": final_test,
        "test_delta": {
            source_id: {
                metric: final_test[source_id][metric] - baseline_test[source_id][metric]
                for metric in ("online_top1", "online_top5", "raster_top1", "raster_top5")
            }
            for source_id in final_test
        },
        "history": history,
        "hwrt_official_test_used": False, "product_validation": False,
    }
    (args.output / "federated_online_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )
    print(json.dumps({key: value for key, value in report.items() if key != "history"}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()