Impressed by Andrei Kapathy’s current YouTube video Let’s try to replicate GPT-2 (124M)we want to rebuild it utilizing most of Jax’s coaching optimizations. Jax is constructed for extremely environment friendly computation velocity and it could be very attention-grabbing to match Pytorch with its newest coaching optimizations and Jax with its associated libraries similar to Flax (Jax’s layer API for neural community coaching) and Optax (JAX’s gradient processing and optimization library). We are going to briefly study what Jax is and rebuild GPT utilizing Jax. Lastly, we are going to examine tokens/sec on multi-GPU coaching between Pytorch and Jax.
What’s Jax?
the Read the documentationJAX is a Python library for accelerator-oriented array computations and program transformations designed for high-performance numerical computation and large-scale machine studying. I want to introduce the title JAX. XLMA (Accelerated Linear Algibra), which I choose to name J(it) A(utograd) X(LA) to mirror its excessive effectivity.
J – Simply-in-time (JIT) compilation. While you run a Python operate, Jax interprets it right into a set of primitive operations referred to as Jaxpr. The Jaxpr expressions are then translated into inputs for XLA, which compiles the low-level script to supply an executable that’s optimized for the goal system (CPU, GPU, or TPU).
A — Autograd. Gradient calculations are a necessary a part of fashionable machine studying strategies. jax.grad() Get hold of the gradients that help you optimize the mannequin.
X — XLA. That is an open supply machine studying compiler for CPUs, GPUs, and ML accelerators. Normally, XLA is: Stable HLO It creates the graph and sends the HLO computation to a backend for additional HLO-level optimizations, which then performs target-specific code technology.
These are a few of the key options of JAX, however there are additionally many extra user-friendly NumPy-like APIs out there. jax.numpy and auto-vectorization jax.vmap after which parallelize the code throughout a number of gadgets. jax.pmap We are going to discover Jax ideas and functions in additional element in future blogs, however for now let’s attempt to replicate NanoGPT utilizing Jax.
From consideration to vary
GPT is a decoder-only Transformer mannequin, and its important constructing block is the Consideration module. First, we outline a mannequin configuration dataclass to retailer the mannequin’s hyperparameters in order that the mannequin module can use it effectively to initialize the mannequin structure. Much like the 124M GPT mannequin, right here we initialize a 12-layer Transformer decoder with 12 heads and 50257 tokens, every with an embedding dimension of 768. The block dimension for consideration computation is 1024.
from dataclasses import dataclass@dataclass
class ModelConfig:
vocab_size: int = 50257
n_head: int = 12
n_embd: int = 768
block_size: int = 1024
n_layer: int = 12
dropout_rate: float = 0.1
Subsequent comes the important thing part of the Transformer mannequin: consideration. The thought is to course of the enter into three weight matrices: key, question, and worth. Right here, flax ,We use the Jax layers and coaching API library to initialize the three weight matrices. flax.linen.Dense As talked about earlier than, Jax has a whole lot of NumPy-like APIs, so we remodel the output after the load matrix as follows: jax.numpy.reshape from [batch_size, sequence_length, embedding_dim] To [batch_size, sequence_length, num_head, embedding_dim / num_head]Since we have to do a matrix multiplication of keys and values, jax requires jax.numpy.matmul API and jax.numpy.transpose (Transpose the multiplication key matrix).
Observe that we have to masks the eye matrix to stop info leakage (stopping earlier tokens from accessing later tokens). jax.numpy.tril This helps to construct the decrease triangular array. jax.numpy.where You may pad with infinite numbers to get 0 after softmax jax.nn.softmax The total code for multi-head consideration could be discovered under.
from flax import linen as nn
import jax.numpy as jnpclass CausalSelfAttention(nn.Module):
config: ModelConfig
@nn.compact
def __call__(self, x, deterministic=True):
assert len(x.form) == 3
b, l, d = x.form
q = nn.Dense(self.config.n_embd)(x)
ok = nn.Dense(self.config.n_embd)(x)
v = nn.Dense(self.config.n_embd)(x)
# q*ok / sqrt(dim) -> softmax -> @v
q = jnp.reshape(q, (b, l, d//self.config.n_head , self.config.n_head))
ok = jnp.reshape(ok, (b, l, d//self.config.n_head , self.config.n_head))
v = jnp.reshape(v, (b, l, d//self.config.n_head , self.config.n_head))
norm = jnp.sqrt(checklist(jnp.form(ok))[-1])
attn = jnp.matmul(q,jnp.transpose(ok, (0,1,3,2))) / norm
masks = jnp.tril(attn)
attn = jnp.the place(masks[:,:,:l,:l], attn, float("-inf"))
probs = jax.nn.softmax(attn, axis=-1)
y = jnp.matmul(probs, v)
y = jnp.reshape(y, (b,l,d))
y = nn.Dense(self.config.n_embd)(y)
return y
It’s possible you’ll discover, __init__ or ahead It is a methodology you’d discover in pytorch. It is a particular characteristic of JAX that lets you explicitly outline layers. setup You may outline a technique or outline it implicitly within the ahead path. nn.compact On high of the __call__ Technique. [ref]
Subsequent, let’s construct the MLP and block layers, which embrace a Dense layer, a Gelu activation operate, a LayerNorm, and a Dropout. Once more, flax.linen has a layer API to assist us construct modules. deterministic Boolean variables that management completely different behaviors throughout coaching or analysis of some layers, similar to Dropout.
class MLP(nn.Module):config: ModelConfig
@nn.compact
def __call__(self, x, deterministic=True):
x = nn.Dense(self.config.n_embd*4)(x)
x = nn.gelu(x, approximate=True)
x = nn.Dropout(fee=self.config.dropout_rate)(x, deterministic=deterministic)
x = nn.Dense(self.config.n_embd)(x)
x = nn.Dropout(fee=self.config.dropout_rate)(x, deterministic=deterministic)
return x
class Block(nn.Module):
config: ModelConfig
@nn.compact
def __call__(self, x):
x = nn.LayerNorm()(x)
x = x + CausalSelfAttention(self.config)(x)
x = nn.LayerNorm()(x)
x = x + MLP(self.config)(x)
return x
Now let’s construct NanoGPT utilizing the blocks above.
Given a sequence token ID enter, flax.linen.Embed We use layers to acquire place embeddings and token embeddings. Then we move them to the block module N instances, the place N is the variety of layers outlined within the mannequin configuration. Lastly, we map the output from the final block to the chance of every token within the vocabulary to foretell the following token. __call__ Let’s additionally create a technique init A way to get dummy inputs to get the mannequin parameters.
class GPT(nn.Module):config: ModelConfig
@nn.compact
def __call__(self, x, deterministic=False):
B, T = x.form
assert T <= self.config.block_size
pos = jnp.arange(0, T)[None]
pos_emb = nn.Embed(self.config.block_size, self.config.n_embd)(pos)
wte = nn.Embed(self.config.vocab_size, self.config.n_embd)
tok_emb = wte(x)
x = tok_emb + pos_emb
for _ in vary(self.config.n_layer):
x = Block(self.config)(x)
x = nn.LayerNorm()(x)
logits = nn.Dense(config.n_embd, config.vocab_size)
# logits = wte.attend(x) # parameter sharing
return logits
def init(self, rng):
tokens = jnp.zeros((1, self.config.block_size), dtype=jnp.uint16)
params = jax.jit(tremendous().init, static_argnums=(2,))(rng, tokens, True)
return params
Now, let’s test the variety of parameters. First, initialize the mannequin configuration knowledge class and random keys, then create dummy inputs and feed them to the GPT mannequin. Subsequent, jax.util.treemap An API for creating rely parameter features. 124439808 (124M) parameters, the identical quantity as Huggingface’s GPT2, BOOM!
DataLoader and the coaching loop
Now let’s attempt overfitting a small dataset, utilizing a toy so we will examine it to Andrej’s Pytorch NanoGPT video. data set I exploit the GPT2 tokenizer that he shared in his video tiktoken Tokenize all of the textual content from the enter file and extract the tokens jax.numpy.array For Jax mannequin coaching.
class DataLoader:
def __init__(self, B, T):
self.current_position = 0
self.B = B
self.T = Twith open("enter.txt","r") as f:
textual content = f.learn()
enc = tiktoken.get_encoding("gpt2")
self.tokens = jnp.array(enc.encode(textual content))
print(f"loaded {len(self.tokens)} tokens within the datasets" )
print(f" 1 epoch = {len(self.tokens)//(B*T)} batches")
def next_batch(self):
B,T = self.B, self.T
buf = self.tokens[self.current_position:self.current_position+B*T+1]
x,y = jnp.reshape(buf[:-1],(B,T)), jnp.reshape(buf[1:],(B,T))
self.current_position += B*T
if self.current_position + B*T+1 > len(self.tokens):
self.current_position = 0
return x,y
Subsequent, let’s neglect about distributed coaching and optimization and create a easy coaching loop for a sanity test. The very first thing we do after initializing the mannequin is Train State,That is the mannequin state the place we will replace parameters and gradients.,TrainState takes three necessary inputs: apply_fn (mannequin,ahead operate), params (mannequin parameters from the init methodology),,and tx (Optax gradient transformation).
Then, use the train_step operate to replace the mannequin state (gradients and parameters) and proceed coaching the mannequin. Optax We offer softmax cross entropy because the loss operate for the following token prediction process, jax.value_and_grad We calculate the gradient of the loss operate and the loss worth. Lastly, we replace the mannequin state with the brand new parameters. apply_gradients API. [ref] Remember to jit your train_step operate to scale back computational overhead.
def init_train_state(key, config) -> TrainState:
mannequin = GPT(config)
params = mannequin.init(key)
optimizer = optax.adamw(3e-4, b1=0.9, b2=0.98, eps=1e-9, weight_decay=1e-1)
train_state = TrainState.create(
apply_fn=mannequin.apply,
params=params,
tx=optimizer)
return train_state@jax.jit
def train_step(state: TrainState, x: jnp.ndarray, y: jnp.ndarray) -> Tuple[jnp.ndarray, TrainState]:
def loss_fn(params: FrozenDict) -> jnp.ndarray:
logits = state.apply_fn(params, x, False)
loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).imply()
return loss
loss, grads = jax.value_and_grad(loss_fn, has_aux=False)(state.params)
new_state = state.apply_gradients(grads=grads)
return loss, new_state
Now all the things is prepared for poorman’s coaching loop. Let’s test the loss worth. Our mannequin’s predictions needs to be higher than random guessing, so the loss needs to be decrease than -ln(1/50257)≈10.825. What we anticipate from overfitting on one batch is that the loss might be near 10.825 at first, after which near 0. Let’s take a batch of (x, y) and run the coaching loop 50 instances. We additionally add an analogous log to calculate the coaching velocity.
As you possibly can see, the loss values are precisely as anticipated, and the coaching throughput is round 400-500k tokens/sec. That is already 40x quicker than the preliminary unoptimized model of Pytorch in Andrej’s video. Observe that I am operating the Jax script on a single A100 GPU, so the velocity comparability eliminates {hardware} variations. .to(system) It is a software for transferring fashions and knowledge from the host CPU to the system GPU, which is without doubt one of the benefits of Jax.
In order that’s it, partly 2 we’ll do some additional optimizations to hurry up coaching by 10x.
part 2: A coaching optimization journey to 1350k tokens/sec on a single GPU!
“All photographs are property of the writer except in any other case famous.”
