# Extending backends 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](../01-getting-started/overview.md): adapters sit under the unified CLI and configuration system ([Configuration system](../02-user-guide/configuration-system.md)), 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: ```bash primus train pretrain --config ``` 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: ```python 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: ```python 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: ```python # 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..__init__` that calls `BackendRegistry.register_adapter(...)` --- ## Minimal backend layout Create a new backend folder under `primus/backends/`: ```text 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` ```python 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/, 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` ```python 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", "") 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` ```python 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): ```yaml modules: pre_trainer: framework: dummy config: dummy_trainer.yaml model: dummy_8B.yaml ``` Run: ```bash ./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//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` ```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` ```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` ```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: ```bash ./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](./adding-models.md). --- ## 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..__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//` and `runner/helpers/hooks/train/posttrain//`, where `` 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: ```bash ./primus-cli direct -- train pretrain --config ``` 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//` (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** ```bash # inside runner/helpers/hooks/train/pretrain//.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** ```bash # 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 `-- ` after hook execution. Typical pattern to install or configure packages for a backend: - Add a script under `runner/helpers/hooks/train/pretrain//-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.