diff --git a/README.md b/README.md index 2530571..617104b 100644 --- a/README.md +++ b/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)
diff --git a/config/tokenizer/default.yaml b/config/tokenizer/default.yaml index 39fe6ae..e789557 100644 --- a/config/tokenizer/default.yaml +++ b/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 diff --git a/config/trainer.yaml b/config/trainer.yaml index f8022a0..4adf8b1 100644 --- a/config/trainer.yaml +++ b/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 diff --git a/config/world_model/default.yaml b/config/world_model/default.yaml index 16769b7..22e8175 100644 --- a/config/world_model/default.yaml +++ b/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 diff --git a/research_journal.md b/research_journal.md index de33a90..5e2b17e 100644 --- a/research_journal.md +++ b/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 + + + diff --git a/src/models/actor_critic.py b/src/models/actor_critic.py index 0fff2b6..75b4261 100644 --- a/src/models/actor_critic.py +++ b/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 diff --git a/src/models/transformer.py b/src/models/transformer.py index 40b4430..86139c4 100644 --- a/src/models/transformer.py +++ b/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: