vLLM Serving for GRPO Training

How to configure and run vLLM as a generation backend for GRPO reinforcement learning in Axolotl.

1 Overview

GRPO (Group Relative Policy Optimization) trains a language model by generating completions, scoring them with reward functions, and updating the policy to favor higher-reward outputs. The generation step is the bottleneck: producing thousands of tokens per training step with the policy model is slow using standard HuggingFace generation.

Axolotl uses vLLM as a high-throughput generation backend, in one of two modes:

  • Server mode: vLLM runs as a separate process on dedicated GPUs (or dedicated nodes) and serves completions over its OpenAI-compatible HTTP API.
  • Colocate mode: the trainer hosts a vLLM engine on its own GPU, so a single GPU is enough.

Axolotl requires vllm>=0.30.0.

┌──────────────────────┐       HTTP        ┌──────────────────────┐
│   Trainer (GPU 1)    │ ───────────────── │   vLLM Server (GPU 0)│
│                      │  prompts/compls   │                      │
│  - Policy model      │ ◄──────────────── │  - Same base model   │
│  - Reward scoring    │                   │  - Fast generation   │
│  - Gradient updates  │  weight sync      │  - LoRA adapter      │
│  - LoRA adapter      │ ─────────────────►│    (periodically     │
│                      │  (every N steps)  │     updated)         │
└──────────────────────┘                   └──────────────────────┘
Important

vLLM must serve the same base model specified in your training config. If the models do not match, weight synchronization will silently produce incorrect results.

2 Server Mode

Server mode runs vLLM as an external process on dedicated GPU(s). This is the recommended configuration for most setups.

2.1 Starting the Server

Use the axolotl vllm-serve command with your training config:

# Terminal 1: Start vLLM on GPU 0
CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve grpo_config.yaml
# Terminal 2: Start training on GPU 1
CUDA_VISIBLE_DEVICES=1 axolotl train grpo_config.yaml

The server reads vLLM settings from the vllm: section of your config and launches vLLM’s own OpenAI-compatible server in dev mode with the NCCL weight-transfer engine (default: http://0.0.0.0:8000). With trl.vllm_lora_sync: true it also passes --enable-lora and --max-lora-rank (the largest of lora_r and lora_rank_pattern) and enables runtime LoRA loading.

If you prefer to run vLLM yourself, the equivalent raw command is:

VLLM_SERVER_DEV_MODE=1 VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_ALLOW_RUNTIME_LORA_UPDATING=True \
  vllm serve <model> --weight-transfer-config '{"backend": "nccl"}' \
  --logprobs-mode processed_logprobs --max-logprobs -1 --enable-scale-out \
  --enable-lora --max-lora-rank <r> --max-loras 2 --api-server-count 1

<r> is lora_r rounded up to a rank vLLM supports (8, 16, 32, 64, …). Drop VLLM_ALLOW_RUNTIME_LORA_UPDATING and the flags after --enable-scale-out when not using LoRA sync. --enable-scale-out is needed for image prompts.

Tip

Use tmux or screen to manage the vLLM server process. Typical startup time is 30-90 seconds depending on model size and whether CUDA graphs are captured.

2.2 Minimal Server Config

base_model: Qwen/Qwen2.5-1.5B-Instruct

vllm:
  host: 0.0.0.0
  port: 8000
  gpu_memory_utilization: 0.85
  dtype: auto
  max_model_len: 4096

rl: grpo
trl:
  use_vllm: true
  vllm_server_host: 0.0.0.0
  vllm_server_port: 8000
  vllm_server_timeout: 300

For a full GRPO config with rewards, dataset, LoRA sync, and async generation, see the GRPO Quick Start.

2.3 Multi-GPU vLLM

For larger models, use tensor parallelism across multiple GPUs:

vllm:
  tensor_parallel_size: 2
  gpu_memory_utilization: 0.85
# vLLM on GPUs 2,3; training on GPUs 0,1
CUDA_VISIBLE_DEVICES=2,3 axolotl vllm-serve grpo_config.yaml
CUDA_VISIBLE_DEVICES=0,1 axolotl train grpo_config.yaml --num-processes 2
Note

Due to how TRL maps vLLM device indices, the vLLM instance should use the last N GPUs (highest device indices), while training uses the first N.

data_parallel_size runs that many replicas of the model behind the same endpoint, for more generation throughput when the model already fits on tensor_parallel_size GPUs.

3 Multi-Node Serving

When one node cannot generate fast enough, or cannot hold the model, axolotl vllm-serve can span several nodes as a single vLLM server. Node 0 runs the API server and dispatches to headless engines on the other nodes. The trainer only talks to node 0, and weight sync reaches every engine through it.

For example, a 16-GPU setup (2 nodes x 8 GPUs) running 8 replicas at TP=2 (4 replicas per node). Use the same config on every node and launch each with its --node-rank:

vllm:
  tensor_parallel_size: 2
  data_parallel_size: 8     # total replicas across all nodes; GPUs = TP x PP x DP = 16
  nnodes: 2
  master_addr: 10.0.0.1     # node 0

trl:
  vllm_server_host: 10.0.0.1
  vllm_server_port: 8000
# Node 0 (10.0.0.1): API server + 4 replicas
axolotl vllm-serve grpo_config.yaml --node-rank 0

# Node 1: 4 headless replicas
axolotl vllm-serve grpo_config.yaml --node-rank 1

# Trainer node (10.0.0.3, must be reachable from the vLLM nodes).
# Start once http://10.0.0.1:8000/health returns 200.
VLLM_HOST_IP=10.0.0.3 axolotl train grpo_config.yaml
Layout Config Link between vLLM nodes
Whole replicas on each node (recommended) data_parallel_size is a multiple of nnodes Ethernet is fine for dense models (replicas are independent). MoE models synchronize expert layers across all replicas every forward pass, so use InfiniBand or RoCE
One replica spanning nodes (model does not fit on one node) data_parallel_size: 1, tensor_parallel_size = GPUs per node, pipeline_parallel_size = nnodes InfiniBand or RoCE

Requirements:

  • Every node needs the same model, packages, and config.
  • NCCL weight sync (full fine-tuning, or LoRA without vllm_lora_sync): trainer rank 0 joins an NCCL group with every vLLM GPU, so each vLLM node must reach the trainer at VLLM_HOST_IP on port 51216 (TRL’s vllm_group_port). Each sync broadcasts the full model across the network, so prefer LoRA sync on Ethernet.
  • LoRA sync: every vLLM node loads the adapter from its own filesystem. The trainer writes it under TMPDIR, so point TMPDIR on the trainer at storage that all vLLM nodes mount at the same path.
  • If NCCL picks the wrong network interface, set NCCL_SOCKET_IFNAME (or NCCL_IB_HCA) on all nodes.

These settings map to vLLM’s data parallel deployment and multi-node deployment. vLLM’s hybrid and external load-balancing modes are not supported: they expose one API server per node, but weight sync has to reach every engine through a single server.

4 Colocate Mode

Colocate mode runs a vLLM engine inside the trainer process, on the same GPU. No axolotl vllm-serve process is needed, so this is the way to run GRPO on a single GPU.

base_model: Qwen/Qwen2.5-1.5B-Instruct
rl: grpo
adapter: lora

trl:
  use_vllm: true
  vllm_mode: colocate
  vllm_enable_sleep_mode: true
  num_generations: 8
  max_completion_length: 512
  reward_funcs:
    - rewards.accuracy_reward

vllm:
  gpu_memory_utilization: 0.3   # fraction reserved for vLLM; the rest is for training
  max_model_len: 2048           # prompt + completion budget for the engine
  enforce_eager: true           # disable CUDA graphs if vLLM's custom all-reduce errors

The vllm: block configures the colocated engine: gpu_memory_utilization (default 0.3 here, not 0.9, since training shares the GPU), tensor_parallel_size (must divide the number of training processes), max_model_len, enforce_eager, enable_prefix_caching, and dtype. Always set max_model_len; otherwise vLLM reserves KV cache for the model’s full context.

With vllm_enable_sleep_mode: true, vLLM offloads its VRAM allocation when not actively generating, freeing memory for training. When the trainer needs new completions, vLLM wakes up and reclaims VRAM.

Warning

Colocate mode is significantly slower than server mode because generation and training cannot overlap. The GPU alternates between the two workloads. This mode is practical only for smaller models (up to ~3B on a 24 GB GPU).

Colocate mode cannot be combined with context_parallel_size > 1, with vllm_lora_sync, or with the async trainer (async_prefetch / use_data_producer) when training an adapter; config validation rejects those combinations.

When to use colocate mode:

  • You have exactly one GPU
  • The model fits in memory with both vLLM and training active (with sleep mode), or is small enough to time-share
  • You accept the performance tradeoff for simpler setup (no separate vLLM process to manage)

When to use server mode:

  • You have two or more GPUs
  • You want maximum throughput (generation overlaps with training via async prefetch)
  • You are running larger models (7B+)

5 LoRA Sync

LoRA sync is the recommended weight synchronization method when training with LoRA adapters. Instead of merging adapter weights into the base model and broadcasting the full merged weights over NCCL, it saves only the LoRA adapter files to the filesystem and tells vLLM to load them natively through its runtime LoRA loading endpoint.

5.1 How It Works

  1. The trainer calls model.save_pretrained() to write the LoRA adapter weights to a temporary directory
  2. The trainer drains in-flight requests (/pause?mode=wait), registers the adapter with /v1/load_lora_adapter as <base model>-v<N>, resumes generation, and addresses subsequent requests to that name; the adapter from two syncs back is unloaded
  3. vLLM loads the adapter using its native LoRA support (Punica kernels)
  4. Generation uses the updated adapter on the next request

5.2 Benefits

  • Smaller sync payload: Transfers ~40 MB of LoRA weights instead of ~1.4 GB+ of merged model weights (for a typical 0.5-3B model)
  • No NCCL communicator: Eliminates the need for a cross-GPU NCCL communication channel, removing GPU contention between vLLM generation and weight sync
  • Faster sync: ~200 ms per sync vs. 350 ms to 5+ seconds for NCCL merge sync
  • Simpler multi-GPU: No need to set up NCCL groups between trainer and vLLM processes

5.3 Configuration

adapter: lora
lora_r: 32
lora_alpha: 64
lora_target_linear: true

trl:
  vllm_lora_sync: true    # Enables LoRA sync mode
  vllm_sync_interval: 5   # Sync every 5 training steps

axolotl vllm-serve enables runtime LoRA loading automatically. A failed adapter load raises an error and stops training.

Important

LoRA sync requires that you are training with a LoRA adapter (adapter: lora or adapter: qlora). It is not applicable to full fine-tuning. The adapter directory must be on a filesystem the vLLM server can read (same host or a shared mount).

6 Weight Synchronization

During GRPO training, the policy model on the trainer is continuously updated via gradient steps. The vLLM server, however, still holds the old weights. Periodically, the trainer must push updated weights to vLLM so that future generations reflect the improved policy.

6.1 Sync Interval

The vllm_sync_interval parameter controls how often weights are synced:

trl:
  vllm_sync_interval: 5   # Sync every 5 optimizer steps

Tradeoffs:

  • Lower interval (e.g., 1-3): Fresher generations, better on-policy data, but more sync overhead per step
  • Higher interval (e.g., 5-10): Less overhead, but generations become increasingly off-policy between syncs
  • Recommended: 3-5 for most setups. Axolotl includes importance sampling correction (vllm_importance_sampling_correction: true) to handle mild distribution mismatch from stale vLLM weights.

6.2 Sync Methods

Method Config Payload Mechanism Typical Time
LoRA sync vllm_lora_sync: true LoRA adapter only (~40 MB) Filesystem + /v1/load_lora_adapter ~200 ms
NCCL merge sync Default (no lora_sync) Full merged weights (~1.4 GB+) Native NCCL weight-transfer engine 350 ms - 5 s
Tip

If you are training with LoRA (which is recommended for GRPO), always enable vllm_lora_sync: true. The performance difference is substantial, especially as training progresses and NCCL contention increases.

NCCL merge sync folds the adapter into the base weights out of place. FP8 base weights are requantized with fresh block scales before being sent. DoRA adapters and bitsandbytes 4-bit bases (QLoRA) cannot be merged this way and raise an error; use vllm_lora_sync: true for them.

6.3 Merged Sync Rounding

Merged sync serves bf16(W + scaling * B @ A), or its FP8 requantization for FP8 bases, while the trainer optimizes the unrounded W + scaling * B @ A. At RL learning rates most entries of the LoRA delta are smaller than half a BF16 step of the base weight (about 0.2% of |w|). The cast rounds those entries back to the base value, so vLLM samples from a policy that has only part of the trained update. On Qwen3-1.7B-Base with Dr.GRPO and rank-16 LoRA on all linear layers, the merged weights at step 100 kept about 24% of the update at learning_rate: 1e-5 and about 2.5% at 1e-6. Full fine-tuning is not affected, because the trainer and vLLM hold the same BF16 weights.

The trainer measures this at the first merged sync and then every trl.vllm_merge_check_interval optimizer steps (default 50; 0 disables). It logs two metrics:

  • lora_merge/retained: the share of the LoRA delta that survives the cast, as the projection of the served weight change onto the delta.
  • lora_merge/changed_fraction: the share of weight entries the merge changes.

When less than half of the update survives, every check after the first logs a warning. Parameters sharded by FSDP or DeepSpeed ZeRO-3 are not measured.

There are two ways to keep the whole update:

Option Config What vLLM serves Use when
LoRA sync trl.vllm_lora_sync: true Base weights plus the unmerged adapter Server mode with the async GRPO trainer (use_data_producer or async_prefetch). It is also the faster sync.
Merge-aware LoRA lora_merge_aware: bf16 Merged BF16 weights, which the trainer now optimizes Weights must be merged: colocate mode, the standard GRPO trainer (neither use_data_producer nor async_prefetch), or a merged BF16 checkpoint

With lora_merge_aware: bf16, each LoRA linear computes x @ bf16(W + scaling * B @ A)^T, bitwise equal to the merged weight that is sent to vLLM, and gradients reach A and B straight through the cast. The forward and backward passes rebuild every LoRA linear’s merged weight, at the cost of a few passes over the targeted base weights per micro-batch, and no more activations are saved than with ordinary LoRA. lora_merge_aware: bf16 requires adapter: lora on an unquantized base. It works on a single GPU and with DDP, DeepSpeed ZeRO-1/2 and FSDP2:

  • Without FSDP, the sync merges in place inside temporarily_merge_adapter, which restores the base weights exactly afterwards.
  • Under FSDP2, lora_fp32_gradients is enabled automatically. FSDP2 then gathers the LoRA factors in FP32 rather than casting them to BF16, so the forward and the served merge compute the same delta. The vLLM sync merges each LoRA linear out of place from its gathered weights instead of merging the sharded weights in place.

Config validation rejects it with:

  • DoRA;
  • the fused LoRA kernels (lora_*_kernel);
  • 4/8-bit bases;
  • DeepSpeed ZeRO-3, because the merge-aware forward reads weights of submodules whose forwards it skips, so ZeRO-3 never gathers them;
  • tensor parallelism;
  • vllm_lora_sync.

LoRA linears whose base weight is FP8 or otherwise quantized keep the ordinary forward, with a warning.

axolotl merge-lora (memory-efficient method) reports the same retained share when it merges into a BF16 or FP16 base, and warns below one half unless the config sets lora_merge_aware.

6.4 Importance Sampling Correction

When vLLM weights are stale (between syncs), the generated data is slightly off-policy. Axolotl can correct for this:

trl:
  vllm_importance_sampling_correction: true
  importance_sampling_level: token          # 'token' or 'sequence'
  off_policy_mask_threshold: 0.5            # KL threshold for masking stale sequences
  • Token-level IS is recommended when using Liger kernel (sequence-level has numerical issues with chunked computation)
  • Off-policy sequence masking (OPSM) drops sequences that have diverged too far from the current policy, providing a safety net against stale data

7 Restart Requirements

Warning

vLLM must be restarted between training runs. Weight syncs from a previous run leave the server in a corrupted state. If you start a new training run against a stale vLLM server, the model may fail to learn.

Restart before every new experiment, after a crashed or interrupted run, or when the base model changes.

A plain kill can leave vLLM’s EngineCore subprocess holding GPU memory. Kill both, then confirm the memory is freed:

pkill -9 -f "vllm|EngineCore"
nvidia-smi
CUDA_VISIBLE_DEVICES=0 axolotl vllm-serve grpo_config.yaml

For multi-node, do this on every node and relaunch each with its --node-rank.

Wait for the health endpoint to return 200 before starting training:

curl http://localhost:8000/health

8 Configuration Reference

8.1 vLLM Server Options (vllm: section)

These control the vLLM server process started by axolotl vllm-serve, and the colocated engine in colocate mode.

Option Type Default Description
host str 0.0.0.0 Host address for the vLLM server
port int 8000 Port for the vLLM server
device str auto Device to use for vLLM
tensor_parallel_size int None Number of GPUs for tensor parallelism
data_parallel_size int None Number of data parallel replicas (total across all nodes)
pipeline_parallel_size int None Number of pipeline stages; for a replica spanning nodes
nnodes int None Number of nodes for multi-node serving; launch each with --node-rank
master_addr str None Address of node 0 when nnodes > 1
gpu_memory_utilization float None Fraction of GPU memory for vLLM (0.0-1.0). Defaults to 0.9 for vllm-serve and 0.3 in colocate mode
dtype str auto Data type (auto, float16, bfloat16)
max_model_len int None Maximum model context length. Set explicitly if the default is too large for your GPU
enable_prefix_caching bool None Enable prefix caching for repeated prompt prefixes
enable_reasoning bool None Enable reasoning mode for models with thinking tokens
reasoning_parser str None Parser for reasoning output
enforce_eager bool None Disable CUDA graph capture (required for some architectures like Qwen3.5 hybrid attention)
serve_module str None Optional custom module exposing main(script_args). Default is the native vLLM server; legacy trl.scripts.vllm_serve and axolotl.scripts.vllm_serve_lora are ignored
worker_extension_cls str None Forwarded to vLLM as --worker-extension-cls; not needed for weight sync. Legacy trl.scripts.vllm_serve.* and axolotl.scripts.* classes are ignored

8.2 Trainer vLLM Options (trl: section)

These control how the trainer interacts with vLLM.

Option Type Default Description
use_vllm bool false Enable vLLM for generation
vllm_mode str None server (external process) or colocate (same GPU)
vllm_server_host str 0.0.0.0 Host of the vLLM server to connect to
vllm_server_port int 8000 Port of the vLLM server to connect to
vllm_server_timeout int None Timeout in seconds for vLLM requests
vllm_lora_sync bool false Sync LoRA adapters via filesystem instead of NCCL merge
vllm_sync_interval int None Sync weights every N optimizer steps
vllm_merge_check_interval int 50 Measure how much of the LoRA update survives merged sync every N steps; 0 disables. See Merged Sync Rounding
vllm_enable_sleep_mode bool None Offload vLLM VRAM when idle (colocate mode)
vllm_guided_decoding_regex str None Regex constraint for structured/guided decoding (mapped to TRL’s vllm_structured_outputs_regex on TRL >= 1.7)

For async pipeline and off-policy correction options, see the GRPO Configuration Reference.

9 Troubleshooting

Problem Likely Cause Solution
Training hangs waiting for vLLM Server not started or wrong port Check curl http://localhost:8000/health and verify vllm_server_host/vllm_server_port match
OOM on vLLM GPU gpu_memory_utilization too high or max_model_len too large Reduce gpu_memory_utilization to 0.7 or set max_model_len explicitly
OOM on training GPU Batch too large for policy logprobs Reduce micro_batch_size or num_generations
Accuracy stays at zero Stale vLLM from previous run Restart vLLM: pkill -9 -f "vllm\|EngineCore", verify with nvidia-smi, restart
ResponseValidationError from vLLM Missing logprobs in response Start the server with axolotl vllm-serve (or the raw command above) so --logprobs-mode processed_logprobs --max-logprobs -1 are set
Multi-node training hangs at startup A vLLM node cannot reach the trainer for the NCCL weight-sync group Set VLLM_HOST_IP on the trainer to an address the vLLM nodes can reach, allow port 51216, and set NCCL_SOCKET_IFNAME if needed
Weight sync takes 5+ seconds NCCL contention with vLLM generation Switch to vllm_lora_sync: true to eliminate NCCL
async_prefetch deadlocks with FSDP Background threads run unsynchronized FSDP collectives Set async_prefetch: false when using FSDP or DeepSpeed multi-GPU