Constitutional AI recipe (#108)

* cai

* add training configuration

* update readme

* Update recipes/cai/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update recipes/cai/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update recipes/cai/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update recipes/cai/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* Update recipes/cai/README.md

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

* rename

* update

* rename

* Quick change

---------

Co-authored-by: lewtun <lewis.c.tunstall@gmail.com>

recipes/constitutional-ai/README.md

@@ -0,0 +1,24 @@
+# Constitutional AI 
+
+This repo includes the recipe for training the following models:
+
+* https://huggingface.co/HuggingFaceH4/mistral-7b-anthropic
+* https://huggingface.co/HuggingFaceH4/mistral-7b-grok
+
+
+## Full training examples
+
+You will require 8 GPUs (80GB of VRAM) to train the full model.
+```shell
+# Step 1 - SFT
+ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_sft.py recipes/constitutional-ai/sft/config_{grok,anthropic}.yaml
+
+# Step 2 - DPO
+ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_dpo.py recipes/constitutional-ai/dpo/config_anthropic.yaml
+# Note that we did not include the DPO recipe for grok, as that model's seems overtrained and too snarky.
+```
+
+
+## Advanced: generating you own dataset
+
+To generate the constitutional AI dataset, see https://github.com/huggingface/llm-swarm/tree/main/examples/constitutional-ai for detailed instructions if you want build or customize the dataset. 

recipes/constitutional-ai/dpo/config_anthropic.yaml

@@ -0,0 +1,41 @@
+# Model arguments
+model_name_or_path: alignment-handbook/mistral-7b-sft-constitutional-ai
+torch_dtype: null
+
+# Data training arguments
+# For definitions, see: src/h4/training/config.py
+dataset_mixer:
+  HuggingFaceH4/ultrafeedback_binarized: 1.0
+  HuggingFaceH4/cai-conversation-harmless: 1.0
+dataset_splits:
+- train_prefs
+- test_prefs
+preprocessing_num_workers: 12
+
+# DPOTrainer arguments
+bf16: true
+beta: 0.1
+do_eval: true
+do_train: true
+evaluation_strategy: steps
+eval_steps: 1000
+gradient_accumulation_steps: 1
+gradient_checkpointing: true
+hub_model_id: mistral-7b-dpo-constitutional-ai
+learning_rate: 5.0e-7
+log_level: info
+logging_steps: 10
+lr_scheduler_type: linear
+max_length: 1024
+max_prompt_length: 512
+num_train_epochs: 3
+optim: rmsprop
+output_dir: data/mistral-7b-dpo-constitutional-ai
+per_device_train_batch_size: 2
+per_device_eval_batch_size: 8
+push_to_hub: true
+save_strategy: "steps"
+save_steps: 100
+save_total_limit: 1
+seed: 42
+warmup_ratio: 0.1

recipes/constitutional-ai/sft/config_anthropic.yaml

@@ -0,0 +1,48 @@
+# Model arguments
+model_name_or_path: mistralai/Mistral-7B-v0.1
+model_revision: main
+torch_dtype: bfloat16
+use_flash_attention_2: true
+
+# Data training arguments
+chat_template: "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
+dataset_mixer:
+  HuggingFaceH4/cai-conversation-harmless: 1.0
+  HuggingFaceH4/ultrachat_200k: 1.0
+dataset_splits:
+- train_sft
+- test_sft
+preprocessing_num_workers: 12
+
+# SFT trainer config
+bf16: true
+do_eval: true
+do_train: true
+evaluation_strategy: epoch # One of ["no", "steps", "epoch"]
+gradient_accumulation_steps: 4
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: False
+hub_model_id: mistral-7b-sft-constitutional-ai
+hub_strategy: every_save
+learning_rate: 2.0e-05
+log_level: info
+logging_steps: 5  
+logging_strategy: steps
+lr_scheduler_type: cosine
+max_seq_length: 2048
+max_steps: -1
+num_train_epochs: 1
+output_dir: data/mistral-7b-sft-constitutional-ai
+overwrite_output_dir: true
+per_device_eval_batch_size: 8
+per_device_train_batch_size: 8
+push_to_hub: true
+remove_unused_columns: true
+report_to:
+- tensorboard
+save_strategy: "steps"
+save_steps: 100
+save_total_limit: 1
+seed: 42
+warmup_ratio: 0.1

recipes/constitutional-ai/sft/config_grok.yaml

@@ -0,0 +1,48 @@
+# Model arguments
+model_name_or_path: mistralai/Mistral-7B-v0.1
+model_revision: main
+torch_dtype: bfloat16
+use_flash_attention_2: true
+
+# Data training arguments
+chat_template: "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
+dataset_mixer:
+  HuggingFaceH4/grok-conversation-harmless: 0.15
+  HuggingFaceH4/ultrachat_200k: 1.0
+dataset_splits:
+- train_sft
+- test_sft
+preprocessing_num_workers: 12
+
+# SFT trainer config
+bf16: true
+do_eval: true
+do_train: true
+evaluation_strategy: epoch # One of ["no", "steps", "epoch"]
+gradient_accumulation_steps: 4
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: False
+hub_model_id: mistral-7b-sft-constitutional-ai
+hub_strategy: every_save
+learning_rate: 2.0e-05
+log_level: info
+logging_steps: 5  
+logging_strategy: steps
+lr_scheduler_type: cosine
+max_seq_length: 2048
+max_steps: -1
+num_train_epochs: 1
+output_dir: data/mistral-7b-sft-constitutional-ai
+overwrite_output_dir: true
+per_device_eval_batch_size: 8
+per_device_train_batch_size: 8
+push_to_hub: true
+remove_unused_columns: true
+report_to:
+- tensorboard
+save_strategy: "steps"
+save_steps: 100
+save_total_limit: 1
+seed: 42
+warmup_ratio: 0.1