eval: checkpoint-interfaced evaluation for fully-async training (dedicated fleet / pause-the-world / external service)

docs/user-guide/cli-reference.md

@@ -200,6 +200,14 @@ Sections mirror the launch-script argument groups.
 | `--eval-max-response-len` | int | – | Max eval response length. Inherits from rollout if unset. |
 | `--eval-temperature` | float | – | Eval temperature. Inherits from rollout if unset. |
 | `--eval-top-p` | float | – | Eval top-p. Inherits from rollout if unset. |
+| `--eval-num-gpus` | int | `0` | Dedicated eval fleet size (required for fully-async eval). `0` = shared-engine eval. |
+| `--eval-num-gpus-per-engine` | int | `1` | Eval engine TP, independent of rollout TP. |
+| `--eval-hf-dir` | str | – | Staging dir for per-eval HF snapshots (tmpfs recommended). Unset + `--save-hf` = reuse mode. |
+| `--eval-model-path` | str | – | Boot checkpoint for eval engines. Defaults to `--hf-checkpoint`. |
+| `--eval-dispatch` | str | `async` | `async` fires eval without blocking training; `blocking` awaits inline. |
+| `--eval-max-in-flight` | int | `2` | Max concurrently pending async evals. |
+| `--eval-overflow-policy` | str | `backpressure` | At the cap: await the oldest eval, or `skip` the new point (logged as `eval/skipped_busy`). |
+| `--eval-keep-snapshots` | int | `2` | GC ring for `--eval-hf-dir`; `--save-hf` output is never deleted. |
 
 ### Performance
 

docs/user-guide/fully-async.md

@@ -100,6 +100,81 @@ signal. Fast [P2P weight transfer](/advanced/p2p-weight-transfer) keeps the
 rollout engines closer to the latest actor weights so fewer groups get recycled by
 `--max-weight-staleness`.
 
+## Evaluation
+
+Pick a posture by two questions: is this a test run or a real run (are checkpoints
+persisted anyway), and does eval get standalone GPU resources?
+
+| | Test run | Real run |
+|---|---|---|
+| **No standalone eval** | Pause-the-world (shared engines) | External service over `--save-hf` output; pause-the-world if eval must be strictly on-time (costs ~eval duration of rollout production per point). The two compose: pause-the-world on a small sanity set for on-time points, the service on the full set for delayed high-fidelity points |
+| **Standalone eval fleet** | Fleet + tmpfs snapshot (`--eval-hf-dir /dev/shm/...`) | Fleet + checkpoint reuse (no `--eval-hf-dir`) |
+
+Without extra GPUs (`--eval-num-gpus` unset), eval **shares the rollout engines**:
+the producer pauses new submissions for the duration of the blocking eval and resumes
+after. This is a gate, not a retract — in-flight rollout requests finish and buffer,
+nothing is aborted, and the `pause_generation` API is never involved; eval requests
+simply share engine capacity with the draining tail.
+
+No extra weight movement happens either: the engines already carry the weights the
+step's `update_weights` broadcast just pushed (in fully-async, only *generation* is
+continuous — weight updates are still driver-scheduled per step, each with its own
+generation pause per `--pause-generation-mode`; eval adds no pause of its own on top).
+Pinning comes from ordering: the driver awaits the eval, so the next
+`update_weights` cannot interleave —
+expect `mixed_version_ratio == 0` with the training fleet's update counter as the
+version label (a constant offset from `eval/step`, unlike the dedicated fleet which
+stamps the rollout_id). This ordering is also why shared-engine eval must stay
+blocking: fired-and-forgotten, the next weight update would rewrite the engines
+mid-eval. The cost is that rollout production stalls for roughly the eval duration,
+which is fine for small debug eval sets.
+
+For eval that never touches training capacity, use a **dedicated eval fleet** synced
+through HF checkpoint snapshots — never by joining training weight updates:
+
+```bash
+--eval-num-gpus 1                        # dedicated eval engines (own router)
+--eval-interval K
+--eval-hf-dir /dev/shm/miles_eval_hf     # snapshot staging; tmpfs = no disk dependency
+--eval-prompt-data aime /path/to/aime.jsonl
+```
+
+Per eval-due step the trainer exports an HF snapshot (seconds to tmpfs), fires the
+eval **fire-and-forget**, and keeps training; the eval fleet pins its weights to the
+snapshot (`weight_version = str(rollout_id)`), runs the standard eval datasets, and the
+point lands at the right x-axis step even when it completes a few steps later
+(`eval/lag_steps` reports how late).
+
+Two production-oriented variants:
+
+- **Reuse mode**: with `--save-hf` set and `--eval-hf-dir` unset, eval reuses the
+  periodic HF checkpoints (requires `eval_interval % save_interval == 0`) — zero extra
+  export cost. Pair with `--eval-overflow-policy skip` so a slow eval set can never
+  stall training.
+- **External service**: `tools/checkpoint_eval_service.py` watches `--save-hf` output
+  with its own sglang server — no GPU carve-out from the training job, restartable,
+  backfills missed points from its ledger. Works for `--colocate` runs too.
+
+Every skipped point is attributable from the dashboard: `eval/skipped_busy`,
+`eval/skipped_ckpt_missing`, `eval/skipped_unhealthy`, `eval/skipped_export_failed`
+are logged at the affected step. `eval/{ds}/weight_version/mean == eval/step` and
+`mixed_version_ratio == 0` confirm every point measured exactly the intended weights.
+
+The eval-engine `weight_version` namespace is the snapshot's `rollout_id` — deliberately
+different from the training fleet's job-local update counter; the two fleets never mix.
+
+Where each posture's weights come from — and what "the weights at step R" means:
+
+| Posture | Weight delivery | Measures |
+|---|---|---|
+| Pause-the-world | none needed — training's own `update_weights` broadcast already put them on the shared engines | the engines' last-broadcast version (equals the actor's current weights when `update_weights_interval` is 1) |
+| Dedicated fleet | `update_weights_from_disk` on a snapshot exported **directly from the actor** | the actor's exact step-R weights, regardless of broadcast schedule |
+| External service | loads `--save-hf` checkpoints itself | the actor's exact step-R weights |
+
+Eval engines are never added to the training broadcast group: collectives cannot skip
+members, so a fleet inside the group would have its weights rewritten by every update
+and asynchronous points could not be pinned.
+
 ## Example implementation
 
 For a complete Qwen3 launch script and worker implementation, see the

examples/fully_async/README.md

@@ -6,6 +6,7 @@ The implementation lives in the core library at `miles/rollout/fully_async_rollo
 
 ## Files
 * `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B.
+* `run-qwen3.5-4b-fully_async-eval.sh`: Qwen3.5‑4B with a dedicated eval fleet (fully-async eval).
 
 ## Prerequisite
 First set up model & environment following the Qwen3-4B example.
@@ -26,8 +27,12 @@ Started fully-async rollout worker
 * Completed groups are pushed into a queue; each step drains until it has `--rollout-batch-size` groups.
 * Aborted or too-stale groups are recycled back into the data source.
 
+## Evaluation
+Without extra GPUs, eval shares the rollout engines (producer pauses during the blocking
+eval). With `--eval-num-gpus`/`--eval-hf-dir`, eval runs on a dedicated fleet synced via
+HF checkpoint snapshots; see `run-qwen3.5-4b-fully_async-eval.sh` and the fully-async docs.
+
 ## Limitations
-* No evaluation mode (`--eval-function-path` must point at the standard `InferenceRolloutFn`).
 * Ordering is best effort (sorted at the end by index).
 
 ## Config Differences (3 Key Points)

examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh

@@ -0,0 +1,164 @@
+#!/bin/bash
+# Fully-async training with a dedicated eval fleet (Qwen3.5-4B, one 8-GPU node).
+# GPU split: 4 actor (TP=2) + 3 rollout engines (TP=1) + 1 eval engine; eval weights
+# are pinned per-eval to an HF snapshot on tmpfs, so eval never pauses training.
+
+# for rerun the task
+pkill -9 sglang
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+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
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/qwen3.5-4B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3.5-4B
+   --ref-load /root/Qwen3.5-4B_torch_dist
+   --load /root/Qwen3.5-4B_miles/
+   --save /root/Qwen3.5-4B_miles/
+   --save-interval 20
+)
+
+PROMPT_SET=/path/to/dapo-math-17k.jsonl
+
+ROLLOUT_ARGS=(
+   --rollout-function-path miles.rollout.fully_async_rollout.FullyAsyncRolloutFn
+   --prompt-data ${PROMPT_SET}
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+
+   --rm-type dapo
+   --reward-key score
+
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 256
+   --balance-data
+
+   # retract (default) can deadlock flush_cache in fully_async under load
+   --pause-generation-mode in_place
+
+   # for staleness control
+   #--max-weight-staleness 2
+)
+
+EVAL_ARGS=(
+   --eval-interval 5
+   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 8
+   --eval-max-response-len 16384
+   --eval-top-p 1
+
+   # dedicated eval fleet: 1 engine, pinned per-eval to a tmpfs HF snapshot
+   --eval-num-gpus 1
+   --eval-num-gpus-per-engine 1
+   --eval-hf-dir /dev/shm/miles_eval_hf
+   --eval-keep-snapshots 2
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 2
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   # --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+
+   --use-tis
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --sglang-mem-fraction-static 0.7
+)
+
+MISC_ARGS=(
+   # default dropout in megatron is 0.1
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   # should be good for model performance
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   # need to comment this when using model with MLA
+   --attention-backend flash
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR\": \"1\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train_async.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --rollout-num-gpus 3 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]}

miles/backends/megatron_utils/actor.py

@@ -162,7 +162,7 @@ def init(
             dict(no_load_optim=False, no_load_rng=False, finetune=False) if recv_ckpt_src_rank is not None else {}
         )
         with inplace_modify_args(args, heal_load_overrides):
-            (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer(
+            self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(
                 args, role, checkpointing_context=checkpointing_context
             )
 
@@ -563,6 +563,21 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
         if self.args.offload_train:
             destroy_process_groups()
 
+    @with_logs
+    @timer
+    def export_hf(self, rollout_id: int, path: str) -> None:
+        """Export current weights as an HF checkpoint to ``path`` (collective).
+
+        Uses the direct megatron->HF converters (the weight updater's machinery), so
+        export coverage matches weight-sync coverage. Unlike the periodic --save-hf
+        path inside save_model, failures propagate to the caller so an eval snapshot
+        that failed to export can be skipped loudly.
+        """
+        self._heartbeat.bump()
+        from miles.backends.megatron_utils.model import save_hf_model
+
+        save_hf_model(self.args, rollout_id, self.model, path=path, raise_on_error=True)
+
     @with_logs
     @timer
     def update_weights(self, info: "EnginesAndLock") -> None: