Adding NanoGPT model definition.

Model definition from repository: https://github.com/cgarciae/nanoGPT-jax

README.md

@@ -55,6 +55,10 @@ The main dependencies are `numpy`, `jax` and `chex` libraries.
 ```bash
 pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html
 ```
+Here are the common JAX libraries compatible with IPU:
+```bash
+pip install chex==0.1.6 flax==0.6.4 equinox==0.7.0 jaxtyping==0.2.8s
+```
 
 ## Documentation
 

experiments/nanogpt/model.py

@@ -0,0 +1,337 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import optax
+from flax import traverse_util
+from flax.core import freeze
+from flax.training import train_state
+from flax.traverse_util import path_aware_map
+
+
+@dataclass
+class GPTConfig:
+ block_size: int = 1024
+ vocab_size: int = 50257
+ n_layer: int = 12
+ n_head: int = 12
+ n_embd: int = 768
+ dropout: float = 0.1
+
+
+class CausalSelfAttention(nn.Module):
+ config: GPTConfig
+
+ def setup(self):
+ config = self.config
+ assert config.n_embd % config.n_head == 0
+ # head_size = config.n_embd // config.n_head
+ # key, query, value projections for all heads, but in a batch
+ self.c_attn = nn.Dense(config.n_embd * 3)
+ # output projection
+ self.c_proj = nn.Dense(config.n_embd)
+ # regularization
+ self.attn_dropout = nn.Dropout(config.dropout)
+ self.resid_dropout = nn.Dropout(config.dropout)
+ # causal mask to ensure that attention is only applied to the left in the input sequence
+ self.n_head = config.n_head
+ self.n_embd = config.n_embd
+
+ def __call__(self, x: jax.Array, *, train: bool) -> jax.Array:
+ B, T, C = x.shape # batch size, sequence length, embedding dimensionality (n_embd)
+
+ # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+ qkv = self.c_attn(x)
+ q, k, v = jnp.split(qkv, 3, axis=-1)
+ q = q.reshape(B, T, self.n_head, C // self.n_head).swapaxes(1, 2) # (B, nh, T, hs)
+ k = k.reshape(B, T, self.n_head, C // self.n_head).swapaxes(1, 2) # (B, nh, T, hs)
+ v = v.reshape(B, T, self.n_head, C // self.n_head).swapaxes(1, 2) # (B, nh, T, hs)
+
+ mask = jnp.tril(jnp.ones((T, T))).reshape((1, 1, T, T))
+
+ # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+ att = (q @ k.swapaxes(-2, -1)) * (1.0 / jnp.sqrt(k.shape[-1]))
+ att = jnp.where(mask == 0, float("-inf"), att)
+ att = nn.softmax(att, axis=-1)
+ att = self.attn_dropout(att, deterministic=not train)
+ y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+ y = y.swapaxes(1, 2).reshape(B, T, C) # re-assemble all head outputs side by side
+ # output projection
+ y = self.resid_dropout(self.c_proj(y), deterministic=not train)
+ return y
+
+
+class MLP(nn.Module):
+ config: GPTConfig
+
+ def setup(self):
+ config = self.config
+ self.c_fc = nn.Dense(4 * config.n_embd)
+ self.c_proj = nn.Dense(config.n_embd)
+ self.dropout = nn.Dropout(config.dropout)
+
+ def __call__(self, x: jax.Array, *, train: bool) -> jax.Array:
+ x = self.c_fc(x)
+ x = nn.gelu(x, approximate=True)
+ x = self.c_proj(x)
+ x = self.dropout(x, deterministic=not train)
+ return x
+
+
+class Block(nn.Module):
+ config: GPTConfig
+
+ def setup(self):
+ config = self.config
+ self.ln_1 = nn.LayerNorm(epsilon=1e-5)
+ self.attn = CausalSelfAttention(config)
+ self.ln_2 = nn.LayerNorm(epsilon=1e-5)
+ self.mlp = MLP(config)
+
+ def __call__(self, x: jax.Array, *, train: bool) -> jax.Array:
+ x = x + self.attn(self.ln_1(x), train=train)
+ x = x + self.mlp(self.ln_2(x), train=train)
+ return x
+
+
+class GPT(nn.Module):
+ config: GPTConfig
+
+ def setup(self):
+ config = self.config
+ assert config.vocab_size is not None
+ assert config.block_size is not None
+
+ self.wte = nn.Embed(config.vocab_size, config.n_embd)
+ self.wpe = nn.Embed(config.block_size, config.n_embd)
+ self.drop = nn.Dropout(config.dropout)
+ self.h = [Block(config) for _ in range(config.n_layer)]
+ self.ln_f = nn.LayerNorm()
+
+ def __call__(self, idx: jax.Array, *, train: bool, targets: Optional[jax.Array] = None) -> jax.Array:
+ b, t = idx.shape
+ assert (
+ t <= self.config.block_size
+ ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+ pos = jnp.arange(0, t, dtype=jnp.int32)[None] # shape (1, t)
+
+ # forward the GPT model itself
+ tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
+ pos_emb = self.wpe(pos) # position embeddings of shape (1, t, n_embd)
+ x = self.drop(tok_emb + pos_emb, deterministic=not train)
+ for block in self.h:
+ x = block(x, train=train)
+ x = self.ln_f(x)
+
+ logits = self.wte.attend(x)
+
+ if targets is not None:
+ # if we are given some desired targets also calculate the loss
+ loss = optax.softmax_cross_entropy_with_integer_labels(logits, targets).mean()
+ else:
+ loss = None
+
+ return logits, loss
+
+ def crop_block_size(self, params: Any, block_size: int) -> Any:
+ # model surgery to decrease the block size if necessary
+ # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+ # but want to use a smaller block size for some smaller, simpler model
+
+ assert block_size <= self.config.block_size
+ self.config.block_size = block_size
+
+ # self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+ def crop_weights(path: Tuple[str, ...], x: Any) -> Any:
+ if path[-2:] == ("wpe", "embedding"):
+ return x[:block_size]
+ return x
+
+ return freeze(path_aware_map(crop_weights, params))
+
+ @classmethod
+ def from_pretrained(cls, model_type, override_args=None):
+ assert model_type in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"}
+ override_args = override_args or {} # default to empty dict
+ # only dropout can be overridden see more notes below
+ assert all(k == "dropout" for k in override_args)
+ from transformers import GPT2LMHeadModel
+
+ print("loading weights from pretrained gpt: %s" % model_type)
+
+ # n_layer, n_head and n_embd are determined from model_type
+ config_args = {
+ "gpt2": dict(n_layer=12, n_head=12, n_embd=768), # 124M params
+ "gpt2-medium": dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+ "gpt2-large": dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+ "gpt2-xl": dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+ }[model_type]
+ # we can override the dropout rate
+ if "dropout" in override_args:
+ config_args["dropout"] = override_args["dropout"]
+ # block_size is always 1024 for GPT model checkpoints
+ # if one wants a lower block_size it has to be done through model surgery
+ # later, by calling crop_block_shape
+
+ # create a from-scratch initialized minGPT model
+ config = GPTConfig(block_size=1024, **config_args)
+ model = GPT(config)
+ variables = jax.eval_shape(
+ lambda: model.init(jax.random.PRNGKey(0), jnp.ones((1, 1), dtype=jnp.int32), train=False)
+ )
+ params = variables["params"]
+ flat_params = traverse_util.flatten_dict(params, sep=".")
+
+ # init a huggingface/transformers model
+ model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+ sd_hf = model_hf.state_dict()
+
+ def copy_from(flax_name, pt_name, transpose=False, add_head_dim=False):
+ pt_tensor = sd_hf[pt_name]
+ jax_array = flat_params[flax_name]
+ if transpose:
+ pt_tensor = pt_tensor.t()
+ pt_array = pt_tensor.detach().cpu().numpy()
+
+ if add_head_dim:
+ # pt_array = pt_array.reshape(*pt_array.shape[:-1], config.n_head, -1, 3)
+ pass
+
+ assert pt_array.shape == jax_array.shape
+
+ flat_params[flax_name] = pt_array
+
+ # transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+ copy_from("wte.embedding", "transformer.wte.weight")
+ copy_from("wpe.embedding", "transformer.wpe.weight")
+
+ for i in range(config.n_layer):
+ copy_from(f"h_{i}.ln_1.scale", f"transformer.h.{i}.ln_1.weight")
+ copy_from(f"h_{i}.ln_1.bias", f"transformer.h.{i}.ln_1.bias")
+ copy_from(f"h_{i}.attn.c_attn.kernel", f"transformer.h.{i}.attn.c_attn.weight", add_head_dim=True)
+ copy_from(f"h_{i}.attn.c_attn.bias", f"transformer.h.{i}.attn.c_attn.bias", add_head_dim=True)
+ copy_from(f"h_{i}.attn.c_proj.kernel", f"transformer.h.{i}.attn.c_proj.weight")
+ copy_from(f"h_{i}.attn.c_proj.bias", f"transformer.h.{i}.attn.c_proj.bias")
+ copy_from(f"h_{i}.ln_2.scale", f"transformer.h.{i}.ln_2.weight")
+ copy_from(f"h_{i}.ln_2.bias", f"transformer.h.{i}.ln_2.bias")
+ copy_from(f"h_{i}.mlp.c_fc.kernel", f"transformer.h.{i}.mlp.c_fc.weight")
+ copy_from(f"h_{i}.mlp.c_fc.bias", f"transformer.h.{i}.mlp.c_fc.bias")
+ copy_from(f"h_{i}.mlp.c_proj.kernel", f"transformer.h.{i}.mlp.c_proj.weight")
+ copy_from(f"h_{i}.mlp.c_proj.bias", f"transformer.h.{i}.mlp.c_proj.bias")
+
+ copy_from("ln_f.scale", "transformer.ln_f.weight")
+ copy_from("ln_f.bias", "transformer.ln_f.bias")
+
+ params = freeze(traverse_util.unflatten_dict(flat_params, sep="."))
+
+ return model, params
+
+ def configure_optimizers(self, params, weight_decay, learning_rate, betas):
+ """
+ This long function is unfortunately doing something very simple and is being very defensive:
+ We are separating out all parameters of the model into two buckets: those that will experience
+ weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
+ We are then returning the PyTorch optimizer object.
+ """
+
+ def get_optimizer(decay):
+ return optax.adamw(learning_rate=learning_rate, b1=betas[0], b2=betas[1], weight_decay=decay)
+
+ def partition_fn(path: Tuple[str, ...], x: Any) -> str:
+ if path[-1] in ("bias", "scale", "embedding"):
+ return "no_decay"
+ elif path[-1] in ("kernel",):
+ return "decay"
+ else:
+ raise ValueError(f"Unrecognized parameter: {path}")
+
+ partition_optimizers = {"decay": get_optimizer(weight_decay), "no_decay": get_optimizer(0.0)}
+ param_partitions = freeze(path_aware_map(partition_fn, params))
+ tx = optax.multi_transform(partition_optimizers, param_partitions)
+
+ return tx
+
+ # @torch.no_grad()
+ def generate(self, key, params, input_tokens, max_new_tokens, temperature=1.0, top_k=None):
+ """
+ Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+ the sequence max_new_tokens times, feeding the predictions back into the model each time.
+ Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+ """
+ B, T = input_tokens.shape
+ padding = jnp.zeros((B, max_new_tokens), dtype=jnp.int32)
+ tokens = jnp.concatenate([input_tokens, padding], axis=-1)
+ indexes = jnp.arange(T, T + max_new_tokens)
+
+ # tokens index -> tokens None
+ def scan_f(tokens, i):
+ # l: x y
+ # t: a b - -
+ # i: 0 1 2 3
+ step_key = jax.random.fold_in(key, i)
+ # if the sequence context is growing too long we must crop it at block_size
+ # idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+ # forward the model to get the logits for the index in the sequence
+ logits, _ = self.apply({"params": params}, tokens, train=False)
+ # pluck the logits at the final step and scale by desired temperature
+ logits = logits[:, i - 1, :] / temperature
+ # optionally crop the logits to only the top k options
+ # sample from the distribution
+ if top_k is not None:
+ top_logits, top_tokens = jax.lax.top_k(logits, min(top_k, logits.shape[-1]))
+ token_idx = jax.random.categorical(step_key, top_logits, axis=-1)
+ next_token = jnp.take_along_axis(top_tokens, token_idx[:, None], axis=-1).squeeze(-1)
+ else:
+ next_token = jax.random.categorical(step_key, logits, axis=-1)
+ # logits = jnp.where(logits < v[:, -1:], float('-inf'), logits)
+ # append sampled index to the running sequence and continue
+ tokens = tokens.at[:, i].set(next_token)
+
+ return tokens, None
+
+ tokens, _ = jax.lax.scan(scan_f, tokens, indexes)
+
+ return tokens
+
+ def create_state(
+ self,
+ learning_rate,
+ weight_decay,
+ beta1,
+ beta2,
+ decay_lr=None,
+ warmup_iters=None,
+ lr_decay_iters=None,
+ min_lr=None,
+ params=None,
+ **kwargs,
+ ):
+ if params is None:
+ variables = self.init(jax.random.PRNGKey(0), jnp.ones((1, 1), dtype=jnp.int32), train=False)
+ params = variables["params"]
+ if decay_lr:
+ assert warmup_iters is not None and lr_decay_iters is not None and min_lr is not None
+ lr_schedule = optax.warmup_cosine_decay_schedule(
+ init_value=0.0,
+ peak_value=learning_rate,
+ warmup_steps=warmup_iters,
+ decay_steps=lr_decay_iters,
+ end_value=min_lr,
+ )
+ else:
+ lr_schedule = learning_rate
+ tx = self.configure_optimizers(
+ params, weight_decay=weight_decay, learning_rate=lr_schedule, betas=(beta1, beta2)
+ )
+ return train_state.TrainState.create(apply_fn=self.apply, params=params, tx=tx)

experiments/nanogpt/model_test.py

@@ -0,0 +1,28 @@
+"""Testing NanoGPT JAX model definition.
+
+Inspired by: https://github.com/cgarciae/nanoGPT-jax/blob/master/train.py
+"""
+
+import jax
+import jax.numpy as jnp
+from model import GPT, GPTConfig
+
+gpt2_tiny = GPTConfig(block_size=128, vocab_size=32000, n_layer=2, n_head=8, n_embd=512)
+train_config = dict(
+ learning_rate=0.001,
+ weight_decay=0.1,
+ beta1=1,
+ beta2=1,
+)
+
+rng_key = jax.random.PRNGKey(0)
+init_value = jnp.ones((1, 1), dtype=jnp.int32)
+
+model = GPT(gpt2_tiny)
+# initialize weights
+# state = model.create_state(**train_config)
+params = model.init(rng_key, init_value, train=False)
+print("Model initialized...")
+
+# Model description
+print(model.tabulate(rng_key, init_value, train=False))