mirror of
https://github.com/wassname/iris_bigvae.git
synced 2026-08-25 11:19:04 +08:00
adding bigvae
This commit is contained in:
wassname
@@ -1,3 +1,10 @@
|
||||
|
||||
Fork of IRIS, where instead of a new transformer as a world model we use an adapter on a pretrained LLM. The hypothesis is that the pretrained LLM will help the world model learn faster and in a more data effecient manner.
|
||||
|
||||
See also:
|
||||
- [AdaVAE](https://github.com/ImKeTT/AdaVAE)
|
||||
- [bigvae](https://github.com/JD-P/minihf/blob/adavae-moe/vae_infer.py)
|
||||
|
||||
# Transformers are Sample-Efficient World Models (IRIS)
|
||||
|
||||
[Transformers are Sample-Efficient World Models](https://openreview.net/forum?id=vhFu1Acb0xb) <br>
|
||||
|
||||
@@ -1,10 +1,2 @@
|
||||
_target_: models.TransformerConfig
|
||||
_target_: models.BigVAEConfig
|
||||
tokens_per_block: 17
|
||||
max_blocks: 20
|
||||
attention: 'causal'
|
||||
num_layers: 10
|
||||
num_heads: 4
|
||||
embed_dim: 256
|
||||
embed_pdrop: 0.1
|
||||
resid_pdrop: 0.1
|
||||
attn_pdrop: 0.1
|
||||
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
# 2023-11-03 11:36:39
|
||||
|
||||
Step 1 trying to get a VAE working
|
||||
|
||||
ah bitsandbytes
|
||||
so I needed to use a `poetry add https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl` to get it to work
|
||||
|
||||
# 2023-11-09 07:29:27
|
||||
|
||||
Where am I up to?
|
||||
|
||||
- [ ] I want to get bigVAE working with Mistral
|
||||
- [ ] run
|
||||
- [ ] Then simplify it
|
||||
- [ ] then understand IRIS, and inser this model
|
||||
|
||||
|
||||
Eror: OOM w minstral. It runs out of mem when set_adapter, with an input
|
||||
- 7.6/24 at first stop point we have
|
||||
- vae: DecoderOnlyTransformerVAE
|
||||
- self.model
|
||||
- and we seem to be running inputs through with grad
|
||||
- which is werid as it's jsut
|
||||
- a frozen model
|
||||
- 2 adaptors (float32)
|
||||
- and a vae head
|
||||
- `next(iter(self.model.parameters()))` is bfloat16, cuda:0. As is vae
|
||||
|
||||
$ self.model
|
||||
PeftModel(
|
||||
(base_model): LoraModel(
|
||||
(model): MistralForCausalLM(
|
||||
(model): MistralModel(
|
||||
(embed_tokens): Embedding(32000, 4096)
|
||||
|
||||
$ self.vae
|
||||
DecoderOnlyTransformerVAE(
|
||||
(model): PeftModel(
|
||||
(base_model): LoraModel(
|
||||
(model): MistralForCausalLM(
|
||||
(model): MistralModel(
|
||||
(embed_tokens): Embedding(32000, 4096)
|
||||
(vae): VAEComponent(
|
||||
(f): Linear(in_features=4096, out_features=1, bias=True)
|
||||
(w_e): Linear(in_features=4096, out_features=768, bias=True)
|
||||
(w_d): Linear(in_features=768, out_features=4096, bias=True)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
Hm JDP said he uses 8xH100 so 64GB*8.
|
||||
- A p5.48xlarge is 80$
|
||||
- a o2.16x is 192GB and $14/h
|
||||
we are using peft
|
||||
|
||||
```py
|
||||
self.model.print_trainable_parameters()
|
||||
# trainable params: 83,886,080 || all params: 7,577,276,416 || trainable%: 1.107074302091819
|
||||
|
||||
|
||||
self.model
|
||||
```
|
||||
|
||||
DecoderOnlyTransformerVAE(
|
||||
(model): PeftModel(
|
||||
(base_model): LoraModel(
|
||||
(model): MistralForCausalLM(
|
||||
(model): MistralModel(
|
||||
(embed_tokens): Embedding(32000, 4096)
|
||||
(layers): ModuleList(
|
||||
(0-31): 32 x MistralDecoderLayer(
|
||||
(self_attn): MistralAttention( )
|
||||
(mlp): MistralMLP()
|
||||
(input_layernorm): MistralRMSNorm()
|
||||
(post_attention_layernorm): MistralRMSNorm()
|
||||
)
|
||||
)
|
||||
(norm): MistralRMSNorm()
|
||||
)
|
||||
(lm_head): Linear(in_features=4096, out_features=32000, bias=False)
|
||||
)
|
||||
)
|
||||
)
|
||||
(vae): VAEComponent(
|
||||
(f): Linear(in_features=4096, out_features=1, bias=True)
|
||||
(w_e): Linear(in_features=4096, out_features=768, bias=True)
|
||||
(w_d): Linear(in_features=768, out_features=4096, bias=True)
|
||||
)
|
||||
)
|
||||
|
||||
How much gpu ram should it take it train a 7B?
|
||||
- 3B -> 14
|
||||
- 12B - 56GB
|
||||
- so 7B should be ~30 :(. or 20 with cpu offloading
|
||||
- batch size of 4... with 1 it takes 20GB then crashes
|
||||
|
||||
# 2023-11-10 11:41:21
|
||||
|
||||
tldr:
|
||||
- I can't use Mistral without a bigger gpu: 30GB+, or maybe I can use the deepspeed gpu offloading (batch=1)
|
||||
- I I can just use GPT2 like adavae. Or stablelm
|
||||
|
||||
Ideally I can use a small one for prototyping, and change to a large one if it works.
|
||||
|
||||
|
||||
Models:
|
||||
- mistral: the BigVAE code is setup for it
|
||||
- gpt2: the AdaVAE code is setup for it... but it's also way messier. Roll your own adapter etc
|
||||
|
||||
|
||||
TODO use no_grad! way better
|
||||
|
||||
|
||||
now
|
||||
- look at training code
|
||||
- what is JP code actually doing? start with generate topic
|
||||
- clean it up, there is so much repeeated code
|
||||
|
||||
|
||||
# 2023-11-11 10:09:50
|
||||
|
||||
## Now how does VAE generate work inside?
|
||||
|
||||
The VAE generate function takes in context, input embeddings, target embeddings, and other parameters in that order. The function performs the following steps:
|
||||
- Encode: The input embeddings are encoded using an adapter on `self.model` and `vaecomponent.encoder`. This step involves using linear and pooling operations, followed by a softmax function to sample and obtain `z`.
|
||||
- Decode/Embed: `z` is decoded to obtain `z_embed`, and the decoder is provided with target ID embeddings. This step is done using a linear operation and the model embeddings.
|
||||
- Model: The model is run on the embeddings to obtain logits. The embeddings can be considered as latent states, and the latent state is expressed in the language of embeddings.
|
||||
|
||||
# 2023-11-12 08:06:25
|
||||
|
||||
Now look at IRIS and TWM world models.
|
||||
|
||||
- https://github.com/eloialonso/iris/blob/main/src/models/world_model.py
|
||||
- Model(x, kv_cache). Where x is B, T, C. Batch, Time, Channels?
|
||||
- if just takes in x and outputs x. The output are logits, from a linear layer.
|
||||
- https://github.com/jrobine/twm/blob/main/twm/world_model.py
|
||||
|
||||
|
||||
|
||||
# 2023-11-12 08:46:46
|
||||
|
||||
Adding BigVAe as transformer layer
|
||||
|
||||
I can't tokenize, then pass in input id's. As I need to to be backpropable. So I need to by pass the embedding layers...
|
||||
- perhaps I can encode actions by things I have previously embedded?
|
||||
|
||||
# 2023-11-12 10:10:01
|
||||
|
||||
What the diff between BigVAERouter and DecoderOnlyTransformerVAE
|
||||
|
||||
- DecoderOnlyTransformerVAE(prefix_ids, input_ids) -> outputs, mean
|
||||
- BigVAERouter(prefix_ids, input_ids, target_ids,),
|
||||
|
||||
OK I want to rename here:
|
||||
- prefix_ids
|
||||
- embed_ids -> input_ds
|
||||
- target_ids
|
||||
- decoder_prefix_ids -> prefix_ids
|
||||
|
||||
Now how to reconcille it with the world model. what does the world model do?
|
||||
- components
|
||||
- transforer x->x: the part we are replacing
|
||||
- embedder: a custom embedder
|
||||
- heads for each output: obs, reward, end_of_episode
|
||||
- when it goes foward
|
||||
- tokens -> (x, obs, rewards, ends)
|
||||
- where is is the output of the transformer/vae
|
||||
- x = transformer(sequences)
|
||||
- where the sequences are tokens embedded with the embedder.
|
||||
- TODO: I will want to use the model embedder if possible?
|
||||
|
||||
## Embedder deep dive
|
||||
|
||||
First the code
|
||||
|
||||
So the embedder, takes in `tokens` then breaks it up using a slices, into a seperate obs and action embedder.
|
||||
We add them, plus positions, and add them into one of embed_dim
|
||||
|
||||
self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
|
||||
|
||||
self.embedder = Embedder(
|
||||
max_blocks=config.max_blocks,
|
||||
block_masks=[act_tokens_pattern, obs_tokens_pattern],
|
||||
embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config.embed_dim), nn.Embedding(obs_vocab_size, config.embed_dim)])
|
||||
)
|
||||
|
||||
sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
|
||||
|
||||
|
||||
Now how does the paper describe it? https://openreview.net/pdf?id=vhFu1Acb0xb
|
||||
|
||||
|
||||
# IDEAS
|
||||
|
||||
Could I just use a single transformer? Dreamer seem use encode and decode to a latent state, but IRIS doesn't?
|
||||
@@ -0,0 +1,486 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
||||
" from .autonotebook import tqdm as notebook_tqdm\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import argparse\n",
|
||||
"from contextlib import contextmanager\n",
|
||||
"from itertools import chain, islice\n",
|
||||
"import json\n",
|
||||
"import math\n",
|
||||
"from pathlib import Path\n",
|
||||
"import random\n",
|
||||
"import os\n",
|
||||
"import sys\n",
|
||||
"import zipfile\n",
|
||||
"\n",
|
||||
"import accelerate\n",
|
||||
"from datasets import load_dataset\n",
|
||||
"import peft\n",
|
||||
"import safetensors.torch as safetorch\n",
|
||||
"import torch\n",
|
||||
"from torch import nn, optim\n",
|
||||
"from torch.nn import functional as F\n",
|
||||
"from torch.utils import data\n",
|
||||
"from tqdm import trange, tqdm\n",
|
||||
"from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
|
||||
"from loguru import logger\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# autoreload import your package\n",
|
||||
"%load_ext autoreload\n",
|
||||
"%autoreload 2\n",
|
||||
"\n",
|
||||
"from vae_llm_worldmodels.models.bigvae.bigvae import set_adapter, DecoderOnlyTransformerVAE, VAERouter\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Params"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"['--rank', '16', '--context=96', '--vae_context=32', '--batch_size=1', '--output=./output/adapter']\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"Namespace(batch_size=1, dropout=0.0, epochs=1, gradient_accumulation_steps=1, gradient_checkpointing=False, lr=0.0001, model='stabilityai/stablelm-3b-4e1t', context=96, vae_context=32, output=PosixPath('output/adapter'), rank=16, save_every=1000, start_from=None, z_dim=768)"
|
||||
]
|
||||
},
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"parser = argparse.ArgumentParser(description=__doc__)\n",
|
||||
"parser.add_argument(\"--batch_size\", type=int, default=2, help=\"microbatch size\")\n",
|
||||
"parser.add_argument(\"--dropout\", type=float, default=0.0, help=\"dropout rate\")\n",
|
||||
"parser.add_argument(\"--epochs\", type=int, default=1, help=\"number of epochs\")\n",
|
||||
"parser.add_argument(\n",
|
||||
" \"--gradient_accumulation_steps\", type=int, default=1, help=\"gradient accumulation steps\"\n",
|
||||
")\n",
|
||||
"parser.add_argument(\n",
|
||||
" \"--gradient_checkpointing\",\n",
|
||||
" action=\"store_true\",\n",
|
||||
" default=False,\n",
|
||||
" help=\"use gradient checkpointing\",\n",
|
||||
")\n",
|
||||
"parser.add_argument(\"--lr\", type=float, default=1e-4, help=\"learning rate\")\n",
|
||||
"parser.add_argument(\n",
|
||||
" \"--model\",\n",
|
||||
" type=str,\n",
|
||||
" # default=\"mistralai/Mistral-7B-v0.1\",\n",
|
||||
" # default=\"yichunkuo/stablelm-3b-4e1t-gptq\",\n",
|
||||
" default=\"stabilityai/stablelm-3b-4e1t\", \n",
|
||||
" # default=\"gpt2\", \n",
|
||||
" # default=\"mlabonne/gpt2-GPTQ-4bit\",\n",
|
||||
" help=\"model name\",\n",
|
||||
")\n",
|
||||
"parser.add_argument(\"--context\", type=int, default=2048, help=\"context window length\")\n",
|
||||
"parser.add_argument(\"--vae_context\", type=int, default=64, help=\"vae embed context\")\n",
|
||||
"parser.add_argument(\"--output\", type=Path, required=True, help=\"path to save adapter\")\n",
|
||||
"parser.add_argument(\"--rank\", type=int, default=32, help=\"the lora rank\")\n",
|
||||
"parser.add_argument(\"--save_every\", type=int, default=1000, help=\"save every n steps\")\n",
|
||||
"parser.add_argument(\"--start_from\", type=str, help=\"start from existing lora\")\n",
|
||||
"parser.add_argument(\"--z_dim\", type=int, default=768, help=\"the latent dimension\")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"argvs = \"\"\"\n",
|
||||
"--rank 16 \n",
|
||||
"--context=96 \n",
|
||||
"--vae_context=32 \n",
|
||||
"--batch_size=1 \n",
|
||||
"--output=./output/adapter \n",
|
||||
"\"\"\"\n",
|
||||
"argvs = argvs.replace('\\n', ' ').strip()\n",
|
||||
"argv = [s.strip() for s in argvs.split(\" \") if s and not s.startswith(\"#\")]\n",
|
||||
"print(argv)\n",
|
||||
"args = parser.parse_args(argv)\n",
|
||||
"args\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"max_length = 32\n",
|
||||
"tokenizer_args = dict(\n",
|
||||
" padding='max_length', max_length=max_length,\n",
|
||||
" truncation=True,\n",
|
||||
")\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from transformers.utils.logging import _get_library_root_logger\n",
|
||||
"\n",
|
||||
"os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
|
||||
"os.environ[\"TRANSFORMERS_VERBOSITY\"] = \"detail\"\n",
|
||||
"\n",
|
||||
"library_root_logger = _get_library_root_logger()\n",
|
||||
"library_root_logger.propagate = True\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Load\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32m2023-11-12 09:44:45.920\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mcontextlib\u001b[0m:\u001b[36minner\u001b[0m:\u001b[36m81\u001b[0m - \u001b[1mLoading model: stabilityai/stablelm-3b-4e1t\u001b[0m\n",
|
||||
"Using pad_token, but it is not set yet.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"ename": "TypeError",
|
||||
"evalue": "VAERouter.__init__() takes from 3 to 4 positional arguments but 5 were given",
|
||||
"output_type": "error",
|
||||
"traceback": [
|
||||
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||||
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
|
||||
"\u001b[1;32m/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb Cell 10\u001b[0m line \u001b[0;36m5\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb#X11sZmlsZQ%3D%3D?line=56'>57</a>\u001b[0m vae_model\u001b[39m.\u001b[39mvae\u001b[39m.\u001b[39mrequires_grad_(\u001b[39mFalse\u001b[39;00m)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb#X11sZmlsZQ%3D%3D?line=57'>58</a>\u001b[0m vae_model\u001b[39m.\u001b[39mvae\u001b[39m.\u001b[39mw_d\u001b[39m.\u001b[39mrequires_grad_()\n\u001b[0;32m---> <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb#X11sZmlsZQ%3D%3D?line=58'>59</a>\u001b[0m router \u001b[39m=\u001b[39m VAERouter(base_model_peft, vae_model, device, peft_config)\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb#X11sZmlsZQ%3D%3D?line=59'>60</a>\u001b[0m \u001b[39mif\u001b[39;00m args\u001b[39m.\u001b[39mstart_from:\n\u001b[1;32m <a href='vscode-notebook-cell:/media/wassname/SGIronWolf/projects5/worldmodels/vae_llm_worldmodel/notebooks/mjc-002-vae_train.ipynb#X11sZmlsZQ%3D%3D?line=60'>61</a>\u001b[0m router\u001b[39m.\u001b[39mload_pretrained(args\u001b[39m.\u001b[39mstart_from, is_trainable\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m)\n",
|
||||
"\u001b[0;31mTypeError\u001b[0m: VAERouter.__init__() takes from 3 to 4 positional arguments but 5 were given"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"accelerator = accelerate.Accelerator(\n",
|
||||
" mixed_precision=\"bf16\", gradient_accumulation_steps=args.gradient_accumulation_steps\n",
|
||||
")\n",
|
||||
"device = accelerator.device if accelerator.num_processes > 1 else \"cuda:0\"\n",
|
||||
"is_main = accelerator.is_main_process\n",
|
||||
"\n",
|
||||
"print = tqdm.external_write_mode()(logger.info)\n",
|
||||
"print0 = accelerator.on_main_process(print)\n",
|
||||
"\n",
|
||||
"if Path(args.model).exists():\n",
|
||||
" model_name = Path(args.model).resolve()\n",
|
||||
"else:\n",
|
||||
" model_name = args.model\n",
|
||||
"\n",
|
||||
"print0(f\"Loading model: {model_name}\")\n",
|
||||
"with accelerator.main_process_first():\n",
|
||||
" tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
|
||||
" tokenizer.padding_side = \"left\"\n",
|
||||
" if tokenizer.pad_token is None:\n",
|
||||
" tokenizer.pad_token = tokenizer.eos_token\n",
|
||||
" bnb_config = BitsAndBytesConfig(\n",
|
||||
" load_in_4bit=True,\n",
|
||||
" bnb_4bit_compute_dtype=torch.bfloat16,\n",
|
||||
" bnb_4bit_quant_type=\"nf4\",\n",
|
||||
" bnb_4bit_use_double_quant=True,\n",
|
||||
" )\n",
|
||||
" base_model = AutoModelForCausalLM.from_pretrained(\n",
|
||||
" model_name,\n",
|
||||
" device_map={\"\": device},\n",
|
||||
" quantization_config=bnb_config,\n",
|
||||
" torch_dtype=torch.bfloat16, \n",
|
||||
" trust_remote_code=True\n",
|
||||
" )\n",
|
||||
" peft_config = peft.LoraConfig(\n",
|
||||
" peft.TaskType.CAUSAL_LM,\n",
|
||||
" inference_mode=False,\n",
|
||||
" r=args.rank,\n",
|
||||
" lora_alpha=8,\n",
|
||||
" lora_dropout=args.dropout,\n",
|
||||
" target_modules=[\n",
|
||||
" \"self_attn.q_proj\",\n",
|
||||
" \"self_attn.k_proj\",\n",
|
||||
" \"self_attn.v_proj\",\n",
|
||||
" \"self_attn.o_proj\",\n",
|
||||
" \"mlp.gate_proj\",\n",
|
||||
" \"mlp.up_proj\",\n",
|
||||
" \"mlp.down_proj\",\n",
|
||||
" ],\n",
|
||||
" )\n",
|
||||
" base_model_peft = peft.get_peft_model(base_model, peft_config)\n",
|
||||
" vae_model = DecoderOnlyTransformerVAE(\n",
|
||||
" base_model_peft, peft_config, device=device, z_dim=args.z_dim,\n",
|
||||
" )\n",
|
||||
" if args.start_from:\n",
|
||||
" vae_model.load_pretrained(args.start_from)\n",
|
||||
" base_model_peft.requires_grad_(False)\n",
|
||||
" vae_model.vae.requires_grad_(False)\n",
|
||||
" vae_model.vae.w_d.requires_grad_()\n",
|
||||
" router = VAERouter(base_model_peft, vae_model, device)\n",
|
||||
" if args.start_from:\n",
|
||||
" router.load_pretrained(args.start_from, is_trainable=True)\n",
|
||||
"accelerator.wait_for_everyone()\n",
|
||||
"\n",
|
||||
"router.train()\n",
|
||||
"if args.gradient_checkpointing:\n",
|
||||
" router.model.gradient_checkpointing_enable()\n",
|
||||
" router.model.enable_input_require_grads()\n",
|
||||
"\n",
|
||||
"if is_main:\n",
|
||||
" router.model.print_trainable_parameters()\n",
|
||||
"\n",
|
||||
"router.model.set_adapter(\"router\")\n",
|
||||
"opt = optim.Adam(router.model.parameters(),\n",
|
||||
" lr=args.lr,\n",
|
||||
" betas=(0.9, 0.99))\n",
|
||||
"accelerator.wait_for_everyone()\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 17,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"https://github.com/JD-P/minihf/blob/adavae-moe/train_vae_router.py#L277\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# prepare dataset\n",
|
||||
"input_ids_all, attention_mask_all = [], []\n",
|
||||
"for shard_name in os.listdir(args.preprocessed):\n",
|
||||
" data_path = os.path.join(args.preprocessed, shard_name)\n",
|
||||
" data_file = safetorch.load_file(data_path)\n",
|
||||
" input_ids = torch.split(data_file[\"input_ids\"], args.context, dim=1)\n",
|
||||
" attention_mask = torch.split(data_file[\"attention_mask\"], args.context, dim=1)\n",
|
||||
" if input_ids[-1].shape[1] != args.context:\n",
|
||||
" input_ids = input_ids[:-1]\n",
|
||||
" attention_mask = attention_mask[:-1]\n",
|
||||
" input_ids_all.extend(input_ids)\n",
|
||||
" attention_mask_all.extend(attention_mask)\n",
|
||||
"del data_file, input_ids, attention_mask\n",
|
||||
"input_ids_all = torch.cat(input_ids_all)\n",
|
||||
"attention_mask_all = torch.cat(attention_mask_all)\n",
|
||||
"valid_indices = attention_mask_all.sum(dim=1) == args.context\n",
|
||||
"input_ids_all = input_ids_all[valid_indices]\n",
|
||||
"attention_mask_all = attention_mask_all[valid_indices]\n",
|
||||
"del valid_indices\n",
|
||||
"\n",
|
||||
"preprocessed = data.TensorDataset(input_ids_all, attention_mask_all)\n",
|
||||
"\n",
|
||||
"dataloader = data.DataLoader(\n",
|
||||
" preprocessed,\n",
|
||||
" batch_size=args.batch_size,\n",
|
||||
" shuffle=True,\n",
|
||||
" drop_last=True,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"router, opt, dataloader = accelerator.prepare(router, opt, dataloader)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from vae_llm_worldmodels.utils import cosine_warmup\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"@torch.no_grad()\n",
|
||||
"@torch.cuda.amp.autocast(dtype=torch.bfloat16)\n",
|
||||
"def demo(model, input_ids, attention_mask, n_tokens):\n",
|
||||
" \"\"\"inference.\"\"\"\n",
|
||||
" bs = min(input_ids.shape[0], 2)\n",
|
||||
" n_outputs = 2\n",
|
||||
" tau = 0.8\n",
|
||||
"\n",
|
||||
" index = random.randrange(args.context - (args.vae_context * 2))\n",
|
||||
" context_ids = input_ids[:,:index]\n",
|
||||
" context_mask = attention_mask[:,:index]\n",
|
||||
" embed_ids = input_ids[:,index:index + args.vae_context]\n",
|
||||
" embed_mask = attention_mask[:,index:index + args.vae_context]\n",
|
||||
" target_ids = input_ids[:,index:index + args.vae_context * 2]\n",
|
||||
" target_mask = input_ids[:,index:index + args.vae_context * 2]\n",
|
||||
"\n",
|
||||
" in_texts = [tokenizer.decode(toks, skip_special_tokens=True)\n",
|
||||
" for toks in torch.cat([context_ids, embed_ids], dim=1)]\n",
|
||||
" mean = model.encode(embed_ids[:bs], embed_mask[:bs])\n",
|
||||
" z = model.vae.vae.sample(mean.repeat_interleave(n_outputs, 0), tau=tau)\n",
|
||||
" context_ids = context_ids[:bs].repeat_interleave(n_outputs, 0)\n",
|
||||
" context_mask = context_mask[:bs].repeat_interleave(n_outputs, 0)\n",
|
||||
" # empty = z.new_zeros([z.shape[0], 0], dtype=torch.long)\n",
|
||||
" output_ids = model.generate(z, context_ids, context_mask, n_tokens, tau=tau)\n",
|
||||
" out_texts = [tokenizer.decode(toks, skip_special_tokens=True) for toks in output_ids]\n",
|
||||
" out_texts = list(batched(out_texts, n_outputs))\n",
|
||||
" print(\"======\")\n",
|
||||
" for in_text, out_batch in zip(in_texts, out_texts):\n",
|
||||
" print(\"=== Input ===\")\n",
|
||||
" print(in_text)\n",
|
||||
" print(\"=== Outputs ===\")\n",
|
||||
" for i, out_text in enumerate(out_batch):\n",
|
||||
" print(out_text)\n",
|
||||
" if i < len(out_batch) - 1:\n",
|
||||
" print(\"===\")\n",
|
||||
" print(\"======\")\n",
|
||||
"\n",
|
||||
"def save():\n",
|
||||
" print0(f\"### Saving model to {args.output}\", file=sys.stderr)\n",
|
||||
" accelerator.wait_for_everyone()\n",
|
||||
" if accelerator.is_main_process:\n",
|
||||
" unwrapped_model = accelerator.unwrap_model(router)\n",
|
||||
" unwrapped_model.save_pretrained(args.output)\n",
|
||||
" state_obj = {\"step\": i, \"last_kl_weight\": kl_sched(i)}\n",
|
||||
" with open(args.output / \"state.json\", \"w\") as f:\n",
|
||||
" json.dump(state_obj, f)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# train\n",
|
||||
"i = 0\n",
|
||||
"kl_sched = cosine_warmup(5000, 0.01)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"accelerator.wait_for_everyone()\n",
|
||||
"for epoch in trange(args.epochs, disable=not is_main):\n",
|
||||
" for input_ids, attention_mask in tqdm(dataloader, disable=not is_main):\n",
|
||||
" input_ids = input_ids.long()\n",
|
||||
" if is_main and i % 100 == 0:\n",
|
||||
" demo(accelerator.unwrap_model(router), input_ids, attention_mask, args.vae_context)\n",
|
||||
" pass\n",
|
||||
" with accelerator.accumulate(router):\n",
|
||||
" index = random.randrange(args.context - (args.vae_context * 2))\n",
|
||||
" context_ids = input_ids[:,:index]\n",
|
||||
" context_mask = attention_mask[:,:index]\n",
|
||||
" embed_ids = input_ids[:,index:index + args.vae_context]\n",
|
||||
" embed_mask = attention_mask[:,index:index + args.vae_context]\n",
|
||||
" target_ids = input_ids[:,index:index + args.vae_context * 2]\n",
|
||||
" target_mask = attention_mask[:,index:index + args.vae_context * 2]\n",
|
||||
"\n",
|
||||
" drop_mask = torch.rand([context_ids.shape[0], 1], device=device) < 0.5\n",
|
||||
" context_ids = torch.where(drop_mask, torch.zeros_like(context_ids), context_ids)\n",
|
||||
" context_mask = torch.where(drop_mask, torch.zeros_like(context_mask), context_mask)\n",
|
||||
" outputs = router(embed_ids, embed_mask,\n",
|
||||
" target_ids[:,:-1], target_mask[:,:-1],\n",
|
||||
" context_ids, context_mask)\n",
|
||||
" rec_losses = F.cross_entropy(\n",
|
||||
" outputs.logits[:, -args.vae_context * 2:].transpose(-1, -2),\n",
|
||||
" target_ids,\n",
|
||||
" reduction=\"none\",\n",
|
||||
" )\n",
|
||||
" n_toks = target_mask.sum()\n",
|
||||
" rec_loss = torch.sum(rec_losses * target_mask, dtype=torch.float32) / n_toks\n",
|
||||
" # kl_loss = torch.sum(mean**2 / 2, dtype=torch.float32) * kl_sched(i) / n_toks\n",
|
||||
" loss = rec_loss # + kl_loss\n",
|
||||
"\n",
|
||||
" # accelerator.backward(loss, inputs=list(p for p in accelerator.unwrap_model(router).model.parameters() if p.requires_grad))\n",
|
||||
" accelerator.backward(loss)\n",
|
||||
" # for n, p in router.named_parameters():\n",
|
||||
" # if p.grad is not None:\n",
|
||||
" # grad_norm = torch.norm(p.grad, dtype=torch.float32)\n",
|
||||
" # if grad_norm != 0:\n",
|
||||
" # print(f\"{n}: {grad_norm:g}\", file=sys.stderr)\n",
|
||||
" opt.step()\n",
|
||||
" opt.zero_grad()\n",
|
||||
"\n",
|
||||
" loss_global, rec_global = accelerator.reduce(\n",
|
||||
" (loss, rec_loss), \"mean\"\n",
|
||||
" )\n",
|
||||
" print0(\n",
|
||||
" f\"epoch: {epoch}, step: {i}, loss: {loss_global.item():g}, rec: {rec_global.item():g}\",\n",
|
||||
" file=sys.stderr,\n",
|
||||
" )\n",
|
||||
" i += 1\n",
|
||||
"\n",
|
||||
" if i % args.save_every == 0:\n",
|
||||
" save()\n",
|
||||
"\n",
|
||||
" save()\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": ".venv",
|
||||
"language": "python",
|
||||
"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.0rc1"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 2
|
||||
}
|
||||
Generated
+3285
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,40 @@
|
||||
[tool.poetry]
|
||||
name = "src"
|
||||
version = "0.1.0"
|
||||
description = "Trying to use a AdaVAE (an LLM VAE) as a world model in an RL agent in a text game"
|
||||
authors = ["wassname <git@wassname.org>"]
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.11,<3.13"
|
||||
torch = {version = "^2.1.0+cu118", source = "pytorch"}
|
||||
simple-parsing = "^0.1.4"
|
||||
tqdm = "^4.66.1"
|
||||
numpy = "^1.26.1"
|
||||
pandas = "^2.1.1"
|
||||
lightning = "^2.1.0"
|
||||
matplotlib = "^3.8.0"
|
||||
loguru = "^0.7.2"
|
||||
einops = "^0.7.0"
|
||||
scikit-learn = "^1.3.1"
|
||||
pytorch-optimizer = "^2.12.0"
|
||||
torchinfo = "^1.8.0"
|
||||
accelerate = "^0.24.1"
|
||||
datasets = "^2.14.6"
|
||||
peft = "^0.5.0"
|
||||
bitsandbytes = {url = "https://github.com/TimDettmers/bitsandbytes/releases/download/0.41.0/bitsandbytes-0.41.0-py3-none-any.whl"}
|
||||
transformers = "4.34.0"
|
||||
|
||||
[[tool.poetry.source]]
|
||||
name = "pytorch"
|
||||
url = "https://download.pytorch.org/whl/cu118"
|
||||
priority = "explicit"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
ipykernel = "^6.25.2"
|
||||
ruff = "^0.1.3"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
@@ -1 +1,2 @@
|
||||
from .transformer import Transformer, TransformerConfig
|
||||
# from .transformer import Transformer, TransformerConfig
|
||||
from .bigvae import BigVAE, BigVAEConfig
|
||||
|
||||
@@ -0,0 +1,419 @@
|
||||
"""
|
||||
Fine-tunes a language model on pre-tokenized data.
|
||||
|
||||
|
||||
From https://raw.githubusercontent.com/JD-P/minihf/adavae-moe/vae_infer.py
|
||||
See https://huggingface.co/jdpressman/BigVAE-Mistral-7B-v0.1/blob/main/README.md
|
||||
|
||||
BigVAE is an [AdaVAE](https://arxiv.org/abs/2205.05862) trained as a pair of LoRa finetunes on [Mistral 7B](https://huggingface.co/mistralai/Mistral-7B-v0.1).
|
||||
It is meant to be used with the [MiniHF VAE inference code](https://github.com/JD-P/minihf/blob/adavae-moe/vae_infer.py) and will not work if you try to load it....
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
import argparse
|
||||
from contextlib import contextmanager
|
||||
from itertools import chain, islice
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
import random
|
||||
import os
|
||||
import sys
|
||||
import zipfile
|
||||
|
||||
import accelerate
|
||||
from datasets import load_dataset
|
||||
import peft
|
||||
import safetensors.torch as safetorch
|
||||
import torch
|
||||
from torch import nn, optim
|
||||
from torch.nn import functional as F
|
||||
from torch.utils import data
|
||||
from tqdm import trange, tqdm
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||||
from loguru import logger
|
||||
from peft import PeftModel, LoraConfig
|
||||
|
||||
logger.add(sys.stderr, format="{time} {level} {message}", filter="my_module", level="INFO")
|
||||
print = tqdm.external_write_mode()(logger.info)
|
||||
|
||||
|
||||
def cosine_warmup(steps, value=1.0):
|
||||
return lambda i: value * math.sin(min(i / steps, 1) * math.pi / 2) ** 2
|
||||
|
||||
|
||||
@contextmanager
|
||||
def set_adapter(model, adapter_name):
|
||||
old_adapter_name = model.active_adapter
|
||||
try:
|
||||
if adapter_name is not None:
|
||||
model.set_adapter(adapter_name)
|
||||
yield model
|
||||
else:
|
||||
with model.disable_adapter():
|
||||
yield model
|
||||
finally:
|
||||
model.set_adapter(old_adapter_name)
|
||||
|
||||
|
||||
def gumbel_like(x):
|
||||
return torch.rand_like(x).log_().nan_to_num_().neg_().log_().neg_()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def disable_causal_mask():
|
||||
raise NotImplementedError("FIXME")
|
||||
import transformers.models.llama.modeling_llama as modeling
|
||||
|
||||
decoder_fn = modeling._make_causal_mask
|
||||
|
||||
def encoder_fn(*args, **kwargs):
|
||||
return torch.zeros_like(decoder_fn(*args, **kwargs))
|
||||
|
||||
try:
|
||||
modeling._make_causal_mask = encoder_fn
|
||||
yield
|
||||
finally:
|
||||
modeling._make_causal_mask = decoder_fn
|
||||
|
||||
|
||||
@contextmanager
|
||||
def disable_causal_mask_mistral():
|
||||
raise NotImplementedError("FIXME")
|
||||
import transformers.models.mistral.modeling_mistral as modeling
|
||||
|
||||
decoder_fn = modeling._make_sliding_window_causal_mask
|
||||
|
||||
def encoder_fn(*args, **kwargs):
|
||||
return torch.zeros_like(decoder_fn(*args, **kwargs))
|
||||
|
||||
try:
|
||||
modeling._make_sliding_window_causal_mask = encoder_fn
|
||||
yield
|
||||
finally:
|
||||
modeling._make_sliding_window_causal_mask = decoder_fn
|
||||
|
||||
|
||||
class VAEHead(nn.Module):
|
||||
def __init__(self, d_model, z_dim):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.z_dim = z_dim
|
||||
self.f = nn.Linear(d_model, 1)
|
||||
self.w_e = nn.Linear(d_model, z_dim)
|
||||
self.w_d = nn.Linear(z_dim, d_model)
|
||||
nn.init.orthogonal_(self.w_e.weight)
|
||||
with torch.no_grad():
|
||||
self.w_d.weight.copy_(self.w_e.weight.T)
|
||||
|
||||
def encode(self, hidden_states, attention_mask):
|
||||
scores = self.f(hidden_states)
|
||||
scores = scores + attention_mask[:, :, None].log().nan_to_num()
|
||||
weights = torch.softmax(scores, dim=1)
|
||||
pooled = torch.sum(hidden_states * weights, dim=1)
|
||||
return self.w_e(pooled)
|
||||
|
||||
def sample(self, mean, tau=1.0):
|
||||
return mean + torch.randn_like(mean) * tau**0.5
|
||||
|
||||
def decode(self, z):
|
||||
return self.w_d(z)
|
||||
|
||||
|
||||
class BigVAE(nn.Module):
|
||||
"""
|
||||
Version of AdaVAE with transformer.
|
||||
"""
|
||||
|
||||
def __init__(self, base_model_peft: PeftModel, peft_config: LoraConfig, z_dim: int=768, device: str = "cuda"):
|
||||
super().__init__()
|
||||
self.model = base_model_peft
|
||||
self.model.add_adapter("encoder", peft_config)
|
||||
self.model.add_adapter("decoder", peft_config)
|
||||
self.model.config.output_hidden_states = True
|
||||
self.vae_head = VAEHead(self.model.config.hidden_size, z_dim).to(device)
|
||||
|
||||
def save_pretrained(self, path):
|
||||
path = Path(path)
|
||||
self.model.save_pretrained(path, safe_serialization=True)
|
||||
safetorch.save_file(self.vae_head.state_dict(), path / "vae.safetensors")
|
||||
|
||||
def load_pretrained(self, path, is_trainable=False):
|
||||
path = Path(path)
|
||||
self.model.delete_adapter("encoder")
|
||||
self.model.load_adapter(path / "encoder", "encoder", is_trainable=is_trainable)
|
||||
self.model.delete_adapter("decoder")
|
||||
self.model.load_adapter(path / "decoder", "decoder", is_trainable=is_trainable)
|
||||
self.vae_head.load_state_dict(safetorch.load_file(path / "vae.safetensors"))
|
||||
|
||||
def encode(self, input_ids, attention_mask):
|
||||
with set_adapter(self.model, "encoder"), disable_causal_mask_mistral():
|
||||
outputs = self.model(
|
||||
input_ids=input_ids, attention_mask=attention_mask, use_cache=False
|
||||
)
|
||||
return self.vae_head.encode(outputs.hidden_states[-1], attention_mask)
|
||||
|
||||
def input_ids_to_embeds(self, input_ids):
|
||||
embed_weight = self.model.get_input_embeddings().weight
|
||||
input_one_hots = F.one_hot(input_ids, num_classes=self.model.config.vocab_size)
|
||||
return input_one_hots.to(embed_weight) @ embed_weight
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, z, input_ids, attention_mask, n_tokens, tau=1.0):
|
||||
"""
|
||||
Takes in a latent vector z from past tokens and generates next n_tokens tokens.
|
||||
|
||||
Used in e.g. https://github.com/JD-P/minihf/blob/adavae-moe/train_vae_overlap.py#L335
|
||||
"""
|
||||
z_embed = self.vae_head.decode(z)[:, None]
|
||||
inputs_embeds = self.input_ids_to_embeds(input_ids)
|
||||
inputs_embeds = torch.cat([z_embed, inputs_embeds], dim=1)
|
||||
attention_mask = torch.cat(
|
||||
[attention_mask.new_ones([attention_mask.shape[0], 1]), attention_mask], dim=1
|
||||
)
|
||||
new_embeds, past = None, None
|
||||
with set_adapter(self.model, "decoder"):
|
||||
for _ in range(n_tokens):
|
||||
outputs = self.model(
|
||||
inputs_embeds=inputs_embeds if past is None else new_embeds,
|
||||
attention_mask=attention_mask,
|
||||
use_cache=True,
|
||||
past_key_values=past,
|
||||
)
|
||||
logits = outputs.logits[:, -1:, :].float()
|
||||
new_input_ids = torch.argmax(logits + gumbel_like(logits) * tau, dim=-1)
|
||||
input_ids = torch.cat([input_ids, new_input_ids], dim=1)
|
||||
new_embeds = self.input_ids_to_embeds(new_input_ids)
|
||||
attention_mask = torch.cat(
|
||||
[attention_mask, attention_mask.new_ones([attention_mask.shape[0], 1])], dim=1
|
||||
)
|
||||
past = outputs.past_key_values
|
||||
return input_ids
|
||||
|
||||
def forward(self, input_ids, attention_mask, prefix_ids, prefix_mask):
|
||||
input_ids_all = torch.cat([prefix_ids, input_ids], dim=1)
|
||||
attn_mask_all = torch.cat([prefix_mask, attention_mask], dim=1)
|
||||
mean = self.encode(input_ids, attention_mask)
|
||||
z = self.vae_head.sample(mean)
|
||||
z_embed = self.vae_head.decode(z)[:, None]
|
||||
inputs_embeds = self.input_ids_to_embeds(input_ids_all)
|
||||
inputs_embeds = torch.cat([z_embed, inputs_embeds], dim=1)
|
||||
attention_mask = torch.cat(
|
||||
[attention_mask.new_ones([attn_mask_all.shape[0], 1]), attn_mask_all], dim=1
|
||||
)
|
||||
with set_adapter(self.model, "decoder"):
|
||||
outputs = self.model(
|
||||
inputs_embeds=inputs_embeds, attention_mask=attention_mask, use_cache=False
|
||||
)
|
||||
return outputs, mean
|
||||
|
||||
class BigVAERouter(nn.Module):
|
||||
def __init__(self, base_model_peft: PeftModel, vae: BigVAE, device: str = "cuda"):
|
||||
super().__init__()
|
||||
peft_config = base_model_peft.peft_config["default"] # debug this
|
||||
self.model = base_model_peft
|
||||
self.model.add_adapter("router", peft_config)
|
||||
self.model.config.output_hidden_states = True
|
||||
self.vae = vae
|
||||
|
||||
def save_pretrained(self, path):
|
||||
path = Path(path)
|
||||
self.model.save_pretrained(path, safe_serialization=True)
|
||||
safetorch.save_file(self.model.state_dict(), path / "router.safetensors")
|
||||
safetorch.save_file(self.vae.vae_head.state_dict(), path / "vae.safetensors")
|
||||
|
||||
def load_pretrained(self, path, is_trainable=False):
|
||||
path = Path(path)
|
||||
self.model.delete_adapter("router")
|
||||
if (path / "router").exists():
|
||||
self.model.load_adapter(path / "router", "router", is_trainable=is_trainable)
|
||||
else:
|
||||
self.model.load_adapter(path / "decoder", "router", is_trainable=is_trainable)
|
||||
|
||||
def encode(self, input_ids, attention_mask):
|
||||
with set_adapter(self.vae.model, "encoder"), disable_causal_mask_mistral():
|
||||
outputs = self.vae.model(
|
||||
input_ids=input_ids, attention_mask=attention_mask, use_cache=False
|
||||
)
|
||||
return self.vae.vae_head.encode(outputs.hidden_states[-1], attention_mask)
|
||||
|
||||
def input_ids_to_embeds(self, input_ids):
|
||||
embed_weight = self.model.get_input_embeddings().weight
|
||||
input_one_hots = F.one_hot(input_ids, num_classes=self.model.config.vocab_size)
|
||||
return input_one_hots.to(embed_weight) @ embed_weight
|
||||
|
||||
def generate(self, z, input_ids, attention_mask, n_tokens, tau=1.0):
|
||||
"""
|
||||
predict next token given a latent vector z and previous tokens as input_ids
|
||||
|
||||
e.g. https://github.com/JD-P/minihf/blob/adavae-moe/vae_infer.py#L428
|
||||
"""
|
||||
z_embed = self.vae.vae_head.decode(z)[:, None]
|
||||
inputs_embeds = self.input_ids_to_embeds(input_ids)
|
||||
inputs_embeds = torch.cat([inputs_embeds, z_embed], dim=1)
|
||||
attention_mask = torch.cat(
|
||||
[attention_mask, attention_mask.new_ones([attention_mask.shape[0], 1])], dim=1
|
||||
)
|
||||
new_embeds, past = None, None
|
||||
with set_adapter(self.vae.model, "router"):
|
||||
for _ in range(n_tokens):
|
||||
outputs = self.model(
|
||||
inputs_embeds=inputs_embeds if past is None else new_embeds,
|
||||
attention_mask=attention_mask,
|
||||
use_cache=True,
|
||||
past_key_values=past,
|
||||
)
|
||||
logits = outputs.logits[:, -1:, :].float()
|
||||
new_input_ids = torch.argmax(logits + gumbel_like(logits) * tau, dim=-1)
|
||||
input_ids = torch.cat([input_ids, new_input_ids], dim=1)
|
||||
new_embeds = self.input_ids_to_embeds(new_input_ids)
|
||||
attention_mask = torch.cat(
|
||||
[attention_mask, attention_mask.new_ones([attention_mask.shape[0], 1])], dim=1
|
||||
)
|
||||
past = outputs.past_key_values
|
||||
return input_ids
|
||||
|
||||
# def generate_cfg(self, z, input_ids, attention_mask, n_tokens, tau=1.0, cfg_scale=1):
|
||||
# """
|
||||
# predict next tokens given a latent vector z and previous tokens as input_ids
|
||||
|
||||
# but this one mixes base and router
|
||||
# was used in topic modelling here https://github.com/JD-P/minihf/blob/adavae-moe/vae_infer.py#L614
|
||||
# I can soon delete it
|
||||
# """
|
||||
# z_embed = self.vae.vae.decode(z)[:, None]
|
||||
# inputs_embeds_base = self.input_ids_to_embeds(input_ids)
|
||||
# inputs_embeds_router = torch.cat([inputs_embeds_base, z_embed], dim=1)
|
||||
# attention_mask = torch.cat(
|
||||
# [attention_mask, attention_mask.new_ones([attention_mask.shape[0], 1])], dim=1
|
||||
# )
|
||||
# new_embeds, base_past, router_past = None, None, None
|
||||
# for _ in range(n_tokens):
|
||||
# with set_adapter(self.vae.model, "router"):
|
||||
# router_outputs = self.model(
|
||||
# inputs_embeds=inputs_embeds_router if router_past is None else new_embeds,
|
||||
# attention_mask=attention_mask,
|
||||
# use_cache=True,
|
||||
# past_key_values=router_past,
|
||||
# )
|
||||
# with set_adapter(self.vae.model, None):
|
||||
# base_outputs = self.model(
|
||||
# inputs_embeds=inputs_embeds_base if base_past is None else new_embeds,
|
||||
# attention_mask=attention_mask[:,:-1],
|
||||
# use_cache=True,
|
||||
# past_key_values=base_past,
|
||||
# )
|
||||
# base_logits = base_outputs.logits[:, -1:, :].float()
|
||||
# router_logits = router_outputs.logits[:, -1:, :].float()
|
||||
|
||||
# # mix base and router prediction based on cfg_scale
|
||||
# logits = base_logits + cfg_scale * (router_logits - base_logits)
|
||||
# new_input_ids = torch.argmax(logits + gumbel_like(logits) * tau, dim=-1)
|
||||
# input_ids = torch.cat([input_ids, new_input_ids], dim=1)
|
||||
# new_embeds = self.input_ids_to_embeds(new_input_ids)
|
||||
# attention_mask = torch.cat(
|
||||
# [attention_mask, attention_mask.new_ones([attention_mask.shape[0], 1])], dim=1
|
||||
# )
|
||||
# base_past = base_outputs.past_key_values
|
||||
# router_past = router_outputs.past_key_values
|
||||
# return input_ids
|
||||
|
||||
def forward(self, input_ids, input_mask, target_ids, target_mask, prefix_ids, prefix_mask):
|
||||
"""Like the decoder only, but with context/prefix."""
|
||||
mean = self.encode(input_ids, input_mask)
|
||||
z = self.vae.vae_head.sample(mean)
|
||||
z_embed = self.vae.vae_head.decode(z)[:, None]
|
||||
prefix_embeds = self.input_ids_to_embeds(prefix_ids)
|
||||
target_embeds = self.input_ids_to_embeds(target_ids)
|
||||
inputs_embeds = torch.cat([prefix_embeds, z_embed, target_embeds], dim=1)
|
||||
attention_mask = torch.cat(
|
||||
[prefix_mask,
|
||||
target_mask.new_ones([prefix_mask.shape[0], 1]),
|
||||
target_mask], dim=1
|
||||
)
|
||||
outputs = self.model(
|
||||
inputs_embeds=inputs_embeds, attention_mask=attention_mask, use_cache=False
|
||||
)
|
||||
return outputs
|
||||
|
||||
def batched(iterable, n):
|
||||
"Batch data into tuples of length n. The last batch may be shorter."
|
||||
# batched('ABCDEFG', 3) --> ABC DEF G
|
||||
if n < 1:
|
||||
raise ValueError("n must be at least one")
|
||||
it = iter(iterable)
|
||||
while batch := tuple(islice(it, n)):
|
||||
yield batch
|
||||
|
||||
|
||||
@dataclass
|
||||
class BigVAEConfig:
|
||||
model_name: str = "stabilityai/stablelm-3b-4e1t"
|
||||
dropout: float = 0
|
||||
rank: int = 32
|
||||
z_dim: int = 768
|
||||
start_from: str = None
|
||||
|
||||
tokens_per_block: int
|
||||
max_blocks: int
|
||||
attention: str
|
||||
|
||||
num_layers: int
|
||||
num_heads: int
|
||||
embed_dim: int
|
||||
|
||||
embed_pdrop: float
|
||||
resid_pdrop: float
|
||||
attn_pdrop: float
|
||||
|
||||
@property
|
||||
def max_tokens(self):
|
||||
return self.tokens_per_block * self.max_blocks
|
||||
|
||||
def load_model(config, device='cuda'):
|
||||
tokenizer = AutoTokenizer.from_pretrained(config.model_name, trust_remote_code=True)
|
||||
tokenizer.padding_side = "left"
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
base_model = AutoModelForCausalLM.from_pretrained(
|
||||
config.model_name,
|
||||
device_map={"": device},
|
||||
quantization_config=bnb_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True
|
||||
)
|
||||
peft_config = peft.LoraConfig(
|
||||
peft.TaskType.CAUSAL_LM,
|
||||
inference_mode=False,
|
||||
r=config.rank,
|
||||
lora_alpha=8,
|
||||
lora_dropout=config.dropout,
|
||||
target_modules=[
|
||||
"self_attn.q_proj",
|
||||
"self_attn.k_proj",
|
||||
"self_attn.v_proj",
|
||||
"self_attn.o_proj",
|
||||
"mlp.gate_proj",
|
||||
"mlp.up_proj",
|
||||
"mlp.down_proj",
|
||||
],
|
||||
)
|
||||
base_model_peft = peft.get_peft_model(base_model, peft_config)
|
||||
vae_model = BigVAE(
|
||||
base_model_peft, device, peft_config, z_dim=config.z_dim,
|
||||
)
|
||||
if config.start_from:
|
||||
vae_model.load_pretrained(config.start_from)
|
||||
base_model_peft.requires_grad_(False)
|
||||
vae_model.vae_head.requires_grad_(False)
|
||||
vae_model.vae_head.w_d.requires_grad_()
|
||||
router = BigVAERouter(base_model_peft, vae_model, device, peft_config)
|
||||
if config.start_from:
|
||||
router.load_pretrained(config.start_from, is_trainable=True)
|
||||
print(router.model.print_trainable_parameters())
|
||||
router.model.set_adapter("router")
|
||||
+95
-95
@@ -1,120 +1,120 @@
|
||||
"""
|
||||
Credits to https://github.com/karpathy/minGPT
|
||||
"""
|
||||
# """
|
||||
# Credits to https://github.com/karpathy/minGPT
|
||||
# """
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
from typing import Optional
|
||||
# from dataclasses import dataclass
|
||||
# import math
|
||||
# from typing import Optional
|
||||
|
||||
from einops import rearrange
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
# from einops import rearrange
|
||||
# import torch
|
||||
# import torch.nn as nn
|
||||
# from torch.nn import functional as F
|
||||
|
||||
from .kv_caching import KeysValues, KVCache
|
||||
# from .kv_caching import KeysValues, KVCache
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransformerConfig:
|
||||
tokens_per_block: int
|
||||
max_blocks: int
|
||||
attention: str
|
||||
# @dataclass
|
||||
# class TransformerConfig:
|
||||
# tokens_per_block: int
|
||||
# max_blocks: int
|
||||
# attention: str
|
||||
|
||||
num_layers: int
|
||||
num_heads: int
|
||||
embed_dim: int
|
||||
# num_layers: int
|
||||
# num_heads: int
|
||||
# embed_dim: int
|
||||
|
||||
embed_pdrop: float
|
||||
resid_pdrop: float
|
||||
attn_pdrop: float
|
||||
# embed_pdrop: float
|
||||
# resid_pdrop: float
|
||||
# attn_pdrop: float
|
||||
|
||||
@property
|
||||
def max_tokens(self):
|
||||
return self.tokens_per_block * self.max_blocks
|
||||
# @property
|
||||
# def max_tokens(self):
|
||||
# return self.tokens_per_block * self.max_blocks
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, config: TransformerConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.drop = nn.Dropout(config.embed_pdrop)
|
||||
self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
|
||||
self.ln_f = nn.LayerNorm(config.embed_dim)
|
||||
# class Transformer(nn.Module):
|
||||
# def __init__(self, config: TransformerConfig) -> None:
|
||||
# super().__init__()
|
||||
# self.config = config
|
||||
# self.drop = nn.Dropout(config.embed_pdrop)
|
||||
# self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
|
||||
# self.ln_f = nn.LayerNorm(config.embed_dim)
|
||||
|
||||
def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
|
||||
device = self.ln_f.weight.device # Assumption that all submodules are on the same device
|
||||
return KeysValues(n, self.config.num_heads, max_tokens, self.config.embed_dim, self.config.num_layers, device)
|
||||
# def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
|
||||
# device = self.ln_f.weight.device # Assumption that all submodules are on the same device
|
||||
# return KeysValues(n, self.config.num_heads, max_tokens, self.config.embed_dim, self.config.num_layers, device)
|
||||
|
||||
def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
|
||||
assert past_keys_values is None or len(past_keys_values) == len(self.blocks)
|
||||
x = self.drop(sequences)
|
||||
for i, block in enumerate(self.blocks):
|
||||
x = block(x, None if past_keys_values is None else past_keys_values[i])
|
||||
# def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
|
||||
# assert past_keys_values is None or len(past_keys_values) == len(self.blocks)
|
||||
# x = self.drop(sequences)
|
||||
# for i, block in enumerate(self.blocks):
|
||||
# x = block(x, None if past_keys_values is None else past_keys_values[i])
|
||||
|
||||
x = self.ln_f(x)
|
||||
return x
|
||||
# x = self.ln_f(x)
|
||||
# return x
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
def __init__(self, config: TransformerConfig) -> None:
|
||||
super().__init__()
|
||||
self.ln1 = nn.LayerNorm(config.embed_dim)
|
||||
self.ln2 = nn.LayerNorm(config.embed_dim)
|
||||
self.attn = SelfAttention(config)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(config.embed_dim, 4 * config.embed_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(4 * config.embed_dim, config.embed_dim),
|
||||
nn.Dropout(config.resid_pdrop),
|
||||
)
|
||||
# class Block(nn.Module):
|
||||
# def __init__(self, config: TransformerConfig) -> None:
|
||||
# super().__init__()
|
||||
# self.ln1 = nn.LayerNorm(config.embed_dim)
|
||||
# self.ln2 = nn.LayerNorm(config.embed_dim)
|
||||
# self.attn = SelfAttention(config)
|
||||
# self.mlp = nn.Sequential(
|
||||
# nn.Linear(config.embed_dim, 4 * config.embed_dim),
|
||||
# nn.GELU(),
|
||||
# nn.Linear(4 * config.embed_dim, config.embed_dim),
|
||||
# nn.Dropout(config.resid_pdrop),
|
||||
# )
|
||||
|
||||
def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
|
||||
x_attn = self.attn(self.ln1(x), past_keys_values)
|
||||
x = x + x_attn
|
||||
x = x + self.mlp(self.ln2(x))
|
||||
return x
|
||||
# def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
|
||||
# x_attn = self.attn(self.ln1(x), past_keys_values)
|
||||
# x = x + x_attn
|
||||
# x = x + self.mlp(self.ln2(x))
|
||||
# return x
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, config: TransformerConfig) -> None:
|
||||
super().__init__()
|
||||
assert config.embed_dim % config.num_heads == 0
|
||||
assert config.attention in ('causal', 'block_causal')
|
||||
self.num_heads = config.num_heads
|
||||
self.key = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
self.query = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
self.value = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
self.attn_drop = nn.Dropout(config.attn_pdrop)
|
||||
self.resid_drop = nn.Dropout(config.resid_pdrop)
|
||||
self.proj = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
# class SelfAttention(nn.Module):
|
||||
# def __init__(self, config: TransformerConfig) -> None:
|
||||
# super().__init__()
|
||||
# assert config.embed_dim % config.num_heads == 0
|
||||
# assert config.attention in ('causal', 'block_causal')
|
||||
# self.num_heads = config.num_heads
|
||||
# self.key = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
# self.query = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
# self.value = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
# self.attn_drop = nn.Dropout(config.attn_pdrop)
|
||||
# self.resid_drop = nn.Dropout(config.resid_pdrop)
|
||||
# self.proj = nn.Linear(config.embed_dim, config.embed_dim)
|
||||
|
||||
causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
|
||||
block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
|
||||
self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
|
||||
# causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
|
||||
# block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
|
||||
# self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
|
||||
|
||||
def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
|
||||
B, T, C = x.size()
|
||||
if kv_cache is not None:
|
||||
b, nh, L, c = kv_cache.shape
|
||||
assert nh == self.num_heads and b == B and c * nh == C
|
||||
else:
|
||||
L = 0
|
||||
# def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
|
||||
# B, T, C = x.size()
|
||||
# if kv_cache is not None:
|
||||
# b, nh, L, c = kv_cache.shape
|
||||
# assert nh == self.num_heads and b == B and c * nh == C
|
||||
# else:
|
||||
# L = 0
|
||||
|
||||
q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
# q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
# k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
# v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2) # (B, nh, T, hs)
|
||||
|
||||
if kv_cache is not None:
|
||||
kv_cache.update(k, v)
|
||||
k, v = kv_cache.get()
|
||||
# if kv_cache is not None:
|
||||
# kv_cache.update(k, v)
|
||||
# k, v = kv_cache.get()
|
||||
|
||||
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
|
||||
att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
|
||||
att = F.softmax(att, dim=-1)
|
||||
att = self.attn_drop(att)
|
||||
y = att @ v
|
||||
y = rearrange(y, 'b h t e -> b t (h e)')
|
||||
# att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
|
||||
# att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
|
||||
# att = F.softmax(att, dim=-1)
|
||||
# att = self.attn_drop(att)
|
||||
# y = att @ v
|
||||
# y = rearrange(y, 'b h t e -> b t (h e)')
|
||||
|
||||
y = self.resid_drop(self.proj(y))
|
||||
# y = self.resid_drop(self.proj(y))
|
||||
|
||||
return y
|
||||
# return y
|
||||
|
||||
@@ -10,7 +10,8 @@ from dataset import Batch
|
||||
from .kv_caching import KeysValues
|
||||
from .slicer import Embedder, Head
|
||||
from .tokenizer import Tokenizer
|
||||
from .transformer import Transformer, TransformerConfig
|
||||
# from .transformer import Transformer, TransformerConfig
|
||||
from .bigvae import BigVAE, BigVAEConfig
|
||||
from utils import init_weights, LossWithIntermediateLosses
|
||||
|
||||
|
||||
@@ -23,11 +24,11 @@ class WorldModelOutput:
|
||||
|
||||
|
||||
class WorldModel(nn.Module):
|
||||
def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: TransformerConfig) -> None:
|
||||
def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: BigVAEConfig) -> None:
|
||||
super().__init__()
|
||||
self.obs_vocab_size, self.act_vocab_size = obs_vocab_size, act_vocab_size
|
||||
self.config = config
|
||||
self.transformer = Transformer(config)
|
||||
self.transformer = BigVAE(config)
|
||||
|
||||
all_but_last_obs_tokens_pattern = torch.ones(config.tokens_per_block)
|
||||
all_but_last_obs_tokens_pattern[-2] = 0
|
||||
|
||||
Reference in New Issue
Block a user