Refactor of models and trainers with base class for common methods

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
 </details>
 
+## 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`. <br />
+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 <type_of_training>
+```
+or
+```bash
+accelerate launch artifacts/main.py artifacts/config/config.yaml --type <type_of_training>
+```
 # License
 
 See the [LICENSE](https://github.com/nebuly-ai/nebullvm/blob/main/apps/accelerate/chatllama/LICENSE) file.

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"

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
     }
   }

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,
+        )
+

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":