Extending backends

July 15, 2026 · View on GitHub

This guide explains how to add a new training backend to Primus using the current runtime architecture. It complements the high-level picture in Primus overview: adapters sit under the unified CLI and configuration system (Configuration system), and each backend plugs in through the same lifecycle and hook points as Megatron-LM, TorchTitan, MaxText, and the other integrated stacks.

The runtime is built around:

  • BackendAdapter – integrates a backend framework
  • BackendRegistry – discovers and instantiates adapters
  • BaseTrainer – defines the minimal training lifecycle that all backends follow
  • PrimusRuntime – orchestrates config loading, environment setup, patches, adapter, and trainer

The examples below use a minimal dummy backend as a template. The dummy files are illustrative and are not checked into this repository; existing backends such as Megatron, TorchTitan, MaxText, Megatron Bridge, and HummingbirdXT show the production pattern.

What happens when you run Primus?

When you run:

primus train pretrain --config <exp.yaml>

the runtime (PrimusRuntime) does roughly:

  1. Load the experiment config—load_primus_config() returns a lightweight SimpleNamespace (not a PrimusConfig)—and select the module_config

  2. Apply CLI overrides to module_config.params

  3. Initialize environment (HF, logging, distributed environment, data directory)

  4. Resolve backend adapter via BackendRegistry.get_adapter(framework)

  5. Call adapter.setup_backend_path(...) to put the backend on sys.path

  6. Call adapter.prepare_backend(module_config) (usually runs backend setup hooks)

  7. Build backend arguments:

    backend_args = adapter.convert_config(module_config.params)
# run "build_args" patches and merge back into module_config.params

  8. Load and construct the trainer:

    TrainerClass = adapter.load_trainer_class(stage=module_config.params.stage or "pretrain")
trainer = TrainerClass(backend_args=backend_args)

  9. Execute the trainer lifecycle (with patches around it). Backend version detection is lazy during patch handling through adapter.detect_backend_version() rather than a separate pre-trainer step:

    # PrimusRuntime (the "build_args" patches from step 7 already ran
# before the trainer was constructed):
run_patches(phase="setup",        backend_args=backend_args)
trainer.setup()

trainer.init()

run_patches(phase="before_train", backend_args=backend_args)
trainer.train()
run_patches(phase="after_train",  backend_args=backend_args)

trainer.cleanup()

So a complete backend must provide:

  • An adapter subclassing BackendAdapter
  • A trainer subclassing BaseTrainer and implementing setup, init, and train (and optionally overriding cleanup, which has a default no-op implementation)
  • A small primus.backends.<name>.__init__ that calls BackendRegistry.register_adapter(...)

Minimal backend layout

Create a new backend folder under primus/backends/:

primus/backends/dummy/
├── __init__.py
├── dummy_adapter.py
└── dummy_pretrain_trainer.py

This mirrors the pattern used by existing backends (for example Megatron, TorchTitan).

Implement the adapter (BackendAdapter)

File: primus/backends/dummy/dummy_adapter.py

from __future__ import annotations

from types import SimpleNamespace
from typing import Any

from primus.core.backend.backend_adapter import BackendAdapter
from primus.core.backend.backend_registry import BackendRegistry
from primus.core.utils.module_utils import log_rank_0


class DummyAdapter(BackendAdapter):
    """Minimal adapter for a 'dummy' backend."""

    def __init__(self, framework: str = "dummy"):
        super().__init__(framework)

    def setup_backend_path(self, backend_path=None) -> str:
        """
        Dummy backend lives inside the Primus tree (no third_party submodule),
        so we don't need to modify sys.path or resolve any external path here.

        For real backends that live under third_party/<backend>, you can rely on
        the default implementation in BackendAdapter instead.
        """
        log_rank_0("[Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend")
        return ""

    def convert_config(self, params: Any) -> Any:
        """
        Convert Primus module params → backend-specific args.

        For a real backend you would build a structured args object. Here we
        just wrap the incoming params in a SimpleNamespace.
        """
        if isinstance(params, dict):
            backend_args = SimpleNamespace(**params)
        else:
            backend_args = params
        log_rank_0("[Primus:DummyAdapter] Converted Primus params -> dummy backend_args")
        return backend_args

    def detect_backend_version(self) -> str:
        """Return a version string used by patch filtering."""
        return "dummy-0.1"

    def load_trainer_class(self, stage: str = "pretrain"):
        """Return the Trainer class for the specified training stage."""
        from primus.backends.dummy.dummy_pretrain_trainer import DummyPretrainTrainer

        log_rank_0("[Primus:DummyAdapter] Loaded trainer class: DummyPretrainTrainer")
        return DummyPretrainTrainer

Key points:

  • Since the dummy backend is implemented directly under primus.backends.dummy (not in third_party/), it overrides setup_backend_path() as a no-op so that the default third-party path resolution is skipped.
  • convert_config() returns whatever your trainer expects as backend_args.
  • load_trainer_class() imports and returns DummyPretrainTrainer directly (similar to MegatronAdapter), without going through a registry lookup.

Implement a runnable trainer (BaseTrainer)

File: primus/backends/dummy/dummy_pretrain_trainer.py

from typing import Any

from primus.core.trainer.base_trainer import BaseTrainer
from primus.core.utils.module_utils import log_rank_0


class DummyPretrainTrainer(BaseTrainer):
    """Minimal runnable trainer for the dummy backend."""

    def __init__(self, backend_args: Any):
        # BaseTrainer stores backend_args and reads torchrun env (rank, world_size, etc.)
        super().__init__(backend_args=backend_args)
        self._initialized = False

    def setup(self):
        # Optional pre-init setup (e.g., logging, sanity checks)
        log_rank_0(f"[DummyPretrainTrainer] setup() on rank={self.rank}")

    def init(self):
        # Build your model / optimizer / dataloader here in a real backend.
        log_rank_0("[DummyPretrainTrainer] init()")
        self._initialized = True

    def train(self):
        if not self._initialized:
            raise RuntimeError("DummyPretrainTrainer.init() must be called before train().")

        log_rank_0("[DummyPretrainTrainer] train()")
        # Example: access a custom param (e.g., 'hello') from backend_args.
        hello_value = getattr(self.backend_args, "hello", "<missing>")
        log_rank_0(f"[DummyPretrainTrainer] hello={hello_value}")
        # Real training loop would go here.
        log_rank_0("[DummyPretrainTrainer] training finished successfully.")

    def cleanup(self, on_error: bool = False):
        # Optional cleanup logic (close files, finalize logging, etc.)
        status = "error" if on_error else "success"
        log_rank_0(f"[DummyPretrainTrainer] cleanup(on_error={status})")

Why this matches the core architecture:

  • BaseTrainer.__init__ reads distributed environment from get_torchrun_env().
  • PrimusRuntime drives the lifecycle: setup → init → train → cleanup and runs patch phases around these steps.
  • Your trainer only needs to implement setup, init, train, and cleanup using backend_args and the resolved environment information.

Register the adapter in BackendRegistry

File: primus/backends/dummy/__init__.py

from primus.backends.dummy.dummy_adapter import DummyAdapter
from primus.core.backend.backend_registry import BackendRegistry


# Register adapter (backend name → adapter class)
BackendRegistry.register_adapter("dummy", DummyAdapter)

At runtime, when framework: dummy is requested:

  • BackendRegistry.get_adapter("dummy") lazily imports primus.backends.dummy (this file), which calls register_adapter("dummy", DummyAdapter).
  • The adapter instance is created and used by PrimusRuntime to set up the backend path, run setup hooks, build backend_args, and load and construct the trainer.

Minimal config example

Create an experiment YAML (simplified; full template in the next section):

modules:
  pre_trainer:
    framework: dummy
    config: dummy_trainer.yaml
    model: dummy_8B.yaml

Run:

./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml

Because this dummy backend is an in-tree template and setup_backend_path() is a no-op, you should see logs similar to:

  • [Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend
  • [Primus:DummyAdapter] Converted Primus params -> dummy backend_args
  • [DummyPretrainTrainer] setup()
  • [DummyPretrainTrainer] init()
  • [DummyPretrainTrainer] train()

Example end-to-end YAML configs

This template mirrors the Megatron pattern. Create these files only when you are actually adding a dummy backend for local development or tests:

  • The top-level experiment config lives under examples/<backend>/configs/...
  • The module config is resolved from primus/configs/modules/{framework}/
  • The model config is resolved from primus/configs/models/{framework}/

Top-level experiment config

File 1: examples/dummy/configs/dummy_8B-pretrain.yaml

work_group: ${PRIMUS_TEAM:local}
user_name: ${PRIMUS_USER:local}
exp_name: ${PRIMUS_EXP_NAME:dummy_8B-pretrain}
workspace: ${PRIMUS_WORKSPACE:./output}

modules:
  pre_trainer:
    framework: dummy
    config: dummy_trainer.yaml

    # model to run
    model: dummy_8B.yaml

    overrides:
      # log / debug
      stderr_sink_level: DEBUG

      # example training overrides (merged into module params)
      train_iters: 100
      global_batch_size: 32
      micro_batch_size: 4
      seq_length: 1024
      hello: world

Module-level trainer config

File 2: primus/configs/modules/dummy/dummy_trainer.yaml

extends:
  - trainer_base.yaml        # optional, if you have a shared base; otherwise omit

train_iters: 1000
global_batch_size: 16
micro_batch_size: 1
seq_length: 512

log_interval: 1
save_interval: 100

This file defines default training hyperparameters for the dummy backend. Fields under modules.pre_trainer.overrides in the top-level config are deep-merged on top of these defaults.

Model-level config

File 3: primus/configs/models/dummy/dummy_8B.yaml

extends: []

model_name: dummy_8B
vocab_size: 32000
hidden_size: 4096
num_layers: 32
num_attention_heads: 32

This file plays the same role as Megatron model configs under primus/configs/models/megatron/. It is loaded via:

  • modules.pre_trainer.model: dummy_8B.yaml
  • resolved as primus/configs/models/{framework}/dummy_8B.yaml

Running the example

Run:

./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml

Primus will:

  • load examples/dummy/configs/dummy_8B-pretrain.yaml
  • resolve modules.pre_trainer.config → primus/configs/modules/dummy/dummy_trainer.yaml
  • resolve modules.pre_trainer.model → primus/configs/models/dummy/dummy_8B.yaml
  • build module_config.params from these sources plus overrides
  • call DummyAdapter.convert_config(params) to build backend_args
  • construct DummyPretrainTrainer(backend_args=...)
  • execute setup → init → train → cleanup.

For adding model YAML for existing backends (Megatron, TorchTitan, and others), see Adding model configurations.

Checklist for a complete backend

Use this as a quick checklist when adding a new backend:

  • Adapter subclass of BackendAdapter implements: - load_trainer_class(stage: str) - convert_config(params) - detect_backend_version() - (optionally) overrides prepare_backend() / third_party_dir_name
  • Trainer subclass of BaseTrainer implements: - setup(), init(), train(), and optional cleanup(on_error: bool)
  • BackendRegistry.register_adapter(backend, AdapterClass) is called in primus.backends.<backend>.__init__
  • At least one unit test is added under tests/unit_tests/backends/

Once these are in place, your backend is fully integrated into the Primus runtime and follows the same lifecycle and patch phases as the built-in backends.

Advanced: Backend-specific setup with train hooks

For more advanced scenarios (for example installing extra Python packages or configuring backend-specific environment variables at runtime), you can use train hooks under runner/helpers/hooks.

  • Hook locations for training:
    • Global hooks (run for all commands): runner/helpers/hooks/*.sh and runner/helpers/hooks/*.py. These are discovered with find ... -maxdepth 1 \( -name "*.sh" -o -name "*.py" \) and executed in lexicographical order of their filenames (see runner/helpers/execute_hooks.sh).
    • Command-specific hooks: runner/helpers/hooks/train/pretrain/*.sh|*.py (and .../posttrain/...), discovered and ordered the same way. For pretrain, this directory contains the dispatcher prepare_experiment.sh.
    • Per-framework hooks: runner/helpers/hooks/train/pretrain/<framework>/ and runner/helpers/hooks/train/posttrain/<framework>/, where <framework> is megatron, torchtitan, dummy, and so on. These are not run directly by execute_hooks; instead prepare_experiment.sh detects the framework from the experiment config, runs that framework folder's *.sh files in lexicographical order, and then invokes the framework's prepare.py dispatcher.

When you run:

./primus-cli direct -- train pretrain --config <exp.yaml>

Primus will:

  • Call execute_hooks train pretrain ..., which:
    • Runs global hooks under runner/helpers/hooks/ (lexicographical order)
    • Then runs command-specific hooks under runner/helpers/hooks/train/pretrain/, including prepare_experiment.sh
    • prepare_experiment.sh resolves the framework from the config and runs the per-framework hooks under runner/helpers/hooks/train/pretrain/<framework>/ (its *.sh files in lexicographical order, then prepare.py)

Each hook script can emit control lines on stdout that Primus parses (the framework hooks' stdout is captured through prepare_experiment.sh):

  • env.*=value → environment variables

    # inside runner/helpers/hooks/train/pretrain/<framework>/<something>.sh
echo "env.MY_BACKEND_FLAG=1"        # becomes: export MY_BACKEND_FLAG=1
echo "env.PYTHONPATH=/opt/mylib:$PYTHONPATH"

    These are exported into the environment of the primus-cli direct process, so downstream backend code and trainers see them.

  • extra.*=value → extra CLI arguments

    # inside the same hook
echo "extra.backend_path=/opt/my-backend"   # becomes: --backend_path /opt/my-backend
echo "extra.train_data_path=/my/data"       # becomes: --train_data_path /my/data

    These extra.* pairs are appended to the Primus CLI invocation as --<name> <value> after hook execution.

Typical pattern to install or configure packages for a backend:

  • Add a script under runner/helpers/hooks/train/pretrain/<framework>/<NNN>-setup.sh (use a numeric prefix such as 000- or 010- to control ordering).
  • In that script:
    • Optionally run python -m pip install ... or other setup commands.
    • Emit env.*=... lines to export any required environment variables.
    • Emit extra.*=... lines if you need to pass additional CLI arguments (for example backend_path) into the Primus runtime for this run.