• config.yaml: config file for model and data set. This allows you to 1) select the model you prefer (LLaMA, OPT, BLOOM, etc) 2) change all the hyperparameters of the training process;
• ds_config.json: config file to define DeepSpeed training parameters;
• templates.json: synthetic data generation templates that can be used to personalize the creation of the dataset. The templates are used for feeding LLMs during the data generation. Note that the templates.json file contains a dictionary having as keys the training steps (actor, reward, rlhf) and as values a string containing the personalization requests of the user. For more details see the dataset preparation section;
• main.py: file to train the model.

wget -O artifacts.zip https://nbllabartifacts.blob.core.windows.net/chatllama/artifacts.zip\?sp\=r\&st\=2023-03-08T14:53:24Z\&se\=2100-03-08T22:53:24Z\&spr\=https\&sv\=2021-06-08\&sr\=b\&sig\=jqr%2B2ZkR0SW9RjV0pDOdQ%2BDulLXLjbZ36vmNd4XxxyQ%3D
unzip artifacts.zip 

Once you have run the command above, you will find the all artificats in the artifacts/ directory. Now you can move on to the next section regarding the dataset preparation.

Before training the model, we need to prepare 3 datasets:

• actor_training_data: this is the JSON dataset used in the supervised fine-tuning. It consists of examples of unlabelled conversations, e.g. collection of prompts and responses;
• rlhf_training_data: this is the JSON dataset used for RLHF training. It consists of a collection of possible input user prompts;
• reward_training_data: this is the JSON dataset used to train the reward model. It consists of responses with associated scores.

In this example, we are using only publicly available dataset and synthetic generation; if you want to use your own data instead, please see the Dataset preparation section.

First, let’s download the actor_training_data and the rlhf_training_data:

python artifacts/download_dataset.py ARLHF --path ./datasets --number_of_samples 200

Finally, let’s create the reward_training_data using davinci-003 for synthetic data generation.

export OPENAI_API_KEY=YOUR_API_KEY
python artifacts/generate_rewards.py ./datasets/reward_training_data.json

> ⚠️ Creating the reward_training_data with davinci-003 is not free, i.e. it costs a few $$. If you prefer avoiding external paid APIs, we suggest using HuggingFace’s models (e.g. flan_t5_xl) as described in more detail in the Supported models section. > > ⚠️ if using OpenAI's API, please be aware of OpenAI's terms of use stating that it is forbidden to "use the Services to develop foundation models or other large scale models that compete with OpenAI".

At this point, we have successfully created the 3 datasets. We can therefore move on to the final section and start the training.

You can train the 3 models in separate steps:

• Train the Reward Model

  python artifacts/main.py artifacts/config/config.yaml --type REWARD
  

• Pre-Train the Actor Model

  python artifacts/main.py artifacts/config/config.yaml --type ACTOR
  

• Training the Actor with reinforcement learning.

  python artifacts/main.py artifacts/config/config.yaml --type RL
  

or, equivantly, the 3 trainings can also be pipelined using the flag ALL.

python artifacts/main.py artifacts/config/config.yaml --type ALL

Note that the path to the datasets and the training hyper-parameters of the training process are specified in the config.yaml file.

We support models that can be run efficiently with a limited amount of compute, such as LLaMA and 🤗 transformers. These are the models with less than 20B parameters currently supported :

• LLaMA: 7B and 13B, please note this library does NOT contain LLaMA’s weights; to access the weights, you need to apply to Meta's form.
• GPTJ: 6B
• GPTNeoX: 1.3B, 20B
• (⚠️WIP) Flan-T5: 80M, 259M, 780M, 3B, 11B
• OPT: 125M, 359M, 1.3B, 2.7B, 6.7B, 13B
• BLOOM: 560M, 1.1B, 1.7B, 3B, 7.1B
• BLOOMZ: 560M, 1.1B, 1.7B, 3B, 7.1B
• Galactica: 125M, 1.3B, 6.7B

We suggest using models under 6B from 🤗 transformers:

• GPT2: 124M, 355M, 774M, 1.5B
• OPT: 125M, 359M, 1.3B, 2.7B
• GPTJ: 6B
• BLOOMZ: 560M, 1.1B, 1.7B, 3B
• (⚠️WIP) OpenAssistant pre-trained reward models

We support both APIs from OpenAI and 🤗 transformers:

• OpenAI: da-vinci-003, gpt-3.5-turbo (⚠️WIP)
• HuggingFace: Flan-T5 (3B and 11B)

> ⚠️ if using OpenAI's API, please be aware of OpenAI's terms of use stating that it is forbidden to "use the Services to develop foundation models or other large scale models that compete with OpenAI".

:watninh

If you need support for different models, please open an issue and we will get to work.

Larger actor models require more powerful hardware. Here is a rough hardware recommendation table, suggesting the right type of hardware for different actor model sizes:

• 125M to 1.3B → 1x Nvidia 3090/4090
• 1.3B to 3B → 1x Nvidia A100 (80Gb)
• 3B with DeepSpeed CPU off-loading → 1x Nvidia 3090/4090
• 3B to 7B with DeepSpeed ZeRO → 4x Nvidia T4
• 3B to 13B → 4x Nvidia A100 (80Gb)
• 13B to 20B with DeepSpeed ZeRO → 4x Nvidia A100 (80Gb)
• 13B to 20B → 8x Nvidia A100 (80Gb)

(⚠️WIP) When it comes to inference optimization, ChatLLaMA will support the following optimization techniques:

• [ ] DeepSpeed ZeRO
• [ ] FlexGen
• [ ] HF Accelerate
• [ ] PyTorch Vanilla

The actor_training_data is a collection of prompts with the associated responses as highlighted below:

[
  {
      "user_input": "here the input of the user",
      "completion": "here the model completion"
  }
]

ChatLLaMA supports 4 different options to prepare the actor_training_data:

• Use 100% synthetic data

  The dataset can be synthetically generated by running the following command:

  python artifacts/generate_actor_dataset.py
  

  > ⚠️ Note that this command will require a subscription to OpenAI. Generating the full dataset with davinci-003 could cost approximately ~200$. > > ⚠️ if using OpenAI's API, please be aware of OpenAI's terms of use stating that it is forbidden to "use the Services to develop foundation models or other large scale models that compete with OpenAI".

  Alternatively, you can generate the dataset for free using 🤗 tranformers as described in the section Supported models.

• Use one of the open source datasets with assistant interactions

  Currently, we support:

  • Anthropic HH RLHF: this dataset consists of structured question/answer pairs with an LLM chatbot that includes selected and rejected answers;
  • Stanford Human Preferences Dataset (SHP): this dataset is curated from selected "ask" subreddits, and includes questions that span a wide range of question/answer pairs based on the most upvoted responses. Please note that, unlike HH RLHF, this dataset is not intended to reduce harassment by selecting the ideal chatbot response, but instead weights the most helpful human responses.

  The datasets can be downloaded running the following command:

  python artifacts/download_dataset.py  --path  --number_of_samples 
  

  Where:

  • `` could be "SHP" or "ARLHF" for the StanfordNLP/SHP dataset or ARLHF for the Anthropic/hh-rlhf dataset respectively;
  • `` is the folder path to where the datasets are going to be created;
  • `` is the number of samples of which the reward_dataset.json is composed.
• Use 100% personalized dataset

  The user provides his own personalized full dataset. Datasets must be JSON files with the following format:

  [
      {
          "user_input": "here the input of the user",
          "completion": "here the model completion"
      }
  ]
  

  Where the list contains multiple dictionaries, and each dictionary corresponds to a data sample. We suggest using more than 1000 data samples to run the actor training.

• (⚠️WIP) Create the full dataset augmenting few custom data samples

  The dataset can be generated synthetically from a few prompt+response examples provided by the user (few =>10).

The dataset for RLHF consists just of prompt examples:

[
    {
        "user_input": "here the example of user input"
    }
]

It can be provided in 2 different ways:

• Few examples provided by the user and dataset synthetically expanded using LLM

  You need to add the key rlhf to the templates.json file with the information about the task you want to perform and extra context needed by the LLM for the generation. Here is an example of template:

  {
    "rlhf": "Here is the template for the generating RLHF prompts. The task we want to perform is ..."
  }
  

  Note that all templates must be saved in a single JSON file named templates.json

• The user provides the full dataset with possible interactions with the model

  The dataset needs to contain more than 1000 prompt examples:

  [
      {
          "user_input": "here the example of user input"
      }
  ]
  

  The file must be named rlhf_training_data.json.

The reward_training_data is a collection of i) prompts, ii) completion and iii) score of the completion assigned accordingly to the user feedback (the Human Feedback in RLHF).

[{
	"user_input": "...",
	"completion": "...",
	"score": 1
},
	...
]

We support 3 different options to prepare the reward_training_data:

• Fully Synthetic Score Generation

  In this case the reward dataset can be synthetically scored using a LLM as Human Feedback. We recommend the reward_training_data having at least 100 data samples.

  [{
  	"user_input": "...",
  	"completion": "...",
  	"score": None
  },
  	...
  ]
  

  A LLM model is used to assign the score to each entry.

  The LLM needs a prompt template containing all the instructions to evaluate the generated text. To do this, you should add the key reward to the templates.json file. Here is an example:

  {
  	"reward": "Here is the template for the reward model. The rules are:\n\n1.Rule 1\n\n2. Rule 2"
  }
  

  If no template is provided the default one is used. You can find the default template in artifacts/generate_rewards.py. Note that all templates must be saved in a single JSON file named templates.json.

  Once you have the unlabelled dataset, you can generate the scores by running the following command:

  python artifacts/generate_rewards.py  --model  --temperature  --max_tokens  --reward_template 
  

  Where:

  • `` path to the reward dataset to be scored;
  • `` model to use for the reward. Default and suggested text-davinci-003 (More to come);
  • `` temperature used to score the model; temperature=0.1;
  • `` max_tokens of the generation;
  • `` is the path to the templates.json file containing the template to be used for generating the reward. If no path is provided, the default template will be used.

• The user provides their personalized full dataset

  Datasets must be JSON files in the following format:

  [
      {
          "user_input": "here type the user input",
          "completion": "here type the completion",
          "score": 4.0
      },
      {
          "user_input": "here type the user input",
          "completion": "random garbage",
          "score": 0.0
      }
  ]
  

  Note that at least 100 data samples are required in this case. The file must be named reward_training_data.json

• (⚠️WIP) Few examples provided by the user and dataset synthetically expanded using LLM