Skip to content

CONFIGURATION.md

Latest commit

 

History

History
304 lines (258 loc) · 13.8 KB

CONFIGURATION.md

File metadata and controls

304 lines (258 loc) · 13.8 KB
Raw

Configuration reference

YAML, env, and CLI for SpecForge on Primus. Copy a recipe under configs/ and change the fields below. Draft architecture, data format, optimizer, export, and SGLang serving are SpecForge's — see the AMD ROCm tutorial (§4 for online disaggregated). SpecForge Hydra keys (specforge_overrides) are documented in SpecForge examples/configs/README.md.

Set a value in the experiment YAML, or override it on the CLI with the same path (specforge_online.server_count=8). ${VAR:default} is expanded when the experiment file loads. Export the variable, or put the literal in YAML.

modules:
  pre_trainer:
    framework: specforge                 # required
    config: offline.yaml                 # offline.yaml | online.yaml
    model: qwen3.5-4b-dflash.yaml        # Primus model preset
    overrides:
      specforge_mode: train              # train | capture  (default train)
      specforge_train_mode: online       # online | offline; required for train
      specforge_config: ...              # SpecForge Hydra YAML; required for train
      specforge_root: ...                # SpecForge checkout (this image: /workspace/SpecForge)
      output_dir: ...                    # checkpoints / logs
      specforge_overrides: {}            # SpecForge Hydra key=value (train)
      specforge_capture: {}              # capture only
      specforge_online: {}               # online train only

Do not set specforge_role. Primus assigns producer/consumer from node rank. Do not pass --role on primus-cli.

Capture

Example config and usage: configs/qwen3.5-4b-dflash-offline-capture.yaml

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

Full config reference. Extra specforge_capture keys are flags for SpecForge scripts/prepare_hidden_states.py.

work_group: ${PRIMUS_TEAM:amd}            # Primus experiment metadata
user_name: ${PRIMUS_USER:root}            # Primus experiment metadata
exp_name: ${PRIMUS_EXP_NAME:qwen3.5-4b-dflash-offline-capture}  # run name
workspace: ${PRIMUS_WORKSPACE:./output}  # Primus workspace

modules:
  pre_trainer:
    framework: specforge                 # required backend
    config: offline.yaml                 # Primus module preset
    model: qwen3.5-4b-dflash.yaml        # Primus model preset
    overrides:
      specforge_mode: capture            # train | capture
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}  # SpecForge checkout
      output_dir: ${OUTPUT_DIR}          # checkpoints / logs

      specforge_capture:
        target_model_path: ${TARGET_MODEL:Qwen/Qwen3.5-4B}  # HF id or local path
        data_path: ${CAPTURE_DATA_PATH}  # ShareGPT-style JSONL
        output_path: ${OUTPUT_DIR}/hidden_states_raw  # raw shards
        nproc_per_node: ${NPROC_PER_NODE:1}  # GPUs for capture
        nnodes: ${NNODES:1}              # capture nodes; slurm -N
        filter_output_path: ${OUTPUT_DIR}/hidden_states  # filtered shards (train on these)
        filter_block_size: ${BLOCK_SIZE:16}  # DFlash block size
        filter_min_kept: ${FILTER_MIN_KEPT:1}  # fail if fewer shards survive
        strategy: dflash                 # SpecForge capture strategy
        draft_model_config: configs/qwen3.5-4b-dflash.json  # relative to specforge_root
        trust_remote_code: true          # allow custom HF code
        cache_dir: ${OUTPUT_DIR}/sf-cache  # download / KV cache
        chat_template: qwen3.5           # chat template name
        max_length: 2048                 # prompt + completion tokens
        tp_size: 1                       # SGLang tensor parallel
        batch_size: ${CAPTURE_BATCH_SIZE:8}  # capture batch
        sglang_mem_fraction_static: 0.8  # SGLang GPU memory fraction
        sglang_context_length: 2560      # SGLang context length

Train later with data.hidden_states_path pointing at filter_output_path (the filtered shards), not output_path.

Offline train

Example config and usage: configs/qwen3.5-4b-dflash-offline.yaml

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

Full config reference. specforge_overrides keys are SpecForge Hydra — see examples/configs/README.md.

work_group: ${PRIMUS_TEAM:amd}            # Primus experiment metadata
user_name: ${PRIMUS_USER:root}            # Primus experiment metadata
exp_name: ${PRIMUS_EXP_NAME:qwen3.5-4b-dflash-offline}  # run name
workspace: ${PRIMUS_WORKSPACE:./output}  # Primus workspace

modules:
  pre_trainer:
    framework: specforge                 # required backend
    config: offline.yaml                 # Primus module preset
    model: qwen3.5-4b-dflash.yaml        # Primus model preset
    overrides:
      specforge_mode: train              # train | capture
      specforge_train_mode: offline      # online | offline
      specforge_config: ${SPECFORGE_CONFIG:/workspace/SpecForge/examples/configs/offline/colocated/qwen3.5-4b-dflash-offline-amd.yaml}  # SpecForge Hydra YAML
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}  # SpecForge checkout
      output_dir: ${OUTPUT_DIR}          # checkpoints / logs

      specforge_overrides:
        training.max_steps: ${MAX_STEPS:20}  # optimizer steps
        training.num_epochs: 1           # epoch cap (max_steps usually wins first)
        training.save_interval: ${MAX_STEPS:20}  # checkpoint every N steps
        training.log_interval: 5         # log every N steps
        model.use_liger_kernel: false    # not shipped in this image
        data.hidden_states_path: ${HIDDEN_STATES_PATH}  # filtered capture output
        deployment.trainer.nproc_per_node: ${NPROC_PER_NODE:1}  # trainer GPUs
        deployment.trainer.nnodes: ${NNODES:1}  # trainer nodes; slurm -N

Online train

Example config and usage: configs/qwen3.5-4b-dflash-online.yaml

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
./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

-N is capture nodes + trainer nodes. --gres=gpu:N is per node: SGLang GPUs on capture nodes (server_count × server_tp), trainer processes on trainer nodes (trainer_nproc). A single GPU id 0 expands to 0..N-1. Rank 0 starts Mooncake and the SpecForge producer; other capture ranks run SGLang only; trainer ranks train.

The SpecForge Hydra file (specforge_config) lists SGLang and Mooncake as http://127.0.0.1:…. Leave those loopback URLs in that file. After the job allocates, Primus rewrites them to the capture nodes. Put run_root on shared storage and consumer_state_dir on node-local disk. Both must be empty for a new RUN_ID.

Full config reference. specforge_overrides keys are SpecForge Hydra — see examples/configs/README.md. Online layout is also in the AMD ROCm tutorial §4.

work_group: ${PRIMUS_TEAM:amd}            # Primus experiment metadata
user_name: ${PRIMUS_USER:root}            # Primus experiment metadata
exp_name: ${PRIMUS_EXP_NAME:qwen3.5-4b-dflash-online}  # run name
workspace: ${PRIMUS_WORKSPACE:./output}  # Primus workspace

modules:
  pre_trainer:
    framework: specforge                 # required backend
    config: online.yaml                  # Primus module preset
    model: qwen3.5-4b-dflash.yaml        # Primus model preset
    overrides:
      specforge_mode: train              # train | capture
      specforge_train_mode: online       # online | offline
      specforge_config: ${SPECFORGE_CONFIG:/workspace/SpecForge/examples/configs/online/disaggregated/external/qwen3.5-4b-dflash-online-amd.yaml}  # SpecForge Hydra YAML
      specforge_root: ${SPECFORGE_ROOT:/workspace/SpecForge}  # SpecForge checkout
      output_dir: ${OUTPUT_DIR}          # checkpoints / logs

      specforge_online:
        run_id: ${RUN_ID}                # required; letters, digits, . _ -
        run_root: ${RUN_ROOT}            # required; shared control dir
        consumer_state_dir: ${CONSUMER_STATE_DIR}  # required; local NVMe
        capture_nnodes: ${CAPTURE_NNODES:1}  # capture ranks
        trainer_nnodes: ${TRAINER_NNODES:1}  # trainer ranks; -N = capture + trainer
        server_count: ${SERVER_COUNT:1}  # SGLang servers per capture node
        server_tp: ${SERVER_TP:1}        # tensor parallel per SGLang server
        server_gpus: ${SERVER_GPUS:0}    # SGLang GPU ids; 0 means 0,1,…,server_count*server_tp-1
        trainer_gpus: ${TRAINER_GPUS:0}  # trainer GPU ids; 0 means 0,1,…,trainer_nproc-1
        trainer_nproc: ${NPROC_PER_NODE:1}  # trainer processes per trainer node
        target_model_path: ${TARGET_MODEL:Qwen/Qwen3.5-4B}  # HF id or local path
        mooncake_protocol: ${MOONCAKE_PROTOCOL:tcp}  # Mooncake transport
        mooncake_lease_ttl_ms: ${MOONCAKE_LEASE_TTL_MS:500}  # KV lease TTL
        # capture_layer_ids: 1,8,15,22,29  # else from SpecForge draft json
        # server_port: 30000             # first SGLang HTTP port
        # server_mem_fraction: 0.85      # SGLang GPU memory fraction
        # mooncake_rpc_port: 35551       # Mooncake RPC
        # mooncake_http_port: 35880      # Mooncake HTTP metadata
        # mooncake_metrics_port: 35903   # Mooncake metrics
        # bind_interface: ens3           # NIC if auto IP is wrong
        # start_timeout_s: 1800          # sidecar start wait
        # peer_timeout_s: 1800           # wait for the other ranks

      specforge_overrides:
        training.max_steps: ${MAX_STEPS:20}  # optimizer steps
        training.num_epochs: 1           # epoch cap (max_steps usually wins first)
        training.save_interval: ${MAX_STEPS:20}  # checkpoint every N steps
        training.log_interval: 5         # log every N steps
        model.use_liger_kernel: false    # not shipped in this image
        data.train_data_path: ${TRAIN_DATA_PATH}  # ShareGPT-style JSONL
        # runtime.in_flight_high_watermark: 64  # SpecForge online queue
        # runtime.in_flight_low_watermark: 32

On Ethernet without IB, pass NCCL_IB_DISABLE=1 and NCCL_SOCKET_IFNAME=<iface>.

Skip the next block unless auto-detected IPs are wrong (Mooncake/SGLang not reachable across nodes). Then set the cluster NIC, or pin IPs:

# NIC name, or:
export PRIMUS_SPECFORGE_BIND_IFACE=ens3
# capture rank 0:
export PRIMUS_SPECFORGE_HEAD_IP=10.0.0.1
# every node:
export PRIMUS_SPECFORGE_LOCAL_IP=10.0.0.2

CLI

direct runs inside the image. container starts the image from the host and then runs direct. slurm allocates, then each node runs container or direct.

Capture and offline train are homogeneous: direct or container on one node, or slurm srun -N <nnodes> for DP across nodes. The same command runs on every node. Point OUTPUT_DIR / HIDDEN_STATES_PATH at shared storage. Multi-GPU on one node is NPROC_PER_NODE. On Ethernet without IB, pass NCCL_IB_DISABLE=1 and NCCL_SOCKET_IFNAME=<iface>.

Online train is slurm only. -N must equal CAPTURE_NNODES + TRAINER_NNODES (default 1+1 so -N 2). Do not use direct or container on a single host. Do not wrap SpecForge managed_local; for one-node online, run the SpecForge CLI inside the image (AMD ROCm tutorial §4). specforge export and data-prep scripts stay on the SpecForge CLI.

Inside the docker image:

./runner/primus-cli direct -- train pretrain --config <experiment.yaml> \
  [dotted.overrides...]

From the host:

./runner/primus-cli container --image primus-specforge:v0.5.14-rocm700-mi35x \
  --shm-size 64g --volume /data:/data --env KEY=VALUE \
  -- train pretrain --config <experiment.yaml> \
  [dotted.overrides...]

On Slurm, from the login node, image already loaded. Capture or offline (-N is the node count):

./runner/primus-cli slurm srun -N <nnodes> --gres=gpu:<gpus-per-node> \
  -- container --image <tag> \
  --volume <host>:<container> \
  --env KEY=VALUE \
  -- train pretrain --config <experiment.yaml> \
  [dotted.overrides...]

Online (-N is capture nodes + trainer nodes):

./runner/primus-cli slurm srun -N <capture+trainer> --gres=gpu:<gpus-per-node> \
  -- container --image <tag> \
  --volume <host>:<container> \
  --env KEY=VALUE \
  -- train pretrain --config <experiment.yaml> \
  [dotted.overrides...]

Dotted keys match the YAML (specforge_overrides.training.max_steps=1000, specforge_online.mooncake_lease_ttl_ms=8000). --exp is an alias for --config. --volume and --env are repeatable. --backend_path points at a SpecForge checkout if it is not /workspace/SpecForge.

Slurm -p, -t, -o, --exclude and launcher --debug / --dry-run are in the CLI reference.