| """Primus / Primus-Turbo backend — pretraining-scale (>70 B params). |
| |
| Subprocess wrapper that expects the rocm/primus:v26.2 container to be active |
| (Primus, AITER, Composable Kernel, hipBLASLt, FP8 transformer-engine support |
| all live there). Lazy availability check. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib.util |
| import os |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import yaml |
|
|
| from mindxtrain.autotune.plan import AutotunePlan |
| from mindxtrain.config.schema import XTrainConfig |
| from mindxtrain.train.axolotl_compile import compile_axolotl_yaml |
|
|
|
|
| def _primus_available() -> bool: |
| return ( |
| importlib.util.find_spec("primus_turbo") is not None |
| or shutil.which("primus") is not None |
| or os.environ.get("PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32") is not None |
| ) |
|
|
|
|
| def run_primus(cfg: XTrainConfig, plan: AutotunePlan, out_dir: Path) -> Path: |
| """Run a Primus pretraining job; return the checkpoint directory.""" |
| if not _primus_available(): |
| msg = ( |
| "primus_turbo not installed and not running inside rocm/primus:v26.2; " |
| "see ops/containerfiles/containerfile_train and pull the container." |
| ) |
| raise RuntimeError(msg) |
|
|
| out_dir = Path(out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| yaml_payload = compile_axolotl_yaml(cfg, plan) |
| yaml_path = out_dir / f"{cfg.meta.run_name}.primus.yaml" |
| yaml_path.write_text(yaml.safe_dump(yaml_payload, sort_keys=False)) |
|
|
| log_path = out_dir / "train.log" |
| env = dict(os.environ) |
| env["PYTORCH_ROCM_ARCH"] = "gfx942" |
| env["PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32"] = "1" |
| if plan.rccl_config == "8gpu_xgmi": |
| env["NCCL_MIN_NCHANNELS"] = "112" |
|
|
| cmd = ["python", "-m", "primus_turbo.train", "--config", str(yaml_path)] |
| with log_path.open("w") as log: |
| proc = subprocess.run(cmd, stdout=log, stderr=subprocess.STDOUT, env=env, check=False) |
| if proc.returncode != 0: |
| sys.stderr.write(f"primus returned {proc.returncode}; see {log_path}\n") |
| raise SystemExit(proc.returncode) |
|
|
| return out_dir / yaml_payload.get("output_dir", "checkpoint") |
|
|
|
|
| __all__ = ["run_primus"] |
|
|
