初始化项目，由ModelHub XC社区提供模型
Model: mnoukhov/pythia410m-sft-tldr Source: Original Platform
This commit is contained in:
ModelHub XC
193
code/run_on_mila.py
Normal file
193
code/run_on_mila.py
Normal file
@@ -0,0 +1,193 @@
|
||||
import argparse
|
||||
import datetime
|
||||
import os
|
||||
import subprocess
|
||||
from copy import deepcopy
|
||||
|
||||
import yaml
|
||||
from accelerate.commands import launch
|
||||
from generate_vllm import generate_relabel_args_dict
|
||||
|
||||
def run_exp(exp_dict, savedir, args):
|
||||
exp_name = exp_dict.pop("name")
|
||||
git_hash = exp_dict.pop("git")
|
||||
print(args)
|
||||
|
||||
if args.wandb:
|
||||
os.environ["WANDB_MODE"] = "online"
|
||||
# os.environ["WANDB_RUN_ID"] = os.path.basename(savedir)
|
||||
os.environ["WANDB_NAME"] = exp_name
|
||||
os.environ["WANDB_RUN_GROUP"] = exp_name + git_hash
|
||||
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("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("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_vllm_args_dict(exp_dict)
|
||||
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("--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['output_dir'] = args.savedir_base
|
||||
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.wandb:
|
||||
timenow = datetime.datetime.now().strftime("%d-%m-%y_%H-%M-%S")
|
||||
exp_list[0]["name"] = exp_list[0]["name"] + f"_local_{timenow}"
|
||||
|
||||
|
||||
# exp_list[0]["save_strategy"] = "no"
|
||||
|
||||
# Run experiments and create results file
|
||||
run_exp(exp_list[0], "output", args)
|
||||
|
||||
Reference in New Issue
Block a user