Nemo Run

This workload (test_template_name is NemoRun) executes NeMo training and fine-tuning tasks using the NeMo Run framework.

Test TOML example:

            

            
name = "my_nemo_test"
description = "Example NeMo Run test"
test_template_name = "NemoRun"

[cmd_args]
recipe = "llama3_8b"
task = "pretrain"
        

Test Scenario example:

            

            
name = "nemo-run-test"

[[Tests]]
id = "nemo.1"
num_nodes = 4
time_limit = "02:00:00"

test_name = "my_nemo_test"
        

Test-in-Scenario example:

            

            
name = "nemo-run-test"

[[Tests]]
id = "nemo.1"
num_nodes = 4
time_limit = "02:00:00"

name = "my_nemo_test"
description = "Example NeMo Run test"
test_template_name = "NemoRun"

  [Tests.cmd_args]
  recipe = "llama3_8b"
  task = "pretrain"
        

Command Arguments

class cloudai.workloads.nemo_run.nemo_run.NeMoRunCmdArgs(*, docker_image_url: str, task: str, recipe_name: str, num_layers: int | None = None, trainer: ~cloudai.workloads.nemo_run.nemo_run.Trainer = <factory>, log: ~cloudai.workloads.nemo_run.nemo_run.Log = <factory>, data: ~cloudai.workloads.nemo_run.nemo_run.Data = <factory>, optim: ~cloudai.workloads.nemo_run.nemo_run.Optim = <factory>, **extra_data: ~typing.Any)[source]

Bases: CmdArgs

NeMoRun test command arguments.

docker_image_url: str

task: str

recipe_name: str

num_layers: int | None

trainer: Trainer

log: Log

data: Data

optim: Optim

Test Definition

class cloudai.workloads.nemo_run.nemo_run.NeMoRunTestDefinition(*, name: str, description: str, test_template_name: str, cmd_args: NeMoRunCmdArgs, extra_env_vars: dict[str, str | List[str]] = {}, extra_cmd_args: dict[str, str] = {}, extra_container_mounts: list[str] = [], git_repos: list[GitRepo] = [], nsys: NsysConfiguration | None = None, predictor: PredictorConfig | None = None, agent: str = 'grid_search', agent_steps: int = 1, agent_metrics: list[str] = ['default'], agent_reward_function: str = 'inverse', script: File = File(src=PosixPath('/home/runner/work/cloudai/cloudai/src/cloudai/workloads/nemo_run/cloudai_nemorun.py')))[source]

Bases: TestDefinition

NeMoRun test definition.

cmd_args: NeMoRunCmdArgs

script: File

property docker_image: DockerImage

property installables: list[Installable]

Get list of installable objects.

constraint_check(tr: TestRun) → bool[source]

Check constraints for NeMoRun.

property update_num_train_samples: int | None

Calculate num_train_samples based on global_batch_size and max_steps.

was_run_successful(tr: TestRun) → JobStatusResult[source]