math_jlens/cli.py

"""Command-line entry point for fitting DAPO Jacobians."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from .corpus import infer_corpus_format, load_fitting_corpus
from .fitting import fit
from .model import load_qwen


def build_parser() -> argparse.ArgumentParser:
    command = argparse.ArgumentParser(description=__doc__)
    command.add_argument("--model", required=True)
    command.add_argument("--data", required=True)
    command.add_argument("--output-dir", default="outputs/jacobians")
    command.add_argument("--num-prompts", type=int, default=20)
    command.add_argument("--offset", type=int, default=0)
    command.add_argument("--seed", type=int, default=17)
    command.add_argument(
        "--corpus-format",
        choices=["auto", "dapo-jsonl", "rollout-parquet"],
        default="auto",
    )
    command.add_argument("--response-window-len", type=int, default=1024)
    command.add_argument("--max-seq-len", type=int)
    command.add_argument("--skip-first", type=int, default=16)
    command.add_argument("--dim-batch", type=int, default=8)
    command.add_argument("--target-layer", type=int)
    command.add_argument("--checkpoint-every", type=int, default=1)
    command.add_argument("--device", default="cuda")
    command.add_argument("--no-resume", action="store_true")
    return command


def main() -> None:
    args = build_parser().parse_args()
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    corpus_format = infer_corpus_format(args.data, args.corpus_format)
    max_seq_len = args.max_seq_len
    if max_seq_len is None:
        max_seq_len = args.response_window_len if corpus_format == "rollout-parquet" else 128
    if max_seq_len < 2:
        raise ValueError("max sequence length must be at least 2")
    prompts = load_fitting_corpus(
        args.data,
        corpus_format=corpus_format,
        count=args.num_prompts,
        seed=args.seed,
        offset=args.offset,
        response_window_length=args.response_window_len,
    )
    if corpus_format == "rollout-parquet":
        lengths = [len(item) for item in prompts]
        print(
            f"loaded {len(prompts)} Parquet response windows: "
            f"token lengths {min(lengths)}..{max(lengths)}, "
            f"seed={args.seed}, offset={args.offset}"
        )
    model = load_qwen(args.model, device=args.device, dtype=torch.bfloat16)
    target = model.n_layers - 1 if args.target_layer is None else args.target_layer
    configuration = vars(args) | {
        "resolved_corpus_format": corpus_format,
        "resolved_max_seq_len": max_seq_len,
        "resolved_target_layer": target, "source_layers": list(range(target)),
        "model_layers": model.n_layers, "d_model": model.d_model,
    }
    (output / "config.json").write_text(json.dumps(configuration, indent=2), encoding="utf-8")
    fit(
        model, prompts, output_path=str(output / "lens-bf16.pt"),
        checkpoint_path=str(output / "fit-checkpoint-fp32.pt"), target_layer=target,
        max_seq_len=max_seq_len, dim_batch=args.dim_batch,
        skip_first=args.skip_first, checkpoint_every=args.checkpoint_every,
        resume=not args.no_resume, export_dtype=torch.bfloat16,
    )


if __name__ == "__main__":
    main()