The original GAN loss is based on Jensen-Shannon (JS) divergence between the real distribution Pr and generated distribution Pg. The Wasserstein GAN is based on Earth Mover distance between these distributions.
-
W(Pr,Pg)=γ∈Π(Pr,Pg)infE(x,y)∼γ∥x−y∥
+
W(Pr,Pg)=γ∈Π(Pr,Pg)infE(x,y)∼γ∥x−y∥
Π(Pr,Pg) is the set of all joint distributions, whose marginal probabilities are γ(x,y).
E(x,y)∼γ∥x−y∥ is the earth mover distance for a given joint distribution (x and y are probabilities).
-
So W(Pr,Pg) is equal to the least earth mover distance for any joint distribution between the real distribution Pr and generated distribution Pg.
+
So W(Pr,Pg) is equal to the least earth mover distance for any joint distribution between the real distribution Pr and generated distribution Pg.
The paper shows that Jensen-Shannon (JS) divergence and other measures for the difference between two probability distributions are not smooth. And therefore if we are doing gradient descent on one of the probability distributions (parameterized) it will not converge.
-
Based on Kantorovich-Rubinstein duality, W(Pr,Pg)=∥f∥L≤1supEx∼Pr[f(x)]−Ex∼Pg[f(x)]
+
Based on Kantorovich-Rubinstein duality, W(Pr,Pg)=∥f∥L≤1supEx∼Pr[f(x)]−Ex∼Pg[f(x)]
where ∥f∥L≤1 are all 1-Lipschitz functions.
That is, it is equal to the greatest difference Ex∼Pr[f(x)]−Ex∼Pg[f(x)] among all 1-Lipschitz functions.
-
For K-Lipschitz functions, W(Pr,Pg)=∥f∥L≤KsupEx∼Pr[K1f(x)]−Ex∼Pg[K1f(x)]
+
For K-Lipschitz functions, W(Pr,Pg)=∥f∥L≤KsupEx∼Pr[K1f(x)]−Ex∼Pg[K1f(x)]
If all K-Lipschitz functions can be represented as fw where f is parameterized by w∈W,
114defforward(self,input_ids):
+115batch_size,input_shape=input_ids.size()
+116
+117token_embeddings=self.token_embedding(input_ids)# B T C
+118position_ids=torch.arange(input_shape,device=config['device'])# T C
+119position_embeddings=self.position_embedding(position_ids)# B T C
+120
+121hidden_states=token_embeddings+position_embeddings
+122
+123forblockinself.blocks:
+124hidden_states=block(hidden_states)
+125
+126hidden_states=self.final_norm(hidden_states)
+127
+128logits=self.lm_head(hidden_states)
+129
+130returnlogits
Low-Rank Adaptation (LoRA) freezes pre-trained model weights and injects trainable rank decomposition matrices into each layer of the transformer. This makes it possible to efficiently fine-tune large langauge models by reducing trainable parameters by a large factor.
+
Here's the training code for training a GPT2 model with LoRA on Tiny Shakespeare dataset.
diff --git a/labml_nn/gan/wasserstein/__init__.py b/labml_nn/gan/wasserstein/__init__.py
index b3c52472..3c5394e0 100644
--- a/labml_nn/gan/wasserstein/__init__.py
+++ b/labml_nn/gan/wasserstein/__init__.py
@@ -26,7 +26,7 @@ marginal probabilities are $\gamma(x, y)$.
$\mathbb{E}_{(x,y) \sim \gamma} \Vert x - y \Vert$ is the earth mover distance for
a given joint distribution ($x$ and $y$ are probabilities).
-So $W(\mathbb{P}_r, \mathbb{P}g)$ is equal to the least earth mover distance for
+So $W(\mathbb{P}_r, \mathbb{P}_g)$ is equal to the least earth mover distance for
any joint distribution between the real distribution $\mathbb{P}_r$ and generated distribution $\mathbb{P}_g$.
The paper shows that Jensen-Shannon (JS) divergence and other measures for the difference between two probability
diff --git a/labml_nn/lora/__init__.py b/labml_nn/lora/__init__.py
new file mode 100644
index 00000000..f5fc197d
--- /dev/null
+++ b/labml_nn/lora/__init__.py
@@ -0,0 +1,151 @@
+"""
+---
+title: Low-Rank Adaptation (LoRA)
+summary: >
+ Annotated implementation of RoRA from paper
+ LoRA: Low-Rank Adaptation of Large Language Models
+---
+
+# Low-Rank Adaptation (LoRA)
+
+This is an implementation of
+[Low-Rank Adaptation (LoRA)](https://arxiv.org/abs/2106.09685)
+in [PyTorch](https://pytorch.org).
+
+Low-Rank Adaptation (LoRA) freezes pre-trained model weights and injects
+ trainable rank decomposition matrices into each layer of the transformer.
+ This makes it possible to efficiently fine-tune large langauge models by
+ reducing trainable parameters by a large factor.
+
+Here's [the training code](experiment.html) for training a GPT2 model with LoRA
+ on Tiny Shakespeare dataset.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class Linear(nn.Module):
+ """
+ ## LoRA Linear Layer
+
+ LoRA linear layer adds a low-rank decomposition to the pre-trained
+ weight matrix ($W_0 \in \mathbb{R}^{d \times k}$)
+ of the linear layer.
+
+ $$W_0 + \Delta W = W_0 + BA$$
+
+ , where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$,
+ and the rank $r \ll min(d, k)$.
+
+ All parameters are frozen except $A$ and $B$.
+
+ $\Delta W$ is initialized to be zero at the beginning of the training.
+
+ They multiple $\Delta W x$ by $\frac{\alpha}{r}$ where $\alpha$ is a hyper-parameter.
+ Once $\alpha$ is tuned it can be kept the same when varying $r$.
+ """
+
+ def __init__(self, in_features: int, out_features: int, bias: bool,
+ r: int, alpha: int = None):
+ """
+ :param in_features: is the number of input features of the linear layer
+ :param out_features: is the number of output features of the linear layer
+ :param bias: is a flag indicating if there is a bias parameter
+ :param r: is the rank of the decomposition $r$
+ :param alpha: is the scaling factor $\alpha$
+ """
+ super().__init__()
+
+ # Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
+ if alpha is None:
+ alpha = r
+
+ # The pre-trained weight $W_0$
+ self.weight = nn.Parameter(torch.empty((out_features, in_features)))
+ # Freeze it
+ self.weight.requires_grad = False
+
+ if bias:
+ # Bias parameter $b_0$ (also frozen)
+ self.bias = nn.Parameter(torch.empty(out_features))
+ self.bias.requires_grad = False
+ else:
+ # No bias parameter
+ self.bias = None
+
+ # scaling factor $\frac{\alpha}{r}$
+ self.scaling = alpha / r
+ # Matrix $A \in \mathbb{R}^{r \times k}$
+ self.lora_a = nn.Parameter(torch.empty((in_features, r)))
+ # Matrix $B \in \mathbb{R}^{d \times r}$, we keep $A$ and $B$ transposed
+ self.lora_b = nn.Parameter(torch.empty((r, out_features)))
+
+ with torch.no_grad():
+ # Initialize $A$ similar to a weight matrix in a normal linear layer
+ nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
+ # Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
+ nn.init.zeros_(self.lora_b)
+
+ def forward(self, x: torch.Tensor):
+ # Compute $W_0 x + b_0$
+ result = nn.functional.linear(x, self.weight, bias=self.bias)
+
+ # Add $\frac{\alpha}{r} \Delta W x = \frac{\alpha}{r} BAx$
+ result += (x @ self.lora_a @ self.lora_b) * self.scaling
+
+ #
+ return result
+
+
+class Embedding(nn.Module):
+ """
+ ## LoRA Embedding Layer
+
+ Similar to LoRA linear layer this adds a low-rank decomposition to the pre-trained
+ embedding weights matrix ($W_0 \in \mathbb{R}^{d \times k}$).
+
+ $$W_0 + \Delta W = W_0 + BA$$
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int,
+ r: int, alpha: int = None):
+ """
+
+ :param num_embeddings: is the number of embeddings
+ :param embedding_dim: is the number embedding dimensions
+ :param r: is the rank of the decomposition $r$
+ :param alpha: is the scaling factor $\alpha$
+ """
+ super().__init__()
+
+ # Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
+ if alpha is None:
+ alpha = r
+
+ # The pre-trained embedding weights $W_0$ (frozen)
+ self.weight = nn.Parameter(torch.empty((num_embeddings, embedding_dim)))
+ self.weight.requires_grad = False
+
+ # scaling factor $\frac{\alpha}{r}$
+ self.scaling = alpha / r
+ # Matrix $A \in \mathbb{R}^{r \times k}$
+ self.lora_a = nn.Parameter(torch.empty((num_embeddings, r)))
+ # Matrix $B \in \mathbb{R}^{d \times r}$
+ self.lora_b = nn.Parameter(torch.empty((r, embedding_dim)))
+
+ with torch.no_grad():
+ # Initialize $A$ with a normal distribution
+ nn.init.normal_(self.lora_a)
+ # Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
+ nn.init.zeros_(self.lora_b)
+
+ def forward(self, x: torch.Tensor):
+ # Compute the embeddings $W_0 \text{onehot}(x)$
+ result = nn.functional.embedding(x, self.weight)
+
+ # Add $\frac{\alpha}{r} \Delta W \text{onehot}(x) = \frac{\alpha}{r} BA \text{onehot}(x_$
+ result += (nn.functional.embedding(x, self.lora_a) @ self.lora_b) * self.scaling
+
+ #
+ return result
diff --git a/labml_nn/lora/gpt2.py b/labml_nn/lora/gpt2.py
new file mode 100644
index 00000000..a83a0276
--- /dev/null
+++ b/labml_nn/lora/gpt2.py
@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer
+from labml_nn.lora import Linear, Embedding
+
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+config = {
+ "layer_norm_epsilon": 1e-05,
+ "n_embd": 768,
+ "n_head": 12,
+ "n_layer": 12,
+ "n_positions": 1024,
+ "vocab_size": 50257,
+ "device": "cuda"
+}
+
+
+class FFN(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.c_fc = Linear(config['n_embd'], dim, r=32, bias=True)
+ self.c_proj = Linear(dim, config['n_embd'], r=32, bias=True)
+ self.act = nn.functional.gelu
+
+ def forward(self, hidden_states):
+ hidden_states = self.c_fc(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.c_proj(hidden_states)
+ return hidden_states
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.embed_dim = config['n_embd']
+ self.num_heads = config['n_head']
+ self.head_dim = self.embed_dim // self.num_heads
+ self.split_size = self.embed_dim
+
+ self.c_att = Linear(config['n_embd'], config['n_embd'] * 3, r=32, bias=True)
+ self.c_proj = Linear(config['n_embd'], config['n_embd'], r=32, bias=True)
+
+ def _split_heads(self, tensor, num_heads, attn_head_size):
+ """
+ Splits hidden_size dim into attn_head_size and num_heads
+ """
+ new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
+ tensor = tensor.view(new_shape)
+ return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
+
+ def forward(self, hidden_states):
+ batch_size, seq_length, _ = hidden_states.size()
+
+ query, key, value = self.c_att(hidden_states).split(self.split_size, dim=2)
+
+ query = self._split_heads(query, self.num_heads, self.head_dim)
+ key = self._split_heads(key, self.num_heads, self.head_dim)
+ value = self._split_heads(value, self.num_heads, self.head_dim)
+
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query,
+ key,
+ value,
+ attn_mask=None,
+ dropout_p=0.0,
+ is_causal=True, # for the triangular mask
+ )
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+ attn_output = attn_output.view(batch_size, seq_length, self.embed_dim)
+
+ attn_output = self.c_proj(attn_output)
+
+ return attn_output
+
+
+class Block(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.pre_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
+ self.attn = MultiHeadAttention()
+ self.post_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
+ self.ffn = FFN(config['n_embd'] * 4)
+
+ def forward(self, hidden_states):
+ residual = hidden_states
+ hidden_states = self.pre_norm(hidden_states)
+
+ attn_output = self.attn(hidden_states)
+
+ hidden_states = attn_output + residual
+ residual = hidden_states
+ hidden_states = self.post_norm(hidden_states)
+ feed_forward_output = self.ffn(hidden_states)
+ hidden_states = feed_forward_output + residual
+
+ return hidden_states
+
+
+class GPTModel(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ self.token_embedding = Embedding(config['vocab_size'], config['n_embd'], r=32)
+ self.position_embedding = Embedding(config['n_positions'], config['n_embd'], r=32)
+
+ self.blocks = nn.ModuleList([Block() for _ in range(config['n_layer'])])
+
+ self.final_norm = nn.LayerNorm(config['n_embd'], eps=config['layer_norm_epsilon'])
+
+ self.lm_head = Linear(config['n_embd'], config['vocab_size'], r=32, bias=False)
+
+ def forward(self, input_ids):
+ batch_size, input_shape = input_ids.size()
+
+ token_embeddings = self.token_embedding(input_ids) # B T C
+ position_ids = torch.arange(input_shape, device=config['device']) # T C
+ position_embeddings = self.position_embedding(position_ids) # B T C
+
+ hidden_states = token_embeddings + position_embeddings
+
+ for block in self.blocks:
+ hidden_states = block(hidden_states)
+
+ hidden_states = self.final_norm(hidden_states)
+
+ logits = self.lm_head(hidden_states)
+
+ return logits
diff --git a/labml_nn/lora/train.ipynb b/labml_nn/lora/train.ipynb
new file mode 100644
index 00000000..68bbb7eb
--- /dev/null
+++ b/labml_nn/lora/train.ipynb
@@ -0,0 +1,217 @@
+{
+ "cells": [
+ {
+ "metadata": {},
+ "cell_type": "code",
+ "outputs": [],
+ "execution_count": null,
+ "source": [
+ "import torch\n",
+ "from torch.optim import Adam\n",
+ "from torch.utils.data import DataLoader, TensorDataset\n",
+ "from torch.utils.data import random_split\n",
+ "from transformers import AutoTokenizer\n",
+ "\n",
+ "from labml import tracker, experiment\n",
+ "from labml_nn.lora.gpt2 import GPTModel"
+ ],
+ "id": "f072832ec9d346e1"
+ },
+ {
+ "cell_type": "code",
+ "id": "initial_id",
+ "metadata": {
+ "collapsed": true,
+ "jupyter": {
+ "outputs_hidden": true
+ }
+ },
+ "source": "# !wget https://raw.github/zusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt",
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "3b1e507015ba6b81",
+ "metadata": {},
+ "source": [
+ "with open('input.txt', 'r', encoding='utf-8') as f:\n",
+ " text = f.read()"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "ac8e51ae5bbfcae7",
+ "metadata": {},
+ "source": [
+ "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
+ "\n",
+ "tokens = tokenizer.encode(text, add_special_tokens=False)"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "aeefcdf813e427e",
+ "metadata": {},
+ "source": [
+ "context_length = 512\n",
+ "batch_size = 2"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "a384b42274f008a2",
+ "metadata": {},
+ "source": [
+ "num_batches = len(tokens) // (batch_size * context_length)\n",
+ "tokens = tokens[:num_batches * batch_size * context_length]"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "5c4cc78ac1a02c1d",
+ "metadata": {},
+ "source": "input_ids = torch.tensor(tokens).view(-1, context_length)",
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "7037fd75e2161382",
+ "metadata": {},
+ "source": [
+ "dataset = TensorDataset(input_ids)\n",
+ "\n",
+ "train_ratio = 0.8\n",
+ "test_ratio = 0.2\n",
+ "\n",
+ "train_size = int(train_ratio * len(dataset))\n",
+ "test_size = len(dataset) - train_size\n",
+ "\n",
+ "train_dataset, test_dataset = random_split(dataset, [train_size, test_size])\n",
+ "\n",
+ "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
+ "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "a98b7baa064b8494",
+ "metadata": {},
+ "source": [
+ "model = GPTModel()\n",
+ "state_dict = torch.load('transformed.pth', weights_only=True)\n",
+ "\n",
+ "_ = model.load_state_dict(state_dict, strict=False)"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "metadata": {},
+ "cell_type": "code",
+ "source": [
+ "device = \"cuda\"\n",
+ "model = model.to(device=\"cuda\")"
+ ],
+ "id": "2e0fa8b3082df716",
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "e2f5076894770740",
+ "metadata": {},
+ "source": [
+ "optimizer = Adam(model.parameters(), lr=5e-5)\n",
+ "criterion = torch.nn.CrossEntropyLoss()\n",
+ "\n",
+ "model.train()\n",
+ "epochs = 3\n",
+ "step = 0\n",
+ "\n",
+ "with experiment.record(name='LoRA.GPT2', app_url='http://localhost:5005/api/v1/track'):\n",
+ " for epoch in range(epochs):\n",
+ " for batch in train_dataloader:\n",
+ " inputs = batch[0]\n",
+ " inputs = inputs.to(device)\n",
+ " labels = inputs.clone()\n",
+ "\n",
+ " outputs = model(inputs)\n",
+ "\n",
+ " shift_logits = outputs[..., :-1, :]\n",
+ " shift_labels = labels[..., 1:]\n",
+ "\n",
+ " loss = criterion(shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1))\n",
+ "\n",
+ " optimizer.zero_grad()\n",
+ " loss.backward()\n",
+ " optimizer.step()\n",
+ "\n",
+ " tracker.save(step, {'loss': loss})\n",
+ " step += 1\n",
+ " print(f'Epoch: {epoch + 1}, Loss: {loss.item()}')\n",
+ "\n",
+ " test_loss = 0\n",
+ " for batch in test_dataloader:\n",
+ " inputs = batch[0]\n",
+ " inputs = inputs.to(device)\n",
+ " labels = inputs.clone()\n",
+ "\n",
+ " outputs = model(inputs)\n",
+ "\n",
+ " shift_logits = outputs[..., :-1, :]\n",
+ " shift_labels = labels[..., 1:]\n",
+ "\n",
+ " loss = criterion(shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1))\n",
+ "\n",
+ " test_loss += loss.item()\n",
+ " test_loss /= len(test_dataloader)\n",
+ " tracker.save(step, {'test_loss': test_loss})\n",
+ "\n",
+ "print(\"Training complete.\")"
+ ],
+ "outputs": [],
+ "execution_count": null
+ },
+ {
+ "cell_type": "code",
+ "id": "da2d4023002648dc",
+ "metadata": {},
+ "source": [],
+ "outputs": [],
+ "execution_count": null
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "base",
+ "language": "python",
+ "name": "base"
+ },
+ "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.12.4"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/labml_nn/lora/transform_hf_model.py b/labml_nn/lora/transform_hf_model.py
new file mode 100644
index 00000000..df53bbf2
--- /dev/null
+++ b/labml_nn/lora/transform_hf_model.py
@@ -0,0 +1,46 @@
+import torch
+from transformers import AutoModelForCausalLM
+
+
+def transform_hf_model():
+ model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+ state_dict = model.state_dict()
+
+ mapping = {
+ 'transformer.wte.weight': 'token_embedding.weight',
+ 'transformer.wpe.weight': 'position_embedding.weight',
+ 'transformer.ln_f.weight': 'final_norm.weight',
+ 'transformer.ln_f.bias': 'final_norm.bias',
+ 'lm_head.weight': 'lm_head.weight'
+ }
+
+ for i in range(12):
+ mapping[f'transformer.h.{i}.ln_1.weight'] = f'blocks.{i}.pre_norm.weight'
+ mapping[f'transformer.h.{i}.ln_1.bias'] = f'blocks.{i}.pre_norm.bias'
+ mapping[f'transformer.h.{i}.attn.c_attn.weight'] = f'blocks.{i}.attn.c_att.weight'
+ mapping[f'transformer.h.{i}.attn.c_attn.bias'] = f'blocks.{i}.attn.c_att.bias'
+ mapping[f'transformer.h.{i}.attn.c_proj.weight'] = f'blocks.{i}.attn.c_proj.weight'
+ mapping[f'transformer.h.{i}.attn.c_proj.bias'] = f'blocks.{i}.attn.c_proj.bias'
+ mapping[f'transformer.h.{i}.ln_2.weight'] = f'blocks.{i}.post_norm.weight'
+ mapping[f'transformer.h.{i}.ln_2.bias'] = f'blocks.{i}.post_norm.bias'
+ mapping[f'transformer.h.{i}.mlp.c_fc.weight'] = f'blocks.{i}.ffn.c_fc.weight'
+ mapping[f'transformer.h.{i}.mlp.c_fc.bias'] = f'blocks.{i}.ffn.c_fc.bias'
+ mapping[f'transformer.h.{i}.mlp.c_proj.weight'] = f'blocks.{i}.ffn.c_proj.weight'
+ mapping[f'transformer.h.{i}.mlp.c_proj.bias'] = f'blocks.{i}.ffn.c_proj.bias'
+
+ new_state_dict = {}
+ for old_key, new_key in mapping.items():
+ if old_key in state_dict:
+ new_state_dict[new_key] = state_dict[old_key]
+
+ # transpose weight matrices of convo 1d layers to use linear layers instead
+ convo_layers = ([f'blocks.{i}.ffn.c_fc.weight' for i in range(12)] +
+ [f'blocks.{i}.ffn.c_proj.weight' for i in range(12)] +
+ [f'blocks.{i}.attn.c_att.weight' for i in range(12)] +
+ [f'blocks.{i}.attn.c_proj.weight' for i in range(12)])
+
+ for layer in convo_layers:
+ new_state_dict[layer] = torch.transpose(new_state_dict[layer], 0, 1)
+
+ torch.save(new_state_dict, 'transformed.pth')
diff --git a/labml_nn/RWKV/__init__.py b/labml_nn/rwkv/__init__.py
similarity index 100%
rename from labml_nn/RWKV/__init__.py
rename to labml_nn/rwkv/__init__.py
diff --git a/labml_nn/RWKV/configs.py b/labml_nn/rwkv/configs.py
similarity index 100%
rename from labml_nn/RWKV/configs.py
rename to labml_nn/rwkv/configs.py
diff --git a/labml_nn/RWKV/experiment.py b/labml_nn/rwkv/experiment.py
similarity index 97%
rename from labml_nn/RWKV/experiment.py
rename to labml_nn/rwkv/experiment.py
index 1f99d66d..983db2c0 100644
--- a/labml_nn/RWKV/experiment.py
+++ b/labml_nn/rwkv/experiment.py
@@ -3,10 +3,10 @@ import math
import torch
import torch.nn as nn
-from labml_nn.RWKV.configs import RWKVConfigs
+from labml_nn.rwkv.configs import RWKVConfigs
-from labml_nn.RWKV import RWKV
-from labml_nn.RWKV import TimeMixing
+from labml_nn.rwkv import RWKV
+from labml_nn.rwkv import TimeMixing
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
diff --git a/translate_cache/transformers/feed_forward.zh.json b/translate_cache/transformers/feed_forward.zh.json
index 66f87871..719c685d 100644
--- a/translate_cache/transformers/feed_forward.zh.json
+++ b/translate_cache/transformers/feed_forward.zh.json
@@ -1,5 +1,5 @@
{
- "
Position-wise Feed-Forward Network (FFN)
\n
This is a PyTorch implementation of position-wise feedforward network used in transformer.
\n
FFN consists of two fully connected layers. Number of dimensions in the hidden layer _^_0_^_, is generally set to around four times that of the token embedding _^_1_^_. So it is sometime also called the expand-and-contract network.
\n
There is an activation at the hidden layer, which is usually set to ReLU (Rectified Linear Unit) activation, _^_2_^_
\n
That is, the FFN function is, _^_3_^_ where _^_4_^_, _^_5_^_, _^_6_^_ and _^_7_^_ are learnable parameters.
\n
Sometimes the GELU (Gaussian Error Linear Unit) activation is also used instead of ReLU. _^_8_^_ where _^_9_^_
\n
Gated Linear Units
\n
This is a generic implementation that supports different variants including Gated Linear Units (GLU). We have also implemented experiments on these:
This is a PyTorch implementation of position-wise feedforward network used in transformer.
\n
FFN consists of two fully connected layers. Number of dimensions in the hidden layer _^_0_^_, is generally set to around four times that of the token embedding _^_1_^_. So it is sometime also called the expand-and-contract network.
\n
There is an activation at the hidden layer, which is usually set to ReLU (Rectified Linear Unit) activation, _^_2_^_
\n
That is, the FFN function is, _^_3_^_ where _^_4_^_, _^_5_^_, _^_6_^_ and _^_7_^_ are learnable parameters.
\n
Sometimes the GELU (Gaussian Error Linear Unit) activation is also used instead of ReLU. _^_8_^_ where _^_9_^_
\n
Gated Linear Units
\n
This is a generic implementation that supports different variants including Gated Linear Units (GLU). We have also implemented experiments on these:
