270 lines
9.9 KiB
Python
270 lines
9.9 KiB
Python
|
import argparse
|
||
|
|
import datetime
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
from copy import deepcopy
|
||
|
|
|
||
|
|
import generate_and_eval
|
||
|
|
import yaml
|
||
|
|
from accelerate.commands import launch
|
||
|
|
from generate_vllm import generate_relabel_args_dict
|
||
|
|
from haven import haven_wizard as hw
|
||
|
|
|
||
|
|
|
||
|
|
def run_exp(exp_dict, savedir, args):
|
||
|
|
exp_name = exp_dict.pop("name")
|
||
|
|
git_hash = exp_dict.pop("git")
|
||
|
|
print(args)
|
||
|
|
print(f"savedir {savedir}")
|
||
|
|
|
||
|
|
exp_dict["output_dir"] = savedir
|
||
|
|
|
||
|
|
os.environ["WANDB_RUN_ID"] = os.path.basename(savedir)
|
||
|
|
os.environ["WANDB_NAME"] = exp_name
|
||
|
|
os.environ["WANDB_RUN_GROUP"] = exp_name + "_" + git_hash
|
||
|
|
|
||
|
|
if args.wandb:
|
||
|
|
os.environ["WANDB_MODE"] = "online"
|
||
|
|
os.environ["WANDB_PROJECT"] = "trl"
|
||
|
|
os.environ["WANDB_ENTITY"] = "mila-language-drift"
|
||
|
|
else:
|
||
|
|
os.environ["WANDB_MODE"] = "disabled"
|
||
|
|
|
||
|
|
if exp_name.startswith("marlhf"):
|
||
|
|
print("MARLHF")
|
||
|
|
accelerate_launch("rl_training_with_ma_value.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("vmrlhf"):
|
||
|
|
print("Separate Value Model RLHF")
|
||
|
|
accelerate_launch("rl_training_value_model.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("rlhf"):
|
||
|
|
print("RLHF")
|
||
|
|
accelerate_launch("rl_training.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("dpo"):
|
||
|
|
print("DPO")
|
||
|
|
accelerate_launch("dpo_training.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("newdpo"):
|
||
|
|
print("DPO")
|
||
|
|
accelerate_launch("dpo.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("rm"):
|
||
|
|
accelerate_launch("reward_modeling.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("gptrm"):
|
||
|
|
accelerate_launch("gpt_reward_modeling.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("sft"):
|
||
|
|
accelerate_launch("supervised_finetuning.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("newsft"):
|
||
|
|
accelerate_launch("sft.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("rouge"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
accelerate_launch("evaluate_rouge.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("pseudo"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
accelerate_launch("inference_pseudolabel.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("create_rlhf"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
accelerate_launch("create_rlhf_dataset.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("vllm"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
exp_dict["num_gpus"] = args.gpus
|
||
|
|
generate_relabel_args_dict(exp_dict)
|
||
|
|
elif exp_name.startswith("geneval"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
exp_dict["num_gpus"] = args.gpus
|
||
|
|
generate_and_eval.main_args_dict(exp_dict)
|
||
|
|
elif exp_name.startswith("scalarrm"):
|
||
|
|
exp_dict.pop("save_strategy", None)
|
||
|
|
accelerate_launch("scalar_rm_model.py", exp_dict, args)
|
||
|
|
elif exp_name.startswith("costa_dpo"):
|
||
|
|
accelerate_launch("costa_dpo.py", exp_dict, args)
|
||
|
|
else:
|
||
|
|
raise Exception(f"Config file {exp_name} does not start with one of the correct prefixes")
|
||
|
|
|
||
|
|
|
||
|
|
def accelerate_launch(training_file, training_args_dict, args):
|
||
|
|
parser = launch.launch_command_parser()
|
||
|
|
training_cmd_args = []
|
||
|
|
if args.accelerate_config is not None and args.accelerate_config != "None":
|
||
|
|
training_cmd_args.extend(["--config_file", args.accelerate_config])
|
||
|
|
# training_cmd_args.extend(["--num_processes", str(args.gpus)])
|
||
|
|
# training_cmd_args.extend(
|
||
|
|
# ["--gradient_accumulation_steps", str(training_args_dict["gradient_accumulation_steps"])]
|
||
|
|
# )
|
||
|
|
elif args.gpus > 1:
|
||
|
|
training_cmd_args.append("--multi_gpu")
|
||
|
|
|
||
|
|
# if training_args_dict.pop("fp16", False):
|
||
|
|
# mixed_precision = "fp16"
|
||
|
|
# elif training_args_dict.pop("bf16", False):
|
||
|
|
# mixed_precision = "bf16"
|
||
|
|
if training_args_dict.get("fp16", False):
|
||
|
|
mixed_precision = "fp16"
|
||
|
|
elif training_args_dict.get("bf16", False):
|
||
|
|
mixed_precision = "bf16"
|
||
|
|
else:
|
||
|
|
mixed_precision = "no"
|
||
|
|
training_cmd_args.extend(["--mixed_precision", mixed_precision])
|
||
|
|
#
|
||
|
|
|
||
|
|
training_cmd_args.extend(["--num_machines", "1"])
|
||
|
|
training_cmd_args.extend(["--num_processes", str(args.gpus)])
|
||
|
|
# if args.gpus > 1:
|
||
|
|
# if args.deepspeed is not None and args.deepspeed != "None":
|
||
|
|
# assert (
|
||
|
|
# "gradient_accumulation_steps" in training_args_dict
|
||
|
|
# ), "Must include gradient_accumulation_steps in config"
|
||
|
|
# training_cmd_args.append("--use_deepspeed")
|
||
|
|
# training_cmd_args.extend(["--zero_stage", str(args.deepspeed)])
|
||
|
|
# training_cmd_args.extend(
|
||
|
|
# ["--gradient_accumulation_steps", str(training_args_dict["gradient_accumulation_steps"])]
|
||
|
|
# )
|
||
|
|
|
||
|
|
training_cmd_args.append(training_file)
|
||
|
|
for key, val in training_args_dict.items():
|
||
|
|
training_cmd_args.append(f"--{key}")
|
||
|
|
if not (isinstance(val, bool) and val is True):
|
||
|
|
training_cmd_args.append(str(val))
|
||
|
|
|
||
|
|
print(" ".join(training_cmd_args))
|
||
|
|
args = parser.parse_args(training_cmd_args)
|
||
|
|
launch.launch_command(args)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
# Specify arguments regarding save directory and job scheduler
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument(
|
||
|
|
"-e",
|
||
|
|
"--exp_group",
|
||
|
|
help="Define the experiment group to run.",
|
||
|
|
nargs="+",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-sb",
|
||
|
|
"--savedir_base",
|
||
|
|
default="/home/toolkit/trl/results",
|
||
|
|
help="Define the base directory where the experiments will be saved.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-r",
|
||
|
|
"--reset",
|
||
|
|
type=int,
|
||
|
|
default=0,
|
||
|
|
help="If true, reset the experiment. Else, resume.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-j",
|
||
|
|
"--job_scheduler",
|
||
|
|
default=None,
|
||
|
|
type=str,
|
||
|
|
help="Run the experiments as jobs in the cluster.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"-p",
|
||
|
|
"--python_binary",
|
||
|
|
default="/home/toolkit/.conda/envs/trl/bin/python",
|
||
|
|
help="path to your python executable",
|
||
|
|
)
|
||
|
|
parser.add_argument("-n", "--gpus", default=1, type=int, help="number of gpus to use for experiment")
|
||
|
|
parser.add_argument("-a", "--accelerate_config", default=None, help="accelerate config")
|
||
|
|
# parser.add_argument("-d", "--deepspeed", default=None, help="ds stage")
|
||
|
|
parser.add_argument("--gpu-mem", default=32, type=int, help="mem of gpus to use for experiment")
|
||
|
|
parser.add_argument("--wandb", action="store_true", help="force enable wandb", default=False)
|
||
|
|
parser.add_argument("--local-save", action="store_true", help="force local save", default=False)
|
||
|
|
parser.add_argument("--search", default=None)
|
||
|
|
# parser.add_argument(
|
||
|
|
# "--exp-id", default=None, help="id used to resume an experiment"
|
||
|
|
# )
|
||
|
|
|
||
|
|
args, extra_args = parser.parse_known_args()
|
||
|
|
|
||
|
|
exp_list = []
|
||
|
|
for exp_file in args.exp_group:
|
||
|
|
with open(exp_file, "r") as fp:
|
||
|
|
exp_dict = yaml.safe_load(fp)
|
||
|
|
|
||
|
|
exp_dict["name"] = os.path.basename(exp_file)
|
||
|
|
exp_dict["git"] = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
|
||
|
|
|
||
|
|
if args.search is not None and args.search != "None":
|
||
|
|
search_key, search_val_str = args.search.split("=")
|
||
|
|
search_vals = search_val_str.split(",")
|
||
|
|
exps = []
|
||
|
|
for val in search_vals:
|
||
|
|
exp_dict_copy = deepcopy(exp_dict)
|
||
|
|
exp_dict_copy[search_key] = val
|
||
|
|
# exp_dict_copy["name"] = exp_dict_copy["name"] + f"/{search_key}={val}"
|
||
|
|
exps.append(exp_dict_copy)
|
||
|
|
# for key, val in vars(extra_args).items():
|
||
|
|
# exp_dict[key] = val
|
||
|
|
# print(exps)
|
||
|
|
else:
|
||
|
|
exps = [exp_dict]
|
||
|
|
|
||
|
|
exp_list.extend(exps)
|
||
|
|
|
||
|
|
args.exp_group = " ".join(args.exp_group)
|
||
|
|
print(args.exp_group)
|
||
|
|
|
||
|
|
if args.job_scheduler == "toolkit":
|
||
|
|
with open("/home/toolkit/wandb_api_key", "r") as f:
|
||
|
|
wandb_api_key = f.read().rstrip()
|
||
|
|
|
||
|
|
job_config = {
|
||
|
|
"account_id": os.environ["EAI_ACCOUNT_ID"],
|
||
|
|
# "image": "registry.console.elementai.com/snow.colab/cuda",
|
||
|
|
# "image": "registry.console.elementai.com/snow.colab_public/ssh",
|
||
|
|
# "image": "registry.console.elementai.com/snow.mnoukhov/rl4lms",
|
||
|
|
"image": "registry.console.elementai.com/snow.interactive_toolkit/default",
|
||
|
|
"data": [
|
||
|
|
"snow.mnoukhov.home:/home/toolkit",
|
||
|
|
"snow.colab.public:/mnt/public",
|
||
|
|
],
|
||
|
|
"environment_vars": [
|
||
|
|
"HOME=/home/toolkit",
|
||
|
|
"HF_HOME=/home/toolkit/huggingface/",
|
||
|
|
f"WANDB_API_KEY={wandb_api_key}",
|
||
|
|
"WANDB_RESUME=allow",
|
||
|
|
"WANDB__SERVICE_WAIT=300",
|
||
|
|
"WANDB_PROJECT=trl",
|
||
|
|
"WANDB_ENTITY=mila-language-drift",
|
||
|
|
],
|
||
|
|
"restartable": True,
|
||
|
|
"resources": {
|
||
|
|
"cpu": 4 * args.gpus,
|
||
|
|
"mem": 64 * args.gpus,
|
||
|
|
"gpu_mem": args.gpu_mem,
|
||
|
|
"gpu": args.gpus,
|
||
|
|
},
|
||
|
|
"interactive": False,
|
||
|
|
"bid": 9999,
|
||
|
|
}
|
||
|
|
job_scheduler = "toolkit"
|
||
|
|
args.wandb = True
|
||
|
|
else:
|
||
|
|
job_config = None
|
||
|
|
job_scheduler = None
|
||
|
|
|
||
|
|
if args.wandb:
|
||
|
|
timenow = datetime.datetime.now().strftime("%d-%m-%y_%H-%M-%S")
|
||
|
|
exp_list[0]["name"] = exp_list[0]["name"] + f"_local_{timenow}"
|
||
|
|
|
||
|
|
if not args.local_save:
|
||
|
|
exp_list[0]["save_strategy"] = "no"
|
||
|
|
|
||
|
|
# Run experiments and create results file
|
||
|
|
# if job_scheduler == "toolkit":
|
||
|
|
from haven import haven_wizard as hw
|
||
|
|
|
||
|
|
hw.run_wizard(
|
||
|
|
func=run_exp,
|
||
|
|
exp_list=exp_list,
|
||
|
|
savedir_base=args.savedir_base,
|
||
|
|
reset=args.reset,
|
||
|
|
job_config=job_config,
|
||
|
|
job_scheduler=job_scheduler,
|
||
|
|
results_fname="results/notebook.ipynb",
|
||
|
|
python_binary_path=args.python_binary,
|
||
|
|
args=args,
|
||
|
|
use_threads=True,
|
||
|
|
save_logs=False,
|
||
|
|
)
|