From 106027a7f98a38ec7c70682f8542c8329edb3cea Mon Sep 17 00:00:00 2001 From: wassname Date: Tue, 14 Nov 2023 10:42:17 +0800 Subject: [PATCH] add 1b model, use contextlib, --- config/trainer.yaml | 4 +-- config/world_model/default.yaml | 2 +- research_journal.md | 16 +++++++++++ src/models/actor_critic.py | 4 +++ src/models/transformer.py | 51 +++++++++++++++++++++++++++------ src/utils.py | 4 ++- 6 files changed, 68 insertions(+), 13 deletions(-) diff --git a/config/trainer.yaml b/config/trainer.yaml index 01af86e..81f5e88 100644 --- a/config/trainer.yaml +++ b/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 diff --git a/config/world_model/default.yaml b/config/world_model/default.yaml index d03a227..cb9fcd3 100644 --- a/config/world_model/default.yaml +++ b/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 diff --git a/research_journal.md b/research_journal.md index 94d442f..bfd295d 100644 --- a/research_journal.md +++ b/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 diff --git a/src/models/actor_critic.py b/src/models/actor_critic.py index c22e641..69dc3de 100644 --- a/src/models/actor_critic.py +++ b/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) diff --git a/src/models/transformer.py b/src/models/transformer.py index 9b28849..40b4430 100644 --- a/src/models/transformer.py +++ b/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 diff --git a/src/utils.py b/src/utils.py index 1f134d4..ed8e31a 100644 --- a/src/utils.py +++ b/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