iris_bigvae/journal.md

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

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

2023-11-12 12:58:27

Hmm Ok so they have a seperate encoder and decoder and the transformer just represents the dynamics model? But where are the encoder and decodr?

Hmm so one image turns into a sequence of tokens, how?

• decoding it just using the tokenizer.embeddings then tokenizer.decode, and it's in the world model

OK it looks like I don't actually need a VAE here! The encode decode is already made and is decoupled from the dynamic model. So lets start from scratch, loading minstral...

IDEAS

Could I just use a single transformer? Dreamer seem use encode and decode to a latent state, but IRIS doesn't?