wassname/alignment-handbook
1
0
Fork 0
mirror of https://github.com/wassname/alignment-handbook.git synced 2026-08-11 11:14:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update Zephyr configs to account for UltraFeedback & TRL fixes (#88)

Browse Source 
* Add files

* Add checkpointing

* Add checkpointing to SFT

* Add loss type

* Fix setup|

* Clean SFT

* Add lora config

* Rename config

* Remove max eval samples

* Add kwargs tp push to hub

* Add DPO configs

* Fix dpo configs

* Extend chat template test to multi-turn

* Add warmup

* Refactor

* Fix LoRA -> QLoRA

* Fix configs

* Specify chat template

* Add sample logging

* Fix push to hub hanging

* Add reentrant

* Fix quality

* Add transformer logging

* Tweak grad acc

* Add null type

* Add doc
This commit is contained in:
lewtun
2024-01-10 17:42:24 +11:00
committed by GitHub
parent c69ae4b8a5
commit f0ffa0d7a6
17 changed files with 266 additions and 187 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/README.md
+9 -7
View File
@@ -9,20 +9,20 @@ In the handbook, we provide three main ways to align LLMs for chat:
- LoRA or QLoRA fine-tuning on a single consumer 24GB GPU (tested on an RTX 4090).
- LoRA fine-tuning on a multi-GPU machine with DeepSpeed ZeRO-3 (tested on a 2 x A100s (80GB)).
In practice, we find comparable performance for both full and LoRA fine-tuning, with the latter having the advantage of producing small adapter weights that are fast to upload and download from the Hugging Face Hub. Here are the general commands to fine-tune your models:
In practice, we find comparable performance for both full and QLoRA fine-tuning, with the latter having the advantage of producing small adapter weights that are fast to upload and download from the Hugging Face Hub. Here are the general commands to fine-tune your models:
```shell
# Full training with ZeRO-3 on 8 GPUs
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_{task}.py recipes/{model_name}/{task}/config_full.yaml
# LoRA training on a single GPU
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml
# QLoRA 4-bit training on a single GPU
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml --load_in_4bit=true
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_qlora.yaml
# LoRA training on a single GPU
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/multi_gpu.yaml --num_processes=1 scripts/run_{task}.py recipes/{model_name}/{task}/config_qlora.yaml --load_in_4bit=false
# LoRA training with ZeRO-3 on two or more GPUs
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml --num_processes={num_gpus} scripts/run_{task}.py recipes/{model_name}/{task}/config_lora.yaml
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml --num_processes={num_gpus} scripts/run_{task}.py recipes/{model_name}/{task}/config_qlora.yaml --load_in_4bit=false
```
Here `{task}` refers to the type of training you wish to run (SFT, DPO, etc), while `{model_name}` refers to the choice of a recipe in the `recipes` directory. For example, to replicate Zephyr-7B-β you can run:
@@ -44,6 +44,8 @@ By default, these scripts will push each model to your Hugging Face Hub username
ACCELERATE_LOG_LEVEL=info accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml scripts/run_{task}.py recipes/{model_name}/{task}/config_full.yaml --per_device_train_batch_size=42 --num_train_epochs=5
```
## Logging with Weights and Biases
By default all training metrics are logged with TensorBoard. If you have a [Weights and Biases](https://wandb.ai/site) account and are logged in, you can view the training metrics by appending `--report_to=wandb`, e.g.
```shell
@@ -58,7 +60,7 @@ If you have access to a Slurm cluster, we provide a `recipes/launch.slurm` scrip
sbatch --job-name=handbook_{task} --nodes=1 recipes/launch.slurm {model_name} {task} {precision} {accelerator}
```
Here `{model_name}` and `{task}` are defined as above, while `{precision}` refers to the type of training (`full` vs `lora`) and `{accelerator}` refers to the choice of 🤗 Accelerate config in `recipes/accelerate_configs`. If you wish to override the default config parameters, you can provide them by appending a space-separated string like `'--arg1=value1 --arg2=value2'. Here's a concrete example to run SFT on 1 node of 8 GPUs:
Here `{model_name}` and `{task}` are defined as above, while `{precision}` refers to the type of training (`full` vs `qlora`) and `{accelerator}` refers to the choice of 🤗 Accelerate config in `recipes/accelerate_configs`. If you wish to override the default config parameters, you can provide them by appending a space-separated string like `'--arg1=value1 --arg2=value2'. Here's a concrete example to run SFT on 1 node of 8 GPUs:
```shell
# Launch on Slurm and override default hyperparameters
scripts/run_dpo.py
+49 -33
View File
@@ -14,19 +14,20 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
import random
import sys
import torch
import transformers
from transformers import AutoModelForCausalLM, set_seed
from accelerate import Accelerator
from alignment import (
DataArguments,
DPOConfig,
H4ArgumentParser,
ModelArguments,
apply_chat_template,
get_checkpoint,
get_datasets,
get_kbit_device_map,
get_peft_config,
@@ -64,12 +65,14 @@ def main():
logger.info(f"Data parameters {data_args}")
logger.info(f"Training/evaluation parameters {training_args}")
# Check for last checkpoint
last_checkpoint = get_checkpoint(training_args)
if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
logger.info(f"Checkpoint detected, resuming training at {last_checkpoint=}.")
# Set seed for reproducibility
set_seed(training_args.seed)
# Increase distributed timeout to 3h to enable push to Hub to complete
accelerator = Accelerator()
###############
# Load datasets
###############
@@ -102,6 +105,12 @@ def main():
{"text_prompt": "prompt", "text_chosen": "chosen", "text_rejected": "rejected"}
)
# Log a few random samples from the training set:
for index in random.sample(range(len(raw_datasets["train"])), 3):
logger.info(f"Prompt sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['prompt']}")
logger.info(f"Chosen sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['chosen']}")
logger.info(f"Rejected sample {index} of the raw training set:\n\n{raw_datasets['train'][index]['rejected']}")
torch_dtype = (
model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
)
@@ -118,10 +127,10 @@ def main():
)
model = model_args.model_name_or_path
if is_adapter_model(model, model_args.model_revision):
# load the model, merge the adapter weights and unload the adapter
# Note: to run QLora, you will need to merge the based model separately as the merged model in 16bit
logger.info(f"Merging peft adapters for {model_args.model_name_or_path=}")
if is_adapter_model(model, model_args.model_revision) is True:
# Load the base model, merge the adapter weights and unload the adapter
# Note: to run QLoRA, you will need to merge the base model separately as the merged model in 16bit
logger.info(f"Merging PEFT adapters for {model_args.model_name_or_path=}")
peft_config = PeftConfig.from_pretrained(model_args.model_name_or_path, revision=model_args.model_revision)
@@ -153,7 +162,7 @@ def main():
#########################
# Instantiate DPO trainer
#########################
dpo_trainer = DPOTrainer(
trainer = DPOTrainer(
model,
ref_model,
model_init_kwargs=model_kwargs,
@@ -166,17 +175,23 @@ def main():
max_length=training_args.max_length,
max_prompt_length=training_args.max_prompt_length,
peft_config=get_peft_config(model_args),
loss_type=training_args.loss_type,
)
###############
# Training loop
###############
train_result = dpo_trainer.train()
checkpoint = None
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
elif last_checkpoint is not None:
checkpoint = last_checkpoint
train_result = trainer.train(resume_from_checkpoint=checkpoint)
metrics = train_result.metrics
metrics["train_samples"] = len(raw_datasets["train"])
dpo_trainer.log_metrics("train", metrics)
dpo_trainer.save_metrics("train", metrics)
dpo_trainer.save_state()
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
logger.info("*** Training complete ***")
@@ -185,35 +200,36 @@ def main():
##########
if training_args.do_eval:
logger.info("*** Evaluate ***")
metrics = dpo_trainer.evaluate()
metrics = trainer.evaluate()
metrics["eval_samples"] = len(raw_datasets["test"])
dpo_trainer.log_metrics("eval", metrics)
dpo_trainer.save_metrics("eval", metrics)
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
##################################
# Save model and create model card
##################################
dpo_trainer.save_model(training_args.output_dir)
logger.info("*** Save model ***")
trainer.save_model(training_args.output_dir)
logger.info(f"Model saved to {training_args.output_dir}")
# Save everything else on main process
if accelerator.is_main_process:
kwargs = {
"finetuned_from": model_args.model_name_or_path,
"dataset": list(data_args.dataset_mixer.keys()),
"dataset_tags": list(data_args.dataset_mixer.keys()),
"tags": ["alignment-handbook"],
}
dpo_trainer.create_model_card(**kwargs)
kwargs = {
"finetuned_from": model_args.model_name_or_path,
"dataset": list(data_args.dataset_mixer.keys()),
"dataset_tags": list(data_args.dataset_mixer.keys()),
"tags": ["alignment-handbook"],
}
if trainer.accelerator.is_main_process:
trainer.create_model_card(**kwargs)
# Restore k,v cache for fast inference
dpo_trainer.model.config.use_cache = True
dpo_trainer.model.config.save_pretrained(training_args.output_dir)
if training_args.push_to_hub is True:
dpo_trainer.push_to_hub()
trainer.model.config.use_cache = True
trainer.model.config.save_pretrained(training_args.output_dir)
# Ensure we don't timeout on model save / push to Hub
logger.info("*** Waiting for all processes to finish ***")
accelerator.wait_for_everyone()
if training_args.push_to_hub is True:
logger.info("Pushing to hub...")
trainer.push_to_hub(**kwargs)
logger.info("*** Run complete! ***")
logger.info("*** Training complete! ***")
if __name__ == "__main__":
scripts/run_sft.py
+23 -15
View File
@@ -26,13 +26,13 @@ import torch
import transformers
from transformers import set_seed
from accelerate import Accelerator
from alignment import (
DataArguments,
H4ArgumentParser,
ModelArguments,
SFTConfig,
apply_chat_template,
get_checkpoint,
get_datasets,
get_kbit_device_map,
get_peft_config,
@@ -52,8 +52,6 @@ def main():
# Set seed for reproducibility
set_seed(training_args.seed)
accelerator = Accelerator()
###############
# Setup logging
###############
@@ -78,6 +76,11 @@ def main():
logger.info(f"Data parameters {data_args}")
logger.info(f"Training/evaluation parameters {training_args}")
# Check for last checkpoint
last_checkpoint = get_checkpoint(training_args)
if last_checkpoint is not None and training_args.resume_from_checkpoint is None:
logger.info(f"Checkpoint detected, resuming training at {last_checkpoint=}.")
###############
# Load datasets
###############
@@ -149,7 +152,12 @@ def main():
# Training loop
###############
logger.info("*** Train ***")
train_result = trainer.train()
checkpoint = None
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
elif last_checkpoint is not None:
checkpoint = last_checkpoint
train_result = trainer.train(resume_from_checkpoint=checkpoint)
metrics = train_result.metrics
metrics["train_samples"] = len(train_dataset)
trainer.log_metrics("train", metrics)
@@ -174,23 +182,23 @@ def main():
logger.info(f"Model saved to {training_args.output_dir}")
# Save everything else on main process
if accelerator.is_main_process:
kwargs = {
"finetuned_from": model_args.model_name_or_path,
"dataset": list(data_args.dataset_mixer.keys()),
"dataset_tags": list(data_args.dataset_mixer.keys()),
"tags": ["alignment-handbook"],
}
kwargs = {
"finetuned_from": model_args.model_name_or_path,
"dataset": list(data_args.dataset_mixer.keys()),
"dataset_tags": list(data_args.dataset_mixer.keys()),
"tags": ["alignment-handbook"],
}
if trainer.accelerator.is_main_process:
trainer.create_model_card(**kwargs)
# Restore k,v cache for fast inference
trainer.model.config.use_cache = True
trainer.model.config.save_pretrained(training_args.output_dir)
if training_args.push_to_hub is True:
logger.info("Pushing to hub...")
trainer.push_to_hub()
if training_args.push_to_hub is True:
logger.info("Pushing to hub...")
trainer.push_to_hub(**kwargs)
accelerator.wait_for_everyone()
logger.info("*** Training complete ***")
if __name__ == "__main__":