Buckets:

EmergentTwilight
/

MMCoT

Files
xet
EmergentTwilight/MMCoT / experiments /src /zebra_sft /evaluate_task_subset.py
EmergentTwilight
20 days ago
download
raw
32.5 kB
"""Distributed direct/interleaved evaluation on an indexed Zebra task subset."""

from __future__ import annotations

import argparse
import fcntl
import gc
import hashlib
import json
import os
import re
import time
import uuid
from collections import Counter, OrderedDict
from datetime import timedelta
from pathlib import Path

import torch
import torch.distributed as dist
from artifacts import is_timestamped_run_name
from PIL import Image
from sft.runtime import claim_run_directory
from peft import PeftModel
from transformers import AutoProcessor
from uniar import UniARForConditionalGeneration, UniARVisualDecoder, configure_processor

from .cached_rollout import CachedInterleavedRollout
from .evaluate_corpus import _load_shard, _readable, _target
from .generated_codes import load_generated_codes, save_generated_codes
from .scoring import score_answer

SUBSET_FIELD = "evaluation_subset_index"


def _write_json(path: Path, value) -> None:
    temporary = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex}")
    try:
        temporary.write_text(
            json.dumps(value, ensure_ascii=False, indent=2, default=str) + "\n",
            encoding="utf-8",
        )
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def _write_text(path: Path, value: str) -> None:
    temporary = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex}")
    try:
        temporary.write_text(value, encoding="utf-8")
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def _validate_sample_commit(samples_dir: Path, row: dict) -> None:
    sample_dir = samples_dir / row["sample_id"]
    for name in ("target.txt", "output.txt", "result.json", "generated_codes.pt"):
        if not (sample_dir / name).is_file():
            raise FileNotFoundError(sample_dir / name)
    result = json.loads((sample_dir / "result.json").read_text(encoding="utf-8"))
    if (
        result.get("sample_id") != row.get("sample_id")
        or result.get("evaluation_index") != row.get("evaluation_index")
    ):
        raise ValueError(f"sample commit provenance mismatch: {sample_dir}")
    load_generated_codes(sample_dir / "generated_codes.pt")


def _wait_for_adapter(args: argparse.Namespace, sleep=time.sleep) -> None:
    """Wait on the atomic SFT adapter marker without changing retry argv."""

    if not args.wait_for_adapter:
        return
    if args.mode != "interleaved" or args.adapter_dir is None:
        raise ValueError("--wait-for-adapter requires interleaved mode and adapter-dir")
    marker = args.adapter_dir / "COMPLETE"
    announced = False
    while not marker.is_file():
        if int(os.environ.get("LOCAL_RANK", "0")) == 0 and not announced:
            print(f"WAITING_FOR_ADAPTER {marker}", flush=True)
            announced = True
        sleep(15.0)


def _load_prediction_rows(
    path: Path, *, resume: bool, samples_dir: Path | None = None
) -> list[dict]:
    if not path.exists():
        return []
    data = path.read_bytes()
    boundary = data.rfind(b"\n") + 1
    changed = boundary != len(data)
    if changed:
        if not resume:
            raise ValueError(f"prediction journal has a torn tail: {path}")
        data = data[:boundary]
    parsed = [
        (line, json.loads(line.decode("utf-8")))
        for line in data.splitlines(keepends=True)
        if line.strip()
    ]
    rows = [row for _, row in parsed]
    if resume and samples_dir is not None:
        committed = []
        committed_lines = []
        for line, row in parsed:
            try:
                _validate_sample_commit(samples_dir, row)
            except (FileNotFoundError, OSError, RuntimeError, ValueError):
                changed = True
                break
            committed.append(row)
            committed_lines.append(line)
        rows = committed
        parsed = list(zip(committed_lines, committed, strict=True))
    if changed:
        temporary = path.with_name(f".{path.name}.recovery-{os.getpid()}")
        try:
            with temporary.open("wb") as handle:
                for line, _ in parsed:
                    handle.write(line)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, path)
        finally:
            temporary.unlink(missing_ok=True)
    return rows


def _acquire_output_lock(output_dir: Path):
    """Hold an advisory lock so only one evaluation writes a run directory."""
    lock_path = output_dir / ".active-writer.lock"
    handle = lock_path.open("a+", encoding="utf-8")
    try:
        fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError as error:
        handle.seek(0)
        holder = handle.read().strip() or "unknown writer"
        handle.close()
        raise RuntimeError(
            f"evaluation output already has an active writer: {holder}"
        ) from error
    handle.seek(0)
    handle.truncate()
    json.dump(
        {
            "pid": os.getpid(),
            "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
            "host": os.uname().nodename,
        },
        handle,
    )
    handle.write("\n")
    handle.flush()
    return handle


def _worker_device_index(
    local_rank: int, visible_devices: int, workers_per_gpu: int
) -> int:
    if local_rank < 0 or visible_devices < 1 or workers_per_gpu < 1:
        raise ValueError("local rank, visible devices, and workers per GPU are invalid")
    capacity = visible_devices * workers_per_gpu
    if local_rank >= capacity:
        raise ValueError(
            f"local rank {local_rank} exceeds capacity {capacity} for "
            f"{visible_devices} visible GPU(s) and {workers_per_gpu} worker(s) per GPU"
        )
    return local_rank // workers_per_gpu


def _distributed(workers_per_gpu: int) -> tuple[int, int, int, torch.device]:
    rank = int(os.environ.get("RANK", "0"))
    world = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    local_world = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
    visible_devices = torch.cuda.device_count()
    expected_local_world = visible_devices * workers_per_gpu
    if local_world != expected_local_world:
        raise ValueError(
            f"LOCAL_WORLD_SIZE={local_world}, expected {expected_local_world} for "
            f"{visible_devices} visible GPU(s) and {workers_per_gpu} worker(s) per GPU"
        )
    device_index = _worker_device_index(
        local_rank, visible_devices, workers_per_gpu
    )
    if world > 1:
        # Rollout lengths are highly imbalanced. A finished rank may wait well
        # over NCCL's default timeout while another completes a long cyclic
        # sample, even though both are healthy.
        backend = "gloo" if workers_per_gpu > 1 else "nccl"
        dist.init_process_group(backend, timeout=timedelta(hours=6))
    torch.cuda.set_device(device_index)
    return rank, world, local_rank, torch.device("cuda", device_index)


def _barrier(world: int) -> None:
    if world > 1:
        dist.barrier()


def _control_collective_device(world: int, device: torch.device) -> torch.device:
    if world > 1 and dist.get_backend() == "gloo":
        return torch.device("cpu")
    return device


def _manifest_rows(cache_root: Path, selection: str = "indexed") -> list[dict]:
    rows = []
    with (cache_root / "manifest.jsonl").open(encoding="utf-8") as handle:
        for line in handle:
            row = json.loads(line)
            if selection == "indexed" and row.get(SUBSET_FIELD) is not None:
                if row["split"] != "test":
                    raise ValueError("evaluation subset contains a non-test row")
                rows.append(row)
            elif selection in {"validation", "test"} and row["split"] == selection:
                rows.append(row)
    if selection == "indexed":
        rows.sort(key=lambda row: row[SUBSET_FIELD])
        actual = [row[SUBSET_FIELD] for row in rows]
        if actual != list(range(len(rows))):
            raise ValueError(f"{SUBSET_FIELD} must be unique and contiguous")
    else:
        rows.sort(key=lambda row: row["sample_id"])
    for index, row in enumerate(rows):
        row["evaluation_index"] = index
    if not rows:
        raise ValueError(f"evaluation selection {selection!r} is empty")
    return rows


def _limit_manifest_rows(
    rows: list[dict], max_samples: int | None, sample_offset: int = 0
) -> list[dict]:
    if sample_offset < 0:
        raise ValueError("sample_offset must be nonnegative")
    if max_samples is None and sample_offset == 0:
        return rows
    if max_samples is None:
        return rows[sample_offset:]
    if max_samples <= 0:
        raise ValueError("max_samples must be positive")
    return rows[sample_offset : sample_offset + max_samples]


def _items(cache_root: Path, rows: list[dict]) -> list[dict]:
    by_shard = {}
    for position, row in enumerate(rows):
        by_shard.setdefault(row["cache_shard"], []).append((position, row))
    items = [None] * len(rows)
    for relative, selected in by_shard.items():
        shard = _load_shard(cache_root / relative)
        for position, row in selected:
            item = shard[row["cache_index"]]
            if item["sample_id"] != row["sample_id"]:
                raise ValueError(f"cache provenance mismatch for {row['sample_id']}")
            items[position] = item
    return items


class _ShardItemLoader:
    """Small mmap-backed shard LRU for dynamically claimed evaluation rows."""

    def __init__(self, cache_root: Path, max_open_shards: int = 2):
        self.cache_root = cache_root
        self.max_open_shards = max_open_shards
        self.shards = OrderedDict()

    def load(self, row: dict) -> dict:
        relative = row["cache_shard"]
        if relative not in self.shards:
            self.shards[relative] = _load_shard(self.cache_root / relative)
            if len(self.shards) > self.max_open_shards:
                self.shards.popitem(last=False)
        self.shards.move_to_end(relative)
        item = self.shards[relative][row["cache_index"]]
        if item["sample_id"] != row["sample_id"]:
            raise ValueError(f"cache provenance mismatch for {row['sample_id']}")
        return item


def _explicit_final(text: str) -> bool:
    return re.search(r"final\s+answer\s*:", text, flags=re.IGNORECASE) is not None


def _result_row(
    item: dict,
    manifest_row: dict,
    result: dict,
    *,
    mode: str,
    model_label: str,
) -> tuple[dict, list[torch.Tensor]]:
    raw_codes = result.pop("raw_codes", [])
    result.pop("target_codes", None)
    result.pop("intervention_images", None)
    sequence_ids = result.pop("sequence_ids").tolist()
    decoded = result["decoded"]
    score = score_answer(decoded, _target(item))
    has_final = _explicit_final(decoded) if mode == "interleaved" else True
    if mode == "understanding":
        events = [
            {
                "type": "text",
                "event": result["stop_reason"],
                "token_ids": sequence_ids,
            }
        ]
    else:
        events = result["events"]
    row = {
        **result,
        "sample_id": item["sample_id"],
        "evaluation_index": manifest_row["evaluation_index"],
        "evaluation_subset_index": manifest_row.get(SUBSET_FIELD),
        "config": item["config"],
        "mode": mode,
        "model_label": model_label,
        "target": _target(item),
        "sequence_ids": sequence_ids,
        "readable_sequence": _readable(decoded),
        "events": events,
        "score": score,
        "has_final_answer": has_final,
        "strict_correct": bool(score["correct"] and has_final),
    }
    return row, raw_codes


def _summarize(rows: list[dict], args: argparse.Namespace, world: int) -> dict:
    strict = sum(row["strict_correct"] for row in rows)
    image_rows = [row for row in rows if "expected_images" in row]
    by_expected_images = {}
    for count in sorted({int(row["expected_images"]) for row in image_rows}):
        selected = [
            row for row in image_rows if int(row["expected_images"]) == count
        ]
        correct = sum(row["strict_correct"] for row in selected)
        by_expected_images[str(count)] = {
            "samples": len(selected),
            "correct": correct,
            "accuracy": correct / len(selected),
            "expected_images": sum(int(row["expected_images"]) for row in selected),
            "generated_images": sum(int(row["generated_images"]) for row in selected),
        }
    complete_grids = sum(
        row.get("expected_grids") == row.get("generated_grids")
        for row in image_rows
    )
    text_tokens = sum(
        len(event.get("token_ids", []))
        for row in rows
        for event in row["events"]
        if event.get("type") == "text"
    )
    visual_tokens = sum(
        int(event.get("serialized_positions", 0))
        for row in rows
        for event in row["events"]
        if event.get("type") == "image"
    )
    return {
        "mode": args.mode,
        "model_label": args.model_label,
        "model_path": str(args.model_path),
        "adapter_dir": str(args.adapter_dir) if args.adapter_dir else None,
        "selection": args.selection,
        "subset_field": (
            SUBSET_FIELD if args.selection == "indexed" else f"split:{args.selection}"
        ),
        "samples": len(rows),
        "world_size": world,
        "workers_per_gpu": args.workers_per_gpu,
        "scheduling": args.scheduling,
        "text_cycle_max_period": args.text_cycle_max_period,
        "text_cycle_min_tokens": args.text_cycle_min_tokens,
        "correct": strict,
        "accuracy": strict / len(rows),
        "raw_scorer_correct": sum(row["score"]["correct"] for row in rows),
        "stop_reasons": dict(Counter(row["stop_reason"] for row in rows)),
        "expected_images": sum(int(row["expected_images"]) for row in image_rows),
        "generated_images": sum(int(row["generated_images"]) for row in image_rows),
        "by_expected_images": by_expected_images,
        "complete_grid_sequences": complete_grids,
        "complete_grid_sequence_rate": (
            complete_grids / len(image_rows) if image_rows else None
        ),
        "assistant_tokens": {
            "text": text_tokens,
            "visual": visual_tokens,
            "total": text_tokens + visual_tokens,
            "mean_text_per_sample": text_tokens / len(rows),
            "mean_visual_per_sample": visual_tokens / len(rows),
        },
        "decoder": {
            "enabled": bool(args.decode_media and args.mode == "interleaved"),
            "steps": args.decoder_steps,
            "guidance": args.decoder_guidance,
            "upsampling_ratio": args.decoder_upsampling,
        },
    }


def _merge_rank_shards(
    output_dir: Path,
    manifest_rows: list[dict],
    world: int,
    scheduling: str = "static",
) -> list[dict]:
    """Merge complete rank shards and enforce their scheduling contract."""
    expected_by_id = {row["sample_id"]: row for row in manifest_rows}
    merged = []
    for shard_rank in range(world):
        path = output_dir / "shards" / f"rank-{shard_rank:02d}" / "predictions.jsonl"
        shard_rows = [
            json.loads(line) for line in path.read_text().splitlines() if line
        ]
        if scheduling == "static":
            expected_indices = [
                row["evaluation_index"] for row in manifest_rows[shard_rank::world]
            ]
            actual_indices = [row["evaluation_index"] for row in shard_rows]
            if actual_indices != expected_indices:
                raise ValueError(
                    f"rank {shard_rank} boundary mismatch: "
                    f"expected={expected_indices}, actual={actual_indices}"
                )
        for row in shard_rows:
            expected = expected_by_id.get(row["sample_id"])
            if expected is None or row["evaluation_index"] != expected["evaluation_index"]:
                raise ValueError(f"manifest mismatch for {row['sample_id']}")
        merged.extend(shard_rows)
    merged.sort(key=lambda row: row["evaluation_index"])
    expected_indices = [row["evaluation_index"] for row in manifest_rows]
    actual_indices = [row["evaluation_index"] for row in merged]
    actual_ids = [row["sample_id"] for row in merged]
    expected_ids = [row["sample_id"] for row in manifest_rows]
    if actual_indices != expected_indices or actual_ids != expected_ids:
        raise ValueError("merged evaluation does not exactly match the manifest")
    return merged


def _publish_evaluation(
    output_dir: Path,
    manifest_rows: list[dict],
    args: argparse.Namespace,
    world: int,
) -> dict:
    merged = _merge_rank_shards(output_dir, manifest_rows, world, args.scheduling)
    with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle:
        for row in merged:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    summary = _summarize(merged, args, world)
    _write_json(output_dir / "summary.json", summary)
    print("EVALUATION_COMPLETE " + json.dumps(summary), flush=True)
    return summary


def _global_completed_ids(local_ids: list[str], world: int) -> set[str]:
    if world == 1:
        gathered = [local_ids]
    else:
        gathered = [None] * world
        dist.all_gather_object(gathered, local_ids)
    completed = [sample_id for rank_ids in gathered for sample_id in rank_ids]
    if len(completed) != len(set(completed)):
        raise ValueError("dynamic resume predictions contain duplicate samples")
    return set(completed)


def _dynamic_indices(
    output_dir: Path,
    pending_indices: list[int],
    rank: int,
    world: int,
):
    """Yield indices claimed from an atomic queue shared by local workers."""
    if world == 1:
        yield from pending_indices
        return

    queue_path = output_dir / ".worker-queue"
    if rank == 0:
        queue_path.unlink(missing_ok=True)
    _barrier(world)
    store = dist.FileStore(str(queue_path), world)
    if rank == 0:
        store.set("next", "0")
    _barrier(world)
    while True:
        position = store.add("next", 1) - 1
        if position >= len(pending_indices):
            break
        yield pending_indices[position]
    _barrier(world)
    del store
    _barrier(world)
    if rank == 0:
        queue_path.unlink(missing_ok=True)
    _barrier(world)


def _decode_media(
    args: argparse.Namespace,
    device: torch.device,
    rows: list[dict],
    shard_dir: Path,
    rank: int,
) -> None:
    if args.mode != "interleaved" or not args.decode_media:
        return
    decoder = None
    total = sum(row.get("generated_images", 0) for row in rows)
    completed = 0
    errors = []
    for row in rows:
        sample_dir = shard_dir / "samples" / row["sample_id"]
        codes_per_step = load_generated_codes(sample_dir / "generated_codes.pt")
        for step, (codes, grid) in enumerate(
            zip(codes_per_step, row.get("generated_grids", []), strict=True), start=1
        ):
            completed += 1
            destination = sample_dir / f"generated_{step:02d}.png"
            temporary = None
            try:
                if destination.exists():
                    try:
                        with Image.open(destination) as existing:
                            existing.verify()
                    except (OSError, ValueError):
                        pass
                    else:
                        print(
                            f"MEDIA rank={rank} {completed}/{total} existing",
                            flush=True,
                        )
                        continue
                if decoder is None:
                    decoder = UniARVisualDecoder.from_pretrained(
                        args.model_path, device=device
                    )
                generator = torch.Generator(device=device).manual_seed(
                    int.from_bytes(
                        hashlib.sha256(
                            f"20260827:{row['sample_id']}:{step}".encode()
                        ).digest()[:8],
                        "big",
                    )
                )
                binary = codes.to(device=device, dtype=torch.bfloat16)
                indices = decoder.bsq_encoder.bsq.codes_to_indexes(binary * 2 - 1)
                image = decoder.decode(
                    indices,
                    ar_height=grid[0] * 32,
                    ar_width=grid[1] * 32,
                    upsampling_ratio=args.decoder_upsampling,
                    num_inference_steps=args.decoder_steps,
                    guidance_scale=args.decoder_guidance,
                    generator=generator,
                )[0]
                temporary = destination.with_name(
                    f".{destination.name}.tmp-{os.getpid()}"
                )
                image.save(temporary, format="PNG")
                os.replace(temporary, destination)
            except (OSError, RuntimeError, ValueError) as error:
                # Preserve completed rollout statistics if RGB decoding fails.
                errors.append(
                    {
                        "sample_id": row["sample_id"],
                        "step": step,
                        "error": repr(error),
                    }
                )
                if temporary is not None:
                    temporary.unlink(missing_ok=True)
                torch.cuda.empty_cache()
            print(f"MEDIA rank={rank} {completed}/{total}", flush=True)
    if errors:
        with (shard_dir / "decode_errors.jsonl").open("w", encoding="utf-8") as handle:
            for error in errors:
                handle.write(json.dumps(error, ensure_ascii=False) + "\n")


@torch.inference_mode()
def evaluate(args: argparse.Namespace) -> None:
    _wait_for_adapter(args)
    args.resume = args.resume or (args.output_dir / ".run-claim.json").is_file()
    rank, world, _, device = _distributed(args.workers_per_gpu)
    if not is_timestamped_run_name(args.output_dir.name):
        raise ValueError("output directory must use a timestamped run name")
    if rank == 0:
        claim_run_directory(args.output_dir, resume=args.resume)
        output_lock = _acquire_output_lock(args.output_dir)
        arguments_path = args.output_dir / "arguments.json"
        current_arguments = {
            key: str(value) if isinstance(value, Path) else value
            for key, value in vars(args).items()
            if key != "resume"
        }
        current_arguments["_world_size"] = world
        current_arguments["_local_world_size"] = int(
            os.environ.get("LOCAL_WORLD_SIZE", "1")
        )
        if args.resume:
            if not arguments_path.exists():
                raise FileNotFoundError("resume output has no arguments.json")
            previous_arguments = json.loads(arguments_path.read_text())
            mismatched = {
                key: (previous_arguments.get(key), value)
                for key, value in current_arguments.items()
                if previous_arguments.get(key) != value
            }
            if mismatched:
                raise ValueError(f"resume arguments differ: {mismatched}")
        else:
            _write_json(arguments_path, current_arguments)
    _barrier(world)
    shard_dir = args.output_dir / "shards" / f"rank-{rank:02d}"
    shard_dir.mkdir(parents=True, exist_ok=args.resume)

    all_manifest_rows = _limit_manifest_rows(
        _manifest_rows(args.cache_root, args.selection),
        args.max_samples,
        args.sample_offset,
    )
    dynamic = args.scheduling == "dynamic"
    manifest_rows = all_manifest_rows if dynamic else all_manifest_rows[rank::world]
    predictions_path = shard_dir / "predictions.jsonl"
    rows = _load_prediction_rows(
        predictions_path,
        resume=args.resume,
        samples_dir=shard_dir / "samples",
    )
    completed_ids = [row["sample_id"] for row in rows]
    if len(completed_ids) != len(set(completed_ids)):
        raise ValueError(f"rank {rank} resume predictions contain duplicate samples")
    expected_ids = {row["sample_id"] for row in manifest_rows}
    unexpected = sorted(set(completed_ids) - expected_ids)
    if unexpected:
        raise ValueError(f"rank {rank} resume has unexpected samples: {unexpected}")
    manifest_by_id = {row["sample_id"]: row for row in manifest_rows}
    for row in rows:
        expected = manifest_by_id[row["sample_id"]]
        if row.get("evaluation_index") != expected["evaluation_index"]:
            raise ValueError(f"rank {rank} resume index mismatch for {row['sample_id']}")
    if dynamic:
        globally_completed = _global_completed_ids(completed_ids, world)
        remaining_indices = sorted(
            [
                index
                for index, row in enumerate(all_manifest_rows)
                if row["sample_id"] not in globally_completed
            ],
            key=lambda index: (
                all_manifest_rows[index]["cache_shard"],
                -int(all_manifest_rows[index]["sequence_length"]),
                all_manifest_rows[index]["sample_id"],
            ),
        )
        remaining_count = len(remaining_indices)
    else:
        completed = set(completed_ids)
        remaining_rows = [
            row for row in manifest_rows if row["sample_id"] not in completed
        ]
        remaining_count = len(remaining_rows)

    collective_device = _control_collective_device(world, device)
    remaining_tensor = torch.tensor(
        remaining_count, device=collective_device, dtype=torch.long
    )
    if world > 1:
        dist.all_reduce(remaining_tensor, op=dist.ReduceOp.MAX if dynamic else dist.ReduceOp.SUM)
    if int(remaining_tensor) == 0:
        _barrier(world)
        if rank == 0:
            _publish_evaluation(args.output_dir, all_manifest_rows, args, world)
        _barrier(world)
        _decode_media(args, device, rows, shard_dir, rank)
        _barrier(world)
        if world > 1:
            dist.destroy_process_group()
        if rank == 0:
            output_lock.close()
        return

    processor = configure_processor(
        AutoProcessor.from_pretrained(args.model_path, use_fast=True)
    )
    base = (
        UniARForConditionalGeneration.from_pretrained(
            args.model_path,
            dtype=torch.bfloat16,
            attn_implementation=args.attn,
            low_cpu_mem_usage=True,
        )
        .to(device)
        .eval()
    )
    model = base
    if args.mode == "interleaved":
        if args.adapter_dir is None:
            raise ValueError("interleaved mode requires --adapter-dir")
        model = PeftModel.from_pretrained(base, args.adapter_dir).eval()
    if dynamic:
        loader = _ShardItemLoader(args.cache_root)
        remaining = (
            (loader.load(all_manifest_rows[index]), all_manifest_rows[index])
            for index in _dynamic_indices(
                args.output_dir, remaining_indices, rank, world
            )
        )
    else:
        remaining_items = _items(args.cache_root, remaining_rows)
        remaining = iter(zip(remaining_items, remaining_rows, strict=True))
    print(
        f"MODEL_READY rank={rank} mode={args.mode} samples={len(manifest_rows)} "
        f"remaining={remaining_count} scheduling={args.scheduling} device={device}",
        flush=True,
    )

    rollout = CachedInterleavedRollout(model, processor)
    file_mode = "a" if args.resume else "w"
    with predictions_path.open(file_mode, encoding="utf-8") as predictions:
        for local_index, (item, manifest_row) in enumerate(remaining, start=1):
            if args.mode == "understanding":
                result = rollout.direct_answer(
                    item, max_new_tokens=args.max_understanding_tokens
                )
            else:
                result = rollout.run(
                    item,
                    max_text_tokens=None,
                    max_scaffold_tokens=None,
                    extra_image_budget=None,
                    cycle_repeats=args.cycle_repeats,
                    text_cycle_max_period=args.text_cycle_max_period,
                    text_cycle_min_tokens=args.text_cycle_min_tokens,
                    include_target_codes=False,
                )
            row, raw_codes = _result_row(
                item,
                manifest_row,
                result,
                mode=args.mode,
                model_label=args.model_label,
            )
            rows.append(row)
            sample_dir = shard_dir / "samples" / item["sample_id"]
            sample_dir.mkdir(parents=True, exist_ok=args.resume)
            _write_text(sample_dir / "target.txt", item["target_render"])
            _write_text(sample_dir / "output.txt", row["readable_sequence"])
            _write_json(sample_dir / "result.json", row)
            save_generated_codes(sample_dir / "generated_codes.pt", raw_codes)
            predictions.write(json.dumps(row, ensure_ascii=False) + "\n")
            predictions.flush()
            os.fsync(predictions.fileno())
            print(
                f"PROGRESS local={local_index} rank={rank} "
                f"index={row['evaluation_index']} correct={row['strict_correct']} "
                f"stop={row['stop_reason']} images={row.get('generated_images', 0)}",
                flush=True,
            )

    _barrier(world)
    if rank == 0:
        _publish_evaluation(args.output_dir, all_manifest_rows, args, world)
    _barrier(world)

    del rollout, model, base
    gc.collect()
    torch.cuda.empty_cache()
    _decode_media(args, device, rows, shard_dir, rank)
    _barrier(world)
    if world > 1:
        dist.destroy_process_group()
    if rank == 0:
        output_lock.close()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--mode", choices=["understanding", "interleaved"], required=True
    )
    parser.add_argument("--model-label", required=True)
    parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT"))
    parser.add_argument("--adapter-dir", type=Path)
    parser.add_argument(
        "--wait-for-adapter",
        action="store_true",
        help="wait for adapter/COMPLETE so a separately queued SFT can finish",
    )
    parser.add_argument("--cache-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--selection",
        choices=["indexed", "validation", "test"],
        default="indexed",
        help="evaluate indexed test100 or a complete validation/test split",
    )
    parser.add_argument(
        "--max-samples",
        type=int,
        default=None,
        help="take this many rows after --sample-offset",
    )
    parser.add_argument(
        "--sample-offset",
        type=int,
        default=0,
        help="skip this many deterministic manifest-order rows",
    )
    parser.add_argument("--attn", default="flash_attention_2")
    parser.add_argument("--max-understanding-tokens", type=int, default=256)
    parser.add_argument("--cycle-repeats", type=int, default=3)
    parser.add_argument("--text-cycle-max-period", type=int, default=128)
    parser.add_argument("--text-cycle-min-tokens", type=int, default=64)
    parser.add_argument(
        "--workers-per-gpu",
        type=int,
        default=1,
        help="share each visible GPU across this many torchrun worker processes",
    )
    parser.add_argument(
        "--scheduling",
        choices=["static", "dynamic"],
        default="static",
        help="use strided rank shards or an atomic work queue",
    )
    parser.add_argument(
        "--decode-media", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument("--decoder-steps", type=int, default=28)
    parser.add_argument("--decoder-guidance", type=float, default=1.5)
    parser.add_argument("--decoder-upsampling", type=float, default=1.0)
    parser.add_argument(
        "--resume",
        action="store_true",
        help="append only missing samples to an interrupted output directory",
    )
    return parser


def main() -> None:
    evaluate(build_parser().parse_args())


if __name__ == "__main__":
    main()
Xet Storage Details

Size:: 32.5 kB
Xet hash:: fd87c685f6b258619040a48096479a162cd8383a3757bfc1145c719713671e24
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.