IQ.Pilot Release Commit @ 0798119

tinygrad_repo/examples/llm.c/.gitignore

@@ -0,0 +1,3 @@
+data
+out.c
+a.out

tinygrad_repo/examples/llm.c/export.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+import os
+if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
+from tinygrad import Device, nn, Tensor, dtypes
+from train_gpt2 import GPT, GPTConfig
+from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
+from tinygrad.engine.realize import get_kernel
+from tinygrad.schedule.memory import memory_planner
+from tinygrad.uop.ops import Ops
+
+DEV.value = "CPU"
+
+TIMING = getenv("TIMING")
+
+if __name__ == "__main__":
+  model = GPT(GPTConfig(n_layer=getenv("NLAYER", 12), n_head=12, n_embd=768))
+  #model.load_pretrained()
+  for p in nn.state.get_parameters(model): p.replace(Tensor.empty(p.shape, dtype=p.dtype)) # fake load pretrained
+
+  #early_sched = create_schedule([x.uop for x in nn.state.get_parameters(model)])
+  #print(f"built model {len(early_sched)}")
+
+  #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
+  B, T = 4, 64
+
+  Tensor.training = True
+  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
+  warmup_count = getenv("WARMUP", 3)
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
+    GlobalCounters.reset()
+    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    _, loss = model(X, Y)
+    optimizer.zero_grad()
+    if getenv("BACKWARD", 1):
+      loss.backward()
+      tensors = optimizer.schedule_step()
+    else:
+      tensors = []
+    sched = loss.schedule(*tensors)
+    print(f"calls {i}:", len(sched))
+    #run_schedule(sched[:])
+  sched = memory_planner(sched)
+  ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
+  srcs = {}
+  for ast in ast_dedup:
+    k = get_kernel(Device["CPU"].renderer, ast)
+    k.linearize()
+    src = Device["CPU"].renderer.render(to_function_name(k.name), k.uops)
+    srcs[ast] = (k.name, src)
+  print("functions:", len(srcs))
+  used_buffers = dedup(flatten([si.bufs for si in sched]))
+  numbered_bufs = {x:i for i,x in enumerate(used_buffers)}
+  print("buffers:", len(numbered_bufs))
+
+  state_dict = nn.state.get_state_dict(model)
+  state_dict.update({'X': X, 'Y': Y, 'loss': loss})
+  grad_state_dict = {}
+  for k,v in state_dict.items():
+    if v.uop.base.buffer not in used_buffers: print(f"UNUSED: {k}")
+    if v.grad is not None: grad_state_dict['grad_'+k] = v.grad
+  state_dict.update(grad_state_dict)
+  state_dict.update({'adam_b1_t': optimizer.b1_t, 'adam_b2_t': optimizer.b2_t, 'adam_lr': optimizer.lr})
+  inverse_state_dict = {v:k for k,v in state_dict.items()}
+  for p,m,v in zip(optimizer.params, optimizer.m, optimizer.v):
+    nm = inverse_state_dict[p]
+    state_dict["adam_m_"+nm] = m
+    state_dict["adam_v_"+nm] = v
+  named_buffers = {v.uop.base.buffer:k.replace(".", "_") for k,v in state_dict.items()}
+
+  c_code = ["#include <stdlib.h>", "#include <tgmath.h>", "#include <stdbool.h>"]
+  if TIMING: c_code += ["#include <stdio.h>", "#include <time.h>"]
+  c_code += [x[1].replace(" restrict ", " ")+"\n" for x in srcs.values()]
+
+  premain = ["int main() {"]
+  if TIMING:
+    premain += ["  struct timespec tm0; clock_gettime(CLOCK_MONOTONIC, &tm0);"]
+  lst = 0
+  main = []
+
+  all_bufs = []
+  for i,si in enumerate(sched):
+    bufs = [(named_buffers.get(b, f"b{numbered_bufs[b]}"), b) for b in si.bufs]
+    all_bufs += bufs
+    if si.ast.op is not Ops.SINK:
+      print(f"// {si.ast.op}", bufs)
+    else:
+      print(f"{srcs[si.ast][0]}({', '.join([x[0] for x in bufs])})")
+      main.append(f"  {to_function_name(srcs[si.ast][0])}({', '.join([x[0] for x in bufs])});")
+      if TIMING:
+        main.append(f"  struct timespec tm{i+1}; clock_gettime(CLOCK_MONOTONIC, &tm{i+1});")
+        main.append(f"  printf(\"%10.2f ms + %7.2f ms @ {to_function_name(srcs[si.ast][0])}\\n\"," +\
+                    f"((tm{i+1}.tv_sec-tm{0}.tv_sec) + (tm{i+1}.tv_nsec-tm{0}.tv_nsec) / 1e9) * 1e3," +\
+                    f"((tm{i+1}.tv_sec-tm{lst}.tv_sec) + (tm{i+1}.tv_nsec-tm{lst}.tv_nsec) / 1e9) * 1e3);")
+      lst = i+1
+      #call = f"{srcs[si.ast][0]}({', '.join(bufs)})"
+      #call += " "*(80-ansilen(call))
+      #print(f"{call} // {i+1}")
+      #print(srcs[si.ast][1])
+  main.append("}")
+
+  mallocs = [f"  {b.dtype.name}* {n} = ({b.dtype.name}*)malloc({b.nbytes});" for n,b in dedup(all_bufs)]
+
+  with open("out.c", "w") as f: f.write('\n'.join(c_code+premain+mallocs+main))

tinygrad_repo/examples/llm.c/train_gpt2.py

@@ -0,0 +1,207 @@
+#!/usr/bin/env python3
+import os, math, time
+import numpy as np
+from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters
+from dataclasses import dataclass
+
+@dataclass
+class GPTConfig:
+  block_size: int = 1024
+  vocab_size: int = 50257
+  padded_vocab_size: int = 50304
+  n_layer: int = 12
+  n_head: int = 12
+  n_embd: int = 768
+
+class CausalSelfAttention:
+  def __init__(self, config:GPTConfig):
+    assert config.n_embd % config.n_head == 0
+    # key, query, value projections for all heads, but in a batch
+    self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+    # output projection
+    self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+    # regularization
+    self.n_head = config.n_head
+    self.n_embd = config.n_embd
+    # not really a 'bias', more of a mask, but following the OpenAI/HF naming though
+    self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
+    self.bias.is_param_(False)
+
+  def __call__(self, x:Tensor):
+    B, T, C = x.shape
+    qkv = self.c_attn(x)
+    q, k, v = qkv.split(self.n_embd, dim=2)
+    k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+    q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+    v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+    # manual implementation of attention
+    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+    att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+    att = att.softmax()
+    y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+    y = y.transpose(1, 2).view(B, T, C) # re-assemble all head outputs side by side
+    # output projection
+    y = self.c_proj(y)
+    return y
+
+class MLP:
+  def __init__(self, config:GPTConfig):
+    self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
+    self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return self.c_proj(self.c_fc(x).gelu())
+
+class Block:
+  def __init__(self, config:GPTConfig):
+    self.ln_1 = nn.LayerNorm(config.n_embd)
+    self.attn = CausalSelfAttention(config)
+    self.ln_2 = nn.LayerNorm(config.n_embd)
+    self.mlp = MLP(config)
+
+  def __call__(self, x:Tensor):
+    x = x + self.attn(self.ln_1(x))
+    x = x + self.mlp(self.ln_2(x))
+    return x
+
+class GPT:
+  def __init__(self, config:GPTConfig):
+    self.config = config
+
+    self.wte = nn.Embedding(config.padded_vocab_size, config.n_embd)
+    self.wpe = nn.Embedding(config.block_size, config.n_embd)
+    self.h = [Block(config) for _ in range(config.n_layer)]
+    self.ln_f = nn.LayerNorm(config.n_embd)
+    self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
+    self.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+  def load_pretrained(self):
+    weights = nn.state.torch_load(fetch(f'https://huggingface.co/gpt2/resolve/main/pytorch_model.bin'))
+    transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
+    for k in weights:
+      if k == "wte.weight":
+        weights[k] = weights[k].pad(((0, self.config.padded_vocab_size-self.config.vocab_size), (0,0))).to(None).contiguous()
+      if k.endswith(transposed):
+        weights[k] = weights[k].to(None).T.contiguous()
+    # lm head and wte are tied
+    weights['lm_head.weight'] = weights['wte.weight']
+    nn.state.load_state_dict(self, weights)
+
+  def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+    for _ in range(max_new_tokens):
+      idx_cond = idx if idx.shape[1] <= self.config.block_size else idx[:, -self.config.block_size:]
+      logits, _ = self(idx_cond)
+      logits = logits[:, -1, :] / temperature
+      idx_next = logits.softmax().multinomial()
+      idx = Tensor.cat(idx, idx_next, dim=1)
+    return idx
+
+  def __call__(self, idx:Tensor, targets=None):
+    b, t = idx.shape
+    pos = Tensor.arange(0, t, device=idx.device)
+
+    tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
+    pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
+    x = tok_emb + pos_emb
+
+    x = self.ln_f(x.sequential(self.h))
+
+    if targets is not None:
+      logits = self.lm_head(x)[:, :, :self.config.vocab_size]
+      loss = logits.sparse_categorical_crossentropy(targets)
+    else:
+      logits = self.lm_head(x[:, [-1], :])[:, :, :self.config.vocab_size]
+      loss = None
+
+    return logits, loss
+
+if __name__ == "__main__":
+  import tiktoken, argparse
+
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--num_iterations", type=int, default=10, help="number of iterations to run")
+  parser.add_argument("--batch_size", type=int, default=4, help="batch size")
+  parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
+  parser.add_argument("--skip_test", action="store_true", help="skip test")
+  parser.add_argument("--gpus", type=int, default=1, help="sequence length")
+  args = parser.parse_args()
+  B, T = args.batch_size, args.sequence_length
+  assert 1 <= T <= 1024
+
+  model = GPT(GPTConfig(n_layer=12, n_head=12, n_embd=768))
+  model.load_pretrained()
+
+  if args.gpus > 1:
+    GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(args.gpus))
+    for x in nn.state.get_parameters(model): x.to_(GPUS)  # we put a copy of the model on every GPU
+
+  # init the tokenizer
+  enc = tiktoken.get_encoding("gpt2")
+  encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
+  decode = lambda l: enc.decode(l)
+
+  # load the tokens
+  # prefer to use tiny_shakespeare if it's available, otherwise use tiny_stories
+  # we're using val instead of train split just because it is smaller/faster
+  tokens_bin = fetch("https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/tiny_shakespeare_val.bin")
+  assert os.path.isfile(tokens_bin)
+  print(f"loading cached tokens in {tokens_bin}")
+  with open(tokens_bin, "rb") as f:
+    f.seek(0x400)
+    tokens = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
+  tokens = Tensor(tokens)
+
+  # lightweight dataloader
+  def get_batch():
+    assert B*T+1 <= len(tokens), "not enough tokens"
+    # for 338,025 tokens. E.g. with B=8 T=1024, this will yield 41 batches before looping
+    i = 0
+    while True:
+      x = tokens[i:i+B*T].view(B, T)
+      y = tokens[i+1:i+B*T+1].view(B, T)
+      yield x, y
+      i += B*T
+      if i + B*T + 1 >= len(tokens):
+        i = 0 # in prod we'd want to randomize the start point a bit
+
+  # forward backward for a few iterations
+  data_iter = iter(get_batch())
+  x, y = next(data_iter) # we'll overfit this batch below
+  optimizer = nn.optim.AdamW(nn.state.get_parameters(model), lr=1e-4, weight_decay=0)
+
+  print(f"model state:     {sum(x.nbytes() for x in nn.state.get_parameters(model))/1e9:.2f} GB")
+  print(f"optimizer state: {sum(x.nbytes() for x in nn.state.get_parameters(optimizer))/1e9:.2f} GB")
+
+  # shard the data on axis 0
+  if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)
+
+  @TinyJit
+  @Tensor.train()
+  def step(x:Tensor, y:Tensor) -> Tensor:
+    _, loss = model(x, y)
+    optimizer.zero_grad()
+    loss.backward()
+    return loss.realize(*optimizer.schedule_step())
+
+  for i in range(args.num_iterations):
+    GlobalCounters.reset()
+    t0 = time.perf_counter()
+    loss = step(x.contiguous(), y.contiguous())
+    Device[Device.DEFAULT].synchronize()
+    t1 = time.perf_counter()
+    print(f"iteration {i}, loss: {loss.item():.6f}, time: {(t1-t0)*1000:.3f}ms, {int(B*T/(t1-t0))} tok/s, {GlobalCounters.global_mem/1e9:.2f} GB")
+
+  if not args.skip_test:
+    # copy back to single gpu for test
+    if args.gpus > 1:
+      for x in nn.state.get_parameters(model): x.to_(Device.DEFAULT)
+    start = "<|endoftext|>"
+    start_ids = encode(start)
+    x = (Tensor(start_ids)[None, ...])
+    max_new_tokens = 16
+    temperature = 1.0
+    top_k = 40
+    y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
+    print(decode(y[0].tolist()))
+

tinygrad_repo/examples/llm.c/ubench/matmul.c

@@ -0,0 +1,65 @@
+// clang -Ofast -Wno-unused-result -march=native matmul.c
+#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+
+float b52[786432];
+float b49[196608];
+float h_0_mlp_c_fc_weight[2359296];
+float h_0_mlp_c_fc_bias[3072];
+
+void matmul_forward(float* out,
+                    float* inp, float* weight, float* bias,
+                    int B, int T, int C, int OC) {
+    // most of the running time is spent here and in matmul_backward
+    // OC is short for "output channels"
+    // inp is (B,T,C), weight is (OC, C), bias is (OC)
+    // out will be (B,T,OC)
+    #pragma omp parallel for collapse(2)
+    for (int b = 0; b < B; b++) {
+        for (int t = 0; t < T; t++) {
+            float* out_bt = out + b * T * OC + t * OC;
+            float* inp_bt = inp + b * T * C + t * C;
+            for (int o = 0; o < OC; o++) {
+                float val = (bias != NULL) ? bias[o] : 0.0f;
+                float* wrow = weight + o*C;
+                for (int i = 0; i < C; i++) {
+                    val += inp_bt[i] * wrow[i];
+                }
+                out_bt[o] = val;
+            }
+        }
+    }
+}
+
+
+void r_256_3072_768(float* restrict data0, const float* restrict data1, const float* restrict data2, const float* restrict data3) {
+  for (int ridx0 = 0; ridx0 < 256; ridx0++) {
+    for (int ridx1 = 0; ridx1 < 3072; ridx1++) {
+      float acc0 = 0.0f;
+      float val0 = data3[ridx1];
+      for (int ridx2 = 0; ridx2 < 768; ridx2++) {
+        float val1 = data1[(ridx0*768)+ridx2];
+        float val2 = data2[(ridx1*768)+ridx2];
+        acc0 = ((val1*val2)+acc0);
+      }
+      data0[(ridx0*3072)+ridx1] = (acc0+val0);
+    }
+  }
+}
+
+
+int main() {
+  for (int i = 0; i < 5; i++) {
+    struct timespec t1, t2, t3;
+    clock_gettime(CLOCK_MONOTONIC, &t1);
+    r_256_3072_768(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias);
+    clock_gettime(CLOCK_MONOTONIC, &t2);
+    matmul_forward(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias, 4, 64, 768, 3072);
+    clock_gettime(CLOCK_MONOTONIC, &t3);
+    double time_gen = (t2.tv_sec - t1.tv_sec) + (t2.tv_nsec - t1.tv_nsec) / 1e9;
+    double time_real = (t3.tv_sec - t2.tv_sec) + (t3.tv_nsec - t2.tv_nsec) / 1e9;
+    printf("%.2f ms gen vs %.2f ms reference\n", time_gen*1e3, time_real*1e3);
+  }
+}
+