# Run SpecForge with Primus

`primus-cli` is the entrypoint for SpecForge **capture** and **draft training**.

| Mode | Values | Meaning |
| --- | --- | --- |
| `specforge_mode` | `train` \| `capture` | Which program Primus runs (`train` if omitted) |
| `specforge_train_mode` | `online` \| `offline` | How train runs. **Required** when mode is `train` (no default). Ignored for capture. |

`offline` train is `specforge train` on pre-captured hidden states. `online`
train is live Mooncake + SGLang capture plus `--role producer` / `--role
consumer`. Capture is SpecForge `scripts/prepare_hidden_states.py`, not a
train mode.

Build the image first: [runtime image](runtime-image.md). Inside the container,
Primus is `/opt/primus` and SpecForge is `/workspace/SpecForge`.

## Launch with `primus-cli`

- **Capture and offline train** use `direct` (inside the image), `container` (from the host), or `slurm`. On Slurm, `-N` is the node count for one SpecForge job: data parallel across nodes, the same command on every node, and a shared `OUTPUT_DIR` / `HIDDEN_STATES_PATH`.
- **Online disaggregated train** is `slurm` only. `-N` is capture nodes plus trainer nodes, not a homogeneous DP job.

`direct` is the in-image command. `container` starts this image from the host
and runs the same command. `slurm` allocates nodes, then each node runs
`container` or `direct`. Multi-GPU on one node is `NPROC_PER_NODE`.

Inside the image:

```bash
./runner/primus-cli direct -- train pretrain --config <experiment.yaml>
```

From the host, capture or offline train:

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

Multi-node capture or offline train (shared output path; `-N` is the node count):

```bash
./runner/primus-cli slurm srun -N 2 --gres=gpu:8 \
  -- container --image primus-specforge:v0.5.14-rocm700-mi35x \
  --volume /shared:/shared \
  -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-offline-capture.yaml
```

Dotted CLI keys override YAML, for example
`specforge_overrides.training.max_steps=1000`. See the
[configuration reference](configuration.md) for the full configuration, CLI,
and env reference.

One-node online (Mooncake + SGLang + producer + consumer on a single host) is
SpecForge `managed_local`: run the SpecForge CLI inside this image, not
`primus-cli`. Follow the
[AMD ROCm tutorial](https://github.com/sgl-project/SpecForge/blob/main/docs/sections/basic_usage/AMD/amd_rocm.md)
§4.

## Offline

Capture, then train on disk hidden states. Example pair for Qwen3.5-4B DFlash
(20-step smoke defaults).

**Capture** — `specforge_mode: capture` runs SpecForge
`scripts/prepare_hidden_states.py`. Recipe:
[qwen3.5-4b-dflash-offline-capture.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/specforge/configs/qwen3.5-4b-dflash-offline-capture.yaml).

```yaml
# examples/specforge/configs/qwen3.5-4b-dflash-offline-capture.yaml
modules:
  pre_trainer:
    framework: specforge
    overrides:
      specforge_mode: capture
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}
      specforge_capture:
        target_model_path: ${TARGET_MODEL:Qwen/Qwen3.5-4B}
        data_path: ${CAPTURE_DATA_PATH}
        output_path: ${OUTPUT_DIR}/hidden_states_raw
        nproc_per_node: ${NPROC_PER_NODE:1}
        nnodes: ${NNODES:1}
```

```bash
cd /opt/primus
export CAPTURE_DATA_PATH=/data/sharegpt.jsonl
export OUTPUT_DIR=/data/runs/capture
export NPROC_PER_NODE=8
./runner/primus-cli direct -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-offline-capture.yaml
```

**Train** — `specforge_mode: train` with `specforge_train_mode: offline`.
`specforge_config` points at SpecForge’s YAML; Primus forwards overrides onto
`specforge train`. Recipe:
[qwen3.5-4b-dflash-offline.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/specforge/configs/qwen3.5-4b-dflash-offline.yaml).

```yaml
# examples/specforge/configs/qwen3.5-4b-dflash-offline.yaml
modules:
  pre_trainer:
    framework: specforge
    overrides:
      specforge_mode: train
      specforge_train_mode: offline
      specforge_config: ${SPECFORGE_CONFIG:/workspace/SpecForge/examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml}
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}
      specforge_overrides:
        training.max_steps: ${MAX_STEPS:20}
        data.hidden_states_path: ${HIDDEN_STATES_PATH}
        deployment.trainer.nproc_per_node: ${NPROC_PER_NODE:1}
        deployment.trainer.nnodes: ${NNODES:1}
```

```bash
cd /opt/primus
export HIDDEN_STATES_PATH=/data/runs/capture/hidden_states
export OUTPUT_DIR=/data/runs/train
export NPROC_PER_NODE=8
export MAX_STEPS=20
./runner/primus-cli direct -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-offline.yaml
```

## Online

`specforge_mode: train` with `specforge_train_mode: online` is live
capture+train. SpecForge's own CLI will not start Mooncake/SGLang across
nodes, so Primus does that and then launches `--role producer` /
`--role consumer`.

Run `./runner/primus-cli slurm` on a Slurm cluster. Training runs **inside**
the docker image, or a bind-mount of this tree at `/opt/primus`. Example:
[qwen3.5-4b-dflash-online.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/specforge/configs/qwen3.5-4b-dflash-online.yaml).

```yaml
# examples/specforge/configs/qwen3.5-4b-dflash-online.yaml
modules:
  pre_trainer:
    framework: specforge
    overrides:
      specforge_mode: train
      specforge_train_mode: online
      specforge_config: ${SPECFORGE_CONFIG:/workspace/SpecForge/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml}
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}
      specforge_online:
        run_id: ${RUN_ID}
        run_root: ${RUN_ROOT}
        consumer_state_dir: ${CONSUMER_STATE_DIR}
        capture_nnodes: ${CAPTURE_NNODES:1}
        trainer_nnodes: ${TRAINER_NNODES:1}
        server_count: ${SERVER_COUNT:1}
        server_tp: ${SERVER_TP:1}
        trainer_nproc: ${NPROC_PER_NODE:1}
```

```bash
export RUN_ID=$(date -u +%Y%m%dT%H%M%SZ)
export RUN_ROOT=/path/to/shared/run/$RUN_ID
export OUTPUT_DIR=$RUN_ROOT/output
export CONSUMER_STATE_DIR=/path/to/local/nvme/specforge/$RUN_ID/consumer-state
export TRAIN_DATA_PATH=/path/to/sharegpt_train.jsonl
export MAX_STEPS=20
cd /opt/primus
./runner/primus-cli slurm srun -N 2 \
  --gres=gpu:1 \
  -- container --image primus-specforge:v0.5.14-rocm700-mi35x \
  --volume /shared:/shared \
  -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-online.yaml
```

The example above is **1 capture GPU + 1 trainer GPU** on 2 nodes. Load the
docker image on **every** node if the scheduler's container store is
node-local. Do not wrap this in `managed_local` or `--role both`. On one
node, use SpecForge `managed_local` inside the image (see Launch above), not
`primus-cli`. SpecForge `deployment.trainer.nnodes` is `TRAINER_NNODES`
(consumer nodes only).
`server_gpus` / `trainer_gpus` default to `0`; a single device id expands to
`0..N-1` (`N = server_count * server_tp` or `trainer_nproc`), so 8+8 only
needs `SERVER_COUNT=8` and `NPROC_PER_NODE=8`. Set an explicit CSV to pin
non-contiguous devices. Raise `-N` with `CAPTURE_NNODES` / `TRAINER_NNODES`
for more nodes.

### How it works

One Slurm allocation, the same `primus-cli` command on every node. Shape is
**C capture nodes + T trainer nodes** (`CAPTURE_NNODES` + `TRAINER_NNODES`,
default 1+1 so `-N 2`):

| Slurm `NODE_RANK` | Role |
| --- | --- |
| `0 .. C-1` | Capture: SGLang on GPU. Rank 0 also starts Mooncake and the SpecForge **producer** (CPU HTTP client to those servers). |
| `C .. C+T-1` | Wait for `inference.ready` → `--role consumer` (`--node-rank` when `T>1`) |

Rank 0 starts Mooncake, then local SGLang, collects every capture rank’s server
URLs, rewrites loopback into the producer and consumer commands, writes
`inference.ready`, and only then starts the producer. Other capture ranks start
SGLang only. Trainer ranks wait for `inference.ready`. Teardown is producer,
then SGLang, then Mooncake, after every `consumer.done`. When `C=1`, there are
no other capture ranks. When `T>1`, each consumer gets `--node-rank` inside
the trainer group.

```mermaid
sequenceDiagram
    participant Cli as primus-cli slurm
    participant Cap0 as Capture rank 0
    participant CapN as Other capture ranks
    participant Train as Trainer ranks
    Cli->>Cap0: NODE_RANK 0
    Cli->>CapN: NODE_RANK 1 to C-1
    Cli->>Train: NODE_RANK C to C+T-1
    Cap0->>Cap0: start Mooncake then local SGLang
    CapN->>Cap0: wait for head.ip
    CapN->>CapN: start SGLang only
    Cap0->>Cap0: collect server URLs and rewrite loopback
    Cap0->>Cap0: write inference.ready
    Cap0->>Cap0: specforge train --role producer
    Train->>Cap0: wait for inference.ready
    Train->>Train: specforge train --role consumer
    Train-->>Cap0: consumer.done
    Cap0->>Cap0: stop producer then SGLang then Mooncake
```

Primus maps `specforge_online.run_root` / `consumer_state_dir` onto SpecForge
`control_dir`, `output_dir`, and `consumer_state_dir`. After allocate it
rewrites the SpecForge loopback endpoints to a routable capture-rank-0 IP.
`PRIMUS_SPECFORGE_HEAD_IP` / `HEAD_IP` apply **only on capture rank 0**
(Mooncake + producer). Capture replicas and trainers advertise
`PRIMUS_SPECFORGE_LOCAL_IP` if set, else the IPv4 on
`PRIMUS_SPECFORGE_BIND_IFACE`, else hostname. A job-wide `HEAD_IP` must not be
the per-node Mooncake hostname. Shared vs local directories, empty
control/SQLite trees, and the producer/consumer contract are SpecForge's — see
[online disaggregated training](https://github.com/sgl-project/SpecForge/blob/main/docs/sections/basic_usage/AMD/amd_rocm.md#4-online-disaggregated-training).

## Reference results on MI355X

Experiment setup and ShareGPT data prep follow the
[SpecForge AMD ROCm tutorial](https://github.com/sgl-project/SpecForge/blob/main/docs/sections/basic_usage/AMD/amd_rocm.md).
Qwen3.5-4B DFlash, one epoch. Prompts were leak-filtered (exact-row and
first-user-turn) and **256 unique rows** were held out for eval. Train labels
are the target’s own **greedy thinking** continuations (`temperature=0`,
reasoning saved) — not original ShareGPT assistant text.

Both runs used the example YAMLs above via `primus-cli`, 8 trainer GPUs,
batch 2, accumulation 1, `log_interval=20`, `save_interval=240`,
`MAX_STEPS=2510` (~40k refs). Offline captured hidden states to disk
(`max_length=2048`) then trained. Online recaptured the same train prompts
live (8 SGLang `TP=1` + 8 trainer GPUs).

Offline train:

```bash
export MAX_STEPS=2510
export NPROC_PER_NODE=8
./runner/primus-cli direct -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-offline.yaml \
  specforge_overrides.training.num_epochs=1 \
  specforge_overrides.training.batch_size=2 \
  specforge_overrides.training.accumulation_steps=1 \
  specforge_overrides.training.save_interval=240 \
  specforge_overrides.training.log_interval=20
```

Online train (8+8):

```bash
export MAX_STEPS=2510
export NPROC_PER_NODE=8
export SERVER_COUNT=8
export SERVER_TP=1
# SERVER_GPUS / TRAINER_GPUS default 0 → 0..7 when the counts above are 8.
./runner/primus-cli slurm srun -N 2 --gres=gpu:8 \
  -- container --image primus-specforge:v0.5.14-rocm700-mi35x \
  --volume /shared:/shared \
  -- train pretrain \
  --config examples/specforge/configs/qwen3.5-4b-dflash-online.yaml \
  specforge_overrides.training.num_epochs=1 \
  specforge_overrides.training.batch_size=2 \
  specforge_overrides.training.accumulation_steps=1 \
  specforge_overrides.training.save_interval=240 \
  specforge_overrides.training.log_interval=20 \
  specforge_overrides.runtime.in_flight_high_watermark=64 \
  specforge_overrides.runtime.in_flight_low_watermark=32 \
  specforge_online.mooncake_lease_ttl_ms=8000
```

![train/loss, offline vs online](assets/train-loss-online-vs-offline.png)

![train/acc, offline vs online](assets/train-acc-online-vs-offline.png)

Held-out SGLang eval on **1 GPU** after `specforge export --to hf` and
normalizing the draft to `DFlashDraftModel` / `block_size=16`: `--tp-size 1`,
`--max-running-requests 1`, `--mem-fraction-static 0.75`,
`--attention-backend aiter`, `--disable-radix-cache`, `--disable-cuda-graph`,
`--reasoning-parser qwen3`. 256 prompts, thinking on, EOS on, 64 new tokens,
concurrency 1.

| | Offline | Online |
| --- | ---: | ---: |
| Target-only tok/s | 52.14 | 52.12 |
| DFLASH tok/s | **151.76** | **156.45** |
| MAL | **3.662** | **3.649** |
| Speedup | 2.911x | 3.002x |