Buckets:
| """LoRA overfit runner for UniAR joint text/visual Zebra-CoT SFT.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import random | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| from peft import LoraConfig, TaskType, get_peft_model | |
| from transformers import AutoProcessor | |
| from artifacts import is_timestamped_run_name | |
| from uniar import UniARForConditionalGeneration, configure_processor | |
| from .dataset import ZebraTrajectoryDataset, image_gen_token_id, prepare_processor_inputs | |
| from .objective import encode_visual_codes, joint_sft_loss, unwrap_ar_model | |
| LORA_TARGETS = [ | |
| "q_proj", | |
| "k_proj", | |
| "v_proj", | |
| "o_proj", | |
| "gate_proj", | |
| "up_proj", | |
| "down_proj", | |
| "embed_tokens", | |
| "lm_head", | |
| ] | |
| def _jsonable(value): | |
| if isinstance(value, Path): | |
| return str(value) | |
| if isinstance(value, torch.Tensor): | |
| return value.detach().float().cpu().item() if value.numel() == 1 else value.tolist() | |
| return value | |
| def _write_json(path: Path, value) -> None: | |
| path.write_text( | |
| json.dumps(value, indent=2, ensure_ascii=False, default=_jsonable) + "\n", | |
| encoding="utf-8", | |
| ) | |
| def _metric_row(metrics: dict) -> dict: | |
| return { | |
| key: (_jsonable(value) if not isinstance(value, float) else value) | |
| for key, value in metrics.items() | |
| if key not in {"loss", "visual_predictions", "visual_targets"} | |
| } | {"loss": float(metrics["loss"].detach().float().cpu())} | |
| def evaluate( | |
| model, | |
| prepared: list[tuple[dict, torch.Tensor]], | |
| marker_id: int, | |
| visual_weight: float, | |
| marker_weight: float, | |
| ): | |
| model.eval() | |
| unwrap_ar_model(model).model.visual.eval() | |
| rows = [] | |
| for batch, codes in prepared: | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| result = joint_sft_loss( | |
| model, | |
| batch, | |
| codes, | |
| image_gen_id=marker_id, | |
| visual_loss_weight=visual_weight, | |
| marker_token_weight=marker_weight, | |
| ) | |
| rows.append(_metric_row(result) | {"sample_id": batch["sample_id"]}) | |
| aggregate = {} | |
| for key in [ | |
| "loss", | |
| "text_loss", | |
| "marker_loss", | |
| "visual_loss", | |
| "text_accuracy", | |
| "visual_accuracy", | |
| "marker_accuracy", | |
| ]: | |
| values = [float(row[key]) for row in rows if not math.isnan(float(row[key]))] | |
| aggregate[key] = sum(values) / len(values) | |
| aggregate["samples"] = len(rows) | |
| aggregate["markers"] = sum(row["markers"] for row in rows) | |
| return aggregate, rows | |
| def _prepare_dataset(model, processor, dataset: ZebraTrajectoryDataset, output_dir: Path): | |
| ar_model = unwrap_ar_model(model) | |
| cache_dir = output_dir / "bsq_cache" | |
| cache_dir.mkdir() | |
| prepared = [] | |
| rendered_dir = output_dir / "rendered_inputs" | |
| rendered_dir.mkdir() | |
| for index, trajectory in enumerate(dataset.records): | |
| batch = prepare_processor_inputs( | |
| processor, trajectory, dataset.root, ar_model.config.image_token_id | |
| ) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| codes = encode_visual_codes( | |
| ar_model, batch.pop("pixel_values"), batch["image_grid_thw"] | |
| ).cpu() | |
| torch.save( | |
| {"codes": codes, "image_grid_thw": batch["image_grid_thw"]}, | |
| cache_dir / f"{batch['sample_id']}.pt", | |
| ) | |
| (rendered_dir / f"{batch['sample_id']}.txt").write_text( | |
| batch.pop("rendered_text"), encoding="utf-8" | |
| ) | |
| prepared.append((batch, codes)) | |
| print( | |
| f"prepared {index + 1}/{len(dataset)} {batch['sample_id']} " | |
| f"seq={batch['input_ids'].shape[1]} visual={codes.shape[0]}", | |
| flush=True, | |
| ) | |
| return prepared | |
| def train(args: argparse.Namespace) -> dict: | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("Zebra UniAR SFT requires CUDA") | |
| if not is_timestamped_run_name(args.output_dir.name): | |
| raise ValueError("output directory must use the canonical timestamped run name") | |
| args.output_dir.mkdir(parents=True, exist_ok=False) | |
| torch.manual_seed(args.seed) | |
| random.seed(args.seed) | |
| processor = configure_processor( | |
| AutoProcessor.from_pretrained(args.model_path, use_fast=True) | |
| ) | |
| marker_id = image_gen_token_id(processor) | |
| model = UniARForConditionalGeneration.from_pretrained( | |
| args.model_path, | |
| torch_dtype=torch.bfloat16, | |
| attn_implementation=args.attn, | |
| low_cpu_mem_usage=True, | |
| ).to("cuda") | |
| model.config.use_cache = False | |
| model.model.visual.requires_grad_(False) | |
| model.model.visual.eval() | |
| lora_config = LoraConfig( | |
| task_type=TaskType.CAUSAL_LM, | |
| r=args.lora_rank, | |
| lora_alpha=args.lora_alpha, | |
| lora_dropout=0.0, | |
| bias="none", | |
| target_modules=LORA_TARGETS, | |
| ) | |
| model = get_peft_model(model, lora_config) | |
| if args.gradient_checkpointing: | |
| model.gradient_checkpointing_enable( | |
| gradient_checkpointing_kwargs={"use_reentrant": False} | |
| ) | |
| model.enable_input_require_grads() | |
| model.print_trainable_parameters() | |
| dataset = ZebraTrajectoryDataset(args.samples_root) | |
| if len(dataset) != 16: | |
| raise ValueError(f"overfit fixture must contain exactly 16 trajectories, got {len(dataset)}") | |
| if not 1 <= args.limit_samples <= 16: | |
| raise ValueError("limit_samples must be between 1 and 16") | |
| dataset.records = dataset.records[: args.limit_samples] | |
| prepared = _prepare_dataset(model, processor, dataset, args.output_dir) | |
| trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] | |
| optimizer = torch.optim.AdamW( | |
| trainable, lr=args.learning_rate, betas=(0.9, 0.95), weight_decay=0.0 | |
| ) | |
| trainable_parameters = sum(parameter.numel() for parameter in trainable) | |
| total_parameters = sum(parameter.numel() for parameter in model.parameters()) | |
| effective = { | |
| "args": vars(args), | |
| "environment": { | |
| "python": sys.version, | |
| "torch": torch.__version__, | |
| "cuda": torch.version.cuda, | |
| "gpu": torch.cuda.get_device_name(), | |
| "visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), | |
| }, | |
| "tokenizer": { | |
| "class": type(processor.tokenizer).__name__, | |
| "image_gen_id": marker_id, | |
| "image_gen_encoding": processor.tokenizer.encode( | |
| "<image_gen>", add_special_tokens=False | |
| ), | |
| }, | |
| "parameters": { | |
| "trainable": trainable_parameters, | |
| "total": total_parameters, | |
| "fraction": trainable_parameters / total_parameters, | |
| }, | |
| "lora_targets": LORA_TARGETS, | |
| } | |
| _write_json(args.output_dir / "config.json", effective) | |
| baseline, baseline_rows = evaluate( | |
| model, | |
| prepared, | |
| marker_id, | |
| args.visual_loss_weight, | |
| args.marker_token_weight, | |
| ) | |
| _write_json(args.output_dir / "baseline_metrics.json", baseline) | |
| _write_json(args.output_dir / "baseline_per_sample.json", baseline_rows) | |
| print("baseline", json.dumps(baseline), flush=True) | |
| log_path = args.output_dir / "train_log.jsonl" | |
| order = list(range(len(prepared))) | |
| model.train() | |
| unwrap_ar_model(model).model.visual.eval() | |
| for step in range(1, args.max_steps + 1): | |
| if (step - 1) % len(order) == 0: | |
| random.shuffle(order) | |
| sample_index = order[(step - 1) % len(order)] | |
| batch, codes = prepared[sample_index] | |
| optimizer.zero_grad(set_to_none=True) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| result = joint_sft_loss( | |
| model, | |
| batch, | |
| codes, | |
| image_gen_id=marker_id, | |
| visual_loss_weight=args.visual_loss_weight, | |
| marker_token_weight=args.marker_token_weight, | |
| ) | |
| result["loss"].backward() | |
| grad_norm = torch.nn.utils.clip_grad_norm_(trainable, args.max_grad_norm) | |
| optimizer.step() | |
| row = _metric_row(result) | { | |
| "step": step, | |
| "sample_id": batch["sample_id"], | |
| "grad_norm": float(grad_norm.detach().float().cpu()), | |
| "learning_rate": args.learning_rate, | |
| } | |
| with log_path.open("a", encoding="utf-8") as handle: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| print( | |
| f"step={step}/{args.max_steps} sample={batch['sample_id']} " | |
| f"loss={row['loss']:.4f} text={row['text_loss']:.4f} " | |
| f"visual={row['visual_loss']:.4f} marker={row['marker_accuracy']:.3f} " | |
| f"v_acc={row['visual_accuracy']:.3f}", | |
| flush=True, | |
| ) | |
| final, final_rows = evaluate( | |
| model, | |
| prepared, | |
| marker_id, | |
| args.visual_loss_weight, | |
| args.marker_token_weight, | |
| ) | |
| _write_json(args.output_dir / "final_metrics.json", final) | |
| _write_json(args.output_dir / "final_per_sample.json", final_rows) | |
| adapter_dir = args.output_dir / "adapter" | |
| model.save_pretrained( | |
| adapter_dir, safe_serialization=True, save_embedding_layers=False | |
| ) | |
| processor.save_pretrained(adapter_dir) | |
| model.eval() | |
| first_batch, first_codes = prepared[0] | |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): | |
| prediction = joint_sft_loss( | |
| model, | |
| first_batch, | |
| first_codes, | |
| image_gen_id=marker_id, | |
| visual_loss_weight=args.visual_loss_weight, | |
| marker_token_weight=args.marker_token_weight, | |
| return_predictions=True, | |
| ) | |
| first_length = first_batch["generated_image_token_counts"][0] | |
| predicted_bits = prediction["visual_predictions"][:first_length].cpu().to(torch.uint8) | |
| target_bits = prediction["visual_targets"][:first_length].cpu().to(torch.uint8) | |
| torch.save( | |
| {"predicted_bits": predicted_bits, "target_bits": target_bits}, | |
| args.output_dir / "teacher_forced_first_image_bits.pt", | |
| ) | |
| summary = { | |
| "baseline": baseline, | |
| "final": final, | |
| "improvement": { | |
| "loss": baseline["loss"] - final["loss"], | |
| "text_loss": baseline["text_loss"] - final["text_loss"], | |
| "marker_loss": baseline["marker_loss"] - final["marker_loss"], | |
| "visual_loss": baseline["visual_loss"] - final["visual_loss"], | |
| "marker_accuracy": final["marker_accuracy"] - baseline["marker_accuracy"], | |
| "visual_accuracy": final["visual_accuracy"] - baseline["visual_accuracy"], | |
| }, | |
| "adapter": str(adapter_dir), | |
| "peak_gpu_memory_gib": torch.cuda.max_memory_allocated() / 1024**3, | |
| } | |
| _write_json(args.output_dir / "summary.json", summary) | |
| print("final", json.dumps(summary, default=_jsonable), flush=True) | |
| return summary | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT")) | |
| parser.add_argument("--samples-root", type=Path, default=Path("data/Zebra-CoT-samples")) | |
| parser.add_argument("--output-dir", type=Path, required=True) | |
| parser.add_argument("--attn", default="flash_attention_2") | |
| parser.add_argument("--max-steps", type=int, default=96) | |
| parser.add_argument("--limit-samples", type=int, default=16) | |
| parser.add_argument("--learning-rate", type=float, default=2e-4) | |
| parser.add_argument("--visual-loss-weight", type=float, default=1.0) | |
| parser.add_argument("--marker-token-weight", type=float, default=20.0) | |
| parser.add_argument("--lora-rank", type=int, default=16) | |
| parser.add_argument("--lora-alpha", type=int, default=32) | |
| parser.add_argument("--max-grad-norm", type=float, default=1.0) | |
| parser.add_argument("--seed", type=int, default=17) | |
| parser.add_argument("--gradient-checkpointing", action=argparse.BooleanOptionalAction, default=True) | |
| return parser | |
| def main() -> None: | |
| train(build_parser().parse_args()) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 12.2 kB
- Xet hash:
- dc79097f5f80818c8c5695bb04d902e3152b114f126ff7a5e36db08231779eca
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.