Configuration system

Configuration system#

Primus experiments are described in YAML. The loader resolves environment variables, extends: inheritance, and module/model/platform presets before training starts. This document focuses on the Python configuration pipeline (primus/core/config/ and primus/core/launcher/parser.py).

Related documentation

Topic

Location

CLI launcher and --config

CLI Reference

Setting env vars and XLA_FLAGS for a run

Environment and XLA flags

Backend parameter references

Megatron parameters, TorchTitan parameters, MaxText parameters

Overview: Three-layer YAML#

A typical experiment ties together:

  1. Experiment YAML—your run: identity, workspace, and a modules section naming framework presets plus overrides.

  2. Module preset—training defaults for a backend (optimizer, schedule, parallelism hooks) under primus/configs/modules/<framework>/.

  3. Model preset—architecture and tokenizer metadata under primus/configs/models/<framework>/.

All of these are deep-merged (see primus/core/config/yaml_loader.py and primus/core/config/merge_utils.py). A platform preset (primus/configs/platforms/) maps distributed environment variable names and logging defaults; if omitted, the parser injects platform_azure.yaml (see PrimusParser.parse_platform in primus/core/launcher/parser.py).

Experiment config structure#

work_group: ${PRIMUS_TEAM:amd}
user_name: ${PRIMUS_USER:root}
exp_name: ${PRIMUS_EXP_NAME:my_experiment}
workspace: ${PRIMUS_WORKSPACE:./output}

env:                             # optional: env vars exported before training starts
  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.96"

modules:
  pre_trainer:
    framework: megatron          # backend name (megatron, torchtitan, maxtext, megatron_bridge, …)
    config: pre_trainer.yaml     # module preset file under primus/configs/modules/<framework>/
    model: llama3_8B.yaml        # model preset file under primus/configs/models/<framework>/
    overrides:                   # training overrides (deep-merged last)
      train_iters: 50
      micro_batch_size: 4

Required top-level keys are validated in PrimusParser.parse_meta_info: work_group, user_name, exp_name, workspace.

Module presets#

  • Location: primus/configs/modules/<framework>/ (for example primus/configs/modules/megatron/pre_trainer.yaml).

  • Purpose: Default training behavior (iterations, batching, optimizer, logging, parallelism-related flags for that backend).

  • Inheritance: Use extends: to compose files in the same directory (or relative paths). Multiple entries merge in order; the current file wins on conflicts (_apply_extends in yaml_loader.py).

Example chain (excerpt): pre_trainer.yaml extends trainer_base.yaml, which extends ../module_base.yaml and other shared fragments (primus/configs/modules/megatron/trainer_base.yaml).

Model presets#

  • Location: primus/configs/models/<framework>/ (for example primus/configs/models/megatron/llama3_8B.yaml).

  • Purpose: Architecture dimensions, tokenizer identifiers, and other model metadata.

  • Inheritance: Same extends: mechanism as modules (for example llama3_8B.yaml → llama3_base.yaml → …).

Platform presets#

  • Location: primus/configs/platforms/ (for example primus/configs/platforms/platform_azure.yaml).

  • Purpose: Names of environment variables used for distributed launch (NNODES, NODE_RANK, MASTER_ADDR, …) and defaults such as master_sink_level and workspace.

  • Default: If the experiment omits platform, the parser sets config: platform_azure.yaml (primus/core/launcher/parser.py).

Environment variable substitution#

primus/core/config/yaml_loader.py expands:

Pattern

Behavior

${VAR}

Required. Raises if VAR is unset.

${VAR:default}

Uses default when VAR is unset.

After substitution, purely numeric strings may be converted to int or float.

The env: block#

An optional top-level env: mapping (sibling of modules:) is exported into the process environment before distributed init and before the backend imports its framework, so JAX, XLA, PyTorch, and RCCL all observe it. It is read from the file separately by PrimusRuntime._apply_config_env, so the ${VAR:default} form above does not apply here — only $VAR and ${VAR} for already-set variables are expanded.

env:
  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.96"
  XLA_FLAGS_APPEND: "--xla_gpu_autotune_level=5"

Quote the values, and keep the block at the top level: an env: nested under a module silently becomes a training parameter on that module instead of being exported. See Environment and XLA flags for the precedence rules, which matter for XLA_FLAGS in particular.

extends: inheritance#

For each YAML file:

  1. Each path in extends: is resolved relative to the directory of the current file.

  2. Presets are loaded recursively (each may have its own extends:).

  3. Merge order: earlier presets in the list are merged first; later presets override earlier ones; the current file overrides all (_apply_extends).

PresetLoader.load (primus/core/config/preset_loader.py) resolves primus/configs/<config_type>/<framework>/<name>.yaml and runs the same parse_yaml pipeline.

CLI overrides (training)#

After the main arguments are parsed, unknown tokens are interpreted as key=value overrides and deep-merged into the active training module namespace (module_cfg.params). The core runtime applies them in PrimusRuntime._apply_overrides (primus/core/runtime/train_runtime.py) using parse_cli_overrides (primus/core/utils/arg_utils.py) followed by deep_merge. Both key=value and --key value forms are accepted. Unknown keys are merged in (not rejected) and forwarded to the backend; the stricter key-existence check in parse_args / _check_keys_exist (primus/core/launcher/parser.py) belongs to a legacy path that the train subcommand does not exercise.

Example (conceptual):

./runner/primus-cli direct -- train pretrain --config exp.yaml \
  --train_iters 100 --micro_batch_size 2

Merge priority (training config)#

The effective ordering of training parameters is:

  1. CLI overrides (key=value after the main train arguments)—highest. Applied last by the runtime (PrimusRuntime._apply_overrides), after preset and experiment merging.

  2. modules.pre_trainer.overrides in the experiment YAML (applied in PrimusParser.parse_trainer_module).

  3. Module preset with model preset additions: the module preset is loaded first, then the model preset is merged in with allow_override=False, so duplicate top-level module keys are preserved while non-duplicate model keys are added (merge_namespace in parse_trainer_module).

  4. Preset chains via extends: inside those files—base layers first, specialized layers later, file body last.

A concise mental model:

CLI overrides > experiment overrides > module preset with model preset additions > each preset’s own extends: chain > shared bases such as module_base.yaml.

Config resolution walkthrough: llama2_7B-BF16-pretrain.yaml#

Example experiment: examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml.

  1. Load experiment—parse_yaml reads the file; ${PRIMUS_*:…} placeholders resolve.

  2. Meta—work_group, user_name, exp_name, workspace are checked.

  3. Platform—If not present, default platform_azure.yaml loads from primus/configs/platforms/.

  4. Module preset—PresetLoader.load("pre_trainer.yaml", "megatron", "modules") loads primus/configs/modules/megatron/pre_trainer.yaml and applies its extends: chain (for example trainer_base.yaml → …).

  5. Model preset—PresetLoader.load("llama2_7B.yaml", "megatron", "models") loads primus/configs/models/megatron/llama2_7B.yaml (which extends llama2_base.yaml → llama_base.yaml → …).

  6. Merge—Module and model namespaces are merged for pre_trainer.

  7. Experiment overrides—Keys under modules.pre_trainer.overrides in the example file (for example mock_data: true, parallelism, LR) are applied on top.

  8. CLI overrides—Any key=value pairs from the command line are merged last.

How to write a new config for a new model#

  1. Add or reuse a model preset under primus/configs/models/<framework>/, using extends: from the closest existing architecture (for example copy llama3_8B.yaml and adjust hidden size, layers, tokenizer).

  2. Point an experiment YAML at it—set modules.pre_trainer.framework, config: <module_preset>.yaml, and model: <your_model>.yaml.

  3. Set overrides in the experiment file for run-specific values (batch sizes, paths, mock_data, parallelism). Prefer small experiment files that reference presets instead of duplicating hundreds of keys.

  4. Validate with --dry-run and by tracing the referenced experiment, module, and model presets (see below).

  5. Optional: add an example under examples/<backend>/configs/<platform>/ for others to copy.

Debugging config issues#

Technique

What it does

--export_config PATH

Parsed by the training config parser, but the default core PrimusRuntime path does not currently write the resolved YAML. Treat this as legacy/future functionality unless your deployment implements it.

./runner/primus-cli --dry-run …

Shows the launcher command without executing (shell layer).

--debug

Enables verbose logging for launcher and Python (PRIMUS_LOG_LEVEL=DEBUG).

Inspect presets

Open the resolved extends: chain under primus/configs/modules/ and primus/configs/models/ for the framework you use.

If ${VAR} substitution fails, set the variable or switch to ${VAR:default} in the YAML.

Cross-references#

  • Default launcher YAML: runner/.primus.yaml

  • YAML loader (env + extends): primus/core/config/yaml_loader.py

  • Parser and merge: primus/core/launcher/parser.py

  • Preset paths: primus/core/config/preset_loader.py

Contents