| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| import argparse |
| import glob |
| import importlib.util |
| import itertools |
| import json |
| import os |
| import sys |
| import time |
|
|
| import torch |
| import torch.nn as nn |
|
|
| INIT_FNS = ("uniform_", "normal_", "trunc_normal_", "kaiming_uniform_", "kaiming_normal_", |
| "xavier_uniform_", "xavier_normal_", "orthogonal_") |
|
|
|
|
| def find_model_file(d): |
| files = sorted(glob.glob(os.path.join(d, "JiRackDeltaNet_*.py"))) |
| if not files: |
| sys.exit(f"❌ no JiRackDeltaNet_*.py in {d}; pass --model-file") |
| return files[-1] |
|
|
|
|
| def load_module(path): |
| spec = importlib.util.spec_from_file_location("jirack_chat_model", path) |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules["jirack_chat_model"] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| def build(mod): |
| """bf16 skeleton without random init (all weights are loaded afterwards) - as in deltanet_gate.py.""" |
| saved = {n: getattr(nn.init, n) for n in INIT_FNS if hasattr(nn.init, n)} |
| for n in saved: |
| setattr(nn.init, n, lambda t, *a, **k: t) |
| torch.set_default_dtype(torch.bfloat16) |
| try: |
| model = mod.JiRackQwen38ForCausalLM() |
| finally: |
| torch.set_default_dtype(torch.float32) |
| for n, f in saved.items(): |
| setattr(nn.init, n, f) |
| return model.eval() |
|
|
|
|
| def load_weights(model, d): |
| from safetensors import safe_open |
| idx = os.path.join(d, "model.safetensors.index.json") |
| files = sorted(set(json.load(open(idx))["weight_map"].values())) if os.path.exists(idx) else \ |
| sorted(os.path.basename(p) for p in glob.glob(os.path.join(d, "*.safetensors"))) |
| sd = model.state_dict() |
| params = {n for n, _ in model.named_parameters()} |
| loaded = set() |
| t0 = time.time() |
| with torch.no_grad(): |
| for i, f in enumerate(files): |
| with safe_open(os.path.join(d, f), framework="pt", device="cpu") as h: |
| for k in h.keys(): |
| if k in sd: |
| sd[k].copy_(h.get_tensor(k)) |
| loaded.add(k) |
| print(f" {i + 1}/{len(files)} {f} {time.time() - t0:.0f} s", flush=True) |
| missing = sorted(params - loaded) |
| if missing: |
| sys.exit(f"❌ {len(missing)} parameters not in {d}: {missing[:5]}") |
|
|
|
|
| class ChatTokenizer: |
| """tokenizer.json (tokenizers) + chat template (jinja2) - no transformers needed, so the |
| tokenizer_config.json format (e.g. extra_special_tokens as a list) does not matter.""" |
|
|
| def __init__(self, d): |
| import jinja2 |
| from jinja2.sandbox import ImmutableSandboxedEnvironment |
| from tokenizers import Tokenizer |
| self.tk = Tokenizer.from_file(os.path.join(d, "tokenizer.json")) |
| cfg = {} |
| if os.path.exists(os.path.join(d, "tokenizer_config.json")): |
| cfg = json.load(open(os.path.join(d, "tokenizer_config.json"), encoding="utf-8")) |
| path = os.path.join(d, "chat_template.jinja") |
| text = open(path, encoding="utf-8").read() if os.path.exists(path) else cfg.get("chat_template") |
| if isinstance(text, list): |
| text = next(t["template"] for t in text if t.get("name") == "default") |
| if not text: |
| sys.exit(f"❌ no chat template in {d}") |
|
|
| def raise_exception(msg): |
| raise jinja2.exceptions.TemplateError(msg) |
|
|
| def tojson(x, ensure_ascii=False, indent=None, separators=None, sort_keys=False): |
| return json.dumps(x, ensure_ascii=ensure_ascii, indent=indent, separators=separators, |
| sort_keys=sort_keys) |
|
|
| env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True) |
| env.filters["tojson"] = tojson |
| env.globals["raise_exception"] = raise_exception |
| self.tmpl = env.from_string(text) |
| s = lambda v: v.get("content") if isinstance(v, dict) else v |
| self.bos, self.eos = s(cfg.get("bos_token")) or "", s(cfg.get("eos_token")) or "" |
| self.special = [t.content for t in self.tk.get_added_tokens_decoder().values() if t.special] |
| self.stop_ids = {i for i in (self.tk.token_to_id("<|im_end|>"), self.tk.token_to_id("<|endoftext|>"), |
| self.tk.token_to_id(self.eos) if self.eos else None) if i is not None} |
|
|
| def prompt_ids(self, msgs, think): |
| text = self.tmpl.render(messages=msgs, add_generation_prompt=True, enable_thinking=think, |
| bos_token=self.bos, eos_token=self.eos) |
| return self.tk.encode(text, add_special_tokens=False).ids |
|
|
| def decode(self, ids): |
| return self.tk.decode(ids, skip_special_tokens=False) |
|
|
|
|
| def _to(x, dev): |
| if torch.is_tensor(x): |
| return x.to(dev) |
| if isinstance(x, (tuple, list)): |
| return type(x)(_to(v, dev) for v in x) |
| if isinstance(x, dict): |
| return {k: _to(v, dev) for k, v in x.items()} |
| return x |
|
|
|
|
| def place(model, devices): |
| """Blocks split evenly over the GPUs (moved there directly - the whole model never sits on one |
| GPU); everything else on the first GPU. Pre-hooks move each module's inputs to its device, |
| so the model code needs no changes.""" |
| blocks = [m for m in model.modules() if type(m).__name__ == "TransformerBlock"] |
| in_block = {id(x) for b in blocks for x in b.modules()} |
| for i, b in enumerate(blocks): |
| dev = devices[i * len(devices) // len(blocks)] |
| b.to(dev) |
| if len(devices) > 1: |
| b.register_forward_pre_hook(lambda m, a, k, dev=dev: (_to(a, dev), _to(k, dev)), with_kwargs=True) |
| for m in model.modules(): |
| if id(m) in in_block: |
| continue |
| m._apply(lambda t: t.to(devices[0]), recurse=False) |
| if len(devices) > 1 and not list(m.children()): |
| m.register_forward_pre_hook(lambda mm, a, k: (_to(a, devices[0]), _to(k, devices[0])), |
| with_kwargs=True) |
| return devices[0] |
|
|
|
|
| def _cast(x, dtype): |
| if torch.is_tensor(x): |
| return x.to(dtype) if x.is_floating_point() and x.dtype != dtype else x |
| if isinstance(x, (tuple, list)): |
| return type(x)(_cast(v, dtype) for v in x) |
| if isinstance(x, dict): |
| return {k: _cast(v, dtype) for k, v in x.items()} |
| return x |
|
|
|
|
| def fp32_blocks(model): |
| """CPU (AVX-only Xeon has no fast bf16 matmul): every TransformerBlock and lm_head runs with |
| fp32 weights/inputs, converted on entry and back on exit - stored weights stay bf16 (RAM), |
| same as the 'fp32blocks' compute of deltanet_gate.py.""" |
| def pre(mod, args, kwargs): |
| saved = [] |
| for t in itertools.chain(mod.parameters(), mod.buffers()): |
| if t.is_floating_point() and t.dtype != torch.float32: |
| saved.append((t, t.data)) |
| t.data = t.data.float() |
| mod._jr_saved = saved |
| return _cast(args, torch.float32), _cast(kwargs, torch.float32) |
|
|
| def post(mod, args, out, back=torch.bfloat16): |
| for t, d in mod._jr_saved: |
| t.data = d |
| mod._jr_saved = [] |
| return _cast(out, back) if back is not None else out |
|
|
| for m in model.modules(): |
| if type(m).__name__ == "TransformerBlock": |
| m.register_forward_pre_hook(pre, with_kwargs=True) |
| m.register_forward_hook(post) |
| head = getattr(model, "lm_head", None) |
| if isinstance(head, nn.Module): |
| head.register_forward_pre_hook(pre, with_kwargs=True) |
| head.register_forward_hook(lambda m, a, o: post(m, a, o, None)) |
|
|
|
|
| def logits_of(out): |
| if isinstance(out, (tuple, list)): |
| return out[0] |
| return out.logits if hasattr(out, "logits") else out |
|
|
|
|
| def sample(logits, temperature, top_p): |
| if temperature <= 0: |
| return int(logits.argmax()) |
| probs = torch.softmax(logits.float() / temperature, dim=-1) |
| sp, si = probs.sort(descending=True) |
| cut = (sp.cumsum(-1) - sp) > top_p |
| sp[cut] = 0 |
| return int(si[torch.multinomial(sp / sp.sum(), 1)]) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--model-dir", default="/mnt/nfs_share/JiRackDeltaNet_70b/expanded") |
| p.add_argument("--model-file", default=None) |
| p.add_argument("--tokenizer-dir", default=None, help="default: --model-dir") |
| p.add_argument("--max-new", type=int, default=2048) |
| p.add_argument("--temperature", type=float, default=0.6, help="0 = greedy") |
| p.add_argument("--top-p", type=float, default=0.95) |
| p.add_argument("--system", default=None, help="file with a system prompt") |
| p.add_argument("--no-think", action="store_true", help="start with thinking off") |
| p.add_argument("--cpu", action="store_true") |
| p.add_argument("--cpu-bf16", action="store_true", help="CPU: plain bf16 instead of fp32 blocks") |
| a = p.parse_args() |
|
|
| tok = ChatTokenizer(a.tokenizer_dir or a.model_dir) |
| stop = tok.stop_ids |
| print(f"stop tokens {sorted(stop)}", flush=True) |
|
|
| mf = a.model_file or find_model_file(a.model_dir) |
| print(f"model file {mf}\nweights {a.model_dir}", flush=True) |
| model = build(load_module(mf)) |
| load_weights(model, a.model_dir) |
| n = sum(q.numel() for q in model.parameters()) |
| if a.cpu or not torch.cuda.is_available(): |
| dev = torch.device("cpu") |
| try: |
| torch.backends.nnpack.set_flags(False) |
| except Exception: |
| pass |
| if not a.cpu_bf16: |
| fp32_blocks(model) |
| print(f"⚠️ CPU: {n / 1e9:.1f}B params, compute {'bf16' if a.cpu_bf16 else 'fp32 blocks'} - " |
| f"every token is a full forward pass, expect minutes per token. Ctrl+C stops the answer.") |
| else: |
| devs = [torch.device("cuda", i) for i in range(torch.cuda.device_count())] |
| dev = place(model, devs) |
| print(f"on {len(devs)} GPU(s): {n / 1e9:.1f}B params, " |
| f"{[round(torch.cuda.memory_allocated(d) / 2 ** 30) for d in devs]} GB") |
|
|
| system = open(a.system, encoding="utf-8").read().strip() if a.system else None |
| think = not a.no_think |
| history = [] |
| print("\nReady. /reset - clear history, /think - thinking on/off, exit - quit") |
| while True: |
| try: |
| user = input("\nYou: ").strip() |
| except (EOFError, KeyboardInterrupt): |
| break |
| if not user: |
| continue |
| if user.lower() in ("exit", "quit"): |
| break |
| if user == "/reset": |
| history = [] |
| print("(history cleared)") |
| continue |
| if user == "/think": |
| think = not think |
| print(f"(thinking {'on' if think else 'off'})") |
| continue |
| history.append({"role": "user", "content": user}) |
| msgs = ([{"role": "system", "content": system}] if system else []) + history |
| ids = torch.tensor([tok.prompt_ids(msgs, think)], dtype=torch.long, device=dev) |
| start = ids.shape[1] |
| print("\nAssistant: ", end="", flush=True) |
| printed, t0 = "", time.time() |
| try: |
| with torch.no_grad(): |
| for _ in range(a.max_new): |
| nxt = sample(logits_of(model(ids))[0, -1], a.temperature, a.top_p) |
| if nxt in stop: |
| break |
| ids = torch.cat([ids, torch.tensor([[nxt]], device=ids.device)], dim=1) |
| text = tok.decode(ids[0, start:].tolist()) |
| if not text.endswith("�"): |
| print(text[len(printed):], end="", flush=True) |
| printed = text |
| except KeyboardInterrupt: |
| print(" [stopped]", end="") |
| gen = ids.shape[1] - start |
| raw = tok.decode(ids[0, start:].tolist()) |
| print(f"\n [{gen} tokens, {gen / max(time.time() - t0, 1e-6):.2f} tok/s]") |
| |
| reply = raw.split("</think>")[-1] |
| for t in tok.special: |
| reply = reply.replace(t, "") |
| history.append({"role": "assistant", "content": reply.strip()}) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
