SpecForge ROCm runtime image

SpecForge ROCm runtime image#

Image for Run SpecForge with Primus. That page is the primus-cli entrypoint and MI355X results; YAML/env/CLI knobs are in configuration. This page is the install reference.

The published primus wheel does not include primus/backends/specforge/. The Dockerfile installs this Primus checkout editable with --no-deps.

What Primus owns#

Piece

Owner

ROCm SGLang + PyTorch

Official SGLang image (do not let pip replace it)

Draft architecture, data scripts, specforge train / specforge export

SpecForge

Capture / train launch, experiment YAML, ROCm stack preflight

Primus framework: specforge

Conflict policy#

Component

Install method

SGLang / torch

Keep the base image versions

SpecForge

pip install -e . --no-deps at a pinned git SHA

Primus

pip install -e /opt/primus --no-deps (this tree, not the PyPI wheel)

Extra imports

Allowlist in install-light-deps.sh, each --no-deps

tensorboard

Exception. The Dockerfile installs tensorboard>=2.14,<2.21 with dependencies, constrained to versions already in the base image, so pip cannot upgrade numpy, protobuf, grpcio, setuptools, or pillow. SummaryWriter imports tensorboard’s own tree, and a --no-deps copy leaves --report-to tensorboard broken. tensorboard 2.21 requires protobuf>=6.31.1 and grpcio>=1.74; the upper bound keeps that release out so the constraint holds instead of upgrading the SGLang stack.

Do not run primus-cli deps sync in this image. That path installs Megatron / CUDA wheels and will clobber the ROCm stack.

Build#

Build from the Primus repository root (this repository). The trailing . is the Docker build context; the Dockerfile copies primus/, runner/, primus_cli.py, and examples/specforge/ from that context into /opt/primus. Do not use a SpecForge checkout as the context.

docker build -f examples/specforge/docker/Dockerfile \
  -t primus-specforge:v0.5.14-rocm700-mi35x .

Pins (BASE_IMAGE digest and SPECFORGE_REF) are Dockerfile ARG defaults. A plain docker build uses them; --build-arg overrides.

The Dockerfile:

  1. Starts from the pinned lmsysorg/sglang:v0.5.14-rocm700-mi35x digest.

  2. Clones SpecForge at SPECFORGE_REF and applies patches/sglang/v0.5.14/spec-capture.patch.

  3. Installs SpecForge editable --no-deps.

  4. Copies this Primus tree and installs it editable --no-deps.

  5. Fills missing light packages (accelerate, wandb, …) --no-deps.

  6. Installs tensorboard>=2.14,<2.21 with dependencies, constrained to the base image’s already-installed versions.

  7. Runs GPU-free smoke checks (enable_spec_capture, import specforge, a SummaryWriter event-file write, primus-cli --help, HIP torch).

Validated on AMD Instinct MI355X (gfx950). Newer SGLang ROCm tags may work if they still ship AITER and the speculative algorithm you will serve.

For MI300X (gfx942), override the base tag to an mi30x image that still has those pieces:

docker build -f examples/specforge/docker/Dockerfile \
  --build-arg BASE_IMAGE=lmsysorg/sglang:v0.5.14-rocm700-mi30x \
  -t primus-specforge:v0.5.14-rocm700-mi30x .

Prefer a digest on BASE_IMAGE if you have one, so the override is pinned too.

Run#

docker run -it --name primus-specforge \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined \
  --ipc=host --shm-size=64g --network host \
  -e HF_HOME=/data/hf-cache \
  -e HF_TOKEN \
  -v /path/to/data:/data \
  primus-specforge:v0.5.14-rocm700-mi35x \
  bash

--device=/dev/kfd --device=/dev/dri --group-add video are required for ROCm. --ipc=host --shm-size=64g is enough shared memory for multi-GPU capture/train. Re-enter with docker exec -it primus-specforge bash.

From the host, primus-cli container is the same as docker run plus primus-cli direct inside. ROCm devices are mapped by default; set --shm-size 64g and this image (not the default rocm/primus tag):

./runner/primus-cli container --image primus-specforge:v0.5.14-rocm700-mi35x \
  --shm-size 64g --volume /path/to/data:/data \
  -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-offline-capture.yaml

Inside the image:

  • Primus: /opt/primus (WORKDIR)

  • SpecForge: /workspace/SpecForge (SPECFORGE_ROOT)

  • Capture patch is already applied to the image SGLang

Set HF_TOKEN if the target model is gated. Point HF_HOME at a durable cache so model weights are not re-downloaded.

Launch capture or train with primus-cli from /opt/primus. See Run SpecForge with Primus.

Capture and online SGLang on this image use AITER with the radix cache off. Primus sets that at launch; you do not put it in YAML or export it.

Contents