# SpecForge ROCm runtime image

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

The published `primus` wheel does **not** include `primus/backends/specforge/`.
The [Dockerfile](https://github.com/AMD-AGI/Primus/blob/main/examples/specforge/docker/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](https://github.com/sgl-project/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`](https://github.com/AMD-AGI/Primus/blob/main/examples/specforge/docker/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.

```bash
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:

```bash
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

```bash
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):

```bash
./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](running.md).

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.