Skip to content
Open
Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
162b110
eval: factor run_eval_datasets out; add eval-fleet retargeting helpers
Zhichenzzz Jul 21, 2026
9fbe31c
fully-async: serve eval from the dedicated eval fleet
Zhichenzzz Jul 21, 2026
2c81c8a
args: dedicated eval fleet flags and validation
Zhichenzzz Jul 21, 2026
3058a58
ray: eval fleet placement, server config synthesis, http client sizing
Zhichenzzz Jul 21, 2026
e29a79f
megatron: on-demand HF export with completeness marker
Zhichenzzz Jul 21, 2026
6aead46
rollout manager: checkpoint-pinned eval controller
Zhichenzzz Jul 21, 2026
3acecf6
train_async: fire-and-forget eval dispatch
Zhichenzzz Jul 21, 2026
b2a1764
tools: standalone checkpoint eval service
Zhichenzzz Jul 21, 2026
c31aa55
docs/examples: fully-async eval
Zhichenzzz Jul 21, 2026
63068ca
megatron: export eval snapshots via the direct converters
Zhichenzzz Jul 21, 2026
1ce152d
tools: fix eval service arg contract and health probe (found in GPU t…
Zhichenzzz Jul 21, 2026
7b68845
rollout manager: bound eval weight loads against zombie engines
Zhichenzzz Jul 21, 2026
eb74bb4
rollout manager: probe eval engines so recover() actually self-heals
Zhichenzzz Jul 21, 2026
75ff11f
eval: survive the post-revival router window and per-sample failures
Zhichenzzz Jul 21, 2026
3706243
tests: e2e for fully-async eval on a dedicated fleet
Zhichenzzz Jul 21, 2026
d58d872
eval: address review — hub-id metadata guard, service errors, formatting
Zhichenzzz Jul 21, 2026
4779ca8
eval: fold hf_export helpers into hf_config
Zhichenzzz Jul 21, 2026
fc5d58f
eval: trim comments to the necessary ones
Zhichenzzz Jul 21, 2026
e045000
eval: move EvalDispatcher into miles/ray/rollout
Zhichenzzz Jul 21, 2026
e499cfc
fully-async: shared-engine eval via producer pause when no fleet is set
Zhichenzzz Jul 21, 2026
460e0af
docs: shared-engine eval version semantics
Zhichenzzz Jul 21, 2026
9d39aa8
eval: apply reward_key only to dict rewards
Zhichenzzz Jul 22, 2026
4e2092d
types: get_reward_value tolerates scalar rewards when reward_key is set
Zhichenzzz Jul 22, 2026
f49a6e8
docs: present eval postures as run-type x standalone-resources matrix
Zhichenzzz Jul 22, 2026
774f2ac
docs: note pause-the-world as the strictly-on-time option for real runs
Zhichenzzz Jul 22, 2026
e448444
eval: never skip the final eval point
Zhichenzzz Jul 22, 2026
3fbe28b
eval: code-quality pass
Zhichenzzz Jul 22, 2026
97e7e04
tests: add eval fleet fields to the ray rollout args fixture
Zhichenzzz Jul 22, 2026
0abdd37
eval: trim narrative comments
Zhichenzzz Jul 22, 2026
ff9b694
eval: drop remaining explanatory comments
Zhichenzzz Jul 22, 2026
e7eb153
service: rename eval_snapshot to evaluate_snapshot
Zhichenzzz Jul 22, 2026
09826bb
docs: clarify weight delivery and pause semantics per eval posture
Zhichenzzz Jul 22, 2026
df1130e
docs: note pause-the-world and external service compose for real runs
Zhichenzzz Jul 22, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions docs/user-guide/cli-reference.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
37 changes: 37 additions & 0 deletions docs/user-guide/fully-async.md
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,43 @@ 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

The rollout fleet never has a quiet window, so fully-async eval runs on a **dedicated
eval fleet** synced through HF checkpoint snapshots — never by joining training weight
updates. Enable it with:

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

## Example implementation

For a complete Qwen3 launch script and worker implementation, see the
Expand Down
6 changes: 5 additions & 1 deletion examples/fully_async/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -26,8 +27,11 @@ 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
Eval runs on a dedicated eval fleet synced via HF checkpoint snapshots (`--eval-num-gpus`,
`--eval-hf-dir`); see `run-qwen3.5-4b-fully_async-eval.sh` and the fully-async docs page.

## 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)
Expand Down
168 changes: 168 additions & 0 deletions examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,168 @@
#!/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. Carving one
# GPU out of rollout costs ~25% rollout throughput vs the 4-GPU baseline; in exchange
# eval runs continuously without ever pausing training. Eval weights are pinned
# per-eval to an HF snapshot exported to tmpfs (--eval-hf-dir on /dev/shm), so eval is
# never blocked by disk and every point measures exactly one weight version.

# 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[@]}
20 changes: 19 additions & 1 deletion miles/backends/megatron_utils/actor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)

Expand Down Expand Up @@ -563,6 +563,24 @@ 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()
if self.args.debug_rollout_only:
return

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:
Expand Down
Loading
Loading