diff --git a/tinygrad/llm/cli.py b/tinygrad/llm/cli.py index 417ffa7..63844c0 100644 --- a/tinygrad/llm/cli.py +++ b/tinygrad/llm/cli.py @@ -49,7 +49,8 @@ class SimpleTokenizer: vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"])) normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1) special_tokens_dict = dict(special_tokens) - return SimpleTokenizer(dict(normal_tokens), special_tokens_dict, kv["tokenizer.ggml.pre"], + pre = kv.get("tokenizer.ggml.pre", "qwen2" if "qwen" in kv.get("general.architecture", "") else "llama3") + return SimpleTokenizer(dict(normal_tokens), special_tokens_dict, pre, bos_id=kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None, eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id', special_tokens_dict.get('<|im_end|>'))) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 5231e8e..3840460 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -437,7 +437,7 @@ class Transformer: n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'], vocab_size=len(kv['tokenizer.ggml.tokens']), head_dim=head_dim, - rope_theta=kv[f'{arch}.rope.freq_base'], + rope_theta=kv.get(f'{arch}.rope.freq_base', 1000000.0 if 'qwen' in arch else 10000.0), rope_dim=rope_dim, v_head_dim=kv.get(f'{arch}.attention.value_length_mla', kv.get(f'{arch}.attention.value_length', head_dim)), max_context=max_context, diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index f814a75..a2b519b 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -1,5 +1,5 @@ from __future__ import annotations -import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref +import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref, atexit assert sys.platform != 'win32' from typing import cast from dataclasses import dataclass diff --git a/tinygrad/runtime/support/nv/ip.py b/tinygrad/runtime/support/nv/ip.py index 58885ac..d54f392 100644 --- a/tinygrad/runtime/support/nv/ip.py +++ b/tinygrad/runtime/support/nv/ip.py @@ -1,5 +1,5 @@ from __future__ import annotations -import ctypes, time, array, struct, itertools, dataclasses +import ctypes, time, array, struct, itertools, dataclasses, contextlib from typing import cast, Any from tinygrad.runtime.autogen import nv, nv_570 as nv_gpu, pci from tinygrad.helpers import lo32, hi32, DEBUG, round_up, round_down, fetch_fw, wait_cond, ceildiv @@ -477,9 +477,10 @@ class NV_GSP(NV_IP): # reserve 512MB for the reserved PDES res_va = self.nvdev.mm.alloc_vaddr(res_sz:=(512 << 20)) - bufs_p = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS(pageSize=res_sz, numLevelsToCopy=3, + pts = list(self.nvdev.mm.page_tables(res_va, size=res_sz)) + bufs_p = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS(pageSize=res_sz, numLevelsToCopy=len(pts), virtAddrLo=res_va, virtAddrHi=res_va + res_sz - 1) - for i,pt in enumerate(self.nvdev.mm.page_tables(res_va, size=res_sz)): + for i,pt in enumerate(pts): bufs_p.levels[i] = nv_gpu.struct_NV90F1_CTRL_VASPACE_COPY_SERVER_RESERVED_PDES_PARAMS_level(physAddress=pt.paddr, size=self.nvdev.mm.pte_cnt[0] * 8 if i == 0 else 0x1000, pageShift=self.nvdev.mm.pte_covers[i].bit_length() - 1, aperture=1) self.rpc_rm_control(hObject=vaspace, cmd=nv_gpu.NV90F1_CTRL_CMD_VASPACE_COPY_SERVER_RESERVED_PDES, params=bufs_p) @@ -513,8 +514,11 @@ class NV_GSP(NV_IP): self.stat_q.wait_resp(nv.NV_VGPU_MSG_EVENT_GSP_INIT_DONE) - self.nvdev.NV_PBUS_BAR1_BLOCK.write(mode=0, target=0, ptr=0) - if self.nvdev.fmc_boot: self.nvdev.NV_VIRTUAL_FUNCTION_PRIV_FUNC_BAR1_BLOCK_LOW_ADDR.write(mode=0, target=0, ptr=0) + if not self.nvdev.fmc_boot: self.nvdev.NV_PBUS_BAR1_BLOCK.write(mode=0, target=0, ptr=0) + else: + with contextlib.suppress(Exception): self.nvdev.NV_VIRTUAL_FUNCTION_PRIV_BAR1_BLOCK.write(mode=0, target=0, ptr=0) + with contextlib.suppress(Exception): self.nvdev.NV_VIRTUAL_FUNCTION_PRIV_FUNC_BAR1_BLOCK_LOW_ADDR.write(mode=0, target=0, ptr=0) + self.nvdev.vram[self.nvdev.mm.root_page_table.paddr : self.nvdev.mm.root_page_table.paddr + 0x1000] = bytes(0x1000) self.priv_root = 0xc1e00004 self.init_golden_image() @@ -611,7 +615,7 @@ class NV_GSP(NV_IP): def rpc_unloading_guest_driver(self): data = nv.rpc_unloading_guest_driver_v(bInPMTransition=0, bGc6Entering=0, newLevel=(__GPU_STATE_FLAGS_FAST_UNLOAD:=1 << 6)) self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_UNLOADING_GUEST_DRIVER, bytes(data)) - self.stat_q.wait_resp(nv.NV_VGPU_MSG_FUNCTION_UNLOADING_GUEST_DRIVER) + with contextlib.suppress(Exception): self.stat_q.wait_resp(nv.NV_VGPU_MSG_FUNCTION_UNLOADING_GUEST_DRIVER, timeout=500) def rpc_set_registry_table(self): table = {'RMForcePcieConfigSave': 0x1, 'RMSecBusResetEnable': 0x1} diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index bddeaf5..e4649dd 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -4,7 +4,7 @@ from tinygrad.helpers import getenv, DEBUG, getbits, round_up from tinygrad.runtime.autogen import pci from tinygrad.runtime.support.memory import TLSFAllocator, MemoryManager, AddrSpace from tinygrad.runtime.support.nv.ip import NV_FLCN, NV_FLCN_COT, NV_GSP -from tinygrad.runtime.support.system import PCIDevice +from tinygrad.runtime.support.system import PCIDevice, RemotePCIDevice from tinygrad.runtime.support.hcq import MMIOInterface NV_DEBUG = getenv("NV_DEBUG", 0) @@ -102,7 +102,7 @@ class NVDev: self.include("dev_fb", "tu102") self.include("dev_gc6_island", "ga102") - if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0: + if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0 and not isinstance(self.pci_dev, RemotePCIDevice): self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2) if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True) self.pci_dev.reset()