| |
| """Exercise the four-prompt T4 inference path on training prompts, never blind data.""" |
| from pathlib import Path |
| import json |
| import sys |
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT / 'src')) |
| from unsloth import FastLanguageModel |
| import torch |
| from bitsandbytes.nn import Linear4bit |
| from generate import UnslothGenerator |
| from vibegame import BASE_MODEL |
|
|
| if not torch.cuda.is_available() or 'T4' not in torch.cuda.get_device_name(): |
| raise RuntimeError('This capacity check requires the already allocated free T4') |
| model, tokenizer = FastLanguageModel.from_pretrained(model_name=BASE_MODEL, |
| revision='4858886f896bd05db823557fc5cbc4ac28342af7', max_seq_length=4096, |
| dtype=torch.float16, load_in_4bit=True, trust_remote_code=False, use_exact_model_name=True) |
| for module in model.modules(): |
| if isinstance(module, Linear4bit): |
| module.compute_dtype = torch.float16 |
| module.compute_type_is_set = True |
| FastLanguageModel.for_inference(model) |
| backend = UnslothGenerator.__new__(UnslothGenerator) |
| backend.model, backend.tokenizer, backend.torch = model, tokenizer, torch |
| prompts = [json.loads(line)['instruction'] for line in (ROOT/'data/teacher_prompts.jsonl').read_text().splitlines()][:4] |
| torch.cuda.reset_peak_memory_stats() |
| results = backend.generate_many(prompts, 3072) |
| report = {'kind':'untrained_base_capacity_smoke_on_training_prompts_not_model_eval', |
| 'batch_size':4, 'max_new_tokens':3072, 'peak_allocated_bytes':torch.cuda.max_memory_allocated(), |
| 'results': [{k:v for k,v in result.items() if k != 'output'} for result in results]} |
| target = Path(sys.argv[1]) |
| target.parent.mkdir(parents=True, exist_ok=True) |
| target.write_text(json.dumps(report, indent=2)+'\n') |
| print(json.dumps(report, indent=2), flush=True) |
|
|
