add 1b model, use contextlib,

config/trainer.yaml

@@ -60,14 +60,14 @@ training:
     start_after_epochs: 5
     steps_per_epoch: 200
   world_model:
-    batch_num_samples: 4
+    batch_num_samples: 8
     grad_acc_steps: 1
     max_grad_norm: 10.0
     weight_decay: 0.01
     start_after_epochs: 25
     steps_per_epoch: 200
   actor_critic:
-    batch_num_samples: 8
+    batch_num_samples: 32
     grad_acc_steps: 1
     max_grad_norm: 10.0
     start_after_epochs: 50

config/world_model/default.yaml

@@ -4,7 +4,7 @@ max_blocks: 20
 attention: 'causal'
 num_layers: 10
 num_heads: 4
-embed_dim: 2560 # change this to whatever the embedding dimension is in your pretrained llm
+embed_dim: 2048 # change this to whatever the embedding dimension is in your pretrained llm 2048 for llama. 2560 for stablelm
 embed_pdrop: 0.1
 resid_pdrop: 0.1
 attn_pdrop: 0.1

research_journal.md

@@ -147,3 +147,19 @@ ok even with a full run I get the error. I think it's a bug in the original repo
     RuntimeError: cannot reshape tensor of 0 elements into shape [8, 0, -1] because the unspecified dimension size -1 can be any value and is ambiguous
 
 Oh maybe it's because we don't keep track of KV cache, but it's actually used to track number of steps!!
+
+# 2023-11-13 20:11:51
+
+I go it working byt ut takes 30 seconds for one one actor critic batch, werird
+
+Experience collection (train_dataset): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:03<00:00, 60.45it/s]
+Training tokenizer: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:17<00:00, 11.26it/s]
+Training world_model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:12<00:00,  1.51it/s]
+Training actor_critic:  82%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                             | 165/200 [1:01:43<13:24, 22.99s/it]
+
+
+hm maybe it's just the face it has to backprop throguh the whole LLM :( damn... is there another way to train it? Daym. How many params did the original have?
+
+well running eval on the transformer brought it down from 100sec to 60, but it's still huge. 
+
+But then why is the model training fast? It makes not sense

src/models/actor_critic.py

@@ -146,6 +146,10 @@ class ActorCritic(nn.Module):
 
             outputs_ac = self(obs)
             action_token = Categorical(logits=outputs_ac.logits_actions).sample()
+            
+            # TODO this is really slow, I guess we need grad? does it help to put it in eval? no
+            # wm_env.world_model.eval()
+            
             obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
 
             all_actions.append(action_token)

src/models/transformer.py

@@ -5,7 +5,7 @@ Credits to https://github.com/karpathy/minGPT
 from dataclasses import dataclass
 import math
 from typing import Optional
-
+from contextlib import contextmanager
 from einops import rearrange
 import torch
 import torch.nn as nn
@@ -30,7 +30,9 @@ class TransformerConfig:
     resid_pdrop: float
     attn_pdrop: float
     
-    model_name: str = "stabilityai/stablelm-3b-4e1t"
+    # model_name: str = "stabilityai/stablelm-3b-4e1t"
+    # https://huggingface.co/PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T
+    model_name: str = "PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T"
     dropout: float = 0
     rank: int = 32
     z_dim: int = 768
@@ -118,13 +120,14 @@ class Transformer(nn.Module):
     # @torch.cuda.amp.autocast(dtype=torch.bfloat16)
     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) == self.config.num_layers
-        sequences = sequences.to(torch.bfloat16)
-        outputs = self.model(
-            inputs_embeds=sequences,
-            return_dict=True,
-            output_hidden_states=True,
-        )
-        x = outputs.logits.to(torch.float32)
+        with set_adapter(self.model, "dynamics"), disable_causal_mask(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            # sequences = sequences.to(torch.bfloat16)
+            outputs = self.model(
+                inputs_embeds=sequences,
+                return_dict=True,
+                output_hidden_states=True,
+            )
+        x = outputs.logits#.to(torch.float32)
         x = self.ln_f(x)
         
         # fake it, since it's used to keep track of steps
@@ -136,6 +139,8 @@ class Transformer(nn.Module):
             past_keys_values[0].update(torch.rand(v_size), torch.rand(v_size))
         return x
 
+        
+
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel, LoraConfig
@@ -179,3 +184,31 @@ def load_pretrained_model(config, device="cuda:0"):
     base_model_peft.add_adapter("dynamics", peft_config)
     print(base_model_peft.print_trainable_parameters())
     return base_model_peft
+
+@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)
+
+@contextmanager
+def disable_causal_mask():
+    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

src/utils.py

@@ -10,6 +10,8 @@ import torch.nn as nn
 
 from episode import Episode
 
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+
 
 def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
     """Credits to https://github.com/karpathy/minGPT"""
@@ -17,7 +19,7 @@ def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_na
     decay = set()
     no_decay = set()
     whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
-    blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+    blacklist_weight_modules = tuple(ALL_LAYERNORM_LAYERS+[torch.nn.Embedding])
     for mn, m in model.named_modules():
         for pn, p in m.named_parameters():
             fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name