[Spike]Validate Qwen3-4B checkpoint save/restore flow (test_qwen3_4B_ckpt)

[CalvinXKY]

Summary

Spike to validate whether the checkpoint save/restore flow in tests/test_qwen3_4B_ckpt.py works on vime + vLLM colocate (Qwen3-4B), in support of #11 (Checkpoint save/restore).

This issue records test procedure and observations from a manual run — not a bug filing.

Relation to #11

#11 tracks checkpoint save/restore for the Qwen3-4B baseline. The upstream CI test already defines a two-job workflow:

Phase	Purpose	Key train.py args
Save	Short colocated run + periodic checkpoint	--save, --save-interval 2, no --load
Load	New Ray job resumes from a fixed step	--load <dir>, --ckpt-step 1

CI today runs this with the SGLang rollout stack. This spike runs the same logic with vLLM colocate.

Optimizer / checkpoint args (same as test_qwen3_4B_ckpt.py)

Besides standard Adam / GRPO settings, the spike uses the same optimizer-related flags as the upstream test:

• --optimizer-cpu-offload — keep optimizer state on CPU to save GPU memory
• --overlap-cpu-optimizer-d2h-h2d — overlap CPU↔GPU copies for the optimizer
• --use-precision-aware-optimizer — Megatron precision-aware optimizer (PAO); affects how optimizer state is stored in torch_dist checkpoints

Optional (CI variant): --async-save on the save phase only.

Topology & vLLM

• 8 GPUs, colocate: --colocate, --actor-num-gpus-per-node 8
• Rollout (Ray resource, not in VLLM_ARGS): --rollout-num-gpus 8 (must match actor GPU count so colocate PG sizing is correct)
• vLLM: --rollout-backend vllm, --rollout-num-gpus-per-engine 2, --vllm-gpu-memory-utilization 0.7
• Weight sync: vLLM colocate IPC (same branch requirement as other Qwen3-4B vLLM colocate smokes, e.g. test/with_ipc)

Spike runner script

RL workspace (outside vime repo): run_scripts/qwen3-4B-vllm-ckpt-verify.sh
Orchestration env vars (script only): CKPT_PHASE=save|load|all, ASYNC_SAVE=0|1, CONCURRENT_RAY=0|1.

Prerequisites

• 8× GPU node (we used 8× A100 80GB)
• Qwen3-4B HF checkpoint + Qwen3-4B_torch_dist
• dapo-math-17k prompt data
• vime on a branch with vLLM colocate IPC weight sync

Test steps

1. Paths

cd <RL-workspace>
git -C vime checkout <branch-with-vllm-colocate-ipc>

export HF_CKPT=/path/to/Qwen3-4B
export REF_LOAD=/path/to/Qwen3-4B_torch_dist
export SAVE_DIR=/path/to/Qwen3-4B_vllm_ckpt
export PROMPT_DATA=/path/to/dapo-math-17k/dapo-math-17k.jsonl

2. Full spike (save → load)

Aligns with test_qwen3_4B_ckpt.py (including the three optimizer flags above):

CONCURRENT_RAY=0 CKPT_PHASE=all \
  bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

Save phase (first Ray job):

• --num-rollout 3, --save-interval 2, --save ${SAVE_DIR}
• --colocate, --actor-num-gpus-per-node 8, --rollout-num-gpus 8
• Optimizer flags listed in the table above

Load phase (second Ray job, after Ray/GPU cleanup):

• --load ${SAVE_DIR}, --ckpt-step 1
• Same training / vLLM / optimizer flags as save phase

3. Optional variants

# Save only
CKPT_PHASE=save CONCURRENT_RAY=0 bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

# Load only (after save succeeded)
CKPT_PHASE=load CKPT_STEP=1 CONCURRENT_RAY=0 bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

# Async save (CI: python tests/test_qwen3_4B_ckpt.py --async-save)
ASYNC_SAVE=1 CKPT_PHASE=all CONCURRENT_RAY=0 bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

4. Log / artifact checks

Save phase

• successfully saved checkpoint from iteration … to ${SAVE_DIR}
• ${SAVE_DIR}/iter_0000001/, iter_0000002/, latest_checkpointed_iteration.txt
• Ray: Job succeeded

Load phase

• loading distributed checkpoint from ${SAVE_DIR} at iteration <CKPT_STEP>
• Training / rollout continues
• Ray: Job succeeded

[CalvinXKY]

Spike results (2025-05-25, branch: test/with_ipc https://github.com/vllm-project/vime/tree/test/with_ipc)

Step	Result	Notes
Save (with --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer)	Pass	Checkpoints at iter 1 and 2; Ray job succeeded
Load (--ckpt-step 1 / 2)	Not pass	Megatron dist ckpt failed merging optimizer param_state (list length 81 vs 52) after restart
vLLM colocate training + save	Pass	Save path independent of load-phase optimizer restore

[CalvinXKY]

When running vime + vLLM colocated Qwen3-4B checkpoint verification aligned with tests/test_qwen3_4B_ckpt.py (including the PAO trio: --optimizer-cpu-offload, --overlap-cpu-optimizer-d2h-h2d, --use-precision-aware-optimizer):

• Save usually succeeds and writes iter_0000001 / iter_0000002
• Load (with optimizer) consistently fails with:

ValueError: Cannot merge two lists with different lengths (81 and 52)
  @ optimizer.param_state / (torch.bfloat16, torch.float32)

Findings from comparison runs:

• Not an NFS issue (same failure with container-local paths under /root/ckpt_validate/)
• Not universal on slime images: slimerl/slime:latest + SGLang + PAO with local ckpt had load-only pass in one run; vime + vLLM + PAO fails on full save→load
• With CKPT_NO_OPTIM_OFFLOAD=1 (PAO/offload disabled), save+load passes → tied to PAO + distributed optimizer checkpoint sharding

---

Related slime / Megatron work

Resource	Notes
THUDM/slime#1221	Save-time shape assertion in distrib_optimizer.py with optimizer-cpu-offload + PAO
PR #1230 (merged)	continue on key == 'step' in sharded_param_state_dp_reshardable — mainly fixes save assertions
Megatron --dist-ckpt-optim-fully-reshardable	Stores optimizer as fully_reshardable; fixes load merge 81/52 in our case

Note: the vime image already includes the #1230 step skip patch, but that alone does not fix vLLM+PAO load shard length mismatch; fully_reshardable is required.

---

Fix and verification

run_scripts/qwen3-4B-vllm-ckpt-verify.sh now defaults on --dist-ckpt-optim-fully-reshardable when PAO+offload is enabled (both save and load must use the same flag).

Verified on A100 / vime container (aosheninferact/vime-vllm:latest, 8 GPUs, local ckpt):

• Save log: Storing distributed optimizer sharded state of type fully_reshardable
• Load log: distrib_optim_sharding_type': 'fully_reshardable'
• Result: CKPT verify OK: save + load succeeded.
• Example log: /data/nfs_87/xky/logs/ckpt_vime_pao_fully_reshardable.log

---

Checklist: running qwen3-4B-vllm-ckpt-verify.sh successfully

1. Environment and code

• vime + vLLM colocated image/branch; colocated IPC needs test/with_ipc or an IPC-merged branch
• Run from the vime repo (script cds to VIME_ROOT and sources scripts/models/qwen3-4B.sh)
• Prefer bridge Docker networking; avoid network=host + stale host Ray/Redis (Session name does not match)
• Script uses dedicated RAY_TMPDIR / ports and runs cleanup_ray_gpu before each phase

2. Paths and data (override via env if needed)

Variable	Purpose
HF_CKPT	HuggingFace weights
REF_LOAD	Megatron torch_dist ref checkpoint
SAVE_DIR	Checkpoint output (wiped on CKPT_PHASE=all or save)
PROMPT_DATA	e.g. dapo-math-17k.jsonl
LOG_ROOT	Log directory (default /data/nfs_87/xky/logs)

Without NFS inside the container, e.g. HF_CKPT=/root/models/Qwen3-4B, REF_LOAD=/root/models/Qwen3-4B_torch_dist, SAVE_DIR=/root/ckpt_validate/xxx.

3. PAO checkpoint requirements (critical)

• PAO trio matches test_qwen3_4B_ckpt.py by default
• --dist-ckpt-optim-fully-reshardable on both save and load (script enables by default with PAO; do not set CKPT_USE_FULLY_RESHARDABLE=0)
• Do not load old checkpoints saved as dp_reshardable — re-save with fully_reshardable first
• For A/B: CKPT_NO_OPTIM_OFFLOAD=1 disables PAO/offload (fully_reshardable off by default)

4. GPUs and training args

• 8 GPUs required; script sets --rollout-num-gpus = NUM_GPUS (avoids colocate counting critic GPUs → 16 vs 8)
• TIGHT_HOST_MEMORY=1 (default): --max-tokens-per-gpu 2048, rollout-max-response-len 1024 (aligned with CI test)
• Do not enable --train-memory-margin-bytes by default (CKPT_USE_TRAIN_MARGIN=1 only) — otherwise post-save torch_memory_saver.pause() may hit cudaError: invalid argument

5. Phases and usage

# Full flow (recommended)
bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

# Save or load only
CKPT_PHASE=save bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh
CKPT_PHASE=load SAVE_DIR=... bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

• CKPT_PHASE=all: save-interval 2, 3 rollouts, load with --ckpt-step 1
• Load requires existing SAVE_DIR with matching iter_*
• Success: log ends with Job '...' succeeded; full run also prints CKPT verify OK: save + load succeeded.

6. Optional env vars

Variable	Default / notes
CKPT_PHASE	all | save | load
CKPT_USE_FULLY_RESHARDABLE	1 with PAO; set 0 to reproduce 81/52
CKPT_NO_OPTIM_OFFLOAD	1 disables PAO (control run)
CKPT_USE_TRAIN_MARGIN	off by default
ASYNC_SAVE	1 for async save test
CONCURRENT_RAY	1 skips aggressive GPU cleanup (parallel jobs)

7. Out of scope for this fix

• Long-run torch_memory_saver / reloadable PG cudaError after save is a separate issue
• Upstream slime test_qwen3_4B_ckpt.py (SGLang) does not yet add fully_reshardable; vLLM path should follow this script

---

Summary: vLLM + PAO + optimizer CPU offload checkpoint load failure is in the same Megatron distributed-optimizer ckpt family as slime #1221. #1230 fixes save assertions; this load failure needs --dist-ckpt-optim-fully-reshardable, now default in the vime verify script and verified end-to-end.

[CalvinXKY]

#!/bin/bash
# Qwen3-4B vLLM colocated checkpoint verify (aligned with tests/test_qwen3_4B_ckpt.py)
# Usage:
#   bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh
#   CKPT_PHASE=save|load|all  ASYNC_SAVE=1  CKPT_NO_OPTIM_OFFLOAD=1

if grep -q $'\r' "$0" 2>/dev/null; then
  exec bash <(sed 's/\r$//' "$0") "$@"
fi

CKPT_PHASE="${CKPT_PHASE:-all}"
ASYNC_SAVE="${ASYNC_SAVE:-0}"
CONCURRENT_RAY="${CONCURRENT_RAY:-0}"
TIGHT_HOST_MEMORY="${TIGHT_HOST_MEMORY:-1}"

export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray_vime_4b_vllm_ckpt}"
export RAY_PORT="${RAY_PORT:-6387}"
export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8273}"
RAY_AGENT_PORT="${RAY_AGENT_PORT:-52383}"
RAY_ADDR="http://127.0.0.1:${RAY_DASHBOARD_PORT}"
unset RAY_ADDRESS RAY_REDIS_ADDRESS || true
mkdir -p "${RAY_TMPDIR}"

cleanup_ray_gpu() {
  pkill -9 sglang 2>/dev/null || true
  pkill -9 vllm 2>/dev/null || true
  ray stop --force 2>/dev/null || true
  pkill -9 ray 2>/dev/null || true
  pkill -9 python 2>/dev/null || true
  sleep 3
  if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | xargs -r kill -9 2>/dev/null || true
  fi
  rm -rf /tmp/ray/session_* /tmp/ray/* 2>/dev/null || true
  rm -rf "${RAY_TMPDIR}"/session_* "${RAY_TMPDIR}"/ray_current_cluster 2>/dev/null || true
  sleep 2
}

if [ "${CONCURRENT_RAY}" = "1" ]; then
  ray stop --force >/dev/null 2>&1 || true
  rm -rf "${RAY_TMPDIR}"/session_* "${RAY_TMPDIR}"/ray_current_cluster 2>/dev/null || true
  sleep 2
else
  cleanup_ray_gpu
fi

set -e

export PYTHONUNBUFFERED=1
LOG_ROOT="${LOG_ROOT:-/tmp/ckpt_verify_logs}"
TS="$(date +%Y%m%d_%H%M%S)"
LOG_SAVE="${LOG_ROOT}/ckpt_verify_save_${TS}.log"
LOG_LOAD="${LOG_ROOT}/ckpt_verify_load_${TS}.log"
mkdir -p "${LOG_ROOT}"

export PYTHONBUFFERED=16

NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then HAS_NVLINK=1; else HAS_NVLINK=0; fi

if command -v nvidia-smi >/dev/null 2>&1; then
  DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ')
else
  DETECTED_GPUS=0
fi
NUM_GPUS=${NUM_GPUS:-${DETECTED_GPUS}}
if [ -z "$NUM_GPUS" ] || [ "$NUM_GPUS" -le 0 ]; then NUM_GPUS=8; fi
ROLLOUT_GPUS_PER_ENGINE="${ROLLOUT_GPUS_PER_ENGINE:-2}"

REPO_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." &>/dev/null && pwd)"
VIME_ROOT="${VIME_ROOT:-${REPO_ROOT}/vime}"
cd "${VIME_ROOT}"
source "${VIME_ROOT}/scripts/models/qwen3-4B.sh"

HF_CKPT="${HF_CKPT:-/root/models/Qwen3-4B}"
REF_LOAD="${REF_LOAD:-/root/models/Qwen3-4B_torch_dist}"
SAVE_DIR="${SAVE_DIR:-/root/ckpt_validate/Qwen3-4B_vllm_ckpt}"
CKPT_STEP="${CKPT_STEP:-1}"
SAVE_INTERVAL="${SAVE_INTERVAL:-2}"
NUM_ROLLOUT="${NUM_ROLLOUT:-3}"
PROMPT_DATA="${PROMPT_DATA:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl}"

if [ "${TIGHT_HOST_MEMORY}" = "1" ]; then
  MAX_TOKENS_PER_GPU=2048
  ROLLOUT_MAX_LEN=1024
else
  MAX_TOKENS_PER_GPU=9216
  ROLLOUT_MAX_LEN=8192
fi

ROLLOUT_ARGS=(
  --prompt-data "${PROMPT_DATA}"
  --input-key prompt
  --label-key label
  --apply-chat-template
  --rollout-shuffle
  --rm-type deepscaler
  --num-rollout "${NUM_ROLLOUT}"
  --rollout-batch-size 4
  --n-samples-per-prompt 8
  --rollout-max-response-len "${ROLLOUT_MAX_LEN}"
  --rollout-temperature 0.8
  --global-batch-size 32
  --balance-data
)

PERF_ARGS=(
  --tensor-model-parallel-size 2
  --sequence-parallel
  --pipeline-model-parallel-size 1
  --context-parallel-size 2
  --expert-model-parallel-size 1
  --expert-tensor-parallel-size 1
  --recompute-granularity full
  --recompute-method uniform
  --recompute-num-layers 1
  --use-dynamic-batch-size
  --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}"
)

GRPO_ARGS=(
  --advantage-estimator grpo
  --kl-loss-coef 0.00
  --kl-loss-type k1
  --kl-coef 0.00
  --entropy-coef 0.00
  --eps-clip 0.2
)

OPTIMIZER_ARGS=(
  --optimizer adam
  --lr 1e-6
  --lr-decay-style constant
  --weight-decay 0.1
  --adam-beta1 0.9
  --adam-beta2 0.98
  --optimizer-cpu-offload
  --overlap-cpu-optimizer-d2h-h2d
  --use-precision-aware-optimizer
)

VLLM_ARGS=(
  --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}"
  --rollout-backend vllm
  --vllm-gpu-memory-utilization 0.7
)
export SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC="${SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC:-900}"

MISC_ARGS=(
  --attention-dropout 0.0
  --hidden-dropout 0.0
  --accumulate-allreduce-grads-in-fp32
  --attention-softmax-in-fp32
  --attention-backend flash
)

if [ "${CKPT_USE_TRAIN_MARGIN:-0}" = "1" ]; then
  MISC_ARGS+=(--train-memory-margin-bytes 2147483648)
fi

_use_fully="${CKPT_USE_FULLY_RESHARDABLE:-}"
if [ -z "${_use_fully}" ]; then
  if [ "${CKPT_NO_OPTIM_OFFLOAD:-0}" = "1" ]; then
    _use_fully=0
  else
    _use_fully=1
  fi
fi
if [ "${_use_fully}" = "1" ]; then
  MISC_ARGS+=(--dist-ckpt-optim-fully-reshardable)
fi

if [ "${CKPT_NO_OPTIM_OFFLOAD:-0}" = "1" ]; then
  OPTIMIZER_ARGS=(
    --optimizer adam
    --lr 1e-6
    --lr-decay-style constant
    --weight-decay 0.1
    --adam-beta1 0.9
    --adam-beta2 0.98
  )
fi

VIME_PYTHONPATH="${VIME_ROOT}:/root/Megatron-LM/"
export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"

run_ckpt_phase() {
  local phase="$1"
  local log_file="$2"
  local -a CKPT_ARGS=()

  CKPT_ARGS=(
    --hf-checkpoint "${HF_CKPT}"
    --ref-load "${REF_LOAD}"
  )

  if [ "${phase}" = "save" ]; then
    CKPT_ARGS+=(
      --save "${SAVE_DIR}"
      --save-interval "${SAVE_INTERVAL}"
    )
    if [ "${ASYNC_SAVE}" = "1" ]; then
      CKPT_ARGS+=(--async-save)
    fi
  elif [ "${phase}" = "load" ]; then
    CKPT_ARGS+=(
      --load "${SAVE_DIR}"
      --ckpt-step "${CKPT_STEP}"
    )
  else
    echo "Unknown phase: ${phase}" >&2
    return 1
  fi

  export TENSORBOARD_DIR="${LOG_ROOT}/tb_qwen3_4b_vllm_ckpt_${phase}_${TS}"
  mkdir -p "${TENSORBOARD_DIR}"

  echo "=== CKPT phase: ${phase} ==="
  echo "SAVE_DIR=${SAVE_DIR} CKPT_STEP=${CKPT_STEP} ASYNC_SAVE=${ASYNC_SAVE}"
  echo "LOG: ${log_file}"

  cleanup_ray_gpu

  set -x
  ray start --head \
    --node-ip-address "${MASTER_ADDR}" \
    --port "${RAY_PORT}" \
    --num-gpus "${NUM_GPUS}" \
    --disable-usage-stats \
    --dashboard-host=0.0.0.0 \
    --dashboard-port="${RAY_DASHBOARD_PORT}" \
    --dashboard-agent-listen-port="${RAY_AGENT_PORT}" \
    --temp-dir "${RAY_TMPDIR}"

  RUNTIME_ENV_JSON="{
    \"env_vars\": {
      \"PYTHONPATH\": \"${VIME_PYTHONPATH}\",
      \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
      \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
      \"TENSORBOARD_DIR\": \"${TENSORBOARD_DIR}\",
      \"SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC\": \"${SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC}\"
    }
  }"

  ray job submit --address="${RAY_ADDR}" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
    -- python3 train.py \
    --train-backend megatron \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node "${NUM_GPUS}" \
    --rollout-num-gpus "${NUM_GPUS}" \
    --colocate \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
    ${GRPO_ARGS[@]} \
    ${PERF_ARGS[@]} \
    ${VLLM_ARGS[@]} \
    ${MISC_ARGS[@]} \
    2>&1 | tee -a "${log_file}"
  local rc=${PIPESTATUS[0]}
  set +x
  return "${rc}"
}

check_job_ok() {
  local log_file="$1"
  if grep -q "Job '.*' succeeded" "${log_file}" 2>/dev/null; then
    return 0
  fi
  return 1
}

case "${CKPT_PHASE}" in
  save)
    rm -rf "${SAVE_DIR}"
    mkdir -p "${SAVE_DIR}"
    run_ckpt_phase save "${LOG_SAVE}"
    exit $?
    ;;
  load)
    if [ ! -d "${SAVE_DIR}" ]; then
      echo "SAVE_DIR not found: ${SAVE_DIR}. Run CKPT_PHASE=save first." >&2
      exit 1
    fi
    run_ckpt_phase load "${LOG_LOAD}"
    exit $?
    ;;
  all)
    rm -rf "${SAVE_DIR}"
    mkdir -p "${SAVE_DIR}"
    run_ckpt_phase save "${LOG_SAVE}" || exit 1
    if ! check_job_ok "${LOG_SAVE}"; then
      echo "SAVE phase failed. See ${LOG_SAVE}" >&2
      exit 1
    fi
    if ! grep -q "successfully saved checkpoint" "${LOG_SAVE}"; then
      echo "WARN: no 'successfully saved checkpoint' in ${LOG_SAVE}; load phase may fail"
    fi
    sleep 10
    run_ckpt_phase load "${LOG_LOAD}" || exit 1
    if ! check_job_ok "${LOG_LOAD}"; then
      echo "LOAD phase failed. See ${LOG_LOAD}" >&2
      exit 1
    fi
    if grep -q "loading distributed checkpoint from ${SAVE_DIR}" "${LOG_LOAD}"; then
      echo "CKPT verify OK: save + load succeeded."
    else
      echo "LOAD succeeded; confirm checkpoint load lines in ${LOG_LOAD}"
    fi
    ;;
  *)
    echo "Invalid CKPT_PHASE=${CKPT_PHASE} (use save|load|all)" >&2
    exit 1
    ;;
esac

[CalvinXKY]

Checkpoint verification update (vime + vLLM colocate, Qwen3-4B)

Spike to validate the save/restore flow aligned with tests/test_qwen3_4B_ckpt.py on 8×A100, container vime, colocated vLLM (--rollout-backend vllm, TP=2). Related: checkpoint save/restore tracking (#11).

Configuration (aligned with CI test)

Item	Value
Optimizer (PAO trio)	--optimizer-cpu-offload, --overlap-cpu-optimizer-d2h-h2d, --use-precision-aware-optimizer
Distributed optimizer ckpt	--dist-ckpt-optim-fully-reshardable on save and load (required for PAO; without it load fails with Cannot merge two lists with different lengths (81 and 52) on optimizer shards — see THUDM/slime#1221)
Async save	--async-save --use-persistent-ckpt-worker (ASYNC_SAVE=1; --async-save alone is warned as disabled without persistent worker)
Scripts	run_scripts/qwen3-4B-vllm-ckpt-verify.sh, run-ckpt-vime-async-on-server.sh, run-ckpt-vime-async-load-only.sh

run-ckpt-vime-async-on-server.sh

#!/bin/bash
set -euo pipefail
LOCK=/tmp/ckpt_vime_verify_host.lock
if [ -f "${LOCK}" ]; then
  echo "Host lock exists: ${LOCK}. Stop other ckpt jobs first." >&2
  exit 1
fi
echo $$ > "${LOCK}"
trap 'rm -f "${LOCK}"' EXIT

TS="$(date +%Y%m%d_%H%M%S)"
LOG="/data/nfs_87/xky/logs/ckpt_vime_pao_async_${TS}.log"
SAVE_DIR="/root/ckpt_validate/vime-pao-async-${TS}"

docker exec vime bash -lc "
  set -e
  export CKPT_RUN_ID=${TS}
  export TIGHT_HOST_MEMORY=1
  export CKPT_PHASE=all
  export ASYNC_SAVE=1
  export CKPT_STEP=1
  export CONCURRENT_RAY=1
  export SAVE_DIR=${SAVE_DIR}
  export LOG_ROOT=/data/nfs_87/xky/logs
  export HF_CKPT=/root/models/Qwen3-4B
  export REF_LOAD=/root/models/Qwen3-4B_torch_dist
  export PROMPT_DATA=/root/datasets/dapo-math-17k/dapo-math-17k.jsonl
  export CKPT_LOCK=/tmp/ckpt_vime_verify.lock
  bash /data/nfs_87/xky/RL/run_scripts/qwen3-4B-vllm-ckpt-verify.sh
" >> "${LOG}" 2>&1

echo "DONE log=${LOG} SAVE_DIR=${SAVE_DIR}"

run-ckpt-vime-async-load-only.sh

#!/bin/bash
# Load-only after a successful async save (reuse SAVE_DIR from prior save run)
set -euo pipefail
SAVE_DIR="${SAVE_DIR:-/root/ckpt_validate/vime-pao-async-real}"
TS="$(date +%Y%m%d_%H%M%S)"
LOG="/data/nfs_87/xky/logs/ckpt_vime_pao_async_load_${TS}.log"

docker exec vime bash -lc "
  export CKPT_RUN_ID=${TS}
  export TIGHT_HOST_MEMORY=1
  export CKPT_PHASE=load
  export ASYNC_SAVE=1
  export CKPT_STEP=1
  export CONCURRENT_RAY=1
  export SAVE_DIR=${SAVE_DIR}
  export LOG_ROOT=/data/nfs_87/xky/logs
  export HF_CKPT=/root/models/Qwen3-4B
  export REF_LOAD=/root/models/Qwen3-4B_torch_dist
  export PROMPT_DATA=/root/datasets/dapo-math-17k/dapo-math-17k.jsonl
  bash /data/nfs_87/xky/RL/run_scripts/qwen3-4B-vllm-ckpt-verify.sh
" >> "${LOG}" 2>&1

echo "DONE log=${LOG}"

Results

Scenario	Save	Load	Status
Sync PAO + fully_reshardable	OK	OK	Verified (CKPT verify OK: save + load succeeded)
Async save (persistent worker)	OK	—	Verified — artifacts at /root/ckpt_validate/vime-pao-async-real (iter_0000001, iter_0000002; log: scheduled an async checkpoint save)
Async load from async save (--ckpt-step 1)	—	OK	Verified (2026-05-25 load-only rerun)
Async CKPT_PHASE=all in one script	—	—	Not re-run in latest pass; save+load covered by separate save artifact + load-only job

Async load evidence (load-only, SAVE_DIR=/root/ckpt_validate/vime-pao-async-real):

• Log: /data/nfs_87/xky/logs/ckpt_vime_pao_async_load_20260525_181037.log
• loading distributed checkpoint from /root/ckpt_validate/vime-pao-async-real at iteration 1
• successfully loaded checkpoint from ... at iteration 1
• Job 'raysubmit_Aa2wz1EU5NUUnPFf' succeeded

Infra / procedure notes

1. Ray job submission needs a healthy dashboard. Failures were traced to:
   • RAY_TMPDIR too long → Ray MetricsHead AF_UNIX path exceeds ~107 bytes → Module MetricsHead failed to start. Received EOF from pipe.
     Fix: short temp dir, e.g. /tmp/r4b_${CKPT_RUN_ID} (in qwen3-4B-vllm-ckpt-verify.sh).
   • --dashboard-host=127.0.0.1 (instead of 0.0.0.0) for stable local API.
2. Ray cleanup aligned with qwen3-4B-vllm-ppo-verify-colocate.sh: default CONCURRENT_RAY=1 (ray stop + clear only this run’s RAY_TMPDIR; no global pkill python/ray between save and load).
3. Single job at a time — host lock /tmp/ckpt_vime_verify_host.lock + container lock /tmp/ckpt_vime_verify.lock; overlapping runs previously killed an in-flight load.

Reproduce

# Sync save+load
bash run_scripts/qwen3-4B-vllm-ckpt-verify.sh

# Async load from existing async save
SAVE_DIR=/root/ckpt_validate/vime-pao-async-real \
  bash run_scripts/run-ckpt-vime-async-load-only.sh

# Full async save+load (new SAVE_DIR)
bash run_scripts/run-ckpt-vime-async-on-server.sh

Conclusion

On vime + vLLM colocate + Qwen3-4B, with PAO and --dist-ckpt-optim-fully-reshardable:

• Sync checkpoint save+load: works
• Async checkpoint save (with persistent worker) + load from that checkpoint: works

No Megatron/vime code change was required for the async load path beyond the existing fully_reshardable requirement and the Ray/script fixes above.