| import argparse |
| import json |
| import logging |
|
|
| import torch |
| import transformers |
|
|
| import jlens |
| from revisions import revision |
|
|
| MODEL = "Qwen/Qwen3.6-27B" |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("worker", type=int) |
| parser.add_argument("n_workers", type=int) |
| parser.add_argument("--prompts", default="prompts_1000.json") |
| parser.add_argument("--limit", type=int, default=None) |
| parser.add_argument("--prefix", default="w") |
| parser.add_argument("--dim-batch", type=int, default=32) |
| args = parser.parse_args() |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") |
|
|
| with open(args.prompts) as f: |
| prompts = json.load(f)[args.worker :: args.n_workers] |
| if args.limit: |
| prompts = prompts[: args.limit] |
| print(f"worker {args.worker}/{args.n_workers}: {len(prompts)} prompts", flush=True) |
|
|
| try: |
| hf = transformers.AutoModelForCausalLM.from_pretrained(MODEL, revision=revision(MODEL), dtype=torch.bfloat16) |
| except (ValueError, KeyError) as exc: |
| print(f"AutoModelForCausalLM failed ({exc}); trying AutoModelForImageTextToText", flush=True) |
| hf = transformers.AutoModelForImageTextToText.from_pretrained(MODEL, revision=revision(MODEL), dtype=torch.bfloat16) |
| hf = hf.cuda() |
| tok = transformers.AutoTokenizer.from_pretrained(MODEL, revision=revision(MODEL)) |
| model = jlens.from_hf(hf, tok) |
| print(model, flush=True) |
|
|
| lens = jlens.fit( |
| model, |
| prompts, |
| dim_batch=args.dim_batch, |
| max_seq_len=128, |
| checkpoint_path=f"ckpt_{args.prefix}{args.worker}.pt", |
| checkpoint_every=5, |
| ) |
| lens.save(f"lens_{args.prefix}{args.worker}.pt") |
| print( |
| f"peak GPU mem: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB; " |
| f"saved lens_{args.prefix}{args.worker}.pt", |
| flush=True, |
| ) |
|
|
