File size: 3,320 Bytes
85b17bd f6158c7 85b17bd f6158c7 85b17bd f6158c7 85b17bd f6158c7 85b17bd f6158c7 85b17bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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()
|