diff --git a/apps/accelerate/chatllama/README.md b/apps/accelerate/chatllama/README.md index 417c087a..9edfdc97 100644 --- a/apps/accelerate/chatllama/README.md +++ b/apps/accelerate/chatllama/README.md @@ -407,6 +407,24 @@ We support 3 different options to prepare the `reward_training_data`: - **(⚠️WIP)** Few examples provided by the user and dataset synthetically expanded using LLM +## Single-Node Multi-GPU Training +Currently chatllama supports [Accelerate](https://github.com/huggingface/accelerate) and [DeepSpeed](https://github.com/microsoft/DeepSpeed) for multi-GPU training. +To run a distributed training you need to enable one of them in your `/artifacts/config/config.yaml` file by setting either +`deepspeed_enable` or `accelerate_enable` to `True` or `False`.
+Each type of training (i.e. reward model training, actor supervised fine-tuning, RLHF) has its own flags to tweak. +Deepspeed settings can be customised using the `/artifacts/config/ds_config.yaml` file, while accelerate can be configured by running +```bash +accelerate config +``` +from the command line. +Once the project is configured, the trainin must be started using: +```bash +deepspeed artifacts/main.py artifacts/config/config.yaml --type +``` +or +```bash +accelerate launch artifacts/main.py artifacts/config/config.yaml --type +``` # License See the [LICENSE](https://github.com/nebuly-ai/nebullvm/blob/main/apps/accelerate/chatllama/LICENSE) file. diff --git a/apps/accelerate/chatllama/artifacts/config/config.yaml b/apps/accelerate/chatllama/artifacts/config/config.yaml index 8f8b58e3..ed36cbf5 100644 --- a/apps/accelerate/chatllama/artifacts/config/config.yaml +++ b/apps/accelerate/chatllama/artifacts/config/config.yaml @@ -14,50 +14,55 @@ trainer_config: # number of episodes and generation performed for each episode # in the train() method num_episodes: 100 - max_timesteps: 32 + max_timesteps: 4 # number of timesteps after which the learn() method is called # (to update the weights) - update_timesteps: 32 + update_timesteps: 4 # number of example sampled at each timestep - num_examples: 1 + num_examples: 32 # batch and epochs for the training batch_size: 1 epochs: 1 # number of episodes after which update the checkpoints in RL training - checkpoint_steps: 1000 + checkpoint_steps: 1 # here specify the name of the actor_rl checkpoint from which resume # during actor RL training. If null load the last one. checkpoint_name: null + n_checkpoints_to_keep: 2 + deepspeed_enable: False + deepspeed_config_path: "artifacts/config/ds_config.json" + accelerate_enable: False actor_config: model: "facebook/opt-1.3b" + load_8bit: False model_folder: "./models" tokenizer_path: "path-to-tokenizer" train_dataset_path: "./datasets/actor_training_data.json" validation_dataset_path: null - # froze model embedding during training + # froze model embedding during training (only for llama) froze_embeddings: True # use fairscale layers to build the model instead of vanilla pytorch # only for llama use_fairscale: False # max sequence length for the actor (i.e. prompt + completion) it depends on # the model used. - max_sequence_length: 2048 + max_sequence_length: 512 # max tokens generated by the actor (completion only) - max_tokens: 2048 + max_tokens: 512 # minimum number of tokens generated by the actor - min_tokens: 100 + min_tokens: 50 # additional prompt tokens to be used for template or as safety additonal_prompt_tokens: 20 # temperature for the actor temperature: 0.1 - batch_size: 2 + batch_size: 1 # number iteration after print iteration_per_print: 1 lr: 0.000009 epochs: 1 # number of backpropagation after saving the checkpoints - checkpoint_steps: 5000 + checkpoint_steps: 100 # number of checkpoints to keep while removing the older # (keep memory consumption of checkpoints reasonable) n_checkpoints_to_keep: 5 @@ -76,14 +81,15 @@ actor_config: reward_config: # model to be chosen are gp2-large, bart-base, longformer-base-4096 # more can be simply added in the reward.py __init__() - model: "facebook/opt-125m" + model: "facebook/opt-1.3b" model_folder: "./models" + load_8bit: False # hidden size of the additional ffw head to produce the scores - model_head_hidden_size: 2048 - max_sequence_length: 2048 + model_head_hidden_size: 4096 + max_sequence_length: 512 train_dataset_path: "./datasets/reward_training_data.json" validation_dataset_path: null - batch_size: 8 + batch_size: 1 epochs: 1 iteration_per_print: 1 # steps after which the checkpoint are saved @@ -97,15 +103,20 @@ reward_config: deepspeed_config_path: "./artifacts/config/ds_config.json" # accelerate settings accelerate_enable: False + peft_enable: False + peft_config_path: "./artifacts/config/peft_config.yaml" critic_config: # model to be chosen are gp2-large, bart-base, longformer-base-4096 # more can be simply added in the reward.py __init__() - model: "facebook/opt-125m" + model: "facebook/opt-1.3b" + load_8bit: False # hidden size of the additional ffw head to produce the scores - model_head_hidden_size: 2048 - max_sequence_length: 2048 + model_head_hidden_size: 4096 + max_sequence_length: 512 model_folder: "./models" # here specify the name of the critic checkpoint from which resume # during critic training. If null load the last one. checkpoint_name: null + peft_enable: False + peft_config_path: "./artifacts/config/peft_config.yaml" diff --git a/apps/accelerate/chatllama/artifacts/config/ds_config.json b/apps/accelerate/chatllama/artifacts/config/ds_config.json index 9986679d..685c2deb 100644 --- a/apps/accelerate/chatllama/artifacts/config/ds_config.json +++ b/apps/accelerate/chatllama/artifacts/config/ds_config.json @@ -48,5 +48,11 @@ "stage3_gather_16bit_weights_on_model_save": true, "ignore_unused_parameters": true, "round_robin_gradients": true + }, + "comms_logger": { + "enabled": false, + "verbose": false, + "prof_all": false, + "debug": false } } \ No newline at end of file diff --git a/apps/accelerate/chatllama/artifacts/datasets/actor_dataset.json b/apps/accelerate/chatllama/artifacts/datasets/actor_training_data.json similarity index 100% rename from apps/accelerate/chatllama/artifacts/datasets/actor_dataset.json rename to apps/accelerate/chatllama/artifacts/datasets/actor_training_data.json diff --git a/apps/accelerate/chatllama/artifacts/datasets/reward_dataset.json b/apps/accelerate/chatllama/artifacts/datasets/reward_training_data.json similarity index 100% rename from apps/accelerate/chatllama/artifacts/datasets/reward_dataset.json rename to apps/accelerate/chatllama/artifacts/datasets/reward_training_data.json diff --git a/apps/accelerate/chatllama/artifacts/datasets/rlhf_dataset.json b/apps/accelerate/chatllama/artifacts/datasets/rlhf_training_data.json similarity index 100% rename from apps/accelerate/chatllama/artifacts/datasets/rlhf_dataset.json rename to apps/accelerate/chatllama/artifacts/datasets/rlhf_training_data.json diff --git a/apps/accelerate/chatllama/artifacts/download_dataset.py b/apps/accelerate/chatllama/artifacts/download_dataset.py index 281c942c..c8cf3399 100644 --- a/apps/accelerate/chatllama/artifacts/download_dataset.py +++ b/apps/accelerate/chatllama/artifacts/download_dataset.py @@ -1,7 +1,11 @@ import argparse import os -from chatllama.rlhf.dataset import AnthropicRLHF, StanfordNLPSHPDataset +from chatllama.rlhf.dataset import ( + AnthropicRLHFDataset, + SelfInstructDataset, + StanfordNLPSHPDataset, +) if __name__ == "__main__": @@ -15,7 +19,7 @@ parser.add_argument( "dataset_name", help="dataset name it can be. SSHP: stanfordnlp/SHP or ", - choices=["SHP", "ARLHF"], + choices=["SHP", "ARLHF", "SI"], ) parser.add_argument( "-p", @@ -44,8 +48,15 @@ dataset.save_dataset(args.path, n_samples) elif args.dataset_name == "ARLHF": - dataset = AnthropicRLHF() + dataset = AnthropicRLHFDataset() dataset.save_dataset( args.path, n_samples, ) + elif args.dataset_name == "SI": + dataset = SelfInstructDataset() + dataset.save_dataset( + args.path, + n_samples, + ) + \ No newline at end of file diff --git a/apps/accelerate/chatllama/artifacts/main.py b/apps/accelerate/chatllama/artifacts/main.py index fc2bf5f8..7ab4792e 100644 --- a/apps/accelerate/chatllama/artifacts/main.py +++ b/apps/accelerate/chatllama/artifacts/main.py @@ -2,7 +2,6 @@ from chatllama.rlhf.actor import ActorTrainer from chatllama.rlhf.config import Config -from chatllama.rlhf.dataset import BaseDataset from chatllama.rlhf.reward import RewardTrainer from chatllama.rlhf.trainer import RLTrainer @@ -31,7 +30,16 @@ parser.add_argument( "-r", "--reward", help="Specify reward model by name", default=None ) -parser.add_argument("--local_rank", help="Local rank parameter for deepspeed", default=None) + +parser.add_argument( + "--local_rank", + type=int, + default=-1, + help="local rank passed from distributed launcher", +) + +# Include DeepSpeed configuration arguments +# parser = deepspeed.add_config_arguments(parser) # parse arguments args = parser.parse_args() @@ -53,15 +61,12 @@ config.critic.max_sequence_length, ) config.actor.max_sequence_length = max_seq - BaseDataset.clean_dataset(config) rlhf_trainer = RLTrainer(config) rlhf_trainer.train() elif args.type == "ACTOR": - BaseDataset.clean_dataset(config.actor) actor_trainer = ActorTrainer(config.actor) actor_trainer.train() elif args.type == "REWARD": - BaseDataset.clean_dataset(config.reward) reward_trainer = RewardTrainer(config.reward) reward_trainer.train() elif args.type == "ALL": diff --git a/apps/accelerate/chatllama/chatllama/rlhf/actor.py b/apps/accelerate/chatllama/chatllama/rlhf/actor.py index f7a53732..decf0b2a 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/actor.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/actor.py @@ -1,183 +1,30 @@ import json -import yaml -import os -import shutil -import deepspeed import torch -from accelerate import Accelerator from beartype import beartype from beartype.typing import Tuple from einops import rearrange -from peft import get_peft_model, LoraConfig, TaskType -from torch.utils.data import DataLoader, Dataset -from transformers import ( - AutoModelForCausalLM, - AutoTokenizer, -) +from torch.utils.data import Dataset +from chatllama.rlhf.base_model import BaseModel, BaseTrainer from chatllama.rlhf.config import ConfigActor from chatllama.rlhf.model_list import ( hf_models_causal_lm, - llama_models, - hf_models, ) - -from chatllama.rlhf.model_loader import ModelLoader -from chatllama.rlhf.utils import TrainingStats +from chatllama.rlhf.utils import my_logger -class ActorModel(torch.nn.Module): +class ActorModel(BaseModel): """Actor model that generates the augmented prompt from the initial user_input. The aim is to train this model to generate better prompts. - Attributes: - model: The model from LLaMA to be used - tokenizer: The LLaMA tokenizer - config (ConfigActor): Configuration for the actor model - Methods: - load: Load the model from a path - save: Save the model to a path forward: Compute the action logits for a given sequence. generate: Generate a sequence from a given prompt """ def __init__(self, config: ConfigActor) -> None: - super().__init__() - - # save config - self.config = config - - # initialize the self.model - if config.model in llama_models: - # llama module might not be present when HF models are used - from chatllama.llama_model import ( - load_model, - setup_model_parallel, - ) # noqa - - local_rank, world_size = setup_model_parallel() - - # use load_model_test for testing - self.model, self.tokenizer = load_model( - ckpt_dir=config.model_folder, - tokenizer_path=config.tokenizer_path, - local_rank=local_rank, - world_size=world_size, - froze_embeddings=config.froze_embeddings, - use_fairscale=config.use_fairscale, - max_batch_size=config.batch_size, - ) - elif config.model in hf_models_causal_lm: - self.tokenizer = self.load_tokenizer(config) - self.model = AutoModelForCausalLM.from_pretrained( - config.model, - ) - - # Setup PEFT model - if config.peft_enable: - - # check that the peft config exist - if os.path.exists(config.peft_config_path): - # Read the peft config from yaml - with open(config.peft_config_path, "r") as c: - config_peft = yaml.safe_load(c) - else: - raise ValueError( - f"PEFT config {config.peft_config_path} not found" - ) - - print(config_peft) - # define lora config for peft - peft_config = LoraConfig( - task_type=TaskType.CAUSAL_LM, **config_peft - ) - - # create peft model - self.model = get_peft_model( - model=self.model, - peft_config=peft_config, - ) - - self.model.to(config.device) - - else: - raise ValueError(f"Model {config.model} not supported") - - # load the model from model_folder - self.load() - - @beartype - def load(self) -> None: - """Load the model from the path""" - # check if there is a model to load - path = ModelLoader.check_model_path( - config=self.config, - is_checkpoint=False, - current_epoch=None, - ) - - # if there is a model to load - if path is not None: - - # load the model - print("Loading ...") - model_dict = torch.load(path) - self.model.load_state_dict(model_dict.get("state_dict") or model_dict.get("model")) - - @beartype - def save(self) -> None: - """Save the model to the path""" - # get the path to save the model - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=False, - current_epoch=None, - ) - - # save the model - print(f"Saving model to {path} ...") - torch.save( - {"state_dict": self.model.state_dict()}, - path, - ) - - @staticmethod - def load_tokenizer(config: ConfigActor): - """Load the tokenizer from the model name""" - if config.model in hf_models: - # load the tokenizer from HF - tokenizer = AutoTokenizer.from_pretrained( - config.model, - padding_side="left", - padding=True, - truncation=True, - model_max_length=config.max_sequence_length, - ) - - # add eos token if not present - if tokenizer.eos_token is None: - tokenizer.eos_token = "" - tokenizer.eos_token_id = 2 # OPT eos-token-id - - # add pad token if not present - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - tokenizer.pad_token_id = tokenizer.eos_token_id - elif config.model in llama_models: - - # llama module might not be present when HF models are used - from chatllama.llama_model import ( - load_tokenizer, - ) # noqa - - tokenizer = load_tokenizer(config.tokenizer_path) - return tokenizer - - def parameters(self): - """Return the parameters of the model""" - return self.model.parameters() + super().__init__(config) @beartype def forward( @@ -187,12 +34,15 @@ def forward( of the actions Args: - sequences (torch.Tensor): Sequences of states and actions used to - compute token logits for the whole list of sequences - attention_mask (torch.Tensor): Mask for the sequences attention + sequences (torch.Tensor) [batch_size, seq_len]: Sequences of + states and actions used to compute token logits + for the whole list of sequences + attention_mask (torch.Tensor) [batch_size, seq_len]: Mask for the + sequences attention Returns: - logits (torch.Tensor): Logits for the actions taken + logits (torch.Tensor) [batch_size, seq_len, vocab_size]: Logits for + the actions taken """ model_output = self.model.forward( sequences, attention_mask=sequences_mask @@ -200,10 +50,7 @@ def forward( # need to return logits for the actions if self.config.model in hf_models_causal_lm: model_output = model_output.logits - if self.config.debug: - print("ActorModel.forward") - print("model_output_logits shape", model_output.shape) - print("model_output logits", model_output) + return model_output @beartype @@ -215,13 +62,15 @@ def generate( (i.e. input of the prompt generator model) Args: - state (torch.Tensor): the input of the user - state_mask (torch.Tensor): Mask for the state input (for padding) + state (torch.Tensor) [batch_size, input_len]: the input of the user + state_mask (torch.Tensor) [batch_size, input_len]: Mask for the + state input (for padding) Returns: - actions (torch.Tensor): Actions generated from the state - sequences (torch.Tensor): Sequences generated from the - state as [states, actions] + actions (torch.Tensor) [batch_size, act_len]: Actions generated + from the state + sequences (torch.Tensor) [batch_size, seq_len]: + Sequences generated from the state as [states, actions] """ # temperature for the actor temperature = self.config.temperature @@ -235,20 +84,22 @@ def generate( # max generation possible given the state and the max sequence length max_generation_possible = max_sequence_length - states.shape[1] - if max_generation_possible < min_tokens: - raise ValueError( + if max_generation_possible <= min_tokens: + raise my_logger.error( + ValueError, f"The prompt is too long w.r.t the " f"model sequence length \n" f"max_sequence_length={max_sequence_length}\n" f"state_length={states.shape[1]}\n" f"min_tokens={min_tokens}\n" f"max_tokens={max_tokens}\n" - f"max_generation_possible={max_generation_possible}\n" + f"max_generation_possible={max_generation_possible}\n", ) # take the minimum the max_tokens and the max_generation_possible max_completion = min(max_tokens, max_generation_possible) + # generate the actions and the sequences sequences = self.model.generate( input_ids=states, attention_mask=state_mask, @@ -256,21 +107,9 @@ def generate( max_new_tokens=max_completion, no_repeat_ngram_size=3, ) + + # take the actions actions = sequences[:, states.shape[1] :] # noqa E203 - if self.config.debug: - print( - f"input length {states.shape[1]} \n" - f"max sequence length {max_sequence_length} \n" - f"max completion {max_completion} \n" - f"generated sequence {sequences.shape[1]} \n" - ) - print("ActorModel.generate") - print("state", states) - print("state shape", states.shape) - print("sequence shape", sequences.shape) - print("sequence", sequences) - print("actions shape", actions.shape) - print("actions", actions) return actions, sequences @@ -307,48 +146,45 @@ def __len__( return len(self.data) -class ActorTrainer: +class ActorTrainer(BaseTrainer): """Used to pre-train the actor model to generate better prompts. Args: config (ConfigActor): Configuration for the actor model Attributes: - config (ConfigActor): Configuration for the actor model model (ActorModel): Actor model loss_function (torch.nn.CrossEntropyLoss): Loss function - optimizer (torch.optim.Adam): Optimizer - validation_flag (bool): Flag to indicate if the validation dataset - is provided - train_dataset (ActorDataset): Training dataset - train_dataloader (DataLoader): Training dataloader - validation_dataset (ActorDataset): Validation dataset - validation_dataloader (DataLoader): Validation dataloader - scheduler (torch.optim.lr_scheduler): Learning rate scheduler - training_stats (TrainingStats): Training statistics - model_engine (ModelEngine): Model engine for deepspeed training - accelerator (Accelerator): Accelerator for accelerate training + optimizer (torch.optim.AdamW): Optimizer + validation_flag (bool): Flag to check if validation dataset is provided + train_dataset (ActorDataset): Dataset for the training + validation_dataset (ActorDataset): Dataset for the validation + scheduler (torch.optim.lr_scheduler): Scheduler for the optimizer + train_dataloader (torch.utils.data.DataLoader): Dataloader for the + training dataset + validation_dataloader (torch.utils.data.DataLoader): Dataloader for the + validation dataset Methods: train: Train the actor model - load_checkpoint: Load a checkpoint - save_checkpoint: Save a checkpoint + add_eos_token: Add the eos token to the end of the sequences + + Known Issues: + - When training with lora, 8bit and accelerate the tensor placement + is a mess and the training crash. """ def __init__(self, config: ConfigActor) -> None: - # store config - self.config = config - # load the model - self.actor = ActorModel(config) + self.model = ActorModel(config) # define loss function self.loss_function = torch.nn.CrossEntropyLoss() # define optimizer self.optimizer = torch.optim.AdamW( - self.actor.parameters(), lr=config.lr, weight_decay=1e-5 + self.model.parameters(), lr=config.lr, weight_decay=1e-5 ) # check if validation dataset is provided @@ -356,16 +192,10 @@ def __init__(self, config: ConfigActor) -> None: if config.validation_dataset_path is not None: self.validation_flag = True - # create dataset and dataloaders + # create dataset self.train_dataset = ActorDataset(config.train_dataset_path) - self.train_dataloader = DataLoader( - self.train_dataset, batch_size=config.batch_size - ) if self.validation_flag: self.eval_dataset = ActorDataset(config.validation_dataset_path) - self.validation_dataloader = DataLoader( - self.eval_dataset, batch_size=config.batch_size - ) # define scheduler for the learning rate # learning rate is decreased until 10% of the initial value @@ -376,188 +206,28 @@ def __init__(self, config: ConfigActor) -> None: eta_min=config.lr * 0.1, ) - # define training statistics - stat_path = ModelLoader.get_training_stats_path(config) - self.training_stats = TrainingStats(stat_path) - - # consistency check between accelerate and deepspeed - if config.accelerate_enable and config.deepspeed_enable: - raise ValueError( - "Both DeepSpeed and Accelerate are enabled for the Actor." - "Please choose one of them." - ) - - # initialize deepspeed - self.model_engine = None - if config.deepspeed_enable is True: - if config.deepspeed_config_path is None: - raise ValueError( - "DeepSpeed config path is None, but deepspeed is enabled" - ) - if os.path.exists(config.deepspeed_config_path) is False: - raise ValueError( - f"DeepSpeed config path {config.deepspeed_config_path}" - f"does not exist" - ) - ( - self.model_engine, - self.optimizer, - self.train_dataloader, - _, - ) = deepspeed.initialize( - args=None, - model=self.actor, - model_parameters=self.actor.parameters(), - training_data=self.train_dataset, - config=self.config.deepspeed_config_path, - ) - print("Training with DeepSpeed") - - # initialize accelerate - self.accelerator = None - if config.accelerate_enable is True: - self.accelerator = Accelerator() - ( - self.actor, - self.optimizer, - self.train_dataloader, - self.scheduler, - ) = self.accelerator.prepare( - self.actor, - self.optimizer, - self.train_dataloader, - self.scheduler, - ) - print("Training with Accelerate") + # super init + super().__init__(config) - @beartype - def save_checkpoint( - self, - current_epoch: int, - current_step: int, - max_epochs: int, - max_steps: int, - ) -> None: - """Save the current checkpoint - - Args: - current_epoch (int): Current epoch - current_step (int): Current step - max_epochs (int): Maximum number of epochs - max_steps (int): Maximum number of steps - """ - - print( - f"Saving checkpoint for epoch {current_epoch + 1}, " - f"step {current_step + 1} ..." - ) - # look for path to save the checkpoint - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=current_epoch, - current_step=current_step, - max_epochs=max_epochs, - max_steps=max_steps, + # create dataloader + self.train_dataloader = self.create_dataloader( + self.train_dataset, batch_size=config.batch_size ) - - # remove the checkpoint if it already exists - if os.path.exists(path): - if self.config.deepspeed_enable: - shutil.rmtree(path) - else: - os.remove(path) - - if self.config.deepspeed_enable: - client_state = { - "epoch": current_epoch, - "step": current_step, - } - self.model_engine.save_checkpoint(path, client_state=client_state) - else: - # save the checkpoint - torch.save( - { - "state_dict": self.actor.model.state_dict(), - "optim_state_dict": self.optimizer.state_dict(), - "training_stats": self.training_stats, - "epoch": current_epoch, - "step": current_step, - }, - path, + if self.validation_flag: + self.validation_dataloader = self.create_dataloader( + self.eval_dataset, batch_size=config.batch_size ) - # remove old checkpoints - n_checkpoints_to_keep = self.config.n_checkpoints_to_keep - ModelLoader.delete_old_checkpoints( - model_folder, model_name, n_checkpoints_to_keep - ) - - @beartype - def load_checkpoint( - self, - ) -> Tuple[int, int]: - """Load a checkpoint from the model folder - - Returns: - Tuple[int, int]: Current epoch and current step to resume - training - """ - - print("Looking for checkpoints...") - # look for a checkpoint - path = ModelLoader.check_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=None, - ) - - # if there is a checkpoint - if path is not None: - print("Loading ...") - - if self.config.deepspeed_enable: - # try to load the checkpoint - try: - _, client_state = self.model_engine.load_checkpoint(path) - except Exception: - print( - "Checkpoint corrupted!" - "Try to remove the last checkpoint." - "Now Starting from epoch 0, step 0" - ) - return 0, 0 - # load epoch and step to resume loops - epoch = client_state["epoch"] - step = client_state["step"] - else: - # try to load the checkpoint - try: - checkpoint = torch.load(path) - except Exception: - print( - "Checkpoint corrupted!" - "Try to remove the last checkpoint." - "Now Starting from epoch 0, step 0" - ) - return 0, 0 - - # assing the checkpoint to the model - epoch = checkpoint["epoch"] - self.actor.model.load_state_dict(checkpoint["state_dict"]) - self.optimizer.load_state_dict(checkpoint["optim_state_dict"]) - self.trainign_stats = checkpoint["training_stats"] - step = checkpoint["step"] - return epoch, step + 1 # return the next episode to train - return 0, 0 - def add_eos_token( self, tokens: torch.Tensor, mask: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: # given tokens and mask, add eos token to the end of each sequence # and update the mask batch_size, seq_len = tokens.shape - eos_token = self.actor.tokenizer.eos_token_id + if self.accelerate_enable: + eos_token = self.model.module.tokenizer.eos_token_id + else: + eos_token = self.model.tokenizer.eos_token_id # see if i can append 1 token n_tokens_to_append = min(self.config.max_sequence_length - seq_len, 1) @@ -568,20 +238,13 @@ def add_eos_token( tokens = torch.cat( [ tokens, - torch.ones(batch_size, n_tokens_to_append) - .long() - .to(tokens.device) + torch.ones(batch_size, n_tokens_to_append).long() * eos_token, ], dim=1, ) mask = torch.cat( - [ - mask, - torch.ones(batch_size, n_tokens_to_append) - .long() - .to(mask.device), - ], + [mask, torch.ones(batch_size, n_tokens_to_append).long()], dim=1, ) return tokens, mask @@ -590,15 +253,21 @@ def train( self, ) -> None: """Train the model""" - print("Start Actor Model Pretraining") - # get config parameters - if self.config.deepspeed_enable: - batch_size = self.train_dataloader.batch_size + my_logger.success("Start Actor Model Pretraining") + + # get batch size + if self.deepspeed_enable: + batch_size = self.model_engine.train_batch_size() + elif self.accelerate_enable: + batch_size = ( + self.config.batch_size * self.accelerator.num_processes + ) else: batch_size = self.config.batch_size + + # get other parameters epochs = self.config.epochs - device = self.config.device checkpoint_steps = self.config.checkpoint_steps # compute the number of iterations @@ -607,6 +276,7 @@ def train( # load model_checkpoint start_epoch, start_step = self.load_checkpoint() + # clean the training stats if we are starting from scratch if start_epoch == 0 and start_step == 0: self.training_stats.clear() @@ -615,34 +285,48 @@ def train( # traing loop for epoch in range(start_epoch, epochs): - self.actor.train() + self.model.train() for i, input_text in enumerate(self.train_dataloader): # skip the first steps if we are resuming training if i < start_step: continue - # tokenize input + # get input and output tensors with torch.no_grad(): - input_tokenized = self.actor.tokenizer( - input_text, - return_tensors="pt", - truncation=True, - padding=True, - ) - # split tokens and mask - input_tokenized_id = input_tokenized["input_ids"] - input_tokenized_mask = input_tokenized["attention_mask"] - - # add eos token - ( - input_tokenized_id, - input_tokenized_mask, - ) = self.add_eos_token( - input_tokenized_id, - input_tokenized_mask, - ) + # tokenize the input + if self.accelerate_enable: + input_tokenized = self.model.module.tokenizer( + input_text, + return_tensors="pt", + truncation=True, + padding=True, + max_length=self.config.max_sequence_length, + ) + else: + input_tokenized = self.model.tokenizer( + input_text, + return_tensors="pt", + truncation=True, + padding=True, + max_length=self.config.max_sequence_length, + ) + + # split tokens and mask + input_tokenized_id = input_tokenized["input_ids"] + input_tokenized_mask = input_tokenized[ + "attention_mask" + ] + + # add eos token at the end of the sequence + ( + input_tokenized_id, + input_tokenized_mask, + ) = self.add_eos_token( + input_tokenized_id, + input_tokenized_mask, + ) # split into input and output training_output = input_tokenized_id[:, 1:] @@ -650,42 +334,78 @@ def train( attention_mask = input_tokenized_mask[:, :-1] # move to device - training_output = training_output.to(device) - training_input = training_input.to(device) - attention_mask = attention_mask.to(device) + if not self.config.load_8bit: + training_output = training_output.to(self.device) + training_input = training_input.to(self.device) + attention_mask = attention_mask.to(self.device) # forward pass - if self.config.deepspeed_enable: + if self.deepspeed_enable: + # deepspeed est_output = self.model_engine( training_input, attention_mask ) + elif self.accelerate_enable: + # accelerate + est_output = self.model(training_input, attention_mask) else: - est_output = self.actor(training_input, attention_mask) + # pytorch mixed precison + with torch.autocast( + device_type=self.config.device_type, + dtype=torch.float16, + ): + est_output = self.model(training_input, attention_mask) # compute loss - est_output = rearrange(est_output, "b s v -> (b s) v") - training_output = rearrange(training_output, "b s -> (b s)") - loss = self.loss_function(est_output, training_output) - self.training_stats.training_loss.append(loss.item()) + if self.deepspeed_enable or self.accelerate_enable: + # deepspeed and accelerate use defualt + est_output = rearrange(est_output, "b s v -> (b s) v") + training_output = rearrange( + training_output, "b s -> (b s)" + ) + loss = self.loss_function(est_output, training_output) + else: + # mixed precision pytorch use autocast + with torch.autocast( + device_type=self.config.device_type, + dtype=torch.float16, + ): + est_output = rearrange(est_output, "b s v -> (b s) v") + training_output = rearrange( + training_output, "b s -> (b s)" + ) + loss = self.loss_function(est_output, training_output) # backward pass - if self.config.deepspeed_enable: + if self.deepspeed_enable: + # deepspeed self.model_engine.backward(loss) self.model_engine.step() - elif self.config.accelerate_enable: + elif self.accelerate_enable: + # accelerate self.optimizer.zero_grad() self.accelerator.backward(loss) self.optimizer.step() self.scheduler.step() else: self.optimizer.zero_grad() - loss.backward() - self.optimizer.step() + if self.config.load_8bit: + # 8 bit from HF + loss.backward() + self.optimizer.step() + else: + # Pytorch mixed precision + self.scaler.scale(loss).backward() + self.scaler.step(self.optimizer) + self.scaler.update() self.scheduler.step() + # save training stats + self.append_training_stats(training_loss=loss.detach().item()) + # print progress if i % self.config.iteration_per_print == 0: - print( + my_logger.info( f"Epoch: {epoch+1}/{epochs}, " f"Iteration: {i+1}/{n_iter}, " f"Training Loss: {loss}" @@ -693,7 +413,12 @@ def train( # save checkpoint periodically if cnt_checkpoint % checkpoint_steps == 0: - self.save_checkpoint(epoch, i, epochs, n_iter) + self.save_checkpoint( + current_epoch=epoch, + max_epochs=epochs, + current_step=i, + max_steps=n_iter, + ) self.training_stats.save() cnt_checkpoint = 1 else: @@ -701,12 +426,12 @@ def train( # Validation if self.validation_flag: - self.actor.eval() + self.model.eval() with torch.no_grad(): for i, input_text in enumerate(self.validation_dataloader): # tokenize input - input_tokenized = self.actor.tokenizer( + input_tokenized = self.model.tokenizer( input_text, return_tensors="pt", padding=True ) validation_output = input_tokenized["input_ids"][:, 1:] @@ -716,7 +441,7 @@ def train( ] # forward pass - est_output = self.actor.forward( + est_output = self.model.forward( validation_input, attention_mask ) validation_output = rearrange( @@ -728,11 +453,11 @@ def train( loss = self.loss_function( est_output, validation_output ) - self.training_stats.validation_loss.append(loss.item()) + self.append_training_stats(validation_loss=loss.item()) # print progress if i % self.config.iteration_per_print == 0: - print( + my_logger.info( f"Epoch: {epoch+1}/{epochs}, " f"Iteration: {i+1}/{n_iter}, " f"Validation Loss: {loss}" @@ -741,5 +466,8 @@ def train( start_step = 0 # save the model - self.actor.save() - print("Training Finished ") + if self.accelerate_enable: + self.model.module.save() + else: + self.model.save() + my_logger.success("Training Finished ") diff --git a/apps/accelerate/chatllama/chatllama/rlhf/base_model.py b/apps/accelerate/chatllama/chatllama/rlhf/base_model.py new file mode 100644 index 00000000..6837c4af --- /dev/null +++ b/apps/accelerate/chatllama/chatllama/rlhf/base_model.py @@ -0,0 +1,945 @@ +import yaml +import os +import shutil + +import deepspeed +import torch +import torch.distributed as dist +from accelerate import Accelerator +from beartype import beartype +from beartype.typing import ( + Tuple, + Union, + Iterable, + Optional, +) +from deepspeed.runtime.dataloader import DeepSpeedDataLoader +from einops.layers.torch import Rearrange +from peft import ( + get_peft_model, + LoraConfig, + TaskType, + prepare_model_for_int8_training, +) +from torch.cuda.amp import GradScaler +from torch.utils.data import DataLoader, Dataset +from transformers import ( + AutoModel, + AutoModelForCausalLM, +) + +from chatllama.rlhf.config import ( + Config, + ConfigActor, + ConfigCritic, + ConfigReward, +) +from chatllama.rlhf.dataset import BaseDataset +from chatllama.rlhf.model_list import ( + hf_models_causal_lm, + llama_models, +) +from chatllama.rlhf.model_loader import ModelLoader +from chatllama.rlhf.utils import ( + IgnoreLabelsWrapper, + TrainingStats, + get_multi_gpu_flags, + my_logger, + load_tokenizer, +) + + +ConfigType = Union[ConfigActor, ConfigReward, ConfigCritic, Config] + + +class BaseModel(torch.nn.Module): + """Base Model for generic methods implementations for Actor, Critic and + Reward models. This class is meant to be inherited by the ActorModel, + CriticModel and RewardModel classes. + + Attributes: + model: The model used + head: The head used if specified directly + tokenizer: The tokenizer used + config (ConfigActor): Configuration for the model + accelerate_enable (bool): Flag to enable Accelerate + deepspeed_enable (bool): Flag to enable DeepSpeed + deepspeed_config_path (str): Path to the DeepSpeed configuration file + is_lora_peft_enable (bool): Flag to signal if LORA PEFT is enabled + + Methods: + load: Load the model from a path + save: Save the model to a path + parameters: Return the parameters of the model + apply_lora_with_peft: Apply LORA with PEFT to the model + """ + + @beartype + def __init__(self, config: ConfigType) -> None: + super().__init__() + + # save config + self.config = config + + # initialize flags for accelerator and deepspeed + ( + self.accelerate_enable, + self.deepspeed_enable, + self.deepspeed_config_path, + ) = get_multi_gpu_flags(config) + + if not isinstance(config, Config): + # Actor, Critic or Reward Model initialization + + # initialize the self.model + if config.model in llama_models: + + # llama is supported only for the actor for NOW + if not isinstance(config, ConfigActor): + raise my_logger.error( + ValueError, + "LLAMA is supported only for the actor as of now", + ) + + # llama module might not be present when HF models are used + from chatllama.llama_model import ( + load_model, + setup_model_parallel, + ) # noqa + + local_rank, world_size = setup_model_parallel() + + # use load_model_test for testing + self.model, self.tokenizer = load_model( + ckpt_dir=config.model_folder, + tokenizer_path=config.tokenizer_path, + local_rank=local_rank, + world_size=world_size, + froze_embeddings=config.froze_embeddings, + use_fairscale=config.use_fairscale, + max_batch_size=config.batch_size, + ) + + elif config.model in hf_models_causal_lm: + + # load tokenizer + self.tokenizer = load_tokenizer(config) + + # check load 8 bit condition + if not config.peft_enable: + config.load_8bit = False + if torch.cuda.device_count() > 1: + config.load_8bit = False + + # load model + if isinstance(config, ConfigActor): + # load model for the actor + if config.load_8bit: + # 8 bit + LoRA + PEFT + self.model = AutoModelForCausalLM.from_pretrained( + config.model, + load_in_8bit=config.load_8bit, + device_map="auto", + ) + else: + # Vanilla HF model + self.model = AutoModelForCausalLM.from_pretrained( + config.model, + ) + elif isinstance(config, ConfigReward) or isinstance( + config, ConfigCritic + ): + # load the model for Critic and Reward + # (i.e. without the LM Head) + if config.load_8bit: + # 8 bit + LoRA + PEFT + self.model = AutoModel.from_pretrained( + config.model, + load_in_8bit=config.load_8bit, + device_map="auto", + ) + else: + # Vanilla HF model + self.model = AutoModel.from_pretrained( + config.model, + ) + + # define the head for the reward and critic + # the head is a ff layer that squash the hidden dimension + # to 1 (i.e. the score) + head_hidden_size = config.model_head_hidden_size + head_dim = self.model.config.hidden_size + if config.model.startswith("gpt2"): + head_dim = self.model.config.n_embd + + self.head = torch.nn.Sequential( + torch.nn.Linear(head_dim, head_hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(head_hidden_size, 1), + Rearrange("... 1 -> ..."), + ) + + # apply LoRA with PEFT + self.apply_lora_with_peft() + + else: + raise my_logger.error( + ValueError, + f"Unsupported model {config.model}.\n" + f"Try to add your model to the model_list.py and if you " + f"want open a new PR " + f"@https://github.com/nebuly-ai/nebullvm/pulls\n" + f"The supported models are:\n\n" + f"- Standard LLaMA:\n{llama_models}\n\n" + f"- HF models:\n{hf_models_causal_lm}", + ) + + # move model to device + if config.load_8bit is False: + self.model.to(config.device) + # move the head for the reward and critic to device + if (isinstance(config, ConfigReward)) or ( + isinstance(config, ConfigCritic) + ): + self.head.to(config.device) + + # load the model from model_folder + self.load() + + if isinstance(config, ConfigActor): + my_logger.success("Actor Model loaded") + elif isinstance(config, ConfigReward): + if config.is_reward: + my_logger.success("Reward Model loaded") + else: + my_logger.success("Critic Model loaded") + + else: + # ActorCritic initialization + pass + + @beartype + def apply_lora_with_peft(self) -> None: + """Apply LoRA with PEFT to the model. + The model is modified in place and the head is not included in the + PEFTmodel beacause we need to train it as well. + """ + + # defualt flag to False + self.is_lora_peft_applied = False + + if self.config.peft_enable: + + # check that the peft config exist + if not os.path.exists(self.config.peft_config_path): + raise my_logger.error( + ValueError, + f"PEFT config {self.config.peft_config_path}" + f" not found. Can't apply LoRA with PEFT.", + ) + + # Read the peft config from yaml + with open(self.config.peft_config_path, "r") as c: + config_peft = yaml.safe_load(c) + + # define lora config for peft + peft_config = LoraConfig( + task_type=TaskType.CAUSAL_LM, **config_peft + ) + + # if the model is a reward or critic model + # needs to be wrapped in a IgnoreLabelsWrapper or lora will pass + # the labels argument that is not present in the reward and critic + if isinstance(self.config, ConfigReward) or isinstance( + self.config, ConfigCritic + ): + self.model = IgnoreLabelsWrapper(self.model) + + # call the prepare model method (as in the lora int8 examples) + prepare_model_for_int8_training(self.model) + + # create peft model + self.model = get_peft_model( + model=self.model, + peft_config=peft_config, + ) + + my_logger.info("LoRA with PEFT applied to the model.") + + # change lora Flag + self.is_lora_peft_applied = True + + @beartype + def load(self) -> None: + """Load the model from the path, if it is an actor model load only the + model otherwise if it is an actor or critic model load also the head. + + - For the Actor the saved dictonary contains only one keyword "model" + corresponding to the whole model weights for the actor. + + - For Critic and Reward the saved dictonary contains two keywords + "model" and "head". + """ + + if not isinstance(self.config, Config): + + # Actor, Critic or Reward Model load() + if ( + (not isinstance(self.config, ConfigActor)) + and (not isinstance(self.config, ConfigReward)) + and (not isinstance(self.config, ConfigCritic)) + ): + raise my_logger.error( + ValueError, + f"Model type not supported: {type(self.config)}", + ) + + # check if there is a model to load + path = ModelLoader.check_model_path( + config=self.config, + is_checkpoint=False, + current_epoch=None, + ) + + # if there is a model to load + if path is not None: + + my_logger.info("Loading ...") + + # load the model + model_dict = torch.load(path) + + # check model_dict["lora_peft"] and self.is_lora_peft_applied + # must be the same i.e. lora with peft must be applied + # to both the model to load and the current model + if "lora_peft" in model_dict: + if model_dict["lora_peft"] != self.is_lora_peft_applied: + raise my_logger.error( + ValueError, + "The model to load is not compatible with the " + "current model. The model to load has " + f"lora_peft={model_dict['lora_peft']} while " + f"the current model has " + f"lora_peft={self.is_lora_peft_applied}.", + ) + + if isinstance(self.config, ConfigActor): + self.model.load_state_dict(model_dict["model"]) + + elif isinstance(self.config, ConfigReward) or isinstance( + self.config, ConfigCritic + ): + self.model.load_state_dict(model_dict["model"]) + self.head.load_state_dict(model_dict["head"]) + + else: + # ActorCritic -- not implemented it relies on the load of the + # actor and critic + pass + + @beartype + def save(self) -> None: + """Save the model to the path, if it is an actor model save only the + model otherwise if it is an actor or critic model save also the head. + In case of ActorCritic model save the actor model as result of RLHF + in the folder actor_rl instead of actor.save() method that saves it + in the actor folder. + + - For the Actor the saved dictonary contains only one keyword "model" + corresponding to the whole model weights for the actor. + + - For Critic and Reward the saved dictonary contains two keywords + "model" and "head". + + """ + + # get the path to save the model + model_folder, model_name, path = ModelLoader.get_model_path( + config=self.config, + is_checkpoint=False, + current_epoch=None, + ) + + # save the model + my_logger.info(f"Saving model to {path} ...") + if isinstance(self.config, ConfigActor): + # Actor Model Save() + torch.save( + { + "model": self.model.state_dict(), + "lora_peft": self.is_lora_peft_applied, + }, + path, + ) + elif isinstance(self.config, ConfigReward) or isinstance( + self.config, ConfigCritic + ): + # Critic or Reward Model save() + torch.save( + { + "model": self.model.state_dict(), + "head": self.head.state_dict(), + "lora_peft": self.is_lora_peft_applied, + }, + path, + ) + elif isinstance(self.config, Config): + # ActorCritic save() + + # get the path to save the actor + model_folder, model_name, path = ModelLoader.get_model_path( + config=self.config, + is_checkpoint=False, + ) + + # save the model + my_logger.info(f"Saving model to {path} ...") + torch.save( + { + "model": self.actor.model.state_dict(), + "lora_peft": self.actor.is_lora_peft_applied, + }, + path, + ) + + # get the path to save the critic model + model_folder, model_name, path = ModelLoader.get_model_path( + config=self.config.critic, + is_checkpoint=False, + ) + + # save the model + my_logger.info(f"Saving model to {path} ...") + torch.save( + { + "model": self.critic.model.state_dict(), + "head": self.critic.head.state_dict(), + "lora_peft": self.critic.is_lora_peft_applied, + }, + path, + ) + + def parameters(self) -> Iterable[torch.nn.Parameter]: + """Return the parameters of the model""" + if isinstance(self.config, Config): + for p in self.actor.parameters(): + yield p + # TODO: check if i am missing some critic parameters. + for p in self.critic.parameters(): + yield p + if ( + isinstance(self.config, ConfigActor) + or isinstance(self.config, ConfigReward) + or isinstance(self.config, ConfigCritic) + ): + for p in self.model.parameters(): + yield p + if isinstance(self.config, ConfigReward) or isinstance( + self.config, ConfigCritic + ): + for p in self.head.parameters(): + yield p + + +class BaseTrainer: + """Base Class for the trainer of the Actor and Reward models. + It contains the common methods for the training of the models. + This class is not meant to be used directly, but to be inherited by + the ActorTrainer and RewardTrainer classes. + + Args: + config (ConfigModel): Config parameters for the model + + Attributes: + config (ConfigModel): Config parameters for the model + device (torch.device): Device to use for the training + trainig_stats (TrainingStats): Training stats + eps (float): Epsilon for numerical stability + accelerator (Accelerator): Accelerator for the training + model_engine (Engine): Engine for the training + deepspeed_enabled (bool): Flag to enable deepspeed + accelerate_enabled (bool): Flag to enable accelerate + + Methods: + save_checkpoints: Save the checkpoints of the model + load_checkpoints: Load the checkpoints of the model + setup_training_stats: Setup the training stats + setup_accelerate: Setup the accelerate library + setup_deepspeed: Setup the deepspeed library + create_dataloader: Create the dataloader + """ + + @beartype + def __init__(self, config: ConfigType) -> None: + """Initialize the trainer to be called after the model initialization + of the child class initialization. + """ + + # save the config + self.config = config + if isinstance(config, Config): + self.device = self.config.trainer.device + else: + self.device = self.config.device + + # initialize trainint stats + self.trainig_stats = self.setup_training_stats() + + # eps for numerical stability + self.eps = 1e-6 + + # attributes for deepspeed and accelerate + self.accelerator = None + self.model_engine = None + ( + self.accelerate_enable, + self.deepspeed_enable, + self.deepspeed_config_path, + ) = get_multi_gpu_flags(config) + + self.setup_accelerate() + self.setup_deepspeed() + + # define the scaler needed for vanilla pytorch with mixed precision + if self.accelerate_enable or self.deepspeed_enable: + self.scaler = None + else: + self.scaler = GradScaler() + + # clean the dataset + if self.accelerate_enable or self.deepspeed_enable: + # TODO fix error for process group when using accelerate + if torch.cuda.device_count() > 1: + if dist.get_rank() == 0: + BaseDataset.clean_dataset(config) + else: + BaseDataset.clean_dataset(config) + else: + BaseDataset.clean_dataset(config) + + @beartype + def setup_training_stats( + self, + ) -> None: + """This method initializes the training stats""" + stats_path = ModelLoader.get_training_stats_path(self.config) + self.training_stats = TrainingStats(stats_path) + + @beartype + def append_training_stats( + self, + training_loss: Optional[float] = None, + training_accuracy: Optional[float] = None, + value_loss: Optional[float] = None, + validation_loss: Optional[float] = None, + validation_accuracy: Optional[float] = None, + ) -> None: + """ + This method appends the training stats to the training stats list + + Args: + training_loss (float): Training loss + training_accuracy (float): Training accuracy + value_loss (float): Value loss + validation_loss (float): Validation loss + validation_accuracy (float): Validation accuracy + """ + if training_loss is not None: + self.training_stats.training_loss.append(training_loss) + elif training_accuracy is not None: + self.training_stats.training_accuracy.append(training_accuracy) + elif value_loss is not None: + self.training_stats.value_loss.append(value_loss) + elif validation_loss is not None: + self.training_stats.validation_loss.append(validation_loss) + elif validation_accuracy is not None: + self.training_stats.validation_accuracy.append(validation_accuracy) + + @beartype + def create_dataloader( + self, + dataset: Dataset, + batch_size: Optional[int] = None, + ) -> Union[DataLoader, DeepSpeedDataLoader]: + """This method creates the dataloader for the training + + Args: + dataset (Dataset): Dataset to be used to create + the dataloader. + + Returns: + DataLoader: Distributed dataloader + """ + + if self.accelerate_enable: + # accelerate + dataloader = DataLoader(dataset, batch_size=batch_size) + return self.accelerator.prepare(dataloader) + elif self.deepspeed_enable: + # deepspeed + return self.model_engine.deepspeed_io(dataset) + else: + # vanilla pytorch + return DataLoader(self.train_dataset, batch_size=batch_size) + + @beartype + def setup_deepspeed( + self, + ) -> None: + """This method initializes the deepspeed engine""" + + # initialize deepspeed + self.model_engine = None + + # create model engine + if self.deepspeed_enable is True: + if isinstance(self.config, Config): + + # RL Training + # here self.optimizer is removed from the arguments to use the + # one from deepspeed. + # the custom optimizer to differentiate between actor and + # critic is not working anymore + ( + self.model_engine, + self.optimizer, + _, + self.scheduler, + ) = deepspeed.initialize( + args=None, + model=self.actorcritic, + model_parameters=self.actorcritic.parameters(), + lr_scheduler=self.scheduler, + config=self.deepspeed_config_path, + ) + else: + + # Actor or Reward Training + # here self.optimizer is removed from the arguments to use the + # one from deepspeed. + ( + self.model_engine, + self.optimizer, + _, + self.scheduler, + ) = deepspeed.initialize( + args=None, + model=self.model, + model_parameters=self.model.parameters(), + lr_scheduler=self.scheduler, + config=self.deepspeed_config_path, + ) + + # assign device + if torch.cuda.device_count() > 1: + self.device = torch.device(f"cuda:{dist.get_rank()}") + + my_logger.info("Training with DeepSpeed") + + @beartype + def setup_accelerate( + self, + ) -> None: + """This method initializes the accelerator""" + + # initialize accelerate + self.accelerator = None + if self.accelerate_enable is True: + self.accelerator = Accelerator() + if isinstance(self.config, Config): + ( + self.actorcritic, + self.optimizer, + self.scheduler, + ) = self.accelerator.prepare( + self.actorcritic, + self.optimizer, + self.scheduler, + ) + else: + ( + self.model, + self.optimizer, + self.scheduler, + ) = self.accelerator.prepare( + self.model, + self.optimizer, + self.scheduler, + ) + + # assign device + # Fix error with process group not initialized when using + if torch.cuda.device_count() > 1: + self.device = torch.device(f"cuda:{dist.get_rank()}") + + my_logger.info("Training with Accelerate") + + @beartype + def save_checkpoint( + self, + current_epoch: int, + max_epochs: int, + current_step: Optional[int] = None, + max_steps: Optional[int] = None, + ) -> None: + """Save the checkpoints of the model + + Args: + current_epoch (int): Current epoch + current_step (int): Current step + max_epochs (int): Maximum number of epochs + max_steps (int): Maximum number of steps + """ + + my_logger.info( + f"Saving checkpoint for epoch {current_epoch + 1}, " + f" step {current_step} ..." + ) + + # get the path to save the checkpoint + model_folder, model_name, path = ModelLoader.get_model_path( + config=self.config, + is_checkpoint=True, + current_epoch=current_epoch, + current_step=current_step, + max_epochs=max_epochs, + max_steps=max_steps, + ) + + # remove the checkpoint if it already exists + if os.path.exists(path): + # check if path is a directory + if os.path.isdir(path): + shutil.rmtree(path) + else: + os.remove(path) + + # if deepspeed is enabled + if self.deepspeed_enable: + + # create client state dictonary + if current_step is None: + client_state = { + "episode": current_epoch, + } + else: + client_state = { + "epoch": current_epoch, + "step": current_step, + } + + # save the checkpoint with deepspeed + self.model_engine.save_checkpoint(path, client_state=client_state) + + else: + if isinstance(self.config, Config): + # save actor for ActorCritic Trainer + # "model" is just for compatibility with load() and save() + if self.accelerate_enable: + self.accelerator.save( + { + "model": self.actorcritic.module.actor.state_dict(), # noqa 501 + "critic": self.actorcritic.module.critic.state_dict(), # noqa 501 + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "training_stats": self.training_stats, + "episode": current_epoch, + "lora_peft": self.actorcritic.module.actor.is_lora_peft_applied, # noqa 501 + "critic_lora_peft": self.actorcritic.module.critic.is_lora_peft_applied, # noqa 501 + }, + path, + ) + else: + torch.save( + { + "model": self.actorcritic.actor.state_dict(), + "critic": self.actorcritic.critic.state_dict(), + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "training_stats": self.training_stats, + "episode": current_epoch, + "lora_peft": self.actorcritic.actor.is_lora_peft_applied, # noqa 501 + "critic_lora_peft": self.actorcritic.critic.is_lora_peft_applied, # noqa 501 + }, + path, + ) + else: + # save the model for other trainers + if self.accelerate_enable: + self.accelerator.save( + { + "model": self.model.module.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "training_stats": self.training_stats, + "epoch": current_epoch, + "step": current_step, + "lora_peft": self.model.module.is_lora_peft_applied, # noqa 501 + }, + path, + ) + else: + torch.save( + { + "model": self.model.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "training_stats": self.training_stats, + "epoch": current_epoch, + "step": current_step, + "lora_peft": self.model.is_lora_peft_applied, + }, + path, + ) + + # remove old checkpoints + if isinstance(self.config, Config): + ModelLoader.delete_old_checkpoints( + model_folder=model_folder, + model_name=model_name, + n_ckp_to_keep=self.config.trainer.n_checkpoints_to_keep, + ) + else: + ModelLoader.delete_old_checkpoints( + model_folder=model_folder, + model_name=model_name, + n_ckp_to_keep=self.config.n_checkpoints_to_keep, + ) + + my_logger.success(f"Checkpoint saved at {path}") + + @beartype + def load_checkpoint( + self, + ) -> Tuple[int, int]: + """Load the checkpoints of the model + + Returns: + Tuple[int, int]: The current epoch and step + from which you should resume the training + """ + + my_logger.info("Looking for checkpoints...") + + # look for the checkpoints + path = ModelLoader.check_model_path( + config=self.config, + is_checkpoint=True, + current_epoch=None, + ) + + # check if a checkpoint exists + if path is not None: + my_logger.info("Loading ...") + + # if deepspeed is enabled + if self.deepspeed_enable: + + # try to load the checkpoint + try: + _, client_state = self.model_engine.load_checkpoint(path) + except Exception: + my_logger.warning( + ( + "Checkpoint corrupted! " + + "Try to remove the last checkpoint. " + + "Now Starting from epoch 0, step 0" + ) + ) + return 0, 0 + + # load epoch and step to resume loops + if "episode" in client_state: + episode = client_state["episode"] + return episode, 0 + else: + epoch = client_state["epoch"] + step = client_state["step"] + + my_logger.success(f"Checkpoint loaded from {path}") + return epoch, step + + else: + + # try to load the checkpoint + try: + checkpoint = torch.load(path) + except Exception: + my_logger.warning( + ( + "Checkpoint corrupted! " + + "Try to remove the last checkpoint. " + + "Now Starting from epoch 0, step 0" + ) + ) + return 0, 0 + + # load model and epochs + if isinstance(self.config, Config): + # ActorCritic Trainer + + # first check for lora_peft compatibility + if "lora_peft" in checkpoint: + if ( + checkpoint["lora_peft"] + != self.actorcritic.actor.is_lora_peft_applied + ): + my_logger.warning( + ( + "Checkpoint is not compatible with the " + + "current lora_peft setting. " + + "Now Starting from epoch 0, step 0" + ) + ) + return 0, 0 + + if "critic_lora_peft" in checkpoint: + if ( + checkpoint["critic_lora_peft"] + != self.actorcritic.critic.is_lora_peft_applied + ): + my_logger.warning( + ( + "Checkpoint is not compatible with the " + + "current lora_peft setting. " + + "Now Starting from epoch 0, step 0" + ) + ) + return 0, 0 + + self.actorcritic.actor.load_state_dict(checkpoint["model"]) + self.actorcritic.critic.load_state_dict( + checkpoint["critic"] + ) + episode = checkpoint["episode"] + my_logger.success(f"Checkpoint loaded from {path}") + return episode, 0 + else: + # Actor and Reward Trainer + + # load optimizer and scheduler state + self.optimizer.load_state_dict(checkpoint["optimizer"]) + self.scheduler.load_state_dict(checkpoint["scheduler"]) + + # first check for lora_peft compatibility + if "lora_peft" in checkpoint: + if ( + checkpoint["lora_peft"] + != self.model.is_lora_peft_applied + ): + my_logger.warning( + ( + "Checkpoint is not compatible with the " + + "current lora_peft setting. " + + "Now Starting from epoch 0, step 0" + ) + ) + return 0, 0 + + self.model.model.load_state_dict(checkpoint["model"]) + self.training_stats = checkpoint["training_stats"] + epoch = checkpoint["epoch"] + step = checkpoint["step"] + my_logger.success(f"Checkpoint loaded from {path}") + return epoch, step + 1 # return the next episode to train + return 0, 0 diff --git a/apps/accelerate/chatllama/chatllama/rlhf/config.py b/apps/accelerate/chatllama/chatllama/rlhf/config.py index baaa5380..fe1d8178 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/config.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/config.py @@ -14,10 +14,13 @@ class ConfigReward: Attributes: device (torch.device): Device to be used for the reward model model (str): Model to be used for the reward model + load_8bit (bool): Load the reward model in 8bit mode model_folder (str): Path to the folder where model are stored (used to load / store finetuned model or checkpoints) model_head_hidden_size (int): Hidden size of the reward model head max_sequence_length (int): Max sequence length of the reward model + peft_enable (bool): Enable peft for the reward model + peft_config_path (str): Path to the peft config file. train_dataset_path (Optional[str]): Path to the training dataset. Default to None. To be specified only for the reward model trainig. validation_dataset_path (Optional[str]): Path to the validation @@ -37,6 +40,7 @@ class ConfigReward: the reward model trainig. checkpoint_name (Optional[str]): Name of the checkpoint. Default to None. + n_checkpoints_to_keep (Optional[int]): Number of checkpoints to keep lr (Optional[float]): Learning rate for the reward model. Default to None. To be specified only for the reward model distillation. llm_enable (bool): Enable reward model distillation. Default to True. @@ -53,20 +57,24 @@ class ConfigReward: Default to None. is_reward (bool): True if the model is a reward model. Default to True. accelerate_enable (bool): Enable accelerate for the reward model - debug (bool): enable prints for Debugging + device_type (str): Device type to be used for the reward model """ device: torch.device model: str + load_8bit: bool model_folder: str model_head_hidden_size: int max_sequence_length: int + peft_enable: bool + peft_config_path: str train_dataset_path: Optional[str] = None validation_dataset_path: Optional[str] = None batch_size: Optional[int] = None epochs: Optional[int] = None iteration_per_print: Optional[int] = None checkpoint_steps: Optional[int] = None + n_checkpoints_to_keep: Optional[int] = None checkpoint_name: Optional[str] = None lr: Optional[float] = None llm_enable: Optional[bool] = False @@ -80,7 +88,7 @@ class ConfigReward: is_reward: bool = True accelerate_enable: bool = False - debug: bool = False + device_type: str = "cuda" # just for naming consistency @@ -93,6 +101,7 @@ class ConfigActor: Attributes: model (str): Model to be used for the actor + load_8bit (bool): Load the actor in 8bit mode model_folder (str): Path to the folder where model are stored (used to load / store finetuned model or checkpoints) tokenizer_path (str): Path to the folder where tokenizer are stored @@ -127,11 +136,12 @@ class ConfigActor: None. peft_enable (bool): Enable peft for the actor peft_config_path (str): Path to the peft config file. - debug (bool): Enable prints for debugging + device_type (str): Device type to be used for the actor """ model: str + load_8bit: bool model_folder: str tokenizer_path: str train_dataset_path: str @@ -159,7 +169,7 @@ class ConfigActor: peft_enable: bool peft_config_path: str checkpoint_name: Optional[str] = None - debug: bool = False + device_type: str = "cuda" @dataclass @@ -194,9 +204,15 @@ class ConfigTrainer: for the actual training of the actor and critic models. epochs (int): Number of epochs to train the actor and critic. checkpoint_steps (int): Number of episodes to interleave checkpoints. - device (torch.device): Device to be used for the actor and critic + device (torch.device): Device to be used for the rl training + deepspeed_enable (bool): Enable deepspeed for rl training. + Default to False. + deepspeed_config_path (str): Path to the deepspeed config file. + accelerate_enable (bool): Enable accelerate for rl training + n_checkpoints_to_keep (int): Number of checkpoints to keep checkpoint_name (Optional[str]): Name of the checkpoint. Default to None. + device_type (str): Device type to be used for the rl training """ actor_lr: int @@ -214,8 +230,12 @@ class ConfigTrainer: epochs: int checkpoint_steps: int device: torch.device + deepspeed_enable: bool + deepspeed_config_path: Optional[str] + accelerate_enable: bool + n_checkpoints_to_keep: int checkpoint_name: Optional[str] = None - debug: bool = False + device_type: str = "cuda" class Config: @@ -255,12 +275,18 @@ def __init__( ) -> None: # if not specified use the device available + if device is not None: + if ":" in str(device): + device_type = str(device).split(":")[0] + else: + device_type = str(device) + if device is None: if torch.cuda.is_available(): device = torch.device("cuda") + device_type = "cuda" else: - raise ValueError("No GPU available") - print(f"Current device used :{str(device)}") + raise ValueError("No GPU available...") if path is None or os.path.exists(path) is False: raise ValueError("Path to the config.yaml is not valid") @@ -276,18 +302,18 @@ def __init__( # Trainer Config trainer_dict["device"] = device - trainer_dict["debug"] = debug + trainer_dict["device_type"] = device_type self.trainer = ConfigTrainer(**trainer_dict) # Actor Config actor_dict["device"] = device - actor_dict["debug"] = debug + actor_dict["device_type"] = device_type self.actor = ConfigActor(**actor_dict) # Critic Config critic_dict["device"] = device - critic_dict["debug"] = debug + critic_dict["device_type"] = device_type self.critic = ConfigCritic(**critic_dict) self.critic.is_reward = False # Reward Config reward_dict["device"] = device - reward_dict["debug"] = debug + reward_dict["device_type"] = device_type self.reward = ConfigReward(**reward_dict) diff --git a/apps/accelerate/chatllama/chatllama/rlhf/dataset.py b/apps/accelerate/chatllama/chatllama/rlhf/dataset.py index 8e1c7c6f..71a29667 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/dataset.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/dataset.py @@ -1,13 +1,13 @@ import json import os +import random import numpy as np - from beartype.typing import Dict, List, Union from datasets import load_dataset + from chatllama.rlhf.config import Config, ConfigActor, ConfigReward -from chatllama.rlhf.reward import RewardModel, CriticModel -from chatllama.rlhf.actor import ActorModel +from chatllama.rlhf.utils import load_tokenizer, my_logger ConfigType = Union[Config, ConfigActor, ConfigReward] @@ -62,11 +62,15 @@ def sort_fun(x): # shuffle if shuffle is True: + # keep the first 128 sorted so that for batch <= 128 + # the out-of-memory error occurs only in the first batch iteration + # (i.e. is the worst case scenario) + keep_sorted = 128 conversations = ( - conversations[:10] + conversations[:keep_sorted] + np.random.choice( - conversations[10:], - size=len(conversations[10:]), + conversations[keep_sorted:], + size=len(conversations[keep_sorted:]), replace=False, ).tolist() ) @@ -110,8 +114,9 @@ def clean_dataset(config: ConfigType): config (Config): config object """ + if isinstance(config, Config): - print("Start cleaning the dataset for RLHF") + my_logger.info("Start cleaning the dataset for RLHF") # constraints r_model_max_seq_len = config.reward.max_sequence_length a_model_max_seq_len = config.actor.max_sequence_length @@ -120,31 +125,31 @@ def clean_dataset(config: ConfigType): # dataset dataset_path = config.trainer.examples_path # tokenizers - r_tokenizer = RewardModel.load_tokenizer(config.reward) - a_tokenizer = ActorModel.load_tokenizer(config.actor) - c_tokenizer = CriticModel.load_tokenizer(config.critic) + r_tokenizer = load_tokenizer(config.reward) + a_tokenizer = load_tokenizer(config.actor) + c_tokenizer = load_tokenizer(config.critic) # safety tokens safety_tokens = config.actor.additonal_prompt_tokens elif isinstance(config, ConfigActor): - print("Start cleaning the dataset for Actor") + my_logger.info("Start cleaning the dataset for Actor") # constraint a_model_max_seq_len = config.max_sequence_length # dataset dataset_path = config.train_dataset_path # tokenizer - a_tokenizer = ActorModel.load_tokenizer(config) + a_tokenizer = load_tokenizer(config) # safety tokens safety_tokens = config.additonal_prompt_tokens elif isinstance(config, ConfigReward): - print("Start cleaning the dataset for Reward") + my_logger.info("Start cleaning the dataset for Reward") # constraint r_model_max_seq_len = config.max_sequence_length # dataset dataset_path = config.train_dataset_path # tokenizer - r_tokenizer = RewardModel.load_tokenizer(config) + r_tokenizer = load_tokenizer(config) # if there is the datasets if os.path.exists(dataset_path): @@ -159,15 +164,29 @@ def clean_dataset(config: ConfigType): conversations, only_input=True, reverse=True, + shuffle=False, ) else: conversations = BaseDataset.sort_conversation( conversations, only_input=False, reverse=True, + shuffle=False, ) - + + # save orginal length of dataset old_len = len(conversations) + + # for reward dataset first remove examples that do not have + # the scores - these check avoids errors in the training + # of the reward model when the score is None + if isinstance(config, ConfigReward): + cnt = 0 + while cnt < len(conversations): + if conversations[cnt]["score"] is None: + conversations.pop(cnt) + cnt = cnt + 1 + # remove too long examples # since datasets are ordered by the length # we can remove the first elements until we find @@ -189,17 +208,17 @@ def clean_dataset(config: ConfigType): r_tokens = r_tokenizer.encode(text, truncation=False) c_tokens = c_tokenizer.encode(text, truncation=False) if ( - len(a_tokens) + min_completion + safety_tokens + len(a_tokens) + min_completion + safety_tokens + 100 > a_model_max_seq_len ): conversations.pop(0) elif ( - len(r_tokens) + min_completion + safety_tokens + len(r_tokens) + min_completion + safety_tokens + 100 > r_model_max_seq_len ): conversations.pop(0) elif ( - len(c_tokens) + min_completion + safety_tokens + len(c_tokens) + min_completion + safety_tokens + 100 > c_model_max_seq_len ): conversations.pop(0) @@ -222,27 +241,164 @@ def clean_dataset(config: ConfigType): else: break - # if the number of examples has changed + # if the number of examples has changed print the cleaning + # stats if len(conversations) != old_len: - print("Number of examples before cleaning: ", old_len) - print( - "Number of examples after cleaning: ", len(conversations) + my_logger.info( + f"Number of examples before cleaning: {old_len}" + ) + my_logger.info( + f"Number of examples after cleaning: {len(conversations)}" ) - - # remove the old dataset - os.remove(dataset_path) # save the new dataset with open(dataset_path, "w") as f: json.dump(conversations, f, indent=4) + my_logger.success("Dataset cleaned") else: - print("Dataset is already clean") + my_logger.success("Dataset is already clean") else: - print( + my_logger.warning( f"Dataset not found at {dataset_path}" f" Skipping cleaning of the dataset" ) + + @staticmethod + def augment_reward_dataset( + reward_conv: List[Dict], + ): + """ + Augment the reward dataset with negative examples + + Args: + reward_conv (List[Dict]): list of reward conversations + """ + + # shuffle question and answer to generate wrong answers and assing + # a low score to them - 20% of the dataset + percentage = 20 + n = int(len(reward_conv) * percentage / 100 / 2) + data_shuffle = [] + for i in range(n): + sample = random.choice(reward_conv) + sample2 = random.choice(reward_conv) + new_sample1 = {"score": 1, + "user_input": sample2["user_input"], + "completion": sample["completion"]} + new_sample2 = {"score": 1, + "user_input": sample["user_input"], + "completion": sample2["completion"]} + data_shuffle.append(new_sample1) + data_shuffle.append(new_sample2) + + # Create blank answers and assing a low score to them + # 10% of the dataset + percentage = 10 + n = int(len(reward_conv) * percentage / 100 / 2) + data_blank = [] + for i in range(n): + sample = random.choice(reward_conv) + new_sample = {"score": 0.5, + "user_input": sample["user_input"], + "completion": ""} + data_blank.append(new_sample) + new_sample = {"score": 0.5, + "user_input": sample["completion"], + "completion": "Assistant:", + } + data_blank.append(new_sample) + + # swap words in the original completion and assing a low score to them + # 10% of the dataset + percentage = 20 + n = int(len(reward_conv) * percentage / 100) + data_swap = [] + for i in range(n): + sample = random.choice(reward_conv) + # swap words order in the completion + completion = sample["completion"].split(" ") + # generate list from 0 to len(completion) + idx = list(range(len(completion))) + random.shuffle(idx) + completion = [completion[i] for i in idx] + completion = " ".join(completion) + # add spaces back in the completion creating a string from the list + new_sample = { + "score": 0.5, + "user_input": sample["user_input"], + "completion": completion + } + data_swap.append(new_sample) + + # add all the augmented data to the reward_conv + reward_conv.extend(data_shuffle) + reward_conv.extend(data_blank) + reward_conv.extend(data_swap) + + return reward_conv + + + @staticmethod + def generate_datasets( + conversations: List[Dict], + dataset_folder: str, + reward_dataset_size: int, + ): + """ Generate the datasets for actor reward and rl trainings + + Args: + conversations (List[Dict]): list of conversations + dataset_folder (str): path to the folder where to save the datasets + reward_dataset_size (int): size of the reward dataset + """ + + # split conversation list in 80 / 20 in two seperate lists + actor_conv = conversations[:int(len(conversations) * 0.8)] + rl_conv = conversations[int(len(conversations) * 0.8):] + + # augment rl_conv with 10% of its size from actor_conv + rl_conv.extend(actor_conv[:int(len(rl_conv) * 0.1)]) + + # create reward_conv sampling randomly from rl_conv reward_dataset_size + # samples + reward_conv = random.sample(rl_conv, reward_dataset_size) + + # augment reward datasets with wrong completions to create negative + # examples + reward_conv = BaseDataset.augment_reward_dataset(reward_conv) + + # sort the lists + actor_conv = BaseDataset.sort_conversation( + actor_conv, + only_input=False, + reverse=True, + shuffle=True, + ) + rl_conv = BaseDataset.sort_conversation( + rl_conv, + only_input=True, + reverse=True, + shuffle=True, + ) + reward_conv = BaseDataset.sort_conversation( + reward_conv, + only_input=False, + reverse=True, + shuffle=True, + ) + + # save actor training data + with open(f"{dataset_folder}/actor_training_data.json", "w") as f: + json.dump(actor_conv, f, indent=4) + + # save actor training data + with open(f"{dataset_folder}/rlhf_training_data.json", "w") as f: + json.dump(rl_conv, f, indent=4) + + # save actor training data + with open(f"{dataset_folder}/reward_training_data.json", "w") as f: + json.dump(reward_conv, f, indent=4) class StanfordNLPSHPDataset(BaseDataset): @@ -251,9 +407,9 @@ class StanfordNLPSHPDataset(BaseDataset): def __init__( self, ) -> None: - print("Download the dataset") + my_logger.info("Download the dataset") self.dataset = load_dataset("stanfordnlp/SHP") - print("Download Completed") + my_logger.success("Download Completed") def reformat_dataset(self, data: List) -> List[Dict]: """Reformat the dataset to the format required by RLHF @@ -291,7 +447,7 @@ def reformat_dataset(self, data: List) -> List[Dict]: return conversations def save_dataset( - self, dataset_folder: str, number_of_samples: int, reverse: bool = True + self, dataset_folder: str, number_of_samples: int, ) -> None: """Save the dataset in the format required by RLHF @@ -304,50 +460,31 @@ def save_dataset( Defaults to True. """ - print("Generate datasets for RLHF") + my_logger.info("Generate datasets for RLHF") # take the train and test dataset to create the finetuning dataset conversations = self.reformat_dataset(self.dataset["train"]) conversations.extend(self.reformat_dataset(self.dataset["test"])) - # sort conversations by length of user_input + completion - conversations = self.sort_conversation(conversations, reverse=reverse) - - # save actor training data - with open(f"{dataset_folder}/actor_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) - - # take N samples and sort them - conversations = self.take_n_samples(conversations, number_of_samples) - conversations = self.sort_conversation(conversations, reverse=reverse) - - # save reward training data - with open(f"{dataset_folder}/reward_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) - - # take the validation dataset for rlhf - conversations = self.reformat_dataset(self.dataset["validation"]) - # sort the validation dataset - conversations = self.sort_conversation( + self.generate_datasets( conversations, - only_input=True, - reverse=reverse, + dataset_folder, + number_of_samples, ) - # save rlhf training data - with open(f"{dataset_folder}/rlhf_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) - print("Generation Completed") + my_logger.success("Generation Completed") + +class AnthropicRLHFDataset(BaseDataset): + """Class for Anthropic RLHF dataset from HuggingFace""" -class AnthropicRLHF(BaseDataset): def __init__( self, ) -> None: - print("Download the dataset") + my_logger.info("Download the dataset") self.dataset = load_dataset("Anthropic/hh-rlhf") - print("Download Completed") + my_logger.success("Download Completed") def reformat_dataset(self, data: List) -> List[Dict]: """Reformat the dataset to the format required by RLHF @@ -374,6 +511,7 @@ def reformat_dataset(self, data: List) -> List[Dict]: # conversation previous_convers = previous_convers.rstrip("\n") user_input = previous_convers + "\n\n##\n\n" + user_input = user_input.lstrip("\n") completion = "Assistant: " + split_answer[-1] conv = { @@ -399,34 +537,91 @@ def save_dataset( Defaults to True. """ - print("Generate datasets for RLHF") + my_logger.info("Generate datasets for RLHF") # generate actor and reward dataset conversations = self.reformat_dataset(self.dataset["train"]) - conversations = self.sort_conversation(conversations, reverse=reverse) + conversations.extend(self.reformat_dataset(self.dataset["test"])) - # save actor training data - with open(f"{dataset_folder}/actor_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) + self.generate_datasets( + conversations, + dataset_folder, + number_of_samples + ) - # sample N number of index from 0 to len(conversations) - conversations = self.take_n_samples(conversations, number_of_samples) - conversations = self.sort_conversation(conversations, reverse=reverse) + my_logger.success("Generation Completed") - # save reward training data - with open(f"{dataset_folder}/reward_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) - # rlhf dataset - conversations = self.reformat_dataset(self.dataset["test"]) +class SelfInstructDataset(BaseDataset): + """Class for SelfInstruct dataset from HuggingFace""" - # sort conversations by length of user_input - conversations = self.sort_conversation( - conversations, only_input=True, reverse=reverse - ) + def __init__( + self, + ) -> None: + my_logger.info("Download the dataset") + self.dataset = load_dataset("HuggingFaceH4/self-instruct") + my_logger.success("Download Completed") - # save rlhf training data - with open(f"{dataset_folder}/rlhf_training_data.json", "w") as f: - json.dump(conversations, f, indent=4) + def reformat_dataset(self, data: List) -> List[Dict]: + """Reformat the dataset to the format required by RLHF + + Args: + data (List): dataset from HuggingFace + + Returns: + List[Dict]: reformatted dataset + """ - print("Generation Completed") + # here do a tiling to reformat the dataset (otherwise is slow) + def reformat_shard(shard_data): + rshard = [ + { + "user_input": "Human: " + shard_data["prompt"][i], + "completion": "Assistant: " + shard_data["completion"][i], + } + for i in range(len(shard_data["prompt"])) + ] + return rshard + + # number of shards + n_split = 100 + + # shard size + shard_size = len(data) // n_split + + # initialize the reformatted dataset list + reformat_data = [] + + # loop over the shards + for i in range(n_split): + current_shard = data[ + i * shard_size : (i + 1) * shard_size # noqa E203 + ] + reformat_data.extend(reformat_shard(current_shard)) + return reformat_data + + def save_dataset( + self, dataset_folder: str, number_of_samples: int, reverse: bool = True + ) -> None: + """Save the dataset in the format required by RLHF + + Args: + dataset_folder (str): path to the folder where the dataset + will be saved + number_of_samples (int): number of samples to take from the + dataset + reverse (bool, optional): sort the dataset in descending order. + Defaults to True. + """ + + my_logger.info("Generate datasets for RLHF") + + conversations = self.reformat_dataset(self.dataset["train"]) + + self.generate_datasets( + conversations, + dataset_folder, + number_of_samples + ) + + my_logger.success("Generation Completed") diff --git a/apps/accelerate/chatllama/chatllama/rlhf/model_list.py b/apps/accelerate/chatllama/chatllama/rlhf/model_list.py index 8237141e..35ee6cfd 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/model_list.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/model_list.py @@ -14,6 +14,7 @@ # decoder only TODO: codegen is still broken hf_models_causal_lm = [ "facebook/opt-125m", + "facebook/opt-350m", "facebook/opt-1.3b", "facebook/opt-2.7b", "facebook/opt-6.7b", @@ -44,6 +45,17 @@ "Salesforce/codegen-2B-mono", "Salesforce/codegen-6B-mono", "Salesforce/codegen-16B-mono", + "cerebras/Cerebras-GPT-111M", + "cerebras/Cerebras-GPT-256M", + "cerebras/Cerebras-GPT-590M", + "cerebras/Cerebras-GPT-1.3B", + "cerebras/Cerebras-GPT-2.7B", + "cerebras/Cerebras-GPT-6.7B", + "cerebras/Cerebras-GPT-13B", + "decapoda-research/llama-7b-hf", + "decapoda-research/llama-13b-hf", + "decapoda-research/llama-30b-hf", + "decapoda-research/llama-65b-hf", ] # create a list of all the models from hf diff --git a/apps/accelerate/chatllama/chatllama/rlhf/model_loader.py b/apps/accelerate/chatllama/chatllama/rlhf/model_loader.py index c5d2fba7..0424a5c2 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/model_loader.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/model_loader.py @@ -10,10 +10,12 @@ ConfigReward, ) from chatllama.rlhf.model_list import hf_models +from chatllama.rlhf.utils import my_logger ConfigType = Union[Config, ConfigActor, ConfigCritic, ConfigReward] + class ModelLoader: """Class to load and save models and their checkpoints during training.""" @@ -87,14 +89,15 @@ def get_checkpoint_name(config: ConfigType) -> str: else: return config.checkpoint_name - @staticmethod - def get_base_model_folder_from_config(config: ConfigType) -> str: + @classmethod + def get_base_model_folder_from_config(cls, config: ConfigType) -> str: if isinstance(config, ConfigActor) or isinstance(config, ConfigReward): return config.model_folder elif isinstance(config, Config): return config.actor.model_folder else: - raise ValueError( + raise my_logger.error( + ValueError, "Config type not recognized during saving or loading" ) @@ -112,8 +115,8 @@ def get_model_type_from_config(config: ConfigType) -> str: elif isinstance(config, Config): return "actor_rl" - @staticmethod - def get_model_name_from_config(config: ConfigType) -> str: + @classmethod + def get_model_name_from_config(cls, config: ConfigType) -> str: model_name = None if isinstance(config, Config): model_name = config.actor.model @@ -124,7 +127,10 @@ def get_model_name_from_config(config: ConfigType) -> str: if model_name in hf_models: return os.path.split(model_name)[-1] if model_name is None: - raise ValueError("Model name not found") + raise my_logger.error( + ValueError, + "Model name not found" + ) return model_name @staticmethod @@ -153,10 +159,12 @@ def delete_old_checkpoints( if len(checkpoints) > n_ckp_to_keep: for c in checkpoints[:-n_ckp_to_keep]: checkpoint_path = os.path.join(model_folder, c) + my_logger.info(f"Deleting {checkpoint_path}") os.remove(checkpoint_path) - @staticmethod + @classmethod def get_model_path( + cls, config: ConfigType, is_checkpoint: bool = False, current_epoch: Optional[int] = None, @@ -214,7 +222,9 @@ def get_model_path( # Make the path if not exists if os.path.exists(model_folder) is False: os.makedirs(model_folder, exist_ok=True) - print(f"Model folder does not exist. Creating it: {model_folder}") + my_logger.info( + f"Model folder does not exist. Creating it: {model_folder}" + ) # Create the model name model_name = ModelLoader.get_model_name_from_config(config) @@ -224,11 +234,11 @@ def get_model_path( # just return the simple model name if is_checkpoint and current_epoch is not None: # number of characters to store the checkpoints - n_char = max(len(str(max_epochs)), len(str(max_steps))) + n_char = max(len(str(max_epochs)), len(str(max_steps))) + 1 # create the string epoch such that it is always the same length # equalt to n_char (i.e. 00000001) necessary for sorting string_epoch = str(current_epoch) - string_epoch = "0" * (n_char - len(string_epoch)) + string_epoch + string_epoch = "0" * (n_char - len(string_epoch)) + string_epoch string_epoch = f"_epoch_{string_epoch}" if current_step is not None: string_step = str(current_step) @@ -249,8 +259,9 @@ def get_model_path( path = os.path.join(model_folder, model_name) return model_folder, model_name, path - @staticmethod + @classmethod def check_model_path( + cls, config: ConfigType, is_checkpoint: bool = False, current_epoch: Optional[int] = None, @@ -273,6 +284,7 @@ def check_model_path( epoch (Optional[int]): the epoch of the checkpoint if an actual checkpoint is found. If no checkpoint is found, return None. """ + model_folder, model_name, path = ModelLoader.get_model_path( config, is_checkpoint, @@ -303,17 +315,17 @@ def check_model_path( if is_checkpoint: checkpoint_name = ModelLoader.get_checkpoint_name(config) if checkpoint_name is not None: - print( + my_logger.info( f"No checkpoint found at {model_folder} " f"with name {config.checkpoint_name}" ) else: - print( + my_logger.info( f"No previous checkpoint found at " f"{model_folder} for {model_name}" ) else: - print( + my_logger.info( f"No previous model found at " f"{model_folder} for model {model_name}" ) @@ -324,24 +336,25 @@ def check_model_path( if "_step_" in path: epoch = int(path.split("_epoch_")[-1].split("_")[0]) step = int(path.split("_step_")[-1].split(".")[0]) - print( + my_logger.info( f"Found checkpoint for epoch {epoch + 1}," f" step {step + 1}..." ) else: epoch = int(path.split("_epoch_")[-1].split(".")[0]) - print(f"Found checkpoint for epoch {epoch + 1} ...") + my_logger.info(f"Found checkpoint for epoch {epoch + 1} ...") else: - print(f"Found model at {path}") + my_logger.info(f"Found model at {path}") return path - def init_critic_from_reward(config: ConfigCritic) -> None: + @classmethod + def init_critic_from_reward(cls, config: ConfigCritic) -> None: """Method to initialize the critic from the reward model. If the critic folder is empty """ - + if config.is_reward is True: - raise ValueError( + raise my_logger.error( "The config should work for the Critic model," "but the config seems to be for the Reward model" ) @@ -350,13 +363,16 @@ def init_critic_from_reward(config: ConfigCritic) -> None: path = ModelLoader.check_model_path(config) _, _, critic_path = ModelLoader.get_model_path(config) if path is None: - print("Initializing Critic from Reward model...") + my_logger.info("Initializing Critic from Reward model...") config.is_reward = True path = ModelLoader.check_model_path(config) if path is not None: _, _, reward_path = ModelLoader.get_model_path(config) # copy the file in reward_path to critic_path shutil.copy(reward_path, critic_path) + my_logger.success( + "Critic Model initialized from Reward model" + ) else: - print("Critic Model remains uninitialized") + my_logger.warning("Critic Model remains uninitialized") config.is_reward = False diff --git a/apps/accelerate/chatllama/chatllama/rlhf/reward.py b/apps/accelerate/chatllama/chatllama/rlhf/reward.py index 4982969f..74ec41d7 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/reward.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/reward.py @@ -1,147 +1,28 @@ import json -import shutil -import os -import deepspeed + import torch -from accelerate import Accelerator from beartype import beartype -from beartype.typing import Iterable, Tuple -from einops.layers.torch import Rearrange -from torch.utils.data import Dataset, DataLoader -from transformers import ( - AutoModel, - AutoTokenizer, -) +from torch.utils.data import Dataset +from chatllama.rlhf.base_model import BaseModel, BaseTrainer from chatllama.rlhf.config import ConfigReward -from chatllama.rlhf.model_list import hf_models -from chatllama.rlhf.model_loader import ModelLoader -from chatllama.rlhf.utils import TrainingStats +from chatllama.rlhf.utils import my_logger -class RewardModel(torch.nn.Module): +class RewardModel(BaseModel): """Model to be trained to predict the reward for RL. or to be used as Critic in RL. It is a Language Model with a head that predicts the reward (a scalar) for a given sequence of tokens. - Attributes: - model (torch.nn.Module): Model to be used for the reward model - tokenizer (torch.nn.Module): Tokenizer to be used for the reward model - head (torch.nn.Module): Head to be used for the reward model - config (ConfigReward): Config parameters for the reward model - Methods: - load_tokenizer: Load the tokenizer for the reward model forward: Forward pass of the model (used by the critic) - save: Save the model - load: Load the model get_reward: Get the reward for a given input (used by the reward model) - parameters: Return the parameters of the reward model """ def __init__(self, config: ConfigReward) -> None: - super().__init__() - - # store config - self.config = config - - # initialize the self.model - head_hidden_size = config.model_head_hidden_size - if config.model in hf_models: - self.tokenizer = self.load_tokenizer(config) - self.model = AutoModel.from_pretrained(config.model) - head_dim = self.model.config.hidden_size - if config.model.startswith("gpt2"): - head_dim = self.model.config.n_embd - self.head = torch.nn.Sequential( - torch.nn.Linear(head_dim, head_hidden_size), - torch.nn.ReLU(), - torch.nn.Linear(head_hidden_size, 1), - Rearrange("... 1 -> ..."), - ) - else: - raise ValueError(f"Model {config.model} not supported") - - # load the model - self.load() - - # freeze model parameters (only train the head) - # for param in self.model.parameters(): - # param.requires_grad = False - - # move model to device - self.model.to(config.device) - self.head.to(config.device) - - @staticmethod - def load_tokenizer(config: ConfigReward): - # load tokenizer from HF - tokenizer = AutoTokenizer.from_pretrained( - config.model, - padding_side="left", - padding=True, - truncation=True, - model_max_length=config.max_sequence_length, - ) - - # add eos token if not present - if tokenizer.eos_token is None: - tokenizer.eos_token = "" - tokenizer.eos_token_id = 2 # OPT eos token id - - # add pad token if not present - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - tokenizer.pad_token_id = tokenizer.eos_token_id - return tokenizer - - @beartype - def load(self) -> None: - """Load the model from the path""" - # look for a pretrained model - path = ModelLoader.check_model_path( - config=self.config, - is_checkpoint=False, - current_epoch=None, - ) - - # check if the model exists - if path is not None: - - # load the model from the path - print("Loading ...") - model_dict = torch.load(path) - self.model.load_state_dict(model_dict.get("state_dict") or model_dict.get("model")) - self.head.load_state_dict(model_dict["head"]) - - @beartype - def save(self) -> None: - """Save the model to the path""" - # get the path to save the model - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=False, - current_epoch=None, - ) - - # save the model - print(f"Saving model to {path} ...") - torch.save( - {"model": self.model.state_dict(), "head": self.head.state_dict()}, - path, - ) - - @beartype - def parameters( - self, - ) -> Iterable[torch.nn.Parameter]: - """Return the parameters of the reward model""" - for p in self.model.parameters(): - yield p - for p in self.head.parameters(): - yield p + super().__init__(config) @beartype def forward( @@ -151,26 +32,21 @@ def forward( what is the quality of the output sequence tokens? Args: - output_sequence (torch.Tensor): The sequence of tokens to be - evaluated - output_sequence_mask (torch.Tensor): Mask for the attention + output_sequence (torch.Tensor) [batch_size, seq_len]: The sequence + of tokens to be evaluated + output_sequence_mask (torch.Tensor) [batch_size, seq_len]: Mask + for the attention Returns: - torch.Tensor: Rewards for the given output sequence + torch.Tensor [batch_size, seq_len]: Rewards for the given output + sequence """ - output = self.model( - output_sequence, attention_mask=output_sequence_mask + output = self.model.forward( + output_sequence, + attention_mask=output_sequence_mask, ) - # What if the output_sequence is longer than the max context of - # the model? rewards = self.head(output.last_hidden_state) - if self.config.debug: - print("RewardModel.forward") - print("output_sequence.shape", output_sequence.shape) - print("output_sequence", output_sequence) - print("reward.shape", rewards.shape) - print("reward", rewards) return rewards @beartype @@ -180,9 +56,13 @@ def get_reward( """Get the reward for the given output sequence Args: - output_sequence (torch.Tensor): The concatenation of initial input - and actor output as tokens - output_sequence_mask (torch.Tensor): Mask for the attention + output_sequence (torch.Tensor) [batch_size, seq_len]: The + concatenation of initial input and actor output as tokens + output_sequence_mask (torch.Tensor) [batch_size, seq_len]: Mask + for the attention + + Returns: + torch.Tensor [batch_size]: The reward for the given output sequence """ if output_sequence.shape[1] > self.config.max_sequence_length: raise ValueError( @@ -215,10 +95,8 @@ class RewardDataset(Dataset): """ def __init__(self, path: str) -> None: - print(f"Loading dataset from {path}") with open(path, "r") as f: self.data = list(json.load(f)) - print(f"Loaded {len(self.data)} samples") def __getitem__(self, idx: int): user_input = self.data[idx]["user_input"] @@ -237,7 +115,7 @@ def __len__( return len(self.data) -class RewardTrainer: +class RewardTrainer(BaseTrainer): """Class to train the reward model Args: @@ -255,31 +133,29 @@ class RewardTrainer: train_dataloader (DataLoader): Dataloader for training validation_dataloader (DataLoader): Dataloader for validation scheduler (torch.optim.lr_scheduler): Scheduler for the optimizer - training_stats (List[Dict]): List of dictionaries with the training - statistics - model_engine (ModelEngine): Model engine to train the model - using deepspeed - accelerator (Accelerator): Accelerator to train the model using - accelerate by HF. - Methods: train: Train the reward model - save_checkpoints: Save the checkpoints of the model - load_checkpoints: Load the checkpoints of the model + + Known Issues: + When using lora and 8 bit precision, the following error is raised: + - result = F.linear( + x, + transpose(self.weight, self.fan_in_fan_out), + bias=self.bias + ) + RuntimeError: self and mat2 must have the same dtype + https://github.com/huggingface/peft/issues/84 """ def __init__(self, config: ConfigReward) -> None: - # save the config - self.config = config - # load the model - self.reward = RewardModel(config) + self.model = RewardModel(config) # optimizer self.optimizer = torch.optim.AdamW( - self.reward.parameters(), lr=config.lr + self.model.parameters(), lr=config.lr ) # loss function @@ -290,16 +166,10 @@ def __init__(self, config: ConfigReward) -> None: if config.validation_dataset_path is not None: self.validation_flag = True - # create dataset and dataloaders + # create dataset self.train_dataset = RewardDataset(config.train_dataset_path) - self.train_dataloader = DataLoader( - self.train_dataset, batch_size=config.batch_size - ) if self.validation_flag: self.eval_dataset = RewardDataset(config.validation_dataset_path) - self.validation_dataloader = DataLoader( - self.eval_dataset, batch_size=config.batch_size - ) # intilize scheduler - learning rate will drop to 10% of the initial # value @@ -311,195 +181,37 @@ def __init__(self, config: ConfigReward) -> None: last_epoch=-1, ) - # initialize training stats - stats_path = ModelLoader.get_training_stats_path(config) - self.training_stats = TrainingStats(stats_path) + # initialize the super class + super().__init__(config) - # consistency check between accelerate and deepspeed - if config.accelerate_enable and config.deepspeed_enable: - raise ValueError( - "Both DeepSpeed and Accelerate are enabled for the Reward." - "Please choose one of them." - ) - - # initialize deepspeed - self.model_engine = None - if config.deepspeed_enable is True: - - if config.deepspeed_config_path is None: - raise ValueError( - "DeepSpeed config path is None, but deepspeed is enabled" - ) - if os.path.exists(config.deepspeed_config_path) is False: - raise ValueError( - f"DeepSpeed config path {config.deepspeed_config_path}" - f"does not exist" - ) - ( - self.model_engine, - self.optimizer, - self.train_dataloader, - self.scheduler, - ) = deepspeed.initialize( - args=None, - model=self.reward, - model_parameters=self.reward.parameters(), - training_data=self.train_dataset, - config=self.config.deepspeed_config_path, - ) - print("Training with DeepSpeed") - - # initialize accelerate - self.accelerator = None - if config.accelerate_enable is True: - self.accelerator = Accelerator() - ( - self.reward, - self.optimizer, - self.train_dataloader, - self.scheduler, - ) = self.accelerator.prepare( - self.reward, - self.optimizer, - self.train_dataloader, - self.scheduler, - ) - print("Training with Accelerate") - - @beartype - def save_checkpoint( - self, - current_epoch: int, - current_step: int, - max_epochs: int, - max_steps: int, - ) -> None: - """Save the checkpoints of the model - - Args: - current_epoch (int): Current epoch - current_step (int): Current step - max_epochs (int): Maximum number of epochs - max_steps (int): Maximum number of steps - """ - - print( - f"Saving checkpoint for epoch {current_epoch + 1}, " - f" step {current_step} ..." - ) - - # get the path to save the checkpoint - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=current_epoch, - current_step=current_step, - max_epochs=max_epochs, - max_steps=max_steps, + # dataloader + self.train_dataloader = self.create_dataloader( + self.train_dataset, batch_size=config.batch_size ) - - # remove the checkpoint if it already exists - if os.path.exists(path): - if self.config.deepspeed_enable: - shutil.rmtree(path) - else: - os.remove(path) - - # save the checkpoint - if self.config.deepspeed_enable: - client_state = { - "epoch": current_epoch, - "step": current_step, - } - self.model_engine.save_checkpoint(path, client_state=client_state) - else: - torch.save( - { - "state_dict": self.reward.model.state_dict(), - "optim_state_dict": self.optimizer.state_dict(), - "scheduler_state_dict": self.scheduler.state_dict(), - "training_stats": self.training_stats, - "epoch": current_epoch, - "step": current_step, - }, - path, + if self.validation_flag: + self.validation_dataloader = self.create_dataloader( + self.eval_dataset, batch_size=config.batch_size ) - @beartype - def load_checkpoint( - self, - ) -> Tuple[int, int]: - """Load the checkpoints of the model - - Returns: - Tuple[int, int]: The current epoch and step - from which you should resume the training - """ - - print("Looking for checkpoints...") - # look for the checkpoints - path = ModelLoader.check_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=None, - ) - - # check if a checkpoint exists - if path is not None: - print("Loading ...") - - if self.config.deepspeed_enable: - # try to load the checkpoint - try: - _, client_state = self.model_engine.load_checkpoint(path) - except Exception: - print( - "Checkpoint corrupted!" - "Try to remove the last checkpoint." - "Now Starting from epoch 0, step 0" - ) - return 0, 0 - # load epoch and step to resume loops - epoch = client_state["epoch"] - step = client_state["step"] - else: - # try to load the checkpoint - try: - checkpoint = torch.load(path) - except Exception: - print( - "Checkpoint corrupted!" - "Try to remove the last checkpoint." - "Now Starting from epoch 0, step 0" - ) - return 0, 0 - - # load the model parameters and optimizer parameters - # from the checkpoint - epoch = checkpoint["epoch"] - self.reward.model.load_state_dict(checkpoint["state_dict"]) - self.optimizer.load_state_dict(checkpoint["optim_state_dict"]) - self.scheduler.load_state_dict( - checkpoint["scheduler_state_dict"] - ) - self.training_stats = checkpoint["training_stats"] - step = checkpoint["step"] - return epoch, step + 1 # return the next episode to train - return 0, 0 - def train( self, ) -> None: """Train the reward model""" - print("Start Training the Reward Model") - # get config parameters - if self.config.deepspeed_enable: - batch_size = self.train_dataloader.batch_size + my_logger.success("Start Training the Reward Model") + + # get batch size + if self.deepspeed_enable: + batch_size = self.model_engine.train_batch_size() + elif self.accelerate_enable: + batch_size = ( + self.config.batch_size * self.accelerator.num_processes + ) else: batch_size = self.config.batch_size + + # get other parameters epochs = self.config.epochs - device = self.config.device iteration_per_print = self.config.iteration_per_print checkpoint_steps = self.config.checkpoint_steps @@ -514,63 +226,115 @@ def train( # traing loop for epoch in range(start_epoch, epochs): - self.reward.train() + self.model.train() for i, inputs in enumerate(self.train_dataloader): # skip the steps if resuming from a checkpoint if i < start_step: continue - # get the inputs - input_text = inputs[0] - score = inputs[1] - - # tokenize the input + # get tensor for forward and loss computation with torch.no_grad(): - input_tokens = self.reward.tokenizer( - input_text, - return_tensors="pt", - truncation=True, - padding=True, - ) - output = torch.as_tensor( - score, dtype=torch.float32, device=device - ) + + # get the inputs + input_text = inputs[0] + score = inputs[1] + + # tokenize the input + if self.accelerate_enable: + input_tokens = self.model.module.tokenizer( + input_text, + return_tensors="pt", + truncation=True, + padding=True, + ) + else: + input_tokens = self.model.tokenizer( + input_text, + return_tensors="pt", + truncation=True, + padding=True, + ) + + # assign the input and output + output = torch.as_tensor(score, dtype=torch.float16) + input_ids = input_tokens["input_ids"] + input_mask = input_tokens["attention_mask"] + + # move to device + if not self.config.load_8bit: + input_ids = input_ids.to(self.device) + input_mask = input_mask.to(self.device) + output = output.to(self.device) # forward pass if self.config.deepspeed_enable: + # deepspeed est_output = self.model_engine( - input_tokens["input_ids"].to(device), - input_tokens["attention_mask"].to(device), + input_ids, + input_mask, + )[:, -1] + elif self.accelerate_enable: + # accelerate + est_output = self.model( + input_ids, + input_mask, )[:, -1] else: - est_output = self.reward.get_reward( - input_tokens["input_ids"].to(device), - input_tokens["attention_mask"].to(device), - ) + # Pytorch + with torch.autocast( + device_type=self.config.device_type, + dtype=torch.float16, + ): + est_output = self.model( + input_ids, + input_mask, + )[:, -1] # compute the loss - loss = self.loss_function(est_output, output) - self.training_stats.training_loss.append(loss.item()) + if self.deepspeed_enable or self.accelerate_enable: + # DeepSpeed and Accelerate + loss = self.loss_function(est_output, output) + + else: + # if mixed precision with pytorch use autocast + with torch.autocast( + device_type=self.config.device_type, + dtype=torch.float16, + ): + loss = self.loss_function(est_output, output) # backward pass - if self.config.deepspeed_enable: + if self.deepspeed_enable: + # DeepSpeed self.model_engine.backward(loss) self.model_engine.step() - elif self.config.accelerate_enable: + elif self.accelerate_enable: + # Accelerate self.optimizer.zero_grad() self.accelerator.backward(loss) self.optimizer.step() self.scheduler.step() else: + # Pytorch self.optimizer.zero_grad() - loss.backward() - self.optimizer.step() + if self.config.load_8bit: + # 8 bit from HF + loss.backward() + self.optimizer.step() + else: + # Pytorch Mixed Precision + self.scaler.scale(loss).backward() + self.scaler.step(self.optimizer) + self.scaler.update() self.scheduler.step() + # save training stats + self.append_training_stats(training_loss=loss.detach().item()) + # print progress if i % iteration_per_print == 0: - print( + my_logger.info( f"Epoch: {epoch+1}/{epochs}, " f"Iteration: {i+1}/{n_iter}, " f"Training Loss: {loss.item()}" @@ -578,52 +342,55 @@ def train( printed_est_output = [ round(float(x), 1) for x in est_output.cpu().tolist() ] - print( - "prediction", - printed_est_output, - "target", - score.cpu().tolist(), + my_logger.info( + f"prediction {printed_est_output} " + f"target {score.cpu().tolist()}" ) # checkpoints saving if cnt_checkpoints % checkpoint_steps == 0: - self.save_checkpoint(epoch, i, epochs, n_iter) + self.save_checkpoint( + current_epoch=epoch, + max_epochs=epochs, + current_step=i, + max_steps=n_iter, + ) cnt_checkpoints = 1 else: cnt_checkpoints += 1 # Validation if self.validation_flag: - self.reward.eval() + self.model.eval() with torch.no_grad(): for i, (text, score) in enumerate( self.validation_dataloader ): # tokenize inputs - input_tokens = self.reward.tokenizer( + input_tokens = self.model.tokenizer( text, return_tensors="pt", padding=True ) - input_tokens = input_tokens.to(device) + input_tokens = input_tokens.to(self.device) # TODO: check on the length of the input tokens if # they are too many it can create problems output = torch.tensor(score, dtype=torch.float32).to( - device + self.device ) # forward pass - est_output = self.reward.get_reward( + est_output = self.model.get_reward( input_tokens["input_ids"], input_tokens["attention_mask"], ) # compute loss loss = self.loss_function(est_output, output) - self.training_stats.validation_loss.append(loss.item()) + self.append_training_stats(validation_loss=loss.item()) # print progress if i % iteration_per_print == 0: - print( + my_logger.info( f"Epoch: {epoch+1}/{epochs}, " f"Iteration: {i+1}/{n_iter}, " f"Validation Loss: {loss.item()}" @@ -632,4 +399,8 @@ def train( start_step = 0 # save the model at the end of the training - self.reward.save() + if self.accelerate_enable: + self.model.module.save() + else: + self.model.save() + my_logger.success("Training is finished") diff --git a/apps/accelerate/chatllama/chatllama/rlhf/trainer.py b/apps/accelerate/chatllama/chatllama/rlhf/trainer.py index 2878f7c8..6e591b41 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/trainer.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/trainer.py @@ -3,17 +3,14 @@ import random from collections import deque, namedtuple -import deepspeed import torch -import torch.distributed as dist -from accelerate import Accelerator from beartype import beartype from beartype.typing import Deque, List, Tuple, Union -from deepspeed.runtime.engine import DeepSpeedEngine -from torch.utils.data import DataLoader, Dataset from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts +from torch.utils.data import Dataset from chatllama.rlhf.actor import ActorModel +from chatllama.rlhf.base_model import BaseModel, BaseTrainer from chatllama.rlhf.config import ( Config, ConfigActor, @@ -23,7 +20,7 @@ from chatllama.rlhf.model_list import hf_models from chatllama.rlhf.model_loader import ModelLoader from chatllama.rlhf.reward import RewardModel, CriticModel -from chatllama.rlhf.utils import TrainingStats, ConversationLog +from chatllama.rlhf.utils import ConversationLog, my_logger """ @@ -138,7 +135,7 @@ def check_model_family(config1: ConfigType, config2: ConfigType) -> bool: return False -class ActorCritic(torch.nn.Module): +class ActorCritic(BaseModel): """Actor Critic class stores both the actor and the critic models and it generates values and action for given sequences during the training of the actor. @@ -146,9 +143,8 @@ class ActorCritic(torch.nn.Module): Attributes: actor (ActorModel): Actor model critic (CriticModel): Critic model - debug (bool): enable prints for Debugging - use_same_tokenizer (bool): if True the actor and critic use the same - tokenizer + use_same_tokenizer (bool): True if actor and critic use the same + tokenizer, False otherwise Methods: forward: given a sequence returns action logits and values (used @@ -159,88 +155,20 @@ class ActorCritic(torch.nn.Module): """ def __init__(self, config: Config) -> None: - super().__init__() - self.config = config + super().__init__(config) + # load the actor self.actor = ActorModel(config.actor) # check if critic must be initialized from reward model ModelLoader.init_critic_from_reward(config.critic) - self.critic = CriticModel(config.critic) - - # if the actor and critic use the same tokenizer is set to True - self.use_same_tokenizer = False - - # debug flag - self.debug = config.actor.debug - - @beartype - def load(self) -> None: - """Load the model from the path. - This method is not implemented since it relies on actor and critic - __init__ methods to perform the loading from their respective paths - then loaded. - - """ - pass - - @beartype - def save(self) -> None: - """Save the model to the path - This method is implemented to save the actor model as result of RLHF - in the folder actor_rl instead of actor.save() method that saves it - in the actor folder. - """ - # get the path to save the actor - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=False, - ) - # save the model - print(f"Saving model to {path} ...") - torch.save( - {"state_dict": self.actor.model.state_dict()}, - path, - ) - - # get the path to save the critic model - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config.critic, - is_checkpoint=False, - ) - - # save the model - print(f"Saving model to {path} ...") - torch.save( - { - "model": self.critic.model.state_dict(), - "head": self.critic.head.state_dict(), - }, - path, - ) - - def save_deepspeed( - self, - model_engine: DeepSpeedEngine, - config: ConfigType, - client_state: dict = None, - ): - """Save the deepspeed model_engine to the path - This method is implemented to save the actor model as result of RLHF - in the folder actor_rl instead of actor.save() method that saves it - in the actor folder. Same goes for the critic model. - """ - # get the path to save the actor - model_folder, model_name, path = ModelLoader.get_model_path( - config=config, - is_checkpoint=False, - ) + # now load the critic + self.critic = CriticModel(config.critic) - # save the model - print(f"Saving model to {path} ...") - model_engine.save_checkpoint( - save_dir=path, client_state=client_state if client_state else {} + # flag to check if actor and critic use the same tokenizer + self.use_same_tokenizer = check_model_family( + config.actor, config.critic ) @beartype @@ -258,23 +186,26 @@ def forward( values for each generation step. Args: - sequences_actor (torch.Tensor): Sequences composed of - [states, actions] for the actor - sequence_mask_actor (torch.Tensor): Mask for the sequences - of the actor - sequences_critic (torch.Tensor): Sequences composed of - [states, actions] for the critic - sequences_mask_critic (torch.Tensor): Mask for the sequences - of the critic + sequences_actor (torch.Tensor) [batch_size, seq_len + composed of [states, actions] for the actor + sequence_mask_actor (torch.Tensor) [batch_size, seq_len]: Mask + for the sequences of the actor + sequences_critic (torch.Tensor) [batch_size, seq_len]: Sequences + composed of [states, actions] for the critic. + Can differ from sequences_actor if the critic uses a different + tokenizer + sequences_mask_critic (torch.Tensor) [batch_size, seq_len]: Mask + for the sequences of the critic. action_len_actor (int): Length of the actions in the sequences for the actor action_len_critic (int): Length of the actions in the sequences for the critic Returns: - action_logits (torch.Tensor): Logits for the actions in the - sequences - values (torch.Tensor): Values for the actions in the sequences + action_logits (torch.Tensor) [batch_size, seq_len, vocab_size]: + Logits for the actions in the sequences + values (torch.Tensor) [batch_size, seq_len]: Values for the actions + in the sequences """ # use a single forward on the whole sequence @@ -290,19 +221,6 @@ def forward( real_actions_logits = actions_logits[:, -action_len_actor:, :] real_values = values[:, -action_len_critic:] - if self.debug: - print("ActorCritic.forward") - print("action_len_actor", action_len_actor) - print("action_len_critic", action_len_critic) - print("sequences_actor.shape", sequences_actor.shape) - print("sequences_actor", sequences_actor) - print("sequences_critic.shape", sequences_critic.shape) - print("sequences_critic", sequences_critic) - print("real_action_logits.shape", actions_logits.shape) - print("real_action_logits", actions_logits) - print("real_values.shape", values.shape) - print("real_values", values) - return ( real_actions_logits, real_values, @@ -319,19 +237,25 @@ def generate( """Generate actions, actions_logits, values and sequences from states Args: - states_actor (torch.Tensor): States for the actor - states_mask_actor (torch.Tensor): Mask for the states for the - actor - states_critic (torch.Tensor): States for the critic + states_actor (torch.Tensor) [batch_size, input_len]: States for the + actor. + states_mask_actor (torch.Tensor) [batch_size, input_len]: Mask for + the states for the actor + states_critic (torch.Tensor) [batch_size, input_len]: States for + the critic. Can differ from states_actor if the critic uses a + different tokenizer Returns: - actions (torch.Tensor): Actions generated from the states - actions_logits (torch.Tensor): Logits for the actions generated - from the states (i.e. pi(y | x)) - values (torch.Tensor): Values generated by the critic model - for the actions generated by the actor (i.e. V(x)) - sequences (torch.Tensor): Sequences generated from the states - as [states, actions] + actions (torch.Tensor) [batch_size, act_len]: Actions generated + from the states. + actions_logits (torch.Tensor) [batch_size, act_len, vocab_size]: + Logits for the actions generated from the states + (i.e. pi(y | x)) + values (torch.Tensor) [batch_size, act_len]: Values generated by + the critic model for the actions generated by the actor + (i.e. V(x)) + sequences (torch.Tensor) [batch_size, seq_len]: Sequences + generated from the states as [states, actions] """ # generate action sequence from the actor @@ -384,16 +308,6 @@ def generate( action_len_actor, action_len_critic, ) - if self.debug: - print("ActorCritic.generate") - print("actions shape", actions.shape) - print("actions", actions) - print("sequence shape", sequences_actor.shape) - print("sequence", sequences_actor) - print("actions_logits shape", actions_logits.shape) - print("actions_logits", actions_logits) - print("values shape", values.shape) - print("values", values) return ( actions, @@ -412,8 +326,6 @@ def generate( Memory = namedtuple( "Memory", [ - "states_actor", - "actions", "values", "rewards", "actions_log_probs", @@ -446,8 +358,6 @@ def __len__( def __getitem__(self, idx) -> Tuple: # return the idx-th memory element as a tuple of tensors on the device item = ( - self.data[idx].states_actor, - self.data[idx].actions, self.data[idx].values, self.data[idx].rewards, self.data[idx].actions_log_probs, @@ -493,30 +403,57 @@ def sample(self, n: int) -> List: return random.sample(self.data, n) -class RLTrainer: +class CustomOptimizer(torch.optim.Optimizer): + """Class to define two different LR for the actor and the critic. + Still not working with distributed trainig. + """ + + def __init__(self, params_actor, params_critic, lr_actor, lr_critic): + self.actor = torch.optim.AdamW(params_actor, lr=lr_actor) + self.critic = torch.optim.AdamW(params_critic, lr=lr_critic) + + self.param_groups = self.actor.param_groups + self.critic.param_groups + + def step(self): + self.actor.step() + self.critic.step() + + def zero_grad(self): + self.actor.zero_grad() + self.critic.zero_grad() + + +class RLTrainer(BaseTrainer): """Train the actor-critic model using RL Attributes: config (Config): Configuration of the trainer - debug (bool): Debug mode actorcritic (ActorCritic): Actor-critic model - actor_optim (torch.optim): Optimizer for the actor - critic_optim (torch.optim): Optimizer for the critic - actor_scheduler (torch.optim.lr_scheduler): Scheduler for the actor - critic_scheduler (torch.optim.lr_scheduler): Scheduler for the critic - reward (RewardModel): Reward model - training_stats (TrainingStats): Class to store training stats - conversation_log (ConversationLog): Class to store the conversation - examples_sampler (ExamplesSampler): Class to sample examples - eps (float): small epsilon to avoid division by zero + reward (Reward): Reward function + optimizer (torch.optim.Optimizer): Optimizer for the actor-critic model + scheduler (torch.optim.lr_scheduler): Scheduler for the optimizer + examples_sampler (ExamplesSampler): Sampler to generate the examples + conversation_log (ConversationLog): List of the conversation logs Methods: train: the training loop that calls the learn function after generating the experiences. learn: Learn from a batch of experiences and update the actor and the critic model. - load_checkpoint: Load the checkpoint of the actor-critic model - save_checkpoint: Save the checkpoint of the actor-critic model + + Known Issues: + - When using load_8bit and peft on the actor model, the following error + is raised: + return torch.layer_norm(input, + normalized_shape, + weight, + bias, + eps, + torch.backends.cudnn.enabled) + RuntimeError: Expected all tensors to be on the same device, + but found at least two devices, cuda:1 and cuda:0! + (when checking argument for argument weight + in method wrapper_CUDA__native_layer_norm) """ def __init__( @@ -524,35 +461,31 @@ def __init__( config: Config, ) -> None: - # save config - self.config = config - - # set debug mode - self.debug = config.trainer.debug - - # initialize agent-critic + # initialize actor-critic self.actorcritic = ActorCritic(config) + # TODO: currently now working with custom optimizer # initialize actor optimizer - self.actor_optimizer = torch.optim.Adam( - self.actorcritic.actor.parameters(), lr=config.trainer.actor_lr - ) - - # initialize critic optimizer - self.critic_optimizer = torch.optim.Adam( - self.actorcritic.critic.parameters(), lr=config.trainer.critic_lr + # self.optimizer = CustomOptimizer( + # self.actorcritic.actor.parameters(), + # self.actorcritic.critic.parameters(), + # config.trainer.actor_lr, + # config.trainer.critic_lr, + # ) + + self.optimizer = torch.optim.AdamW( + self.actorcritic.parameters(), + lr=config.trainer.actor_lr, ) - # scheduler (defined in the learn() method (i need dataset size)) - self.actor_scheduler = None - self.critic_scheduler = None + # scheduler init to None (need the dataset info to initialize it) + self.scheduler = None # initialize reward model self.reward = RewardModel(config.reward) - # initialize class to store training stats - path = ModelLoader.get_training_stats_path(config) - self.training_stats = TrainingStats(path) + # initialize class to store conversations logs + # initialize class to store conversations logs model_folder, _, _ = ModelLoader.get_model_path( config, is_checkpoint=True, @@ -563,247 +496,151 @@ def __init__( # initialize examples sampler self.example_sampler = ExamplesSampler(config.trainer.examples_path) - # check if actor and critic use the same tokenizer - self.actorcritic.use_same_tokenizer = check_model_family( - config.actor, config.critic - ) + # Base Trainer Initialization: + super().__init__(config) - # check if actor and reward use the same tokenizer - self.use_same_tokenizer = check_model_family( - config.actor, config.reward - ) - - # eps - self.eps = 1e-8 - - # deepspeed initialization - self.actor_model_engine = None - self.critic_model_engine = None - self.is_deepspeed_init = None - - if ( - self.config.actor.deepspeed_enable - or self.config.critic.deepspeed_enable - or self.config.critic.deepspeed_enable - ): - deepspeed.init_distributed("nccl") - self.is_deepspeed_init = True - os.environ["TOKENIZERS_PARALLELISM"] = "False" + def ppo_loss( + self, + actions_logits: torch.Tensor, + old_actions_log_probs: torch.Tensor, + old_values: torch.Tensor, + values: torch.Tensor, + rewards: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute the PPO loss - else: - self.is_deepspeed_init = False - - if self.config.actor.deepspeed_enable: - ( - self.actor_model_engine, - self.actorcritic.actor, - self.actor_optimizer, - ) = self.initialize_deepspeed_model( - config=self.config.actor, model=self.actorcritic.actor - ) + Args: + actions_logits (torch.Tensor): Logits of the actions + old_actions_log_probs (torch.Tensor): Log prob of the actions + old_values (torch.Tensor): Values of the actions + values (torch.Tensor): Values of the actions + rewards (torch.Tensor): Rewards of the actions - if self.config.critic.deepspeed_enable: - ( - self.critic_model_engine, - self.actorcritic.critic, - self.critic_optimizer, - ) = self.initialize_deepspeed_model( - config=self.config.critic, model=self.actorcritic.critic - ) + Returns: + torch.Tensor: PPO loss - if self.config.reward.deepspeed_enable: - ( - _, - self.reward, - _, - ) = self.initialize_deepspeed_model( - config=self.config.reward, model=self.reward - ) + """ - @staticmethod - def initialize_deepspeed_model( - config: Union[ConfigActor, ConfigCritic, ConfigReward], - model: torch.nn.Module, - ): - - if config.deepspeed_config_path is None: - raise ValueError("DeepSpeed config path is None, but deepspeed is enabled") - if os.path.exists(config.deepspeed_config_path) is False: - raise ValueError( - f"DeepSpeed config path" - f"{config.deepspeed_config_path}" - f"does not exist" - ) - (model_engine, ds_optimizer, _, _,) = deepspeed.initialize( - args=None, - model=model, - model_parameters=model.parameters(), - config=config.deepspeed_config_path, - ) - # model_engine.module has to be returned to make custom methods - # of Module accessible - return model_engine, model_engine.module, ds_optimizer + # get hyperparameters + actor_eps_clip = self.config.trainer.actor_eps_clip + critic_eps_clip = self.config.trainer.critic_eps_clip + beta_s = self.config.trainer.beta_s - @beartype - def save_checkpoint( - self, - current_episode: int, - max_episode: int, - ) -> None: + # get action log prob + actions_prob = torch.softmax(actions_logits, dim=-1).max(dim=-1).values + actions_log_prob = torch.log(actions_prob + self.eps) - print(f"Saving checkpoint for episode {current_episode+1}..") + # compute entropy + entropies = (actions_prob * actions_log_prob).sum(dim=-1) - # get the path to save the checkpoint for the critic - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config.critic, - is_checkpoint=True, - current_epoch=current_episode, - max_epochs=max_episode, - max_steps=0, + # compute KL divergence + kl_div_loss = ( + (actions_prob * (actions_log_prob - old_actions_log_probs)) + .sum(dim=-1) + .mean() ) - # if the checkpoint already exists remove it. - # Deepspeed checkpoints are already directories and will be overwritten - if os.path.exists(path) and not self.is_deepspeed_init: - os.remove(path) - - # save the checkpoint - actor_checkpoint_dict = { - "episode": current_episode, - "critic_state_dict": self.actorcritic.critic.state_dict(), - "critic_optim_state_dict": self.critic_optimizer.state_dict(), - } - - if self.config.actor.deepspeed_enable: - # The model and optimizer state dicts are actually already saved - # In the deepspeed model engine. But to make sure no depending - # methods fail, the states are included in actor_checkpoint_dict. - # ATTENTION: If you use deepspeed zero optimization, the client_state - # will not be saved - self.actor_model_engine.save_checkpoint( - save_dir=path, client_state=actor_checkpoint_dict - ) - else: - torch.save(actor_checkpoint_dict, path) + # compute ratios + ratios = (actions_log_prob - old_actions_log_probs).exp() - # get the path to save the checkpoint for the actor - model_folder, model_name, path = ModelLoader.get_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=current_episode, - max_epochs=max_episode, - max_steps=0, - ) + # compute policy loss + if check_model_family(self.config.actor, self.config.critic): - # if the checkpoint already exists remove it. - # Deepspeed checkpoints are already directories and will be overwritten - if os.path.exists(path) and not self.is_deepspeed_init: - os.remove(path) - - # save the checkpoint - critic_checkpoint_dict = { - "episode": current_episode, - "actor_state_dict": self.actorcritic.actor.state_dict(), - "actor_optim_state_dict": self.actor_optimizer.state_dict(), - "training_stats": self.training_stats, - } - - if self.config.critic.deepspeed_enable: - # The model and optimizer state dicts are actually already saved - # In the deepspeed model engine. But to make sure no depending - # methods fail, the states are included in critic_checkpoint_dict. - # ATTENTION: If you use deepspeed zero optimization, the client_state - # will not be saved - self.critic_model_engine.save_checkpoint( - save_dir=path, client_state=critic_checkpoint_dict + # compute discounted rewards as in TRL + gamma = self.config.trainer.gamma_discounted + discounted_rewards = torch.zeros_like(old_values) + gamma_power = torch.arange( + 0, + discounted_rewards.shape[1], + device=rewards.device, + dtype=rewards.dtype, ) + gamma_vect = gamma**gamma_power + for i in range(discounted_rewards.shape[1]): + discounted_rewards[:, i] = torch.matmul( + rewards[:, i:], gamma_vect[: rewards[:, i:].shape[1]] + ) + advantages = ( + discounted_rewards - old_values + ) # TRL has opposite sign for old values + if advantages.std() > self.eps: + advantages = (advantages - advantages.mean(dim=-1)) / ( + advantages.std() + self.eps + ) + surr1 = advantages * ratios else: - torch.save(critic_checkpoint_dict, path) + advantages = rewards - old_values[:, -1] + surr1 = advantages * ratios - @beartype - def load_checkpoint( - self, - ) -> int: - - critic_episode = -1 - actor_episode = -1 - - # check if there are some checkpoint for the critic - print("Looking for checkpoints...") - path = ModelLoader.check_model_path( - config=self.config.critic, - is_checkpoint=True, - current_epoch=None, + surr2 = ( + torch.clamp(ratios, 1 - actor_eps_clip, 1 + actor_eps_clip) + * advantages ) - # if there are checkpoint - if path is not None: - - # load the critic checkpoint - print("Loading ...") - try: - checkpoint = torch.load(path) - except Exception: - print( - "Checkpoint of critic corrupted!" - "Try to remove the last checkpoint." - "Now Starting from episode 0" - ) - return 0 + policy_loss = -torch.min(surr1, surr2) - beta_s * entropies + policy_loss = policy_loss.mean() + policy_loss = policy_loss + kl_div_loss - # load checkpoint into model - critic_episode = checkpoint["episode"] - self.actorcritic.critic.load_state_dict( - checkpoint["critic_state_dict"] - ) - self.critic_optimizer.load_state_dict( - checkpoint["critic_optim_state_dict"] - ) - - # check if there are checkpoints for the actor - print("Looking for checkpoints...") - path = ModelLoader.check_model_path( - config=self.config, - is_checkpoint=True, - current_epoch=None, + # compute value loss + value_loss_clipped = old_values + (values - old_values).clamp( + -critic_eps_clip, critic_eps_clip ) - - # if there are some checkpoints - if path is not None: - - # load the actor checkpoint - print("Loading ...") - try: - checkpoint = torch.load(path) - except Exception: - print( - "Checkpoint of actor corrupted!" - "Try to remove the last checkpoint." - "Now Starting from episode 0" + value_loss1 = (value_loss_clipped - rewards) ** 2 + value_loss2 = (values - rewards) ** 2 + value_loss = torch.max(value_loss1, value_loss2).mean() + + # Sum the two losses + loss = policy_loss + value_loss + + # check the losses + if torch.isnan(loss): + if torch.isnan(policy_loss): + if torch.isnan(kl_div_loss): + raise my_logger.error( + ValueError, + "KL div Loss is nan", + ) + else: + if torch.isnan(entropies.mean()): + raise my_logger.error( + ValueError, + "Entropies Loss is nan", + ) + elif torch.isnan(surr1.mean()) or torch.isnan( + surr2.mean() + ): + if torch.isnan(advantages.mean()): + if torch.isnan(rewards.mean()): + raise my_logger.error( + ValueError, + "Rewards are nan", + ) + elif torch.isnan(old_values.mean()): + raise my_logger.error( + ValueError, + "Old Values are nan", + ) + elif torch.isnan(gamma_vect.mean()): + raise my_logger.error( + ValueError, + "Gamma Vector is nan", + ) + else: + raise my_logger.error( + ValueError, + "Advantages are nan", + ) + elif torch.isnan(ratios.mean()): + raise my_logger.error( + ValueError, + "Ratios are nan", + ) + else: + raise my_logger.error( + ValueError, + "Value Loss is nan", ) - return 0 - - # load checkpoint into the model - actor_episode = checkpoint["episode"] - self.actorcritic.actor.load_state_dict( - checkpoint["actor_state_dict"] - ) - self.actor_optimizer.load_state_dict( - checkpoint["actor_optim_state_dict"] - ) - self.training_stats = checkpoint["training_stats"] - - # check if there are some discrepancies between the checkpoints - if critic_episode == actor_episode: - # all ok start from next episode - return critic_episode + 1 - else: - print( - f"There are some discrepancies between the checkpoints" - f"of actor and critic \nactor episode: {actor_episode}" - f"\n critic episode: {critic_episode}\n" - ) - return min(critic_episode, actor_episode) + 1 + return loss, policy_loss, value_loss @beartype def learn(self, memories: Deque[Memory]) -> None: @@ -812,84 +649,35 @@ def learn(self, memories: Deque[Memory]) -> None: - then compare action logits probs with memories one and values with rewards to compute the PPO loss and update the actor-critic model """ - print("Start to Learn...") + my_logger.info("Start to Learn...") # get parameters epochs = self.config.trainer.epochs - actor_eps_clip = self.config.trainer.actor_eps_clip - critic_eps_clip = self.config.trainer.critic_eps_clip - beta_s = self.config.trainer.beta_s + + # TODO: check this with distributed training batch_size = self.config.trainer.batch_size - device = ( - torch.device(f"cuda:{dist.get_rank()}") - if self.is_deepspeed_init - else self.config.trainer.device - ) # create dataset from memories - dataset = ExperienceDataset(memories, device) - if self.is_deepspeed_init: - engine = self.actor_model_engine or self.critic_model_engine - dataloader = engine.deepspeed_io(dataset) - else: - dataloader = DataLoader(dataset, batch_size=batch_size) - - # initialize scheduler for actor - actor_lr = self.config.trainer.actor_lr - # This lr_scheduler is not available in deepspeed - # see https://deepspeed.readthedocs.io/en/latest/schedulers.html - if not self.is_deepspeed_init: - self.actor_scheduler = CosineAnnealingWarmRestarts( - self.actor_optimizer, T_0=len(dataset), eta_min=actor_lr * 0.1 - ) - - # initialize scheduler for critic - critic_lr = self.config.trainer.critic_lr - # This lr_scheduler is not available in deepspeed - # see https://deepspeed.readthedocs.io/en/latest/schedulers.html - if not self.is_deepspeed_init: - self.critic_scheduler = CosineAnnealingWarmRestarts( - self.critic_optimizer, T_0=len(dataset), eta_min=critic_lr * 0.1 - ) + self.train_dataset = ExperienceDataset(memories, self.device) + self.train_dataloader = self.create_dataloader( + self.train_dataset, batch_size + ) - # initialize actor accelerate - if self.config.actor.accelerate_enable is True: - actor_accelerator = Accelerator() - ( - actor_model, - self.actor_optimizer, - self.train_dataloader, - self.actor_scheduler, - ) = actor_accelerator.prepare( - self.actorcritic.actor, - self.actor_optimizer, - self.train_dataloader, - self.actor_scheduler, - ) - self.actorcritic.actor = actor_model - - # initialize critic accelerate - if self.config.critic.accelerate_enable is True: - critic_accelerator = Accelerator() - ( - critic_model, - self.critic_optimizer, - self.critic_scheduler, - ) = critic_accelerator.prepare( - self.actorcritic.critic, - self.critic_optimizer, - self.critic_scheduler, - ) - self.actorcritic.critic = critic_model + # assign a scheduler to the optimizer + self.scheduler = CosineAnnealingWarmRestarts( + self.optimizer, + T_0=len(self.train_dataset) // batch_size, + T_mult=1, + eta_min=self.config.trainer.actor_lr * 0.1, + ) # train agent-critic self.actorcritic.train() for epoch in range(epochs): - for k, batch in enumerate(dataloader): + for k, batch in enumerate(self.train_dataloader): + # TODO: now first two elements are not used can be removed? ( - states_actor, - old_actions, old_values, rewards, old_actions_log_probs, @@ -899,165 +687,104 @@ def learn(self, memories: Deque[Memory]) -> None: sequences_mask_critic, action_len_actor, action_len_critic, - ) = [tensor.to(device) for tensor in batch] - - if self.debug: - print( - f"#########################################" - f" batch from memories {k} \n " - f"#########################################" - f"states_actor {states_actor.shape} \n" - f"old_actions {old_actions.shape} \n" - f"old_values {old_values.shape} \n" - f"rewards {rewards.shape} \n" - f"old_actions_log_probs " - f"{old_actions_log_probs.shape}\n" - f"sequences_actor {sequences_actor.shape} \n" - f"sequences_mask_actor " - f"{sequences_mask_actor.shape} \n" - f"sequences_critic {sequences_critic.shape} \n" - f"sequences_mask_critic " - f"{sequences_mask_critic.shape} \n" - f"action_len_actor {action_len_actor} \n" - f"action_len_critic {action_len_critic} \n" - f"#########################################" - ) + ) = [tensor.to(self.device) for tensor in batch] # get actor critic new probabilities and values - actions_logits, values = self.actorcritic.forward( - sequences_actor, - sequences_mask_actor, - sequences_critic, - sequences_mask_critic, - action_len_actor.item(), - action_len_critic.item(), - ) - - # get action log prob - actions_prob = ( - torch.softmax(actions_logits, dim=-1).max(dim=-1).values - ) - actions_log_prob = torch.log(actions_prob + self.eps) - - # compute entropy - entropies = (actions_prob * actions_log_prob).sum(dim=-1) - - # compute KL divergence - kl_div_loss = ( - (actions_prob * (old_actions_log_probs - actions_log_prob)) - .sum(dim=-1) - .mean() - ) - - # compute ratios - ratios = (actions_log_prob - old_actions_log_probs).exp() - - # compute PPO loss - if check_model_family(self.config.actor, self.config.critic): - # compute discounted rewards as in TRL - gamma = self.config.trainer.gamma_discounted - discounted_rewards = torch.zeros_like(old_values) - for i in range(discounted_rewards.shape[1]): - for j in range(i, discounted_rewards.shape[1]): - discounted_rewards[:, i] += ( - gamma ** (j - i) * rewards[:, j] - ) - - advantages = ( - discounted_rewards - old_values - ) # TRL has opposite sign for old values - advantages = (advantages - advantages.mean(dim=-1)) / ( - advantages.std() + self.eps + if self.deepspeed_enable: + actions_logits, values = self.model_engine.forward( + sequences_actor, + sequences_mask_actor, + sequences_critic, + sequences_mask_critic, + action_len_actor.item(), + action_len_critic.item(), + ) + elif self.accelerate_enable: + actions_logits, values = self.actorcritic.forward( + sequences_actor, + sequences_mask_actor, + sequences_critic, + sequences_mask_critic, + action_len_actor.item(), + action_len_critic.item(), ) - - surr1 = advantages * ratios else: - advantages = rewards - old_values[:, -1] - surr1 = advantages * ratios - - surr2 = ( - torch.clamp(ratios, 1 - actor_eps_clip, 1 + actor_eps_clip) - * advantages - ) - - policy_loss = -torch.min(surr1, surr2) - beta_s * entropies - policy_loss = policy_loss.mean() - loss = policy_loss + kl_div_loss - - # check if loss item is NaN - if torch.isnan(loss): - raise ValueError("Loss is nan") - - # update actor with loss - if self.config.actor.deepspeed_enable: - self.actor_model_engine.backward(loss) - self.actor_model_engine.step() - elif self.config.actor.accelerate_enable: - self.actor_optimizer.zero_grad() - actor_accelerator.backward(loss) - self.actor_optimizer.step() - self.actor_scheduler.step() + with torch.autocast( + device_type=self.config.trainer.device_type, + dtype=torch.float16, + ): + actions_logits, values = self.actorcritic.forward( + sequences_actor, + sequences_mask_actor, + sequences_critic, + sequences_mask_critic, + action_len_actor.item(), + action_len_critic.item(), + ) + # compute the loss + if self.accelerate_enable or self.deepspeed_enable: + (loss, policy_loss, value_loss,) = self.ppo_loss( + actions_logits, + old_actions_log_probs, + old_values, + values, + rewards, + ) else: - self.actor_optimizer.zero_grad() - loss.backward() - self.actor_optimizer.step() - self.actor_scheduler.step() - - # compute value loss - # the loss is the distance between the rewards and the values - # I want this distance to be small so that values are - # representative of the rewards, for this reason i took the - # maximum between the two. - # The clip is limiting the slew-rate of values_loss_clipped - value_loss_clipped = old_values + (values - old_values).clamp( - -critic_eps_clip, critic_eps_clip - ) - value_loss1 = (value_loss_clipped - rewards) ** 2 - value_loss2 = (values - rewards) ** 2 - value_loss = torch.max(value_loss1, value_loss2).mean() - - if torch.isnan(value_loss): - raise ValueError("Value loss is nan") - - # upate critic - if self.config.critic.deepspeed_enable: - self.critic_model_engine.backward(value_loss) - self.critic_model_engine.step() - elif self.config.critic.accelerate_enable: - self.critic_optimizer.zero_grad() - critic_accelerator.backward(loss) - self.critic_optimizer.step() - self.critic_scheduler.step() + with torch.autocast( + device_type=self.config.trainer.device_type, + dtype=torch.float16, + ): + (loss, policy_loss, value_loss,) = self.ppo_loss( + actions_logits, + old_actions_log_probs, + old_values, + values, + rewards, + ) + + # backward pass + if self.deepspeed_enable: + # DeepSpeed backward pass + self.model_engine.backward(loss) + self.model_engine.step() + elif self.accelerate_enable: + # Accelerate backward pass + self.optimizer.zero_grad() + self.accelerator.backward(loss) + self.optimizer.step() + self.scheduler.step() else: - self.critic_optimizer.zero_grad() - value_loss.backward() - self.critic_optimizer.step() - self.critic_scheduler.step() + # PyTorch mixed precision backward pass + self.optimizer.zero_grad() + self.scaler.scale(loss).backward() + self.scaler.step(self.optimizer) + self.scaler.update() + self.scheduler.step() # append the losses to the training stats - self.training_stats.training_loss.append( - loss.detach().cpu().item() - ) - self.training_stats.value_loss.append( - value_loss.detach().cpu().item() + self.append_training_stats( + training_loss=policy_loss.detach().cpu().item(), + value_loss=value_loss.detach().cpu().item(), ) # print iteration info - print( - f"Epoch {epoch+1}/{epochs}", - f"Step {k+1}/{int(len(dataloader) / batch_size)}", - f"Loss {loss.detach().cpu().item():.4f}", - f"Value Loss {value_loss.detach().cpu().item():.4f}", + my_logger.info( + f"Epoch {epoch+1}/{epochs} " + f"Step " + f"{k+1}/{int(len(self.train_dataloader) / batch_size)} " + f"Policy Loss {policy_loss.detach().cpu().item():.4f} " + f"Value Loss {value_loss.detach().cpu().item():.4f} " ) self.actorcritic.eval() - print("End Learning") + my_logger.info("End Learning") def train( self, ) -> None: - print("Start RL Training") + my_logger.success("Start RL Training") # initialize settings num_episodes = self.config.trainer.num_episodes @@ -1066,11 +793,6 @@ def train( update_timesteps = self.config.trainer.update_timesteps batch_size = self.config.trainer.batch_size checkpoint_steps = self.config.trainer.checkpoint_steps - device = ( - torch.device(f"cuda:{dist.get_rank()}") - if self.is_deepspeed_init - else self.config.trainer.device - ) # number of elements that the memories should contain when learning number_of_memories_per_learn_iteration = ( @@ -1095,7 +817,7 @@ def train( memories = deque([]) # load checkpoint - start_episode = self.load_checkpoint() + start_episode, _ = self.load_checkpoint() # if it is a new training from the start clear the conversation log if start_episode == 0: @@ -1105,171 +827,283 @@ def train( cnt_timesteps = 0 cnt_learn_iter = 0 + # move reward model to correct device + self.reward.to(self.device) + # loop over episodes and timesteps self.actorcritic.eval() for episode in range(start_episode, num_episodes): for timestep in range(max_timesteps): - # print the iteration info - print( - f"Episode: {episode + 1}/{num_episodes}, " - f"Timestep: {timestep + 1}/{max_timesteps}", - f"Learning Cnt: {cnt_timesteps + 1}/{update_timesteps}", - ) + # ensure that no gradients are computed during this step + with torch.no_grad(): - # counter used to count timesteps into memory - cnt_timesteps += 1 + # print the iteration info + my_logger.info( + f"Episode: {episode + 1}/{num_episodes}, " + f"Timestep: {timestep + 1}/{max_timesteps}, " + f"Learning Cnt: " + f"{cnt_timesteps + 1}/{update_timesteps}" + ) - # sample num_examples examples from example dataset - inputs = self.example_sampler.sample(num_examples) + # counter used to count timesteps into memory + cnt_timesteps += 1 - # tokenize examples for the actor - tok_inputs_act = self.actorcritic.actor.tokenizer( - inputs, padding=True, return_tensors="pt", truncation=True - ) + # sample num_examples examples from example dataset + inputs = self.example_sampler.sample(num_examples) - # states are [batch_size, seq_len_of_states] - states_actor = tok_inputs_act["input_ids"].to(device) - states_mask_actor = tok_inputs_act["attention_mask"].to(device) + # tokenize examples for the actor + if self.accelerate_enable: + tok_inputs_act = self.actorcritic.module.actor.tokenizer( # noqa 501 + inputs, + padding=True, + return_tensors="pt", + truncation=True, + ) + else: + tok_inputs_act = self.actorcritic.actor.tokenizer( + inputs, + padding=True, + return_tensors="pt", + truncation=True, + ) - # tokenize examples for the critic - tok_inputs_crt = self.actorcritic.critic.tokenizer( - inputs, padding=True, return_tensors="pt", truncation=True - ) + states_actor = tok_inputs_act["input_ids"] + states_mask_actor = tok_inputs_act["attention_mask"] - # states are [batch_size, seq_len_of_states] - states_critic = tok_inputs_crt["input_ids"].to(device) + # tokenize examples for the critic + if self.accelerate_enable: + tok_inputs_crt = self.actorcritic.module.critic.tokenizer( # noqa 501 + inputs, + padding=True, + return_tensors="pt", + truncation=True, + ) + else: + tok_inputs_crt = self.actorcritic.critic.tokenizer( + inputs, + padding=True, + return_tensors="pt", + truncation=True, + ) - # generate sequences of actions and values - ( - actions, - actions_logits, - values, - sequences_actor, - sequences_mask_actor, - sequences_critic, - sequences_mask_critic, - action_len_actor, - action_len_critic, - ) = self.actorcritic.generate( - states_actor, states_mask_actor, states_critic - ) + states_critic = tok_inputs_crt["input_ids"] + + # move to device + if not self.config.critic.load_8bit: + states_critic = states_critic.to(self.device) + states_actor = states_actor.to(self.device) + states_mask_actor = states_mask_actor.to(self.device) + + # generate sequences of actions and values + if self.accelerate_enable: + ( + actions, + actions_logits, + values, + sequences_actor, + sequences_mask_actor, + sequences_critic, + sequences_mask_critic, + action_len_actor, + action_len_critic, + ) = self.actorcritic.module.generate( + states_actor, states_mask_actor, states_critic + ) + else: + with torch.autocast( + device_type=self.config.trainer.device_type, + dtype=torch.float16, + ): + ( + actions, + actions_logits, + values, + sequences_actor, + sequences_mask_actor, + sequences_critic, + sequences_mask_critic, + action_len_actor, + action_len_critic, + ) = self.actorcritic.generate( + states_actor, states_mask_actor, states_critic + ) - # compute action log probs - action_prob = ( - torch.softmax(actions_logits, dim=-1).max(dim=-1).values - ) - actions_log_probs = torch.log(action_prob + self.eps) - - # get tokenized sequence for the reward models - if self.use_same_tokenizer: - reward_sequence = sequences_actor - reward_mask = sequences_mask_actor - elif check_model_family( - self.config.critic, self.config.reward - ): - reward_sequence = sequences_critic - reward_mask = sequences_mask_critic - else: - tokenized_responses = change_tokenization( - sequences_actor, - self.actorcritic.actor.tokenizer, - self.reward.tokenizer, - ) - # get tokens and mask - reward_sequence = tokenized_responses["input_ids"].to( - device - ) - reward_mask = tokenized_responses["attention_mask"].to( - device + # compute action log probs + action_prob = ( + torch.softmax(actions_logits, dim=-1) + .max(dim=-1) + .values ) + actions_log_probs = torch.log(action_prob + self.eps) + + # get tokenized sequence for the reward models + # if the reward model uses the same tokenizer as the actor + # then the sequences are already tokenized + # otherwise the sequences need to be converted + if check_model_family( + self.config.actor, self.config.reward + ): + + # they can be directly used since the tokenizer is the + # same + reward_sequence = sequences_critic + reward_mask = sequences_mask_critic + else: + + # convert tokenization + if self.accelerate_enable: + tokenized_responses = change_tokenization( + sequences_actor, + self.actorcritic.module.actor.tokenizer, + self.reward.tokenizer, + ) + else: + tokenized_responses = change_tokenization( + sequences_actor, + self.actorcritic.actor.tokenizer, + self.reward.tokenizer, + ) - # compute rewards - rewards = self.reward.forward( - reward_sequence, - reward_mask, - ) + # get tokens and mask + reward_sequence = tokenized_responses["input_ids"] + reward_mask = tokenized_responses["attention_mask"] - rewards = rewards[:, -action_len_critic:] - reward = rewards[:, -1] - - # store memories of the episode / timestep - for i in range(states_actor.shape[0]): - memories.append( - Memory( - states_actor[i, :].detach().cpu(), - actions[i, :].detach().cpu(), - values[i, :].detach().cpu(), - rewards[i, :].detach().cpu(), - actions_log_probs[i, :].detach().cpu(), - sequences_actor[i, :].detach().cpu(), - sequences_mask_actor[i, :].detach().cpu(), - sequences_critic[i, :].detach().cpu(), - sequences_mask_critic[i, :].detach().cpu(), - int(action_len_actor), - int(action_len_critic), - ) - ) + # move to device + if not self.config.reward.load_8bit: + reward_sequence = reward_sequence.to(self.device) + reward_mask = reward_mask.to(self.device) - # decode completions to be logged in the conversation log - completions = [ - self.actorcritic.actor.tokenizer.decode(action) - for action in actions - ] - # remove pad tokens from completions - completions = [ - c.replace(self.actorcritic.actor.tokenizer.pad_token, "") - for c in completions - ] - # remove eos tokens from completions - completions = [ - c.replace(self.actorcritic.actor.tokenizer.eos_token, "") - for c in completions - ] - # strange i need to force this? - completions = [c.replace("", "") for c in completions] - - # log the memories in the conversation log - for i in range(states_actor.shape[0]): - self.conversation_log.append( - inputs[i], - completions[i], - reward[i].detach().cpu().item(), - cnt_learn_iter, + # compute rewards + rewards = self.reward( + reward_sequence, + reward_mask, ) + # the interesting rewards are only for the action + # TODO: need to use action_len_reward and compute it + # before just in case + rewards = rewards[:, -action_len_critic:] + + # the scalar "reward" is the last reward (see reward model) + reward = rewards[:, -1] + + # store memories of the episode and timestep + for i in range(states_actor.shape[0]): + memories.append( + Memory( + values[i, :].detach().cpu(), + rewards[i, :].detach().cpu(), + actions_log_probs[i, :].detach().cpu(), + sequences_actor[i, :].detach().cpu(), + sequences_mask_actor[i, :].detach().cpu(), + sequences_critic[i, :].detach().cpu(), + sequences_mask_critic[i, :].detach().cpu(), + int(action_len_actor), + int(action_len_critic), + ) + ) + + # decode completions to be logged in the conversation log + if self.accelerate_enable: + completions = [ + self.actorcritic.module.actor.tokenizer.decode( + action + ) # noqa 501 + for action in actions + ] + else: + completions = [ + self.actorcritic.actor.tokenizer.decode(action) + for action in actions + ] + + # remove pad tokens from completions + if self.accelerate_enable: + completions = [ + c.replace( + self.actorcritic.module.actor.tokenizer.pad_token, # noqa 501 + "", + ) + for c in completions + ] + else: + completions = [ + c.replace( + self.actorcritic.actor.tokenizer.pad_token, "" + ) + for c in completions + ] + + # remove eos tokens from completions + if self.accelerate_enable: + completions = [ + c.replace( + self.actorcritic.module.actor.tokenizer.eos_token, # noqa 501 + "", + ) + for c in completions + ] + else: + completions = [ + c.replace( + self.actorcritic.actor.tokenizer.eos_token, "" + ) + for c in completions + ] + + # strange i need to force this? + # TODO check how to remove this + completions = [c.replace("", "") for c in completions] + + # log the memories in the conversation log + for i in range(states_actor.shape[0]): + self.conversation_log.append( + inputs[i], + completions[i], + reward[i].detach().cpu().item(), + cnt_learn_iter, + ) + # learn from memories if (cnt_timesteps % update_timesteps == 0) and ( cnt_timesteps != 0 ): - print("len memories", len(memories)) - if not self.is_deepspeed_init or (dist.get_rank() == 0): - self.conversation_log.save() - self.learn(memories) + + my_logger.info(f"Len memories {len(memories)}") + # self.conversation_log.show(cnt_learn_iter) mean_reward = sum([m.rewards[-1] for m in memories]) / len( memories ) - print(f"Mean Reward: {mean_reward}") + my_logger.success(f"Mean Reward: {mean_reward}") + self.learn(memories) memories.clear() cnt_timesteps = 0 cnt_learn_iter += 1 - if not self.is_deepspeed_init or (dist.get_rank() == 0): + + # save conversations for now works only with 1 gpu + if (not self.deepspeed_enable) and ( + not self.accelerate_enable + ): self.conversation_log.save() # save checkpoints if (episode % checkpoint_steps == 0) and (episode != 0): self.save_checkpoint( - current_episode=episode, max_episode=num_episodes + current_epoch=episode, + max_epochs=num_episodes, ) - if not self.is_deepspeed_init or (dist.get_rank() == 0): + + # save conversations for now works only with 1 gpu + if (not self.deepspeed_enable) and ( + not self.accelerate_enable + ): self.conversation_log.save() # save the models - if self.is_deepspeed_init: - self.actorcritic.save_deepspeed(self.actor_model_engine, self.config) - self.actorcritic.save_deepspeed( - self.critic_model_engine, self.config.critic - ) + if self.accelerate_enable: + self.actorcritic.module.save() else: self.actorcritic.save() - print("End RL Training") + + my_logger.success("End RL Training") diff --git a/apps/accelerate/chatllama/chatllama/rlhf/utils.py b/apps/accelerate/chatllama/chatllama/rlhf/utils.py index 50b123fd..421f8059 100644 --- a/apps/accelerate/chatllama/chatllama/rlhf/utils.py +++ b/apps/accelerate/chatllama/chatllama/rlhf/utils.py @@ -1,7 +1,288 @@ import json import os +import re +import sys + +import deepspeed +import logging +import torch from beartype import beartype +from beartype.typing import Union, Tuple, Optional +from loguru import logger from plotly import graph_objects as go +from transformers import AutoTokenizer + +from chatllama.rlhf.config import ( + Config, + ConfigActor, + ConfigCritic, + ConfigReward, +) +from chatllama.rlhf.model_list import hf_models, llama_models + + +ConfigType = Union[Config, ConfigActor, ConfigCritic, ConfigReward] + + +class Singleton: + """Singleton class to ensure only one instance of a class is created""" + + __instance = None + + def __new__(cls, *args, **kwargs): + if cls.__instance is None: + cls.__instance = super().__new__(cls) + return cls.__instance + + +class LogMessages(Singleton): + """ + Class to handle all logging messages + """ + + def __init__( + self, + ): + + set_global_logging_level() + + self.log_config = { + "handlers": [ + { + "sink": sys.stdout, + "format": ( + " {time:YY-MM-DD HH:mm:ss.S}" + " | " + "{message} " + ), + "colorize": True, + "level": "INFO", + }, + { + "sink": "file.log", + "serialize": True, + }, + ], + } + logger.configure(**self.log_config) + + # flag to enable multi gpu logging + self.is_multi_gpu = False + + # rank to print when multi gpu + self.log_rank = -1 + + # local rank + self.local_rank = -2 + + def setup_logger(self, accelerate_enable: bool, deepspeed_enable: bool): + """Setup logger for multi gpu training + + Args: + accelerate_enable (bool): flag to signal if using accelerate + deepspeed_enable (bool): flag to signal if using deepspeed + """ + if accelerate_enable or deepspeed_enable: + self.is_multi_gpu = True + self.log_rank = 0 + self.local_rank = int(os.environ.get("RANK", "0")) + else: + self.is_multi_gpu = False + + def error(self, error_type, text: str): + """Log error message + + Args: + error_type (Exception): type of error to raise + text (str): error message to log + """ + if self.is_multi_gpu: + if self.local_rank == self.log_rank: + logger.error(f"[Rank {self.local_rank}] {text}") + else: + logger.error(text) + return error_type("") + + def warning(self, text: str): + """Log warning message + + Args: + text (str): warning message to log + """ + if self.is_multi_gpu: + if self.local_rank == self.log_rank: + logger.warning(f"[Rank {self.local_rank}] {text}") + else: + logger.warning(text) + + def info(self, text: str): + """Log info message + + Args: + text (str): info message to log + """ + if self.is_multi_gpu: + if self.local_rank == self.log_rank: + logger.info(f"[Rank {self.local_rank}] {text}") + else: + logger.info(text) + + def success(self, text: str): + """Log success message + + Args: + text (str): success message to log + """ + if self.is_multi_gpu: + if self.local_rank == self.log_rank: + logger.success(f"[Rank {self.local_rank}] {text}") + else: + logger.success(text) + + def debug(self, text: str): + """Log debug message + + Args: + text (str): debug message to log + """ + + if self.is_multi_gpu: + if self.local_rank == self.log_rank: + logger.debug(f"[Rank {self.local_rank}] {text}") + else: + logger.debug(text) + + +# To control logging level for various modules used in the application: +def set_global_logging_level(level=logging.ERROR, prefices=[""]): + """ + Override logging levels of different modules based on their name as a + prefix. + It needs to be invoked after the modules have been loaded so that their + loggers have been initialized. + + Args: + - level: desired level. e.g. logging.INFO. Optional. + Default is logging.ERROR + - prefices: list of one or more str prefices to match + (e.g. ["transformers", "torch"]). Optional. Default is `[""]` + to match all active loggers. + The match is a case-sensitive `module_name.startswith(prefix)` + """ + prefix_re = re.compile(rf'^(?:{ "|".join(prefices) })') + for name in logging.root.manager.loggerDict: + if re.match(prefix_re, name): + logging.getLogger(name).setLevel(level) + + # force deepspeed + deepspeed.utils.logging.logger.setLevel(level) + + +# configure logger and logging level +my_logger = LogMessages() + + +@beartype +def get_multi_gpu_flags( + config: ConfigType, +) -> Tuple[bool, bool, Optional[str]]: + """Setup for multi-gpu training + + Args: + config (ConfigType): Config object + + Returns: + Tuple[bool, bool, Optional[str]]: Tuple of flags for multi-gpu training + """ + + # flags for training + if isinstance(config, Config): + accelerate_enable = config.trainer.accelerate_enable + deepspeed_enable = config.trainer.deepspeed_enable + deepspeed_config_path = config.trainer.deepspeed_config_path + else: + accelerate_enable = config.accelerate_enable + deepspeed_enable = config.deepspeed_enable + deepspeed_config_path = config.deepspeed_config_path + + # check consistency of flags + if accelerate_enable and deepspeed_enable: + raise my_logger.error( + ValueError, + ( + "Both DeepSpeed and Accelerate are enabled" + + "Please choose one of them." + ), + ) + + # check deepspeed config + if deepspeed_enable: + if deepspeed_config_path is None: + raise my_logger.error( + ValueError, + "DeepSpeed config path is None, but deepspeed is enabled", + ) + if os.path.exists(deepspeed_config_path) is False: + raise my_logger.error( + ValueError, + f"DeepSpeed config path" + f" {deepspeed_config_path} " + f"does not exist", + ) + + # setup the logger consequently + my_logger.setup_logger(accelerate_enable, deepspeed_enable) + + return accelerate_enable, deepspeed_enable, deepspeed_config_path + + +def load_tokenizer(config: ConfigType): + """Load the tokenizer from the model name + placed in utils to avoid circular imports from dataset and model class + + Args: + config (ConfigType): config object + """ + + # disable tokenizer parallelization (Avoid warnings in HF) + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + if config.model in hf_models: + + # load the tokenizer from HF + tokenizer = AutoTokenizer.from_pretrained( + config.model, + padding_side="left", + padding=True, + truncation=True, + model_max_length=config.max_sequence_length, + ) + + # add eos token if not present + if tokenizer.eos_token is None: + tokenizer.eos_token = "" + tokenizer.eos_token_id = 2 # OPT eos-token-id + + # add pad token if not present + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + tokenizer.pad_token_id = tokenizer.eos_token_id + + elif config.model in llama_models: + + if not isinstance(config, ConfigActor): + raise my_logger.error( + ValueError, + "LLaMA models can only be used as actor", + ) + + # llama module might not be present when HF models are used + from chatllama.llama_model import ( + load_tokenizer, + ) # noqa + + tokenizer = load_tokenizer(config.tokenizer_path) + return tokenizer class TrainingStats: @@ -144,11 +425,13 @@ def append( ) def save(self): - print("Saving conversations log") - if os.path.exists(self.path): - with open(self.path, "r") as f: - conversation = json.load(f) - self.conversation.extend(conversation) + """Save the conversation log""" + my_logger.info("Saving conversations log") + # load previous conversations - commented out + # if os.path.exists(self.path): + # with open(self.path, "r") as f: + # conversation = json.load(f) + # self.conversation.extend(conversation) self.conversation = sorted( self.conversation, key=lambda x: float(x["learn_counter"]) ) @@ -156,11 +439,13 @@ def save(self): json.dump(self.conversation, f, indent=4) def load(self): + """Load the conversation log""" with open(self.path, "r") as f: self.conversation = json.load(f) def clear(self): - print("Clearing conversations log") + """Clear the conversation log""" + my_logger.info("Clearing conversations log") self.conversation = [] # remove the file in path exists if os.path.exists(self.path): @@ -174,6 +459,8 @@ def show(self, current_iteration: int = None): if not None, print only the conversations that happened at """ + + # TODO: use logger to show messages... for i, c in enumerate(self.conversation): if current_iteration is None: print( @@ -196,3 +483,39 @@ def show(self, current_iteration: int = None): f"## Model Output:\n\n{c['model_output']}\n\n" f"## Reward: {c['reward']}\n\n" ) + + +class IgnoreLabelsWrapper(torch.nn.Module): + """Wrapper to ignore labels arguments when using lora models + with AutoModel HF models. + """ + + def __init__(self, base_model: torch.nn.Module): + super().__init__() + self.base_model = base_model + self.config = base_model.config + + def forward( + self, + input_ids=None, + attention_mask=None, + inputs_embeds=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + # Remove labels, which are unused by the base model + return self.base_model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + def prepare_inputs_for_generation(self, *args, **kwargs): + return self.model.prepare_inputs_for_generation(*args, **kwargs) diff --git a/apps/accelerate/chatllama/setup.py b/apps/accelerate/chatllama/setup.py index 6dba4d1e..99716209 100644 --- a/apps/accelerate/chatllama/setup.py +++ b/apps/accelerate/chatllama/setup.py @@ -3,6 +3,9 @@ REQUIREMENTS = [ + "bitsandbytes", + "loguru", + "peft", "accelerate", "beartype", "deepspeed", @@ -23,7 +26,7 @@ setup( name="chatllama-py", - version="0.0.4", + version="0.0.3", packages=find_packages(), install_requires=REQUIREMENTS, long_description=long_description,