refactor(runner): retire the examples launchers in favour of primus-cli

benchmark/megatron/checkpoint/ckpt_launch.py

@@ -135,7 +135,7 @@ def run_training_subprocess(shared_fs_path="."):
         with open(new_yaml_config_path, "w") as f:
             yaml.dump(yaml_config, f)
 
-        command = "bash examples/run_slurm_pretrain.sh"
+        command = "bash runner/helpers/launch/slurm_pretrain.sh"
         result = subprocess.run(command, shell=True, capture_output=True, text=True, env=env)
         logger.debug(f"training subprocess stdout : {result.stdout}")
         logger.debug(f"training subprocess stderr : {result.stderr}")

benchmark/megatron/model/benchmark_model.py

@@ -30,7 +30,7 @@ def run_benchmark(config_path, log_dir: Path, nnodes: int):
 
     with open(log_file, "w") as f:
         subprocess.run(
-            ["bash", "examples/run_slurm_pretrain.sh"],
+            ["bash", "runner/helpers/launch/slurm_pretrain.sh"],
             env=env,
             stdout=f,
             stderr=subprocess.STDOUT,

docs/02-user-guide/megatron-lm-training.md

@@ -579,7 +579,7 @@ PRIMUS_TRAIN_RUNTIME=legacy ./runner/primus-cli container \
 
 ### 3.2 Multi-node training
 
-To run training on multiple nodes, you can use `primus-cli` (recommended) or the [`run_slurm_pretrain.sh`](https://github.com/AMD-AGI/Primus/blob/main/examples/run_slurm_pretrain.sh) script to launch multi-node workloads. Below are the multi-node setup and examples to run multi-node tests.
+To run training on multiple nodes, you can use `primus-cli` (recommended) or the shared [`runner/helpers/launch/slurm_pretrain.sh`](../../runner/helpers/launch/slurm_pretrain.sh) helper, which turns the `EXP` / `NNODES` environment contract into the same `primus-cli slurm srun ... -- container -- train pretrain` invocation. Below are the multi-node setup and examples to run multi-node tests.
 
 **Multi-node setup**
 
@@ -619,62 +619,62 @@ Notes:
 
 ```bash
 # Adjust the training parameters. For example, `global_batch_size: 8 * #single_node_bs` for 8 nodes in this case
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml bash ./examples/run_slurm_pretrain.sh --global_batch_size 1024
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --global_batch_size 1024
 ```
 
 - **Llama2-7B FP8, 8 nodes:**
 
 ```bash
 # Adjust the training parameters. For example, `global_batch_size: 8 * #single_node_bs` for 8 nodes in this case
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml bash ./examples/run_slurm_pretrain.sh --global_batch_size 2048
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_7B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --global_batch_size 2048
 ```
 
 - **Llama3.1-70B FP8, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 4 --global_batch_size 256 --recompute_num_layers 80
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_70B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 4 --global_batch_size 256 --recompute_num_layers 80
 ```
 
 - **Llama3.1-70B BF16, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --recompute_num_layers 12
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.1_70B-BF16-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --recompute_num_layers 12
 ```
 
 - **Llama2-70B FP8, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 10 --global_batch_size 640 --recompute_num_layers 80
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_70B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 10 --global_batch_size 640 --recompute_num_layers 80
 ```
 
 - **Llama2-70B BF16, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml bash ./examples/run_slurm_pretrain.sh --micro_batch_size 2 --global_batch_size 1536 --recompute_num_layers 12
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama2_70B-BF16-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 2 --global_batch_size 1536 --recompute_num_layers 12
 ```
 
 - **Llama3.3-70B FP8, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.3_70B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 4 --global_batch_size 256 --recompute_num_layers 80
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.3_70B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 4 --global_batch_size 256 --recompute_num_layers 80
 ```
 
 - **Llama3.3-70B BF16, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.3_70B-BF16-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --recompute_num_layers 12
+NNODES=8 EXP=examples/megatron/configs/MI300X/llama3.3_70B-BF16-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --recompute_num_layers 12
 ```
 
 - **Mixtral 8x7B BF16, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 2 --global_batch_size 256
+NNODES=8 EXP=examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 2 --global_batch_size 256
 ```
 
 - **Qwen2.5-72B FP8, 8 nodes:**
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 8 --global_batch_size 512 --recompute_num_layers 80
+NNODES=8 EXP=examples/megatron/configs/MI300X/qwen2.5_72B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 8 --global_batch_size 512 --recompute_num_layers 80
 ```
 
 - **Mixtral-8x22B BF16, 4 nodes, MI355X**
@@ -686,10 +686,10 @@ Launch the training using `primus-cli` (recommended):
 ./runner/primus-cli slurm srun -N 4 -- train pretrain --config examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml --micro_batch_size 1 --global_batch_size 512 --num_virtual_stages_per_pipeline_rank 2 --pipeline_model_parallel_size 4 --expert_model_parallel_size 8 --recompute_num_layers 1 --moe_use_legacy_grouped_gemm True --gradient_accumulation_fusion True
 ```
 
-Launch the training using the legacy script:
+Or with the shared `EXP` / `NNODES` helper:
 
 ```bash
-NNODES=4 EXP=examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 512 --num_virtual_stages_per_pipeline_rank 2 --pipeline_model_parallel_size 4 --expert_model_parallel_size 8 --recompute_num_layers 1 --moe_use_legacy_grouped_gemm True --gradient_accumulation_fusion True
+NNODES=4 EXP=examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 512 --num_virtual_stages_per_pipeline_rank 2 --pipeline_model_parallel_size 4 --expert_model_parallel_size 8 --recompute_num_layers 1 --moe_use_legacy_grouped_gemm True --gradient_accumulation_fusion True
 ```
 
 - **Llama3.1-405B FP8, 8 nodes, MI325X**
@@ -703,10 +703,10 @@ Launch the training using `primus-cli` (recommended):
 
 We use TP=8 for the Llama 3.1 405B model on 8 nodes. Because it has 126 layers, which is not divisible by 8, you need to set `decoder_first_pipeline_num_layers` and `decoder_last_pipeline_num_layers`.
 
-Launch the training using the legacy script:
+Or with the shared `EXP` / `NNODES` helper:
 
 ```bash
-NNODES=8 EXP=examples/megatron/configs/MI325X/llama3.1_405B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --decoder_first_pipeline_num_layers 15 --decoder_last_pipeline_num_layers 15
+NNODES=8 EXP=examples/megatron/configs/MI325X/llama3.1_405B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --micro_batch_size 1 --global_batch_size 256 --decoder_first_pipeline_num_layers 15 --decoder_last_pipeline_num_layers 15
 ```
 
 ---

docs/02-user-guide/pretraining.md

@@ -257,21 +257,19 @@ Run on a JAX base image (for example `rocm/jax-training`) or a bare-metal JAX en
 
 ### Quick start (run from a bare Primus checkout)
 
-Use `run_pretrain.sh` with `BACKEND=MaxDiffusion`. When `PRIMUS_SKIP_PIP` is unset, the launcher runs `setup_maxdiffusion_env.sh` for you (installs the stack + applies the patches), sets `NVTE_FRAMEWORK=jax` and `MAXDIFFUSION_PATH`, then launches:
+Use `primus-cli direct` with `BACKEND=MaxDiffusion`. When `PRIMUS_SKIP_PIP` is unset, the prepare hooks run `setup_maxdiffusion_env.sh` for you (installs the stack + applies the patches), set `NVTE_FRAMEWORK=jax` and `MAXDIFFUSION_PATH`, then launch:
 
 ```bash
-BACKEND=MaxDiffusion \
-EXP=examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml \
-  bash ./examples/run_pretrain.sh
+BACKEND=MaxDiffusion ./primus-cli direct -- train pretrain \
+  --config examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml
 ```
 
 To run the environment setup once by itself (e.g. to warm an image or a shared venv), invoke the script directly, then launch with `PRIMUS_SKIP_PIP=1`:
 
 ```bash
 bash examples/maxdiffusion/setup_maxdiffusion_env.sh
-PRIMUS_SKIP_PIP=1 BACKEND=MaxDiffusion \
-EXP=examples/maxdiffusion/configs/MI355X/flux_dev-pretrain.yaml \
-  bash ./examples/run_pretrain.sh
+PRIMUS_SKIP_PIP=1 BACKEND=MaxDiffusion ./primus-cli direct -- train pretrain \
+  --config examples/maxdiffusion/configs/MI355X/flux_dev-pretrain.yaml
 ```
 
 ### Quick start (container mode)
@@ -350,7 +348,7 @@ The tables above in the Megatron, TorchTitan, and MaxText sections are curated M
 | TorchTitan | `examples/torchtitan/configs/MI300X/` | `parallelism.*` (e.g. `tensor_parallel_degree`, `pipeline_parallel_degree`, `expert_parallel_degree`, FSDP shard settings). |
 | MaxText | `examples/maxtext/configs/MI300X/` | `ici_fsdp_parallelism`, `ici_data_parallelism`, `dcn_fsdp_parallelism`, `dcn_data_parallelism`. |
 
-For scripting patterns that predate `primus-cli`, the repository still documents `examples/run_local_pretrain.sh` and `examples/run_slurm_pretrain.sh` in `examples/README.md`; equivalent launches are shown above using `./runner/primus-cli`.
+`./runner/primus-cli` is the only entry point. The packaged launchers under `examples/customer_package/` and `examples/moe_package/` reach it through the shared helper `runner/helpers/launch/slurm_pretrain.sh`, which translates their `EXP` / `NNODES` / `DATA_PATH` environment contract into a `primus-cli slurm` invocation; call the CLI directly as shown above.
 
 ---
 

docs/02-user-guide/torchtitan-training.md

@@ -261,10 +261,10 @@ Launch the training using `primus-cli` (recommended):
 ./runner/primus-cli slurm srun -N 4 -- train pretrain --config examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml --training.local_batch_size 6 --training.global_batch_size 192 --training.mock_data True
 ```
 
-Launch the training using the legacy script:
+Or with the shared `EXP` / `NNODES` helper:
 
 ```bash
-NNODES=4 EXP=examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --training.local_batch_size 6 --training.global_batch_size 192 --training.mock_data True
+NNODES=4 EXP=examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --training.local_batch_size 6 --training.global_batch_size 192 --training.mock_data True
 ```
 
 - **Llama3.1-405B FP8, 8 nodes, MI355X**
@@ -276,10 +276,10 @@ Launch the training using `primus-cli` (recommended):
 ./runner/primus-cli slurm srun -N 8 -- train pretrain --config examples/torchtitan/configs/MI355X/llama3.1_405B-FP8-pretrain.yaml --training.local_batch_size 3 --training.global_batch_size 192 --training.mock_data True
 ```
 
-Launch the training using the legacy script:
+Or with the shared `EXP` / `NNODES` helper:
 
 ```bash
-NNODES=8 EXP=examples/torchtitan/configs/MI355X/llama3.1_405B-FP8-pretrain.yaml bash examples/run_slurm_pretrain.sh --training.local_batch_size 3 --training.global_batch_size 192 --training.mock_data True
+NNODES=8 EXP=examples/torchtitan/configs/MI355X/llama3.1_405B-FP8-pretrain.yaml bash ./runner/helpers/launch/slurm_pretrain.sh --training.local_batch_size 3 --training.global_batch_size 192 --training.mock_data True
 ```
 
 ---