File size: 4,613 Bytes
31f7037 | 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 86 87 88 89 90 91 92 | """Resident CUDA inference for native Julia and CUDA INT8 checkpoints."""
import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from .cuda import configure, move
from .data import Collator, validate_row
from .probabilities import display_probabilities
def context_length(checkpoint, requested):
config = json.loads((Path(checkpoint) / 'encoder/config.json').read_text())
limit = config['max_position_embeddings']
value = limit if requested is None else requested
if type(value) is not int or not 1 <= value <= limit:
raise ValueError(f'max_length must be an integer between 1 and {limit}')
return value
class TransformerEngine:
def __init__(self, checkpoint, device='cuda', max_length=None, head_length=256, *, memory_map=True):
from transformers import AutoModel, AutoTokenizer
from .model import JuliaDecisionModel
self.device = configure(device)
root = Path(checkpoint)
max_length = context_length(root, max_length)
if (root / 'INCOMPLETE').exists():
raise ValueError('Refusing to load an incomplete INT8 export')
self.tokenizer = AutoTokenizer.from_pretrained(root / 'tokenizer', trust_remote_code=False)
if (root / 'quantization.json').exists():
if self.device.type != 'cuda':
raise ValueError('This INT8 checkpoint requires CUDA')
encoder = AutoModel.from_pretrained(root / 'encoder', device_map={'': self.device.index or 0},
attn_implementation='sdpa', trust_remote_code=False)
config = json.loads((root / 'julia_config.json').read_text())
self.model = JuliaDecisionModel(encoder, **{k: config[k] for k in ('head_layers', 'n_act', 'dropout')})
heads = load_file(str(root / 'heads.safetensors'))
expected = {k for k in self.model.state_dict() if not k.startswith('encoder.')}
if set(heads) != expected:
raise ValueError('INT8 decision-head tensor keys do not match checkpoint architecture')
self.model.load_state_dict(heads, strict=False)
for name, child in self.model.named_children():
if name != 'encoder':
child.to(self.device)
else:
self.model = JuliaDecisionModel.from_pretrained(root, memory_map=memory_map).to(self.device)
self.model.eval()
self.collate = Collator(self.tokenizer, max_length, head_length)
@torch.inference_mode()
def logits(self, rows):
if not rows:
return []
for i, row in enumerate(rows):
validate_row(row, i + 1)
batch = move(self.collate(rows, include_targets=False), self.device)
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=self.device.type == 'cuda'):
logits = self.model(**batch)
logits = logits.cpu()
if not torch.isfinite(logits).all():
raise FloatingPointError('Inference returned nonfinite logits')
return [values[:len(row['options'])].tolist() for values, row in zip(logits, rows)]
def predict(self, rows=None, questions=None, *, state=None):
if questions is not None:
from .typed import predict_typed
if rows is not None and state is not None:
raise ValueError('Pass state either positionally or by keyword, not both')
return predict_typed(self, state if rows is None else rows, questions)
if state is not None or rows is None:
raise ValueError('Provide legacy rows or state with questions')
result = []
for values in self.logits(rows):
probabilities = torch.tensor(values).softmax(-1)
result.append(dict(index=int(probabilities.argmax()),
probabilities=display_probabilities(probabilities.tolist())))
return result
def load_model(checkpoint, device='cpu', max_length=None, head_length=256, *, backend=None, **kwargs):
"""Load a Julia checkpoint through the supported resident inference runtime."""
from .router.engine import FastEngine
if checkpoint is None:
raise ValueError('checkpoint is required; pass a local model directory')
if backend not in (None, 'torch', 'bend', 'bend-dense'):
raise ValueError('backend must be torch, bend, or bend-dense')
return FastEngine(checkpoint, device, max_length, head_length,
transformer_backend=backend, **kwargs)
Engine = load_model # Backward compatible entry point.
|