Skip to content

Instantly share code, notes, and snippets.

@HudsonGraeme

HudsonGraeme/OOR.ipynb

Created November 20, 2023 21:10
Show Gist options
  • Save HudsonGraeme/90ee8f59cd20a33805ba7469577f3948 to your computer and use it in GitHub Desktop.
Save HudsonGraeme/90ee8f59cd20a33805ba7469577f3948 to your computer and use it in GitHub Desktop.
Download ZIP
Value OOR during proof
Display the source blob
Display the rendered blob
Raw
not_OOR_const.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Install required dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "m-28wslFI2wX",
"outputId": "42aedcf7-3ae6-4df0-b891-c3c18895c883"
},
"outputs": [],
"source": [
"%pip install dataclasses\n",
"%pip install matplotlib\n",
"%pip install torch\n",
"%pip install numpy\n",
"%pip install pickle\n",
"%pip install requests\n",
"%pip install onnxruntime\n",
"%pip install onnx\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Model Definition"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-61wEcCP1OZ5"
},
"outputs": [],
"source": [
"\"\"\"\n",
"Full definition of a GPT Language Model, all of it in this single file.\n",
"References:\n",
"1) the official GPT-2 TensorFlow implementation released by OpenAI:\n",
"https://github.com/openai/gpt-2/blob/master/src/model.py\n",
"2) huggingface/transformers PyTorch implementation:\n",
"https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py\n",
"\n",
"This implementation is a modified version of model.py found in the NanoGPT repository. It supports an output of N characters rather than just a single character at a time.\n",
"\"\"\"\n",
"\n",
"import math\n",
"import inspect\n",
"from dataclasses import dataclass\n",
"\n",
"import torch\n",
"import torch.nn as nn\n",
"from torch.nn import functional as F\n",
"\n",
"\n",
"class LayerNorm(nn.Module):\n",
" \"\"\"LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False\"\"\"\n",
"\n",
" def __init__(self, ndim, bias):\n",
" super().__init__()\n",
" self.weight = nn.Parameter(torch.ones(ndim))\n",
" self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None\n",
"\n",
" def forward(self, input):\n",
" return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)\n",
"\n",
"\n",
"class CausalSelfAttention(nn.Module):\n",
" def __init__(self, config):\n",
" super().__init__()\n",
" assert config.n_embd % config.n_head == 0\n",
" # key, query, value projections for all heads, but in a batch\n",
" self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)\n",
" # output projection\n",
" self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)\n",
" # regularization\n",
" self.attn_dropout = nn.Dropout(config.dropout)\n",
" self.resid_dropout = nn.Dropout(config.dropout)\n",
" self.n_head = config.n_head\n",
" self.n_embd = config.n_embd\n",
" self.dropout = config.dropout\n",
" # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0\n",
" self.flash = hasattr(torch.nn.functional, \"scaled_dot_product_attention\")\n",
" if not self.flash:\n",
" print(\n",
" \"WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0\"\n",
" )\n",
" # causal mask to ensure that attention is only applied to the left in the input sequence\n",
" self.register_buffer(\n",
" \"bias\",\n",
" torch.tril(torch.ones(config.block_size, config.block_size)).view(\n",
" 1, 1, config.block_size, config.block_size\n",
" ),\n",
" )\n",
"\n",
" def forward(self, x):\n",
" (\n",
" B,\n",
" T,\n",
" C,\n",
" ) = x.size() # batch size, sequence length, embedding dimensionality (n_embd)\n",
"\n",
" # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n",
" q, k, v = self.c_attn(x).split(self.n_embd, dim=2)\n",
" k = k.view(B, T, self.n_head, C // self.n_head).transpose(\n",
" 1, 2\n",
" ) # (B, nh, T, hs)\n",
" q = q.view(B, T, self.n_head, C // self.n_head).transpose(\n",
" 1, 2\n",
" ) # (B, nh, T, hs)\n",
" v = v.view(B, T, self.n_head, C // self.n_head).transpose(\n",
" 1, 2\n",
" ) # (B, nh, T, hs)\n",
"\n",
" # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n",
" if self.flash:\n",
" # efficient attention using Flash Attention CUDA kernels\n",
" y = torch.nn.functional.scaled_dot_product_attention(\n",
" q,\n",
" k,\n",
" v,\n",
" attn_mask=None,\n",
" dropout_p=self.dropout if self.training else 0,\n",
" is_causal=True,\n",
" )\n",
" else:\n",
" # manual implementation of attention\n",
" att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n",
" att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float(\"-inf\"))\n",
" att = F.softmax(att, dim=-1)\n",
" att = self.attn_dropout(att)\n",
" y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n",
" y = (\n",
" y.transpose(1, 2).contiguous().view(B, T, C)\n",
" ) # re-assemble all head outputs side by side\n",
"\n",
" # output projection\n",
" y = self.resid_dropout(self.c_proj(y))\n",
" return y\n",
"\n",
"\n",
"class MLP(nn.Module):\n",
" def __init__(self, config):\n",
" super().__init__()\n",
" self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)\n",
" self.gelu = nn.GELU()\n",
" self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)\n",
" self.dropout = nn.Dropout(config.dropout)\n",
"\n",
" def forward(self, x):\n",
" x = self.c_fc(x)\n",
" x = self.gelu(x)\n",
" x = self.c_proj(x)\n",
" x = self.dropout(x)\n",
" return x\n",
"\n",
"\n",
"class Block(nn.Module):\n",
" def __init__(self, config):\n",
" super().__init__()\n",
" self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)\n",
" self.attn = CausalSelfAttention(config)\n",
" self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)\n",
" self.mlp = MLP(config)\n",
"\n",
" def forward(self, x):\n",
" x = x + self.attn(self.ln_1(x))\n",
" x = x + self.mlp(self.ln_2(x))\n",
" return x\n",
"\n",
"\n",
"@dataclass\n",
"class GPTConfig:\n",
" block_size: int = 1024\n",
" vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency\n",
" n_layer: int = 12\n",
" n_head: int = 12\n",
" n_embd: int = 768\n",
" dropout: float = 0.0\n",
" bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster\n",
"\n",
"\n",
"class GPT(nn.Module):\n",
" def __init__(self, config, train_mode=False):\n",
" super().__init__()\n",
" assert config.vocab_size is not None\n",
" assert config.block_size is not None\n",
" self.config = config\n",
" self.train_mode = train_mode\n",
" self.transformer = nn.ModuleDict(\n",
" dict(\n",
" wte=nn.Embedding(config.vocab_size, config.n_embd),\n",
" wpe=nn.Embedding(config.block_size, config.n_embd),\n",
" drop=nn.Dropout(config.dropout),\n",
" h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),\n",
" ln_f=LayerNorm(config.n_embd, bias=config.bias),\n",
" )\n",
" )\n",
" self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
" # with weight tying when using torch.compile() some warnings get generated:\n",
" # \"UserWarning: functional_call was passed multiple values for tied weights.\n",
" # This behavior is deprecated and will be an error in future versions\"\n",
" # not 100% sure what this is, so far seems to be harmless. TODO investigate\n",
" self.transformer.wte.weight = (\n",
" self.lm_head.weight\n",
" ) # https://paperswithcode.com/method/weight-tying\n",
"\n",
" # init all weights\n",
" self.apply(self._init_weights)\n",
" # apply special scaled init to the residual projections, per GPT-2 paper\n",
" for pn, p in self.named_parameters():\n",
" if pn.endswith(\"c_proj.weight\"):\n",
" torch.nn.init.normal_(\n",
" p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)\n",
" )\n",
"\n",
" # report number of parameters\n",
" print(\"number of parameters: %.2fM\" % (self.get_num_params() / 1e6,))\n",
"\n",
" def get_num_params(self, non_embedding=True):\n",
" \"\"\"\n",
" Return the number of parameters in the model.\n",
" For non-embedding count (default), the position embeddings get subtracted.\n",
" The token embeddings would too, except due to the parameter sharing these\n",
" params are actually used as weights in the final layer, so we include them.\n",
" \"\"\"\n",
" n_params = sum(p.numel() for p in self.parameters())\n",
" if non_embedding:\n",
" n_params -= self.transformer.wpe.weight.numel()\n",
" return n_params\n",
"\n",
" def _init_weights(self, module):\n",
" if isinstance(module, nn.Linear):\n",
" torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
" if module.bias is not None:\n",
" torch.nn.init.zeros_(module.bias)\n",
" elif isinstance(module, nn.Embedding):\n",
" torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
"\n",
" # @torch.no_grad()\n",
" def forwards(self, idx, targets=None):\n",
" device = idx.device\n",
" b, t = idx.size()\n",
" assert (\n",
" t <= self.config.block_size\n",
" ), f\"Cannot forward sequence of length {t}, block size is only {self.config.block_size}\"\n",
" pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)\n",
"\n",
" # forward the GPT model itself\n",
" tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)\n",
" pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)\n",
" x = self.transformer.drop(tok_emb + pos_emb)\n",
" for block in self.transformer.h:\n",
" x = block(x)\n",
" x = self.transformer.ln_f(x)\n",
"\n",
" if targets is not None:\n",
" # if we are given some desired targets also calculate the loss\n",
" logits = self.lm_head(x)\n",
" loss = F.cross_entropy(\n",
" logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1\n",
" )\n",
" else:\n",
" # inference-time mini-optimization: only forward the lm_head on the very last position\n",
" logits = self.lm_head(\n",
" x[:, [-1], :]\n",
" ) # note: using list [-1] to preserve the time dim\n",
" loss = None\n",
" return logits, loss\n",
"\n",
" def crop_block_size(self, block_size):\n",
" # model surgery to decrease the block size if necessary\n",
" # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)\n",
" # but want to use a smaller block size for some smaller, simpler model\n",
" assert block_size <= self.config.block_size\n",
" self.config.block_size = block_size\n",
" self.transformer.wpe.weight = nn.Parameter(\n",
" self.transformer.wpe.weight[:block_size]\n",
" )\n",
" for block in self.transformer.h:\n",
" if hasattr(block.attn, \"bias\"):\n",
" block.attn.bias = block.attn.bias[:, :, :block_size, :block_size]\n",
"\n",
" @classmethod\n",
" def from_pretrained(cls, model_type, override_args=None):\n",
" assert model_type in {\"gpt2\", \"gpt2-medium\", \"gpt2-large\", \"gpt2-xl\"}\n",
" override_args = override_args or {} # default to empty dict\n",
" # only dropout can be overridden see more notes below\n",
" assert all(k == \"dropout\" for k in override_args)\n",
" from transformers import GPT2LMHeadModel\n",
"\n",
" print(\"loading weights from pretrained gpt: %s\" % model_type)\n",
"\n",
" # n_layer, n_head and n_embd are determined from model_type\n",
" config_args = {\n",
" \"gpt2\": dict(n_layer=12, n_head=12, n_embd=768), # 124M params\n",
" \"gpt2-medium\": dict(n_layer=24, n_head=16, n_embd=1024), # 350M params\n",
" \"gpt2-large\": dict(n_layer=36, n_head=20, n_embd=1280), # 774M params\n",
" \"gpt2-xl\": dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params\n",
" }[model_type]\n",
" print(\"forcing vocab_size=50257, block_size=1024, bias=True\")\n",
" config_args[\"vocab_size\"] = 50257 # always 50257 for GPT model checkpoints\n",
" config_args[\"block_size\"] = 1024 # always 1024 for GPT model checkpoints\n",
" config_args[\"bias\"] = True # always True for GPT model checkpoints\n",
" # we can override the dropout rate, if desired\n",
" if \"dropout\" in override_args:\n",
" print(f\"overriding dropout rate to {override_args['dropout']}\")\n",
" config_args[\"dropout\"] = override_args[\"dropout\"]\n",
" # create a from-scratch initialized minGPT model\n",
" config = GPTConfig(**config_args)\n",
" model = GPT(config)\n",
" sd = model.state_dict()\n",
" sd_keys = sd.keys()\n",
" sd_keys = [\n",
" k for k in sd_keys if not k.endswith(\".attn.bias\")\n",
" ] # discard this mask / buffer, not a param\n",
"\n",
" # init a huggingface/transformers model\n",
" model_hf = GPT2LMHeadModel.from_pretrained(model_type)\n",
" sd_hf = model_hf.state_dict()\n",
"\n",
" # copy while ensuring all of the parameters are aligned and match in names and shapes\n",
" sd_keys_hf = sd_hf.keys()\n",
" sd_keys_hf = [\n",
" k for k in sd_keys_hf if not k.endswith(\".attn.masked_bias\")\n",
" ] # ignore these, just a buffer\n",
" sd_keys_hf = [\n",
" k for k in sd_keys_hf if not k.endswith(\".attn.bias\")\n",
" ] # same, just the mask (buffer)\n",
" transposed = [\n",
" \"attn.c_attn.weight\",\n",
" \"attn.c_proj.weight\",\n",
" \"mlp.c_fc.weight\",\n",
" \"mlp.c_proj.weight\",\n",
" ]\n",
" # basically the openai checkpoints use a \"Conv1D\" module, but we only want to use a vanilla Linear\n",
" # this means that we have to transpose these weights when we import them\n",
" assert len(sd_keys_hf) == len(\n",
" sd_keys\n",
" ), f\"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}\"\n",
" for k in sd_keys_hf:\n",
" if any(k.endswith(w) for w in transposed):\n",
" # special treatment for the Conv1D weights we need to transpose\n",
" assert sd_hf[k].shape[::-1] == sd[k].shape\n",
" with torch.no_grad():\n",
" sd[k].copy_(sd_hf[k].t())\n",
" else:\n",
" # vanilla copy over the other parameters\n",
" assert sd_hf[k].shape == sd[k].shape\n",
" with torch.no_grad():\n",
" sd[k].copy_(sd_hf[k])\n",
"\n",
" return model\n",
"\n",
" def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n",
" # start with all of the candidate parameters\n",
" param_dict = {pn: p for pn, p in self.named_parameters()}\n",
" # filter out those that do not require grad\n",
" param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}\n",
" # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.\n",
" # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.\n",
" decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]\n",
" nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]\n",
" optim_groups = [\n",
" {\"params\": decay_params, \"weight_decay\": weight_decay},\n",
" {\"params\": nodecay_params, \"weight_decay\": 0.0},\n",
" ]\n",
" num_decay_params = sum(p.numel() for p in decay_params)\n",
" num_nodecay_params = sum(p.numel() for p in nodecay_params)\n",
" print(\n",
" f\"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters\"\n",
" )\n",
" print(\n",
" f\"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters\"\n",
" )\n",
" # Create AdamW optimizer and use the fused version if it is available\n",
" fused_available = \"fused\" in inspect.signature(torch.optim.AdamW).parameters\n",
" use_fused = fused_available and device_type == \"cuda\"\n",
" extra_args = dict(fused=True) if use_fused else dict()\n",
" optimizer = torch.optim.AdamW(\n",
" optim_groups, lr=learning_rate, betas=betas, **extra_args\n",
" )\n",
" print(f\"using fused AdamW: {use_fused}\")\n",
"\n",
" return optimizer\n",
"\n",
" def estimate_mfu(self, fwdbwd_per_iter, dt):\n",
" \"\"\"estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS\"\"\"\n",
" # first estimate the number of flops we do per iteration.\n",
" # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311\n",
" N = self.get_num_params()\n",
" cfg = self.config\n",
" L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size\n",
" flops_per_token = 6 * N + 12 * L * H * Q * T\n",
" flops_per_fwdbwd = flops_per_token * T\n",
" flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter\n",
" # express our flops throughput as ratio of A100 bfloat16 peak flops\n",
" flops_achieved = flops_per_iter * (1.0 / dt) # per second\n",
" flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS\n",
" mfu = flops_achieved / flops_promised\n",
" return mfu\n",
"\n",
" # @torch.no_grad()\n",
" def forward(self, idx, targets=None):\n",
" \"\"\"\n",
" Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete\n",
" the sequence max_new_tokens times, feeding the predictions back into the model each time.\n",
" Most likely you'll want to make sure to be in model.eval() mode of operation for this.\n",
" \"\"\"\n",
" if self.train_mode:\n",
" return self.forwards(\n",
" idx, targets\n",
" ) # When training don't override the forward method\n",
" for _ in range(1):\n",
" # if the sequence context is growing too long we must crop it at block_size\n",
" idx_cond = (\n",
" idx\n",
" if idx.size(1) <= self.config.block_size\n",
" else idx[:, -self.config.block_size :]\n",
" )\n",
" # forward the model to get the logits for the index in the sequence\n",
" logits, _ = self.forwards(idx_cond)\n",
" # pluck the logits at the final step and scale by desired temperature\n",
" logits = logits[:, -1, :] / 0.8\n",
" # optionally crop the logits to only the top k options\n",
"\n",
" # apply softmax to convert logits to (normalized) probabilities\n",
" probs = F.softmax(logits, dim=-1)\n",
"\n",
" # append sampled index to the running sequence and continue\n",
" idx = torch.cat(\n",
" (idx.to(\"cpu\"), torch.full_like(idx, fill_value=9).to(\"cpu\")),\n",
" dim=1,\n",
" )\n",
"\n",
" return idx\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Prepare Dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "C-SNhvuVJtyr",
"outputId": "fd5b3b8a-2bb3-4aa8-e549-a6bccb3a354f"
},
"outputs": [],
"source": [
"\"\"\"\n",
"Prepare the Shakespeare dataset for character-level language modeling.\n",
"So instead of encoding with GPT-2 BPE tokens, we just map characters to ints.\n",
"Will save train.bin, val.bin containing the ids, and meta.pkl containing the\n",
"encoder and decoder and some other related info.\n",
"\"\"\"\n",
"import os\n",
"import pickle\n",
"import requests\n",
"import numpy as np\n",
"\n",
"# download the tiny shakespeare dataset\n",
"\n",
"DATA_DIR=\"./data\"\n",
"FILE_PATH=os.path.join(DATA_DIR, 'shakespeare.txt')\n",
"if not os.path.exists(DATA_DIR):\n",
" os.mkdir(DATA_DIR)\n",
"data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'\n",
"with open(FILE_PATH, 'w') as f:\n",
" f.write(requests.get(data_url).text)\n",
"\n",
"with open(FILE_PATH, 'r') as f:\n",
" data = f.read()\n",
"print(f\"length of dataset in characters: {len(data):,}\")\n",
"\n",
"# get all the unique characters that occur in this text\n",
"chars = sorted(list(set(data)))\n",
"vocab_size = len(chars)\n",
"print(\"all the unique characters:\", ''.join(chars))\n",
"print(f\"vocab size: {vocab_size:,}\")\n",
"\n",
"# create a mapping from characters to integers\n",
"stoi = { ch:i for i,ch in enumerate(chars) }\n",
"itos = { i:ch for i,ch in enumerate(chars) }\n",
"def encode(s):\n",
" return [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
"def decode(l):\n",
" return ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
"\n",
"# create the train and test splits\n",
"n = len(data)\n",
"train_data = data[:int(n*0.9)]\n",
"val_data = data[int(n*0.9):]\n",
"\n",
"# encode both to integers\n",
"train_ids = encode(train_data)\n",
"val_ids = encode(val_data)\n",
"print(f\"train has {len(train_ids):,} tokens\")\n",
"print(f\"val has {len(val_ids):,} tokens\")\n",
"\n",
"# export to bin files\n",
"train_ids = np.array(train_ids, dtype=np.uint16)\n",
"val_ids = np.array(val_ids, dtype=np.uint16)\n",
"train_ids.tofile(os.path.join(DATA_DIR, 'train.bin'))\n",
"val_ids.tofile(os.path.join(DATA_DIR, 'val.bin'))\n",
"\n",
"# save the meta information as well, to help us encode/decode later\n",
"meta = {\n",
" 'vocab_size': vocab_size,\n",
" 'itos': itos,\n",
" 'stoi': stoi,\n",
"}\n",
"with open(os.path.join(DATA_DIR, 'meta.pkl'), 'wb') as f:\n",
" pickle.dump(meta, f)\n",
"\n",
"# length of dataset in characters: 1115394\n",
"# all the unique characters:\n",
"# !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
"# vocab size: 65\n",
"# train has 1003854 tokens\n",
"# val has 111540 tokens\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Train the model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 524
},
"id": "EN1ItHTBLxSi",
"outputId": "b1c6379f-03c7-497f-9e4b-3d2e1787215b"
},
"outputs": [],
"source": [
"\n",
"\"\"\"\n",
"This training script can be run both on a single gpu in debug mode,\n",
"and also in a larger training run with distributed data parallel (ddp).\n",
"\n",
"To run on a single GPU, example:\n",
"$ python train.py --batch_size=32 --compile=False\n",
"\n",
"To run with DDP on 4 gpus on 1 node, example:\n",
"$ torchrun --standalone --nproc_per_node=4 train.py\n",
"\n",
"To run with DDP on 4 gpus across 2 nodes, example:\n",
"- Run on the first (master) node with example IP 123.456.123.456:\n",
"$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=123.456.123.456 --master_port=1234 train.py\n",
"- Run on the worker node:\n",
"$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=123.456.123.456 --master_port=1234 train.py\n",
"(If your cluster does not have Infiniband interconnect prepend NCCL_IB_DISABLE=1)\n",
"\"\"\"\n",
"\n",
"import os\n",
"import time\n",
"import math\n",
"import pickle\n",
"from contextlib import nullcontext\n",
"from matplotlib import pyplot as plt\n",
"from threading import Thread\n",
"\n",
"import numpy as np\n",
"import torch\n",
"from torch.nn.parallel import DistributedDataParallel as DDP\n",
"from torch.distributed import init_process_group, destroy_process_group\n",
"\n",
"train_losses = []\n",
"val_losses = []\n",
"\n",
"# -----------------------------------------------------------------------------\n",
"# default config values designed to train a gpt2 (124M) on OpenWebText\n",
"# I/O\n",
"out_dir = 'pretrained'\n",
"eval_interval = 200\n",
"log_interval = 1\n",
"eval_iters = 20\n",
"eval_only = False # if True, script exits right after the first eval\n",
"always_save_checkpoint = True # if True, always save a checkpoint after each eval\n",
"init_from = 'scratch' # 'scratch' or 'resume' or 'gpt2*'\n",
"# wandb logging\n",
"wandb_log = False # disabled by default\n",
"wandb_project = 'owt'\n",
"wandb_run_name = 'gpt2' # 'run' + str(time.time())\n",
"# data\n",
"dataset = 'openwebtext'\n",
"gradient_accumulation_steps = 5 * 8 # used to simulate larger batch sizes\n",
"batch_size = 12 # if gradient_accumulation_steps > 1, this is the micro-batch size\n",
"block_size = 64\n",
"# model\n",
"n_layer = 4\n",
"n_head = 4\n",
"n_embd = 128\n",
"dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+\n",
"bias = False # do we use bias inside LayerNorm and Linear layers?\n",
"# adamw optimizer\n",
"learning_rate = 6e-4 # max learning rate\n",
"max_iters = 600000 # total number of training iterations\n",
"weight_decay = 1e-1\n",
"beta1 = 0.9\n",
"beta2 = 0.95\n",
"grad_clip = 0.0 # clip gradients at this value, or disable if == 0.0\n",
"# learning rate decay settings\n",
"decay_lr = True # whether to decay the learning rate\n",
"warmup_iters = 100 # how many steps to warm up for\n",
"lr_decay_iters = 600000 # should be ~= max_iters per Chinchilla\n",
"min_lr = 6e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla\n",
"# DDP settings\n",
"backend = 'nccl' # 'nccl', 'gloo', etc.\n",
"# system\n",
"device = 'mps' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks\n",
"dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler\n",
"compile = True # use PyTorch 2.0 to compile the model to be faster\n",
"# -----------------------------------------------------------------------------\n",
"config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]\n",
"config = {k: globals()[k] for k in config_keys} # will be useful for logging\n",
"# -----------------------------------------------------------------------------\n",
"# train a miniature character-level shakespeare model\n",
"# good for debugging and playing on macbooks and such\n",
"\n",
"out_dir = 'pretrained'\n",
"eval_interval = 200 # keep frequent because we'll overfit\n",
"eval_iters = 200\n",
"log_interval = 1 # don't print too too often\n",
"\n",
"# we expect to overfit on this small dataset, so only save when val improves\n",
"always_save_checkpoint = True\n",
"\n",
"wandb_log = False # override via command line if you like\n",
"wandb_project = 'shakespeare-char'\n",
"wandb_run_name = 'mini-gpt'\n",
"\n",
"dataset = 'data'\n",
"gradient_accumulation_steps = 1\n",
"batch_size = 10\n",
"block_size = 256 # context of up to 256 previous characters\n",
"\n",
"# baby GPT model :)\n",
"n_layer = 2\n",
"n_head = 2\n",
"n_embd = 128\n",
"dropout = 0.0\n",
"\n",
"learning_rate = 1e-3 # with baby networks can afford to go a bit higher\n",
"max_iters = 200\n",
"lr_decay_iters = 200 # make equal to max_iters usually\n",
"min_lr = 1e-4 # learning_rate / 10 usually\n",
"beta2 = 0.99 # make a bit bigger because number of tokens per iter is small\n",
"\n",
"warmup_iters = 100 # not super necessary potentially\n",
"\n",
"# on macbook also add\n",
"device = 'mps' # run on cpu only\n",
"compile = False # do not torch compile the model\n",
"\n",
"\n",
"# various inits, derived attributes, I/O setup\n",
"ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?\n",
"if ddp:\n",
" init_process_group(backend=backend)\n",
" ddp_rank = int(os.environ['RANK'])\n",
" ddp_local_rank = int(os.environ['LOCAL_RANK'])\n",
" ddp_world_size = int(os.environ['WORLD_SIZE'])\n",
" device = f'cuda:{ddp_local_rank}'\n",
" torch.cuda.set_device(device)\n",
" master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.\n",
" seed_offset = ddp_rank # each process gets a different seed\n",
" # world_size number of processes will be training simultaneously, so we can scale\n",
" # down the desired gradient accumulation iterations per process proportionally\n",
" assert gradient_accumulation_steps % ddp_world_size == 0\n",
" gradient_accumulation_steps //= ddp_world_size\n",
"else:\n",
" # if not ddp, we are running on a single gpu, and one process\n",
" master_process = True\n",
" seed_offset = 0\n",
" ddp_world_size = 1\n",
"tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size\n",
"print(f\"tokens per iteration will be: {tokens_per_iter:,}\")\n",
"\n",
"if master_process:\n",
" os.makedirs(out_dir, exist_ok=True)\n",
"torch.manual_seed(1337 + seed_offset)\n",
"torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul\n",
"torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn\n",
"device_type = 'cuda' if 'cuda' in device else 'mps' # for later use in torch.autocast\n",
"# note: float16 data type will automatically use a GradScaler\n",
"ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]\n",
"ctx = nullcontext()\n",
"\n",
"# poor man's data loader\n",
"data_dir = os.path.join(dataset)\n",
"train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')\n",
"val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')\n",
"def get_batch(split):\n",
" data = train_data if split == 'train' else val_data\n",
" ix = torch.randint(len(data) - block_size, (batch_size,))\n",
" x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])\n",
" y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])\n",
" if device_type == 'cuda':\n",
" # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)\n",
" x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)\n",
" else:\n",
" x, y = x.to(device), y.to(device)\n",
" return x, y\n",
"\n",
"# init these up here, can override if init_from='resume' (i.e. from a checkpoint)\n",
"iter_num = 0\n",
"best_val_loss = 1e9\n",
"\n",
"# attempt to derive vocab_size from the dataset\n",
"meta_path = os.path.join(data_dir, 'meta.pkl')\n",
"meta_vocab_size = None\n",
"if os.path.exists(meta_path):\n",
" with open(meta_path, 'rb') as f:\n",
" meta = pickle.load(f)\n",
" meta_vocab_size = meta['vocab_size']\n",
" print(f\"found vocab_size = {meta_vocab_size} (inside {meta_path})\")\n",
"\n",
"# model init\n",
"model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,\n",
" bias=bias, vocab_size=None, dropout=dropout) # start with model_args from command line\n",
"if init_from == 'scratch':\n",
" # init a new model from scratch\n",
" print(\"Initializing a new model from scratch\")\n",
" # determine the vocab size we'll use for from-scratch training\n",
" if meta_vocab_size is None:\n",
" print(\"defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)\")\n",
" model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304\n",
" gptconf = GPTConfig(**model_args)\n",
" model = GPT(gptconf, True)\n",
"elif init_from == 'resume':\n",
" print(f\"Resuming training from {out_dir}\")\n",
" # resume training from a checkpoint.\n",
" ckpt_path = os.path.join(out_dir, 'ckpt.pt')\n",
" checkpoint = torch.load(ckpt_path, map_location=device)\n",
" checkpoint_model_args = checkpoint['model_args']\n",
" # force these config attributes to be equal otherwise we can't even resume training\n",
" # the rest of the attributes (e.g. dropout) can stay as desired from command line\n",
" for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:\n",
" model_args[k] = checkpoint_model_args[k]\n",
" # create the model\n",
" gptconf = GPTConfig(**model_args)\n",
" model = GPT(gptconf)\n",
" state_dict = checkpoint['model']\n",
" # fix the keys of the state dictionary :(\n",
" # honestly no idea how checkpoints sometimes get this prefix, have to debug more\n",
" unwanted_prefix = '_orig_mod.'\n",
" for k,v in list(state_dict.items()):\n",
" if k.startswith(unwanted_prefix):\n",
" state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
" model.load_state_dict(state_dict)\n",
" iter_num = checkpoint['iter_num']\n",
" best_val_loss = checkpoint['best_val_loss']\n",
"elif init_from.startswith('gpt2'):\n",
" print(f\"Initializing from OpenAI GPT-2 weights: {init_from}\")\n",
" # initialize from OpenAI GPT-2 weights\n",
" override_args = dict(dropout=dropout)\n",
" model = GPT.from_pretrained(init_from, override_args)\n",
" # read off the created config params, so we can store them into checkpoint correctly\n",
" for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:\n",
" model_args[k] = getattr(model.config, k)\n",
"# crop down the model block size if desired, using model surgery\n",
"if block_size < model.config.block_size:\n",
" model.crop_block_size(block_size)\n",
" model_args['block_size'] = block_size # so that the checkpoint will have the right value\n",
"model.to(device)\n",
"\n",
"# initialize a GradScaler. If enabled=False scaler is a no-op\n",
"scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))\n",
"\n",
"# optimizer\n",
"optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)\n",
"if init_from == 'resume':\n",
" optimizer.load_state_dict(checkpoint['optimizer'])\n",
"checkpoint = None # free up memory\n",
"\n",
"# compile the model\n",
"if compile:\n",
" print(\"compiling the model... (takes a ~minute)\")\n",
" unoptimized_model = model\n",
" model = torch.compile(model) # requires PyTorch 2.0\n",
"\n",
"# wrap model into DDP container\n",
"if ddp:\n",
" model = DDP(model, device_ids=[ddp_local_rank])\n",
"\n",
"# helps estimate an arbitrarily accurate loss over either split using many batches\n",
"@torch.no_grad()\n",
"def estimate_loss():\n",
" out = {}\n",
" model.eval()\n",
" for split in ['train', 'val']:\n",
" losses = torch.zeros(eval_iters)\n",
" for k in range(eval_iters):\n",
" X, Y = get_batch(split)\n",
" with ctx:\n",
" logits, loss = model(X, Y)\n",
" losses[k] = loss.item()\n",
" out[split] = losses.mean()\n",
" train_losses.append(out['train'])\n",
" val_losses.append(out['val'])\n",
" model.train()\n",
" return out\n",
"\n",
"# learning rate decay scheduler (cosine with warmup)\n",
"def get_lr(it):\n",
" # 1) linear warmup for warmup_iters steps\n",
" if it < warmup_iters:\n",
" return learning_rate * it / warmup_iters\n",
" # 2) if it > lr_decay_iters, return min learning rate\n",
" if it > lr_decay_iters:\n",
" return min_lr\n",
" # 3) in between, use cosine decay down to min learning rate\n",
" decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)\n",
" assert 0 <= decay_ratio <= 1\n",
" coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1\n",
" return min_lr + coeff * (learning_rate - min_lr)\n",
"\n",
"# logging\n",
"if wandb_log and master_process:\n",
" import wandb\n",
" wandb.init(project=wandb_project, name=wandb_run_name, config=config)\n",
"\n",
"# training loop\n",
"X, Y = get_batch('train') # fetch the very first batch\n",
"t0 = time.time()\n",
"local_iter_num = 0 # number of iterations in the lifetime of this process\n",
"raw_model = model.module if ddp else model # unwrap DDP container if needed\n",
"running_mfu = -1.0\n",
"while True:\n",
"\n",
" # determine and set the learning rate for this iteration\n",
" lr = get_lr(iter_num) if decay_lr else learning_rate\n",
" for param_group in optimizer.param_groups:\n",
" param_group['lr'] = lr\n",
"\n",
" # evaluate the loss on train/val sets and write checkpoints\n",
" if iter_num % eval_interval == 0 and master_process:\n",
" losses = estimate_loss()\n",
"\n",
" print(f\"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
" if wandb_log:\n",
" wandb.log({\n",
" \"iter\": iter_num,\n",
" \"train/loss\": losses['train'],\n",
" \"val/loss\": losses['val'],\n",
" \"lr\": lr,\n",
" \"mfu\": running_mfu*100, # convert to percentage\n",
" })\n",
" if losses['val'] < best_val_loss or always_save_checkpoint:\n",
" best_val_loss = losses['val']\n",
" if iter_num > 0:\n",
" checkpoint = {\n",
" 'model': raw_model.state_dict(),\n",
" 'optimizer': optimizer.state_dict(),\n",
" 'model_args': model_args,\n",
" 'iter_num': iter_num,\n",
" 'best_val_loss': best_val_loss,\n",
" 'config': config,\n",
" }\n",
" print(f\"saving checkpoint to {out_dir}\")\n",
" torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))\n",
" if iter_num == 0 and eval_only:\n",
" break\n",
"\n",
" # forward backward update, with optional gradient accumulation to simulate larger batch size\n",
" # and using the GradScaler if data type is float16\n",
" for micro_step in range(gradient_accumulation_steps):\n",
" if ddp:\n",
" # in DDP training we only need to sync gradients at the last micro step.\n",
" # the official way to do this is with model.no_sync() context manager, but\n",
" # I really dislike that this bloats the code and forces us to repeat code\n",
" # looking at the source of that context manager, it just toggles this variable\n",
" model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)\n",
" with ctx:\n",
" logits, loss = model(X, Y)\n",
" loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation\n",
" # immediately async prefetch next batch while model is doing the forward pass on the GPU\n",
" X, Y = get_batch('train')\n",
" # backward pass, with gradient scaling if training in fp16\n",
" scaler.scale(loss).backward()\n",
" # clip the gradient\n",
" if grad_clip != 0.0:\n",
" scaler.unscale_(optimizer)\n",
" torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
" # step the optimizer and scaler if training in fp16\n",
" scaler.step(optimizer)\n",
" scaler.update()\n",
" # flush the gradients as soon as we can, no need for this memory anymore\n",
" optimizer.zero_grad(set_to_none=True)\n",
"\n",
" # timing and logging\n",
" t1 = time.time()\n",
" dt = t1 - t0\n",
" t0 = t1\n",
" if iter_num % log_interval == 0 and master_process:\n",
" # get loss as float. note: this is a CPU-GPU sync point\n",
" # scale up to undo the division above, approximating the true total loss (exact would have been a sum)\n",
" lossf = loss.item() * gradient_accumulation_steps\n",
" if local_iter_num >= 5: # let the training loop settle a bit\n",
" mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)\n",
" running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu\n",
" loss = Thread(target=estimate_loss)\n",
" print(f\"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%\")\n",
" iter_num += 1\n",
" local_iter_num += 1\n",
"\n",
" # termination conditions\n",
" if iter_num > max_iters:\n",
" break\n",
"\n",
"\n",
"def plot_losses(train_losses, val_losses):\n",
" plt.figure(figsize=(10, 5))\n",
" plt.plot(train_losses, label='Training Loss', color='blue')\n",
" plt.plot(val_losses, label='Validation Loss', color='red')\n",
" plt.xlabel('Epochs')\n",
" plt.ylabel('Loss')\n",
" plt.legend()\n",
" plt.title('Training & Validation Losses')\n",
" plt.show()\n",
"\n",
"plot_losses(train_losses, val_losses)\n",
"\n",
"\n",
"if ddp:\n",
" destroy_process_group()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Export the torch model to ONNX"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"import torch\n",
"import pickle\n",
"import json\n",
"import os\n",
"import onnxruntime as ort\n",
"\n",
"# Sample text to be fed into the model for input.json creation. This can be any string.\n",
"EXAMPLE_TEXT = 'four in let '\n",
"\n",
"# Paths for the export directory (out_dir), the pretrained model ckpt directory and the training data directory (containing meta.pkl)\n",
"out_dir = 'model'\n",
"train_data_dir = 'data'\n",
"pretrained_dir = 'pretrained'\n",
"\n",
"# Joined paths\n",
"checkpoint_path = os.path.join(pretrained_dir, \"ckpt.pt\")\n",
"meta_path = os.path.join(train_data_dir, \"meta.pkl\")\n",
"onnx_path = os.path.join(out_dir, \"network.onnx\")\n",
"input_path = os.path.join(out_dir, \"input.json\")\n",
"\n",
"# Create out dir if not exists\n",
"os.makedirs(out_dir, exist_ok=True)\n",
"\n",
"# Load the model\n",
"checkpoint = torch.load(checkpoint_path)\n",
"gptconf = GPTConfig(**checkpoint[\"model_args\"])\n",
"model = GPT(gptconf)\n",
"state_dict = checkpoint['model']\n",
"\n",
"# Remove unwanted prefix from some of the state dict params (from NanoGPT)\n",
"unwanted_prefix = '_orig_mod.'\n",
"for k,v in list(state_dict.items()):\n",
" if k.startswith(unwanted_prefix):\n",
" state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
"\n",
"# Export the model to ONNX and create input.json\n",
"\n",
"def export():\n",
" # Have tried using torch.compile, removing eval, moving eval below the load state dict... All with the same result.\n",
" model.eval()\n",
" # Load the trained state dict into the model\n",
" model.load_state_dict(state_dict, strict=True)\n",
"\n",
" # Open the metadata pickle and load string to int and vice versa mappings so we can preview decoded text\n",
" with open(meta_path, 'rb') as f:\n",
" meta = pickle.load(f)\n",
" stoi, itos = meta['stoi'], meta['itos']\n",
" encode = lambda s: [stoi[c] for c in s]\n",
" decode = lambda l: ''.join([itos[i] for i in l])\n",
"\n",
" # Convert sample text into tokens for passing into the model\n",
" tokenized_input = (torch.tensor(encode(EXAMPLE_TEXT), dtype=torch.long)[None, ...])\n",
"\n",
" torch.onnx.export(\n",
" model,\n",
" tokenized_input,\n",
" onnx_path,\n",
" export_params=True,\n",
" opset_version=14,\n",
" do_constant_folding=True,\n",
" input_names=['input'],\n",
" output_names=['output'],\n",
" dynamic_axes={'input': {0: 'batch_size'},'output': {0: 'batch_size'}},\n",
" )\n",
"\n",
"export()\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Install EZKL CLI"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import subprocess\n",
"\n",
"subprocess.run([\"curl\", \"https://hub.ezkl.xyz/install_ezkl_cli.sh\", \"|\", \"bash\"])\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Generate input.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import json\n",
"\n",
"# Convert the input string into tokens using the encode method\n",
"input_data = encode(EXAMPLE_TEXT)\n",
"\n",
"# Convert input data into tensor\n",
"input_tensor = torch.tensor(input_data, dtype=torch.long)[None, ...]\n",
"\n",
"# Run the model\n",
"output_tensor = model(input_tensor)\n",
"\n",
"# Convert output tensor to list\n",
"output_data = output_tensor.tolist()[0]\n",
"print(\"Decoded output\", decode(output_data))\n",
"\n",
"# Create a dictionary to hold the data\n",
"data = {\"input_data\": [input_data], \"output_data\": [output_data]}\n",
"\n",
"# Define the path to the input.json file\n",
"input_json_path = \"./model/input.json\"\n",
"\n",
"# Write the data to the input.json file\n",
"with open(input_json_path, 'w') as f:\n",
" json.dump(data, f)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Circuitize the model\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import subprocess\n",
"import os\n",
"\n",
"# Define the base directory\n",
"model_dir = \"./model/\"\n",
"\n",
"# Constructing paths using os.path.join\n",
"onnx_path = os.path.join(model_dir, \"network.onnx\")\n",
"input_path = os.path.join(model_dir, \"input.json\")\n",
"settings_path = os.path.join(model_dir, \"settings.json\")\n",
"srs_path = os.path.join(model_dir, \"kzg.srs\")\n",
"ezkl_path = os.path.join(model_dir, \"network.ezkl\")\n",
"pk_path = os.path.join(model_dir, \"pk.key\")\n",
"vk_path = os.path.join(model_dir, \"vk.key\")\n",
"witness_path = os.path.join(model_dir, \"witness.json\")\n",
"proof_path = os.path.join(model_dir, \"proof.proof\")\n",
"sol_path = os.path.join(model_dir, \"verif.sol\")\n",
"abi_path = os.path.join(model_dir, \"verif.abi\")\n",
"\n",
"# Printing initial message\n",
"print(f\"Creating circuit for {onnx_path}...\")\n",
"\n",
"# Executing the commands\n",
"subprocess.run(\n",
" [\n",
" \"ezkl\",\n",
" \"gen-settings\",\n",
" \"-M\",\n",
" onnx_path,\n",
" \"--param-visibility\",\n",
" \"fixed\",\n",
" \"--input-visibility\",\n",
" \"public\",\n",
" \"--output-visibility\",\n",
" \"public\",\n",
" \"-O\",\n",
" settings_path,\n",
" ]\n",
")\n",
"subprocess.run(\n",
" [\n",
" \"ezkl\",\n",
" \"calibrate-settings\",\n",
" \"-M\",\n",
" onnx_path,\n",
" \"-D\",\n",
" input_path,\n",
" \"-O\",\n",
" settings_path,\n",
" \"--target\",\n",
" \"resources\",\n",
" \"--scales\",\n",
" \"0\",\n",
" ]\n",
")\n",
"subprocess.run([\"ezkl\", \"get-srs\", \"-S\", settings_path, \"--srs-path\", srs_path])\n",
"subprocess.run(\n",
" [\n",
" \"ezkl\",\n",
" \"compile-circuit\",\n",
" \"-M\",\n",
" onnx_path,\n",
" \"-S\",\n",
" settings_path,\n",
" \"--compiled-circuit\",\n",
" ezkl_path,\n",
" ]\n",
")\n",
"subprocess.run(\n",
" [\n",
" \"ezkl\",\n",
" \"setup\",\n",
" \"-M\",\n",
" ezkl_path,\n",
" \"--vk-path\",\n",
" vk_path,\n",
" \"--pk-path\",\n",
" pk_path,\n",
" \"--srs-path\",\n",
" srs_path,\n",
" ]\n",
")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Reproduction of the issue\n",
"\n",
"When attempting to generate a proof for this model, an OOR error is thrown."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"subprocess.run(\n",
" [\"ezkl\", \"gen-witness\", \"-D\", input_path, \"-M\", ezkl_path, \"-O\", witness_path]\n",
")\n",
"\n",
"subprocess.run(\n",
" [\n",
" \"ezkl\",\n",
" \"prove\",\n",
" \"-M\",\n",
" ezkl_path,\n",
" \"-W\",\n",
" witness_path,\n",
" \"--pk-path\",\n",
" pk_path,\n",
" \"--proof-path\",\n",
" proof_path,\n",
" \"--srs-path\",\n",
" srs_path,\n",
" ]\n",
")\n"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.6"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
Display the source blob
Display the rendered blob
Raw
OOR.ipynb
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment