3min epoch breakout

https://wandb.ai/wassname/iris/runs/1zjsneyu?workspace=user-wassname

README.md

@@ -1,3 +1,16 @@
+
+A fork of IRIS where I use a pretrained LLM as the tranformer (with LoRa). 
+
+My hypothesis: Pretrained LLM's make good world models by including a lot of world information!
+
+details:
+- for speed and cost I use a small 1.5B model. But it would be interesting to try a 7B one
+- for speed I use a smaller actor critic than in IRIS
+- max_blocks 20->10
+- batch smaller because of my small machine
+- actor_critic.steps_per_epoch 200->20
+- world_model.batch_num_sampler; 64->8 because the forzen transformer uses lots of gpu ram
+
 # Transformers are Sample-Efficient World Models (IRIS)
 
 [Transformers are Sample-Efficient World Models](https://openreview.net/forum?id=vhFu1Acb0xb) <br>

config/tokenizer/default.yaml

@@ -1,14 +1,14 @@
 _target_: models.tokenizer.Tokenizer
 
-vocab_size: 512
-embed_dim: 512
+vocab_size: 2048
+embed_dim: 2048
 encoder:
   _target_: models.tokenizer.Encoder
   config:
     _target_: models.tokenizer.EncoderDecoderConfig
     resolution: 64
     in_channels: 3
-    z_channels: 512
+    z_channels: 2048
     ch: 64
     ch_mult: [1, 1, 1, 1, 1]
     num_res_blocks: 2

config/trainer.yaml

@@ -60,7 +60,7 @@ training:
     start_after_epochs: 5
     steps_per_epoch: 200
   world_model:
-    batch_num_samples: 8
+    batch_num_samples: 8 # pretrained models use lots of 
     grad_acc_steps: 1
     max_grad_norm: 10.0
     weight_decay: 0.01
@@ -71,7 +71,7 @@ training:
     grad_acc_steps: 1
     max_grad_norm: 10.0
     start_after_epochs: 50
-    steps_per_epoch: 20
+    steps_per_epoch: 40
     imagine_horizon: ${common.sequence_length}
     burn_in: 20
     gamma: 0.995

config/world_model/default.yaml

@@ -1,10 +1,9 @@
 _target_: models.TransformerConfig
+max_blocks: 10  # this is the rollout length when training policy
+num_layers: 1
+num_heads: 1
+embed_dim: 2048  # change this to whatever the embedding dimension is in your pretrained llm 2048 for llama. 2560 for stablelm
+dropout: 0.1
+model_name: "PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T"
+rank: 32
 tokens_per_block: 17
-max_blocks: 10
-attention: 'causal'
-num_layers: 10
-num_heads: 4
-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

@@ -199,3 +199,27 @@ Training world_model: 100%|█████████████████
 Training tokenizer: 55sec
 Training world_model 72 sec
 train actor_critic 3min. It looks like it scales with lstm size!
+
+
+new changes 10mins
+lets try no lstm?
+Right now it will take 41 hours for on epoch lol
+
+
+
+ram during stages:
+- actor critic 20G/24
+
+
+how big does my actor critic need to be?
+- IRIS: large 512 lstm on 64,64,3 obs
+  - We ran our experiments with 8 Nvidia A100 40GB GPUs. With two Atari environments running on the same GPU, training takes around 7 days, resulting in an average of 3.5 days per environment.
+- twm: mlp 512
+
+
+How long to train?
+`600*10//6/24` = 41 days
+- 600 epochs * 10 minutes / 6 to get hours, 24 to get days
+
+
+

src/models/actor_critic.py

@@ -36,18 +36,19 @@ class ImagineOutput:
 class ActorCritic(nn.Module):
     def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
         super().__init__()
-        shrink = 4
+        shrink = 8
+        s = 2
         self.use_original_obs = use_original_obs
-        self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
+        self.conv1 = nn.Conv2d(3, 32//s, 3, stride=1, padding=1)
         self.maxp1 = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(32, 32, 3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(32//s, 32//s, 3, stride=1, padding=1)
         self.maxp2 = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(32//s, 64//s, 3, stride=1, padding=1)
         self.maxp3 = nn.MaxPool2d(2, 2)
-        self.conv4 = nn.Conv2d(64, 64//shrink, 3, stride=1, padding=1)
+        self.conv4 = nn.Conv2d(64//s, 64//shrink, 3, stride=1, padding=1)
         self.maxp4 = nn.MaxPool2d(2, 2)
 
-        self.lstm_dim = 64
+        self.lstm_dim = 16
         self.lstm = nn.LSTMCell(1024//shrink, self.lstm_dim)
         self.hx, self.cx = None, None
 

src/models/transformer.py

@@ -17,94 +17,23 @@ from .kv_caching import KeysValues, KVCache
 @dataclass
 class TransformerConfig:
 
-        
-    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
+    tokens_per_block: int
     
     # 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
-    start_from: str = None
 
     @property
     def max_tokens(self):
         return self.tokens_per_block * self.max_blocks
 
-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
-
-
-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)
-
-    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)
-
-        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)')
-
-        y = self.resid_drop(self.proj(y))
-
-        return y
 
 class Transformer(nn.Module):
     def __init__(self, config: TransformerConfig) -> None: