3min epoch breakout

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

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