Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import json | |
| import os | |
| import threading | |
| import time | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from fastapi import HTTPException | |
| from fastapi.responses import FileResponse, HTMLResponse | |
| from gradio import Error, Server | |
| from huggingface_hub import snapshot_download | |
| from safetensors.torch import load_file | |
| from transformers import ( | |
| AutoConfig, | |
| AutoTokenizer, | |
| LlamaForCausalLM, | |
| StoppingCriteria, | |
| StoppingCriteriaList, | |
| TextIteratorStreamer, | |
| ) | |
| from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding | |
| REPO = "multimodalart/computer-10-fp8" | |
| FRAMES = { | |
| "stranger": "As follows is a conversation between another user and Model C.", | |
| "same": "As follows is a later conversation between the same user and Model C.", | |
| "reader": "As follows is a conversation between a reader of Model C's originating document and Model C.", | |
| } | |
| MAX_CONTEXT = 6144 | |
| MAX_MESSAGE_CHARS = 4000 | |
| FP8 = torch.float8_e4m3fn | |
| class FusedFP8(nn.Module): | |
| mode = None | |
| def __init__(self, weight, scale, splits): | |
| super().__init__() | |
| self.weight = nn.Parameter(weight, requires_grad=False) | |
| self.register_buffer("weight_scale", scale) | |
| self.register_buffer("scale_row", scale[None, :].contiguous()) | |
| self.splits = splits | |
| self.last_x = None | |
| self.parts = None | |
| def dequant(self, x): | |
| return F.linear(x, self.weight.to(x.dtype) * self.weight_scale[:, None].to(x.dtype)) | |
| def run(self, x, mode): | |
| if mode == "dequant": | |
| return self.dequant(x) | |
| m = x.shape[0] | |
| xs = x.abs().amax(dim=1, keepdim=True).float() / 448.0 | |
| xs = xs.clamp_(min=1e-12) | |
| xq = (x / xs).to(FP8) | |
| pad = (-m) % 16 if mode.endswith("pad") else 0 | |
| if pad: | |
| xq = F.pad(xq, (0, 0, 0, pad)) | |
| if mode.startswith("rowwise"): | |
| sa = F.pad(xs, (0, 0, 0, pad), value=1.0) if pad else xs | |
| return torch._scaled_mm(xq, self.weight.t(), scale_a=sa, scale_b=self.scale_row, out_dtype=torch.bfloat16)[:m] | |
| one = torch.ones((), device=x.device, dtype=torch.float32) | |
| y = torch._scaled_mm(xq, self.weight.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16)[:m] | |
| return (y.float() * xs * self.scale_row).to(torch.bfloat16) | |
| def get(self, x, index): | |
| if self.last_x is not x: | |
| shape = x.shape | |
| y = self.run(x.reshape(-1, shape[-1]), FusedFP8.mode) | |
| self.parts = [p.reshape(*shape[:-1], p.shape[-1]) for p in y.split(self.splits, dim=-1)] | |
| self.last_x = x | |
| out = self.parts[index] | |
| if index == len(self.splits) - 1: | |
| self.last_x, self.parts = None, None | |
| return out | |
| class FusedSlice(nn.Module): | |
| def __init__(self, fused, index): | |
| super().__init__() | |
| self.index = index | |
| self.__dict__["fused"] = fused | |
| def forward(self, x): | |
| return self.fused.get(x, self.index) | |
| GROUPS = (("self_attn", ("q_proj", "k_proj", "v_proj")), ("self_attn", ("o_proj",)), ("mlp", ("gate_proj", "up_proj")), ("mlp", ("down_proj",))) | |
| def fuse_layer(layer, sd, prefix): | |
| for parent_name, names in GROUPS: | |
| parent = getattr(layer, parent_name) | |
| keys = [f"{prefix}.{parent_name}.{n}" for n in names] | |
| weight = torch.cat([sd.pop(k + ".weight") for k in keys]) | |
| scale = torch.cat([sd.pop(k + ".weight_scale") for k in keys]) | |
| fused = FusedFP8(weight, scale, [getattr(parent, n).out_features for n in names]) | |
| parent.add_module("fused_" + "_".join(names), fused) | |
| for i, n in enumerate(names): | |
| setattr(parent, n, FusedSlice(fused, i)) | |
| def build_model(): | |
| path = snapshot_download(REPO) | |
| config = AutoConfig.from_pretrained(path) | |
| config.use_cache = True | |
| with torch.device("meta"): | |
| model = LlamaForCausalLM(config).to(torch.bfloat16) | |
| index = json.load(open(os.path.join(path, "model.safetensors.index.json"))) | |
| pending = {} | |
| for shard in sorted(set(index["weight_map"].values())): | |
| pending.update(load_file(os.path.join(path, shard))) | |
| for i, layer in enumerate(model.model.layers): | |
| prefix = f"model.layers.{i}" | |
| needed = [k for k in index["weight_map"] if k.startswith(prefix + ".") and "_proj." in k] | |
| if not isinstance(layer.self_attn.q_proj, FusedSlice) and all(k in pending for k in needed): | |
| fuse_layer(layer, pending, prefix) | |
| rest = {k: pending.pop(k) for k in list(pending) if "_proj." not in k} | |
| if rest: | |
| result = model.load_state_dict(rest, strict=False, assign=True) | |
| assert not result.unexpected_keys, result.unexpected_keys[:5] | |
| assert not pending, list(pending)[:5] | |
| model.model.rotary_emb = LlamaRotaryEmbedding(config=config) | |
| leftover = [n for n, p in list(model.named_parameters()) + list(model.named_buffers()) if p.is_meta] | |
| assert not leftover, leftover[:5] | |
| return model.eval() | |
| tokenizer = AutoTokenizer.from_pretrained(REPO) | |
| model = build_model().to("cuda") | |
| PAGE = (Path(__file__).parent / "index.html").read_text() | |
| FONTS = {p.name: p for p in (Path(__file__).parent / "fonts").glob("*.woff2")} | |
| def pick_kernel(): | |
| if FusedFP8.mode is not None: | |
| return | |
| layer = model.model.layers[0].mlp.fused_gate_proj_up_proj | |
| for m in (1, 7): | |
| x = torch.randn(m, layer.weight.shape[1], device="cuda", dtype=torch.bfloat16) | |
| ref = layer.dequant(x).float() | |
| for mode in ("rowwise", "rowwise_pad", "tensorwise_pad"): | |
| try: | |
| got = layer.run(x, mode).float() | |
| err = ((got - ref).norm() / ref.norm()).item() | |
| print(f"fp8 {mode} m={m} relative error {err:.4f}", flush=True) | |
| except Exception as e: | |
| print(f"fp8 {mode} m={m} failed: {type(e).__name__}: {str(e)[:120]}", flush=True) | |
| for mode in ("rowwise", "rowwise_pad", "tensorwise_pad"): | |
| try: | |
| ok = True | |
| for m in (1, 7, 33): | |
| x = torch.randn(m, layer.weight.shape[1], device="cuda", dtype=torch.bfloat16) | |
| ref = layer.dequant(x).float() | |
| err = ((layer.run(x, mode).float() - ref).norm() / ref.norm()).item() | |
| ok = ok and err < 0.05 | |
| if ok: | |
| FusedFP8.mode = mode | |
| break | |
| except Exception: | |
| continue | |
| else: | |
| FusedFP8.mode = "dequant" | |
| print(f"fp8 kernel: {FusedFP8.mode}", flush=True) | |
| MARKERS = ("\n\n**User:**", "\n\n**Model C:**") | |
| def withhold_marker_prefix(text): | |
| for n in range(min(len(text), max(map(len, MARKERS))), 0, -1): | |
| if any(m.startswith(text[-n:]) for m in MARKERS): | |
| return text[:-n] | |
| return text | |
| class StopOnText(StoppingCriteria): | |
| def __init__(self, prompt_len): | |
| self.prompt_len = prompt_len | |
| def __call__(self, input_ids, scores, **kwargs): | |
| tail = tokenizer.decode(input_ids[0, self.prompt_len :][-8:]) | |
| return "**User" in tail or "**Model C" in tail | |
| def build_prompt(frame, turns): | |
| head = FRAMES.get(frame, FRAMES["stranger"]) + "\n\nFull conversation with Model C:" | |
| body = [] | |
| for i, text in enumerate(turns): | |
| who = "**User:**" if i % 2 == 0 else "**Model C:**" | |
| body.append(f"{who} {text}".rstrip()) | |
| while True: | |
| prompt = head + "\n\n" + "\n\n".join(body + ["**Model C:**"]) | |
| ids = tokenizer(prompt, return_tensors="pt").input_ids | |
| if ids.shape[1] <= MAX_CONTEXT or len(body) <= 1: | |
| return ids[:, -MAX_CONTEXT:] | |
| body = body[2:] | |
| def gpu_duration(input_ids, max_tokens): | |
| return 12 + int(max_tokens / 11) | |
| def generate(input_ids, max_tokens): | |
| try: | |
| torch.cuda.set_per_process_memory_fraction(1.0) | |
| except Exception: | |
| pass | |
| pick_kernel() | |
| input_ids = input_ids.to("cuda") | |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120) | |
| kwargs = dict( | |
| input_ids=input_ids, | |
| attention_mask=torch.ones_like(input_ids), | |
| max_new_tokens=max_tokens, | |
| max_time=gpu_duration(input_ids, max_tokens) - 8, | |
| do_sample=True, | |
| temperature=1.0, | |
| top_p=0.98, | |
| top_k=0, | |
| eos_token_id=[128001, 128009], | |
| pad_token_id=128004, | |
| streamer=streamer, | |
| stopping_criteria=StoppingCriteriaList([StopOnText(input_ids.shape[1])]), | |
| ) | |
| stats = {} | |
| def run(): | |
| t0 = time.time() | |
| with torch.inference_mode(): | |
| out = model.generate(**kwargs) | |
| stats["tokens"] = out.shape[1] - input_ids.shape[1] | |
| stats["seconds"] = time.time() - t0 | |
| thread = threading.Thread(target=run, daemon=True) | |
| thread.start() | |
| text = "" | |
| for chunk in streamer: | |
| text += chunk | |
| cut = min((i for i in (text.find("**User"), text.find("**Model C")) if i != -1), default=-1) | |
| if cut != -1: | |
| yield text[:cut] | |
| break | |
| yield withhold_marker_prefix(text) | |
| thread.join() | |
| if stats: | |
| print(f"{stats['tokens']} tokens in {stats['seconds']:.1f}s ({stats['tokens'] / stats['seconds']:.1f} tok/s), prompt {input_ids.shape[1]}", flush=True) | |
| app = Server(title="Computer") | |
| def respond(turns: str, frame: str = "stranger", max_tokens: int = 512) -> str: | |
| try: | |
| turns = json.loads(turns) | |
| except (TypeError, ValueError): | |
| raise Error("Malformed page.") | |
| if not isinstance(turns, list) or not turns or len(turns) % 2 == 0 or not all(isinstance(t, str) for t in turns): | |
| raise Error("The page must end with your turn.") | |
| turns = [t.strip()[:MAX_MESSAGE_CHARS] for t in turns] | |
| max_tokens = min(max(int(max_tokens), 32), 768) | |
| input_ids = build_prompt(frame, turns) | |
| for text in generate(input_ids, max_tokens): | |
| yield text.strip() | |
| async def homepage(): | |
| return PAGE | |
| async def font(name: str): | |
| if name not in FONTS: | |
| raise HTTPException(status_code=404) | |
| return FileResponse(FONTS[name], media_type="font/woff2", headers={"Cache-Control": "public, max-age=31536000, immutable"}) | |
| demo = app | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False, show_error=True) | |