Skip to content

Switch trtllm slurm scripts to use mpirun directly for better stability #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Kipok merged 2 commits into from
Feb 21, 2024
Merged

Switch trtllm slurm scripts to use mpirun directly for better stability #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d67f684
Switch to direct mpirun for trtllm server
Kipok Feb 20, 2024
ec68b53
Correct trt parameters
Kipok Feb 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions pipeline/launcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,16 +40,23 @@ def fill_env_vars(format_dict, env_vars):


def get_server_command(server_type, num_gpus):
num_tasks = num_gpus
if server_type == 'nemo':
server_start_cmd = (
f"(python /code/nemo_skills/inference/server/serve_nemo.py gpt_model_file=/model trainer.devices={num_gpus} "
f"tensor_model_parallel_size={num_gpus} > /tmp/server_logs.txt &)"
)
# somehow on slurm nemo needs multiple tasks, but locally only 1
if CLUSTER_CONFIG["cluster"] == "local":
num_tasks = 1
else:
server_start_cmd = (
"(python /code/nemo_skills/inference/server/serve_trt.py --model_path /model > /tmp/server_logs.txt &)"
f"(mpirun -np {num_gpus} --allow-run-as-root --oversubscribe python /code/nemo_skills/inference/server/serve_trt.py "
"--model_path /model > /tmp/server_logs.txt &)"
)
return server_start_cmd
num_tasks = 1 # we launch via mpirun directly

return server_start_cmd, num_tasks


SLURM_HEADER = """
Expand Down
7 changes: 1 addition & 6 deletions pipeline/run_eval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ def get_sampling_cmd(benchmark, random_seed, extra_arguments=""):
args.model_path = Path(args.model_path).absolute()
args.output_dir = Path(args.output_dir).absolute()

server_start_cmd = get_server_command(args.server_type, args.num_gpus)
server_start_cmd, num_tasks = get_server_command(args.server_type, args.num_gpus)

format_dict = {
"model_path": args.model_path,
Expand Down Expand Up @@ -131,11 +131,6 @@ def get_sampling_cmd(benchmark, random_seed, extra_arguments=""):
# splitting eval cmds equally across num_nodes nodes
eval_cmds = [" ".join(eval_cmds[i :: args.num_nodes]) for i in range(args.num_nodes)]

num_tasks = format_dict["num_gpus"]
# somehow on slurm nemo needs multiple tasks, but locally only 1
if args.server_type == "nemo" and CLUSTER_CONFIG["cluster"] == "local":
num_tasks = 1

for idx, eval_cmd in enumerate(eval_cmds):
extra_sbatch_args = ["--parsable", f"--output={args.output_dir}/slurm_logs_eval{idx}.log"]
launch_job(
Expand Down
10 changes: 3 additions & 7 deletions pipeline/run_labeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,15 +67,10 @@ def run_script(format_dict, seed, extra_arguments, partition=None, dependency=No
if dependency is not None:
extra_sbatch_args.append(f"--dependency=afterany:{dependency}")

num_tasks = format_dict["num_gpus"]
# somehow on slurm nemo needs multiple tasks, but locally only 1
if format_dict["server_type"] == "nemo" and CLUSTER_CONFIG["cluster"] == "local":
num_tasks = 1

job_id = launch_job(
cmd=SLURM_CMD.format(**format_dict),
num_nodes=1,
tasks_per_node=num_tasks,
tasks_per_node=format_dict["num_tasks"],
gpus_per_node=format_dict["num_gpus"],
job_name=JOB_NAME.format(**format_dict),
container=CLUSTER_CONFIG["containers"][format_dict["server_type"]],
Expand Down Expand Up @@ -116,14 +111,15 @@ def run_script(format_dict, seed, extra_arguments, partition=None, dependency=No

extra_arguments = f'{" ".join(unknown)}'

server_start_cmd = get_server_command(args.server_type, args.num_gpus)
server_start_cmd, num_tasks = get_server_command(args.server_type, args.num_gpus)

format_dict = {
"model_path": args.model_path,
"model_name": args.model_path.name,
"output_dir": args.output_dir,
"num_gpus": args.num_gpus,
"server_start_cmd": server_start_cmd,
"num_tasks": num_tasks,
"server_type": args.server_type,
}
fill_env_vars(format_dict, ["NEMO_SKILLS_CODE"])
Expand Down
7 changes: 1 addition & 6 deletions pipeline/start_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@

args.model_path = Path(args.model_path).absolute()

server_start_cmd = get_server_command(args.server_type, args.num_gpus)
server_start_cmd, num_tasks = get_server_command(args.server_type, args.num_gpus)

format_dict = {
"model_path": args.model_path,
Expand All @@ -75,11 +75,6 @@
}
fill_env_vars(format_dict, ["NEMO_SKILLS_CODE"])

num_tasks = format_dict["num_gpus"]
# somehow on slurm nemo needs multiple tasks, but locally only 1
if args.server_type == "nemo" and CLUSTER_CONFIG["cluster"] == "local":
num_tasks = 1

job_id = launch_job(
cmd=SLURM_CMD.format(**format_dict),
num_nodes=1,
Expand Down