Skip to content

Refactor training launchers around primus-cli - #976

Open
WangLingxun wants to merge 4 commits into
mainfrom
refactor/primus-cli-launchers
Open

Refactor training launchers around primus-cli#976
WangLingxun wants to merge 4 commits into
mainfrom
refactor/primus-cli-launchers

Conversation

@WangLingxun

@WangLingxun WangLingxun commented Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

Retires the hand-rolled examples/run_* launcher scripts and makes runner/primus-cli the single entry point for training.

  • Launchers removed. examples/run_pretrain.sh, run_local_pretrain.sh, run_slurm_pretrain.sh, run_k8s_pretrain.sh, the *_cli.sh shims, and the examples/scripts/ helpers are gone. Everything they did — env seeding, JIT/kernel cache layout, hipBLASLt tuning, AINIC/ANP selection, deterministic mode — already exists in runner/helpers/, so this deletes duplicated logic rather than reimplementing it.
  • Downstream callers migrated. examples/customer_package/, examples/moe_package/, benchmark/, tools/auto_benchmark/, tools/hybrid/, ODC, and the MaxText E2E test now go through primus-cli. Packaged launchers that depend on the EXP / NNODES / DATA_PATH environment contract share a new helper, runner/helpers/launch/slurm_pretrain.sh, which translates it into a primus-cli slurm srun invocation.
  • Spur srun flags normalised. primus-cli slurm srun now drops --export and rewrites --ntasks-per-node=<k> into -n <nodes * k> when a spur binary is on PATH. Spur's srun rejects both with unexpected argument, so the incompatibility is pre-existing, not a regression: examples/run_slurm_pretrain.sh on main hands the same --exclusive --export ALL --ntasks-per-node=1 straight to srun, and has therefore never run on Spur either. It is fixed here because routing through primus-cli slurm gives these flags a single normalisation point for the first time — the old script invoked srun directly, so there was nowhere central to put the knowledge. Standard SLURM is byte-for-byte unchanged and pinned by tests.
  • Prepare-hook behaviour restored. Moving examples/*/prepare.py under runner/helpers/hooks/ had silently dropped three behaviours, now back and covered by tests: the stage: sft short-circuit, the held-out validation split with rank-0-only preparation, and NeMo AutoModel prepare dispatch. The split seed is also renamed so it no longer collides with the Megatron training seed.
  • FLA KDA Triton workaround rewritten. tools/hybrid/patch_fla_triton_autotune_hang.sh, which rewrote files inside the installed fla package in site-packages, is replaced by an in-process patch that narrows the Triton autotune space at startup. Nothing on disk is modified.
  • Documentation rewritten across docs/, examples/, and config header comments to describe primus-cli usage.

Behaviour removal: Kubernetes launcher

examples/run_k8s_pretrain.sh is deleted, along with its section in examples/README.md and the Kubernetes paragraph in docs/05-operations/deployment.md. It was a client script that talked to a bespoke REST API to create and manage training workloads; there is no replacement in this repository, and no Helm chart or operator ships here either.

Kubernetes users are not blocked: Primus only needs the standard distributed variables (MASTER_ADDR, MASTER_PORT, NNODES, NODE_RANK, GPUS_PER_NODE) set consistently across pods, after which every pod runs the same primus-cli direct command. The docs now say that instead of pointing at a script that no longer exists. Please flag if anyone still depends on the old script.

Split seed decoupled from the training seed

primus-cli direct hands one argument list to the prepare hooks and to training. The Megatron prepare hook called its own knob --seed, a name Megatron already owns (trainer_base.yaml defaults seed: 1234). Both consumers read the same flag with different defaults, so passing --seed 7 to reseed training also changed the split cache key — which re-downloaded BookCorpus, re-tokenized the corpus, and deleted the tokenized files already on disk. A seed sweep paid for a full data preparation per seed.

The hook's knob is now --split_seed, default 42. Because the hook parses with parse_known_args, --seed stays in the unknown arguments and reaches training untouched, so it is once again purely the training RNG. The default is unchanged and the on-disk payload keeps its "seed" key, so existing caches stay valid and the rename triggers no rebuild.

FLA KDA autotune narrowing

On Triton 3.6 (ROCm 7.2) the num_stages=4 candidates of FLA's KDA intra-chunk kernels fail to compile — the old shell workaround was misnamed patch_fla_triton_autotune_hang.sh, but this is a compiler abort, not a hang. Triton 3.7+ compiles the full sweep and the autotuner settles on num_stages=2 anyway, and both FLA 0.5.2 (pinned) and upstream main still declare the [2, 3, 4] space, so there is no fixed release to upgrade to. The patch is gated on the ROCm platform plus FLA KDA actually being in use, not on an architecture allowlist, because the faulty pass is chip-independent and an allowlist would silently stop protecting architectures added later; over-applying it costs ~1% at the KDA op level (5.4% on a single kernel in the worst shape measured, within noise end to end). Unsafe entries in FLA's JSON config cache are clamped rather than dropped, since dropping them re-autotunes on every kernel call under FLA_CACHE_MODE=ALWAYS. Delete the patch once the supported toolchain floor reaches Triton >= 3.7.

Test plan

  • pre-commit run --from-ref origin/main --to-ref HEAD — clean.
  • tests/runner/run_all_tests.sh — 10/10 suites.
  • pytest tests/runner/ tests/unit_tests/ci/test_select_tests.py tests/unit_tests/backends/megatron/test_fla_kda_autotune_patches.py — pass.
  • pytest tests/unit_tests/backends/megatron/ — no new failures; the 14 pre-existing failures and 2 collection errors are diffusion/FSDP2 on a driverless host.
  • Spur srun normalisation, via a fake spur on PATH in tests/runner/test_primus_cli_slurm.sh: both --export spellings dropped, -N 4 --ntasks-per-node=1 collapsing to a single -n 4, --ntasks-per-node 2 with -N 4 becoming -n 8, an explicit -n never duplicated, a non-integer node count left alone with a warning, and standard SLURM still receiving --export ALL and --ntasks-per-node=1 verbatim with no -n injected.
  • MI355X smoke on tasimage/primus:latest. Targeted pytest 69 passed; tests/runner/run_all_tests.sh 10/10 suites; the direct-launcher sub-suite 77/77; single-GPU Megatron Llama3 8B BF16 (4 layers, 3 steps) trains with finite loss and grad norm at roughly 995 TFLOP/s.
  • Megatron pretrain on real hardware via primus-cli direct and primus-cli container.
  • SFT config on 2 GPUs — stage: sft skips BookCorpus and needs no HF_TOKEN.
  • Held-out validation split generated, metadata keyed by split seed and test_size.
  • Split seed / training seed independence — --seed 1234 passes through to training, split still uses 42, cache key unchanged.
  • FLA_CACHE_MODE=always does not re-autotune; clamping preserves the other config fields.
  • FLA KDA patch active in the log (kept 4/12 autotune configs); run completes on Triton 3.6, which previously failed to compile.
  • UEP preflight variables forwarded into the container.
  • Cache-root fallback with a root-owned /workspace; explicit PRIMUS_CACHE_ROOT respected; layout unchanged when /workspace is writable.
  • slurm_pretrain.sh propagates the child exit code (0 before the fix, 42 after, against a stubbed primus-cli).
  • BNXT tar mount and MLflow/Databricks passthrough via primus-cli container --dry-run.

Known follow-ups

Tokenized cache key omits tokenizer_model (pre-existing, #515). Switching tokenizer_model within the same tokenizer_type silently reuses the previous tokenizer's token IDs. Noted only because this PR's coordination-marker hash does include tokenizer_model and so reads like protection — it is not, as dataset_cache_ready() never validates the marker's contents.

@WangLingxun
WangLingxun force-pushed the refactor/primus-cli-launchers branch 10 times, most recently from 5d583ef to 3016b30 Compare August 14, 2026 09:14
@WangLingxun
WangLingxun marked this pull request as ready for review August 14, 2026 10:24
@WangLingxun
refactor(runner): retire the examples launchers in favour of primus-cli
24de88c 
The examples/ launch scripts had grown into a second, divergent copy of the
launcher: they re-implemented the cluster env setup, cache layout, NCCL tuning
and hook dispatch that runner/primus-cli already owns, and the parallel
run_*_cli.sh variants meant every fix had to be made twice. Keeping the two
copies in sync was a recurring source of drift, and the K8s submission client
that justified keeping run_pretrain.sh alive is no longer needed.

- Move the behaviour that only lived under examples/ into the runner (the
  persistent kernel/JIT cache layout, TRAIN_LOG honoured by direct mode) and
  delete the entry scripts and networking helpers outright, so there is one
  launcher interface to learn and one implementation to maintain.
- Turn the srun + container assembly that the packaged launchers still share
  into an explicit helper under runner/helpers/launch/ rather than a
  user-facing entry point. The packaged customer/MoE launchers and the
  benchmark runners were relying on the EXP/NNODES contract for real value, so
  they call that helper instead of open-coding the same srun flags nine times;
  everything that only needed a single-node launch calls primus-cli directly.
- Preserve the packaged-launcher contracts while consolidating: Slurm GPU
  counts remain scheduler-derived unless explicitly overridden, UCCL-EP and
  clean-container controls cross the container boundary, and
  PRIMUS_SKIP_PIP consistently skips per-run dependency hooks.
- Drop the backend prepare.py copies under examples/: they were only ever
  invoked by the deleted dispatchers, and the hooks the CLI actually runs own
  that logic now. The two behaviours the megatron copy still had over its hook
  are restored in the hook itself in a follow-up.
- Repoint the docstrings, hook headers and CI comments that named
  examples/run_pretrain.sh at primus-cli or at the runner hook that owns the
  behaviour, since they now describe a file that does not exist.

- Probe the persistent cache root before exporting it. The /workspace default
  is only writable inside the container, so bare-metal direct runs, which now
  source the same base_env.sh, fall back to a per-user cache with a warning
  instead of failing at the first Triton JIT compile. An explicitly configured
  PRIMUS_CACHE_ROOT is still taken as-is.
- Mount PATH_TO_BNXT_TAR_PACKAGE into the container the way the deleted
  launcher did. Forwarding the variable alone left the rebuild hook unable to
  see the tar it names, and that hook now warns rather than informs when it
  skips a rebuild the user explicitly asked for.
- Keep the MLflow and Databricks credentials in the container passthrough
  allowlist. The deleted launcher forwarded them explicitly, the automatic
  prefix passthrough does not cover them, and the logging guide already
  documents them as allowlisted.
- Take the launcher's exit code from PIPESTATUS[0] in the shared srun helper.
  tee was reporting success for failed training runs, which SLURM, CI and the
  benchmark runners all read as a passing job.

cleanup_rs.sh matters beyond the rename: its pkill pattern matched the old
script name, so it had stopped killing the process it is meant to reap.

The container env passthrough list already carried HSA_KERNARG_POOL_SIZE twice
before this change; the duplicate is dropped while the list is being edited
here, leaving the order that the environment-variables reference documents.
@WangLingxun
fix(runner): restore the prepare-hook behaviour lost with the launchers
2720721 
Folding the examples/ launchers into primus-cli routed every run through
runner/helpers/hooks/train/pretrain/, and three preparation behaviours that
only the old dispatcher and its prepare.py had did not survive the move. Each
one fails late and quietly, so restore them where the CLI actually looks.

- Move the NeMo AutoModel prepare entry into the hook directory searched by
  primus-cli and let it consume only its own arguments, so training overrides
  continue to the trainer instead of being rejected by the hook parser.
- Return early for SFT, and only create a held-out validation split when
  evaluation needs one and the config sets neither valid_data_path nor
  test_data_path. Either one means the run already has its own evaluation
  data, and the generated split would replace it.
- Advertise the generated split only through the extra.* keys the config
  left unset. Those lines become CLI arguments, so emitting both keys
  unconditionally overrode an explicitly configured test_data_path.
- Honor custom train/eval tokenized prefixes exactly. Key split caches by the
  split seed and test size, validate their metadata, and coordinate nodes with
  atomic completion/failure markers plus a bounded wait on shared storage.
- Keep exposing the resolved Megatron backend and generated split paths through
  the hook extra-argument protocol used by primus-cli direct mode.

The split seed is spelled --split_seed, with a PRIMUS_SPLIT_SEED environment
entry, rather than --seed. primus-cli direct mode hands the same argument list
to the hooks and to training, and seed is already a Megatron training parameter
(trainer_base.yaml defaults it to 1234). Sharing one flag between two consumers
meant a --seed the user passed to reseed training also changed the split cache
key, which re-downloaded BookCorpus, re-tokenized the corpus and deleted the
tokenized files already on disk -- so a seed sweep paid for a full data
preparation per seed. The hook parser uses parse_known_args, so under the new
name --seed is left in the unknown arguments and reaches the config overrides
untouched. The default stays 42 and the on-disk payload keeps its "seed" key,
so existing bookcorpus_{train,valid}_<key>.json caches and .split.json metadata
remain valid. PRIMUS_SPLIT_SEED is the entry to prefer through primus-cli,
since it keeps the knob out of the argument list training also parses.
@WangLingxun
fix(megatron): narrow the FLA KDA autotune space in-process on ROCm
27550b0 
AMD Triton cannot compile the num_stages=4 variants of flash-linear-attention's
KDA intra-chunk kernels. On rocm/primus:v26.3 (torch 2.10, ROCm 7.2, Triton
3.6.0) the first autotune sweep of a KDA run dies in the MLIR pass pipeline:

    error: 'tt.load' op operation destroyed but still has uses
    error: Failures have been detected while processing an MLIR pass pipeline
    RuntimeError: PassManager::run failed

The failure is in AMD's software pipeliner --
tritonamdgpu-schedule-loops{num_stages=4} feeding tritonamdgpu-pipeline. On
Triton 3.7.0 and 3.7.1 the full sweep compiles and the autotuner settles on
num_stages=2. The [2, 3, 4] space is still declared in FLA 0.5.2 (the pinned
version) and in upstream main, so there is no fixed FLA release to move to.

The workaround used to be a prepare-time shell hook that shelled out to a sed
script and rewrote chunk_intra.py inside site-packages. That mutated a shared
image, needed a separate PrimusParser subprocess just to decide whether the run
used KDA at all, and left the installed package differing from its wheel. Do it
inside the training process instead: a before_train patch drops the unsafe
candidates from the two kernels' autotune space, keeping the num_warps sweep so
the autotuner still picks the best remaining config. Nothing on disk changes.

- Add megatron.fla.kda_safe_autotune. It walks Triton's decorator chain to the
  object owning .configs -- @triton.heuristics wraps FLA's CachedAutotuner, so
  the config list is one .fn hop below the module attribute -- and filters it.
- Also guard CachedAutotuner.maybe_load_cached_config, which writes a config
  read from a best-config JSON straight into self.cache without consulting
  .configs. FLA_CACHE_MODE plus FLA_CONFIG_DIR could otherwise reintroduce an
  unsafe config behind the filter. The guard lowers num_stages on the entries it
  finds rather than deleting them, because deleting loops under
  FLA_CACHE_MODE=ALWAYS: should_check_fla_cache returns True unconditionally in
  that mode, so run() reloads the JSON before every launch, and an emptied cache
  sends Triton back to benchmarking only for the next reload to overwrite the
  result with the unsafe config again -- a full autotune per kernel call for the
  duration of the run. The other cache modes gate on `key not in self.cache` and
  were never exposed to it. num_stages is adjusted in place so num_ctas,
  maxnreg, pre_hook and ir_override -- which FLA populates on Triton >= 3.5.1 --
  survive; rebuilding the config would silently reset them.
- Gate on ROCm, and on whether the run reaches FLA's KDA kernels at all --
  use_fla_triton_kda, kda_backend and use_kimi_k3_attention_backend are all
  merged onto args before before_train runs. The second condition is there to
  avoid importing fla into a run that has no other use for it; trimming the
  config list of a kernel that never launches would itself be harmless.
- Retire PRIMUS_FLA_TRITON_AUTOTUNE_WORKAROUND without a replacement knob,
  rather than carrying a tri-state override for something this cheap.
- Fail soft everywhere: a missing FLA, a renamed kernel or a restructured
  .configs logs and returns rather than blocking training.

There is deliberately no architecture gate. What breaks is a compiler pass,
which is essentially chip-independent, and gfx94* only ever described where the
failure happened to be seen first. 18071e1 records the same lesson from the
buffer-store miscompile: a gfx942-only gate was later found to skip real
kernels carrying the same pattern on gfx950, and "until someone remembers to,
that architecture is silently unprotected". kimi_k3 is KDA on MI355X (gfx950),
so that is a live configuration here rather than a hypothetical.

Over-applying it is cheap, and measured. Letting the full sweep run on gfx942
and comparing its winner against the narrowed space: zebra 300M (K=32, V=64)
and a k3 curve shape (K=V=128, B=2, T=2048, H=16) both pick num_warps=2,
num_stages=2, so narrowing is free; the k3 full shape (K=V=128, B=1, T=4096,
H=96) picks num_stages=3, where the narrowed kernel is 5.4% slower in isolation
and 0.96% slower at the KDA op level. An end-to-end A/B on 4x MI300X (60
iterations, four interleaved repeats each) stayed inside run-to-run noise, at
+0.10% and -0.83% in opposite directions. So ~1% of a single op bounds the cost
of applying this too widely, against a hard compile failure when it is missed.

chunk_kda_fwd_kernel_intra_sub_chunk is narrowed even though it is dead code
today: it is only reached with safe_gate=True, which nothing in Primus sets,
and training logs only ever show autotune records for the bwd kernel. It stays
in so the protection is already in place if safe_gate is ever enabled, with a
comment noting that it is also the one kernel where num_stages=4 actually wins
(by 4.9% at safe_gate shapes), so that call needs re-measuring at that point.

Delete this patch once the supported toolchain floor reaches Triton 3.7.

Validated on 8x MI300X with zebra_llama_300M_kda_pure: the patch narrows both
kernels from 12 to 4 configs and training runs to completion.
@WangLingxun
docs: describe training launches in terms of primus-cli
25ade5c 
Every guide, config header and example still told users to export EXP and call
examples/run_*.sh. Those scripts are gone, so the instructions now use the
runner modes and reference the hooks that own the behavior.

- Rewrite launch snippets to direct/container/slurm forms and keep container
  environment options at the container layer, with a smoke test preventing the
  `-- --env DATA_PATH` parsing regression.
- Document the HipBLASLt tuning master switch and align the container
  passthrough reference with runner/.primus.yaml, including UCCL-EP controls.
- Describe the exact TOKENIZED_DATA_PATH, TOKENIZED_TRAIN_DATA_PATH, and
  TOKENIZED_EVAL_DATA_PATH contracts, split-cache metadata, bounded rank waits,
  and the shared-storage requirement. Spell out that the split seed is
  PRIMUS_SPLIT_SEED / --split_seed and not the Megatron seed training parameter,
  since the whole point of keeping them apart is that reseeding training must not
  invalidate a tokenized corpus.
- Document the FLA KDA autotune narrowing as gated on ROCm rather than on a
  specific GPU architecture, matching the patch.
- Point the remaining multi-node benchmark commands at the shared
  runner/helpers/launch/slurm_pretrain.sh helper, which keeps the EXP and
  NNODES contract those snippets were written against, and add a docs test
  that fails when a Markdown command invokes a script the tree does not
  contain. The test matches only bash/sh/source invocations, so prose that
  merely names a script stays out of scope.
@WangLingxun
WangLingxun force-pushed the refactor/primus-cli-launchers branch from 3016b30 to 25ade5c Compare August 17, 2026 03:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Xiaoming-AMD Xiaoming-AMD Awaiting requested review from Xiaoming-AMD Xiaoming-AMD is a code owner
@wenxie-amd wenxie-amd Awaiting requested review from wenxie-amd wenxie-amd is a code owner
@limou102 limou102 Awaiting requested review from limou102 limou102 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@WangLingxun