让我们用 JAX 重建 NanoGPT！（第一部分）

原文：towardsdatascience.com/lets-reproduce-nanogpt-with-jax-part-1-95bec4630eb4?source=collection_archive---------2-----------------------#2024-07-21

fromdataclassesimportdataclass@dataclassclassModelConfig:vocab_size:int=50257n_head:int=12n_embd:int=768block_size:int=1024n_layer:int=12dropout_rate:float=0.1

接下来是变换器模型的关键构建模块——注意力机制（Attention）。其思想是将输入处理成三个权重矩阵：Key、Query 和 Value。在这里，我们依赖于flax，这是一个 Jax 层和训练 API 库，用来初始化这三个权重矩阵，只需要调用[flax.linen.Dense](https://flax.readthedocs.io/en/v0.5.3/_autosummary/flax.linen.Dense.html)。如前所述，Jax 有许多类似 numpy 的 API，因此我们使用[jax.numpy.reshape](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.reshape.html)将权重矩阵后的输出从[batch_size, sequence_length, embedding_dim]重塑为[batch_size, sequence_length, num_head, embedding_dim / num_head]。由于我们需要对 Key 和 Value 矩阵执行矩阵乘法，jax 还提供了[jax.numpy.matmul](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.matmul.html)和[jax.numpy.transpose](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.transpose.html)API（用于转置 Key 矩阵以进行乘法运算）。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/8f2916bfb42338ef17e1526a677e4f85.png

多头注意力（Multihead Attention）

请注意，我们需要在注意力矩阵上加上一个掩码，以避免信息泄漏（防止之前的 tokens 访问到后面的 tokens），[jax.numpy.tril](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.tril.html)帮助构建一个下三角数组，而[jax.numpy.where](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.where.html)可以为我们填充无限大的数值，以便在 softmax[jax.nn.softmax](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softmax.html)后得到 0。多头注意力的完整代码如下所示。

fromflaximportlinenasnnimportjax.numpyasjnpclassCausalSelfAttention(nn.Module):config:ModelConfig@nn.compactdef__call__(self,x,deterministic=True):assertlen(x.shape)==3b,l,d=x.shape q=nn.Dense(self.config.n_embd)(x)k=nn.Dense(self.config.n_embd)(x)v=nn.Dense(self.config.n_embd)(x)# q*k / sqrt(dim) -> softmax -> @vq=jnp.reshape(q,(b,l,d//self.config.n_head,self.config.n_head))k=jnp.reshape(k,(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(list(jnp.shape(k))[-1])attn=jnp.matmul(q,jnp.transpose(k,(0,1,3,2)))/norm mask=jnp.tril(attn)attn=jnp.where(mask[:,:,: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)returny

你可能会注意到，在 Pytorch 中常见的__init__或forward方法在这里并不存在。这是 jax 的特点，在 jax 中你可以显式地通过setup方法定义层，或者通过在__call__方法上添加nn.compact来隐式定义它们。[参考]

接下来让我们构建 MLP 和 Block 层，包括 Dense 层、Gelu 激活函数、LayerNorm 和 Dropout。再次，flax.linen 提供了层的 API，帮助我们构建模块。请注意，我们会传递一个deterministic布尔变量来控制某些层（如 Dropout）在训练或评估期间的不同行为。

classMLP(nn.Module):config:ModelConfig@nn.compactdef__call__(self,x,deterministic=True):x=nn.Dense(self.config.n_embd*4)(x)x=nn.gelu(x,approximate=True)x=nn.Dropout(rate=self.config.dropout_rate)(x,deterministic=deterministic)x=nn.Dense(self.config.n_embd)(x)x=nn.Dropout(rate=self.config.dropout_rate)(x,deterministic=deterministic)returnxclassBlock(nn.Module):config:ModelConfig@nn.compactdef__call__(self,x):x=nn.LayerNorm()(x)x=x+CausalSelfAttention(self.config)(x)x=nn.LayerNorm()(x)x=x+MLP(self.config)(x)returnx

现在让我们使用上述模块来构建 NanoGPT：

给定一个序列的 token ids 输入，我们使用[flax.linen.Embed](https://flax.readthedocs.io/en/v0.5.3/_autosummary/flax.linen.Embed.html)层来获取位置嵌入和 token 嵌入。然后，我们将它们传入 Block 模块 N 次，其中 N 是模型配置中定义的层数。最后，我们将来自最后一个 Block 的输出映射到每个词汇表 token 的概率，以预测下一个 token。除了前向__call__方法之外，我们还需要创建一个init方法来获取虚拟输入并获得模型的参数。

classGPT(nn.Module):config:ModelConfig@nn.compactdef__call__(self,x,deterministic=False):B,T=x.shapeassertT<=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_embfor_inrange(self.config.n_layer):x=Block(self.config)(x)x=nn.LayerNorm()(x)logits=nn.Dense(config.n_embd,config.vocab_size)(x)# logits = wte.attend(x) # parameter sharingreturnlogitsdefinit(self,rng):tokens=jnp.zeros((1,self.config.block_size),dtype=jnp.uint16)params=jax.jit(super().init,static_argnums=(2,))(rng,tokens,True)returnparams

现在让我们验证一下参数的数量：我们首先初始化模型配置的数据类和随机密钥，然后创建一个虚拟输入并将其输入到 GPT 模型中。接着，我们利用jax.util.treemapAPI 创建一个计数参数函数。我们得到了124439808（124M）个参数，与 Huggingface 的 GPT2 相同，哇！

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/513bfdf90096dcb03e3cd4a76910d2d7.png

Colab 结果：参数数量

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/576fc03b5cb1912c0ba770b510de7b73.png

验证 Huggingface 的 GPT2 参数数量

数据加载器和训练循环

现在让我们在一个小数据集上进行过拟合。为了与 Andrej 的 Pytorch NanoGPT 视频中进行对比，我们使用他在视频中分享的玩具 dataset。我们使用tiktoken库的 GPT2 分词器对输入文件中的所有文本进行分词，并将这些 token 转换为jax.numpy.array以便 Jax 的模型训练。

classDataLoader:def__init__(self,B,T):self.current_position=0self.B=B self.T=Twithopen("input.txt","r")asf:text=f.read()enc=tiktoken.get_encoding("gpt2")self.tokens=jnp.array(enc.encode(text))print(f"loaded{len(self.tokens)}tokens in the datasets")print(f" 1 epoch ={len(self.tokens)//(B*T)}batches")defnext_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*Tifself.current_position+B*T+1>len(self.tokens):self.current_position=0returnx,y

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/da59795faf894d9f4bd6f4d1ed782ece.png

Colab 结果：简单的数据加载器，批量大小为 4，序列长度为 128

接下来，让我们暂时忽略分布式训练和优化，先创建一个简单的训练循环进行基本检查。初始化模型后的第一件事是创建一个TrainState，这是一个可以更新参数和梯度的模型状态。TrainState 接受三个重要输入：apply_fn（模型前向函数）、params（来自初始化方法的模型参数）和 tx（一个 Optax 梯度变换）。

然后我们使用 train_step 函数来更新模型状态（梯度和参数），以继续模型训练。Optax提供了用于下一个令牌预测任务的 softmax 交叉熵作为损失函数，jax.value_and_grad用于计算损失函数的梯度和损失值。最后，我们使用apply_gradientsAPI 更新模型的状态和新参数。[ref] 别忘了对 train_step 函数进行 JIT 编译，以减少计算开销！

definit_train_state(key,config)->TrainState:model=GPT(config)params=model.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=model.apply,params=params,tx=optimizer)returntrain_state@jax.jitdeftrain_step(state:TrainState,x:jnp.ndarray,y:jnp.ndarray)->Tuple[jnp.ndarray,TrainState]:defloss_fn(params:FrozenDict)->jnp.ndarray:logits=state.apply_fn(params,x,False)loss=optax.softmax_cross_entropy_with_integer_labels(logits,y).mean()returnloss loss,grads=jax.value_and_grad(loss_fn,has_aux=False)(state.params)new_state=state.apply_gradients(grads=grads)returnloss,new_state