From 0e05a35abfa97d038ac7fa408fa05f95af33fdc0 Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Wed, 21 Aug 2024 18:35:25 -0700 Subject: [PATCH] Fixes Signed-off-by: Hemil Desai --- examples/llm/run/llama_pretraining.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/llm/run/llama_pretraining.py b/examples/llm/run/llama_pretraining.py index 8527d9d746c3..ef4d563b044e 100644 --- a/examples/llm/run/llama_pretraining.py +++ b/examples/llm/run/llama_pretraining.py @@ -3,6 +3,7 @@ from typing import Any, Optional import nemo_run as run + from nemo.collections import llm @@ -30,7 +31,7 @@ def get_parser(): parser.add_argument( "--local", action="store_true", - help="Do a dryrun and exit", + help="Run locally using run.LocalExecutor", default=False, ) return parser @@ -139,7 +140,8 @@ def main(): ckpt_dir=f"/{exp_name}/checkpoints", ) - # dataloader = + # TODO: Overwrite the dataloader in the recipe. + # dataloader = set_your_custom_dataloader # pretrain.data = dataloader pretrain.trainer.val_check_interval = 400 @@ -153,6 +155,7 @@ def main(): if args.local: executor = local_executor_torchrun(nodes=pretrain.trainer.num_nodes, devices=pretrain.trainer.devices) else: + # TODO: Set your custom parameters for the Slurm Executor. executor = slurm_executor( user="", host="",