From b4f45e2e5da14fe3bf24e989f472eaafd48948e2 Mon Sep 17 00:00:00 2001 From: Tao Lin Date: Wed, 12 Aug 2026 18:26:43 -0400 Subject: [PATCH] docs: mirror examples/ READMEs into the Examples tab (#2481) Co-authored-by: Claude Fable 5 --- .gitattributes | 3 + .pre-commit-config.yaml | 7 + docs/README.md | 7 +- docs/developer/debug.md | 4 +- docs/developer/index.md | 2 +- docs/docs.json | 52 +- docs/examples/fully-async.md | 222 ++------ docs/examples/geo3k-vlm.md | 106 ++++ docs/examples/geo3k-vlm/multi-turn.md | 53 ++ docs/examples/index.md | 95 ++-- docs/examples/infra-features.md | 8 + docs/examples/infra-features/low-precision.md | 132 +++++ .../infra-features/p2p-weight-transfer.md | 90 +++ docs/examples/infra-features/random-async.md | 36 ++ .../train-infer-mismatch-helper.md | 221 ++++++++ .../examples/infra-features/true-on-policy.md | 77 +++ docs/examples/multi-agent.md | 197 ------- docs/examples/multi-lora.md | 149 +++++ docs/examples/on-policy-distillation.md | 22 + .../qwen3-5-35b-selfdistill.md | 213 ++++++++ docs/examples/openhermes-sft.md | 163 ------ docs/examples/ppo.md | 95 ++++ docs/examples/reproducibility.md | 144 ----- docs/examples/retool-v2.md | 36 ++ docs/examples/search-r1.md | 275 ---------- docs/examples/swe-agent-harbor-docker.md | 142 +++++ docs/user-guide/rollout-endpoints.md | 3 +- examples/README.md | 16 + scripts/tools/sync_example_docs.py | 514 ++++++++++++++++++ tests/fast/doc/test_sync_example_docs.py | 165 ++++++ 30 files changed, 2217 insertions(+), 1032 deletions(-) create mode 100644 .gitattributes create mode 100644 docs/examples/geo3k-vlm.md create mode 100644 docs/examples/geo3k-vlm/multi-turn.md create mode 100644 docs/examples/infra-features.md create mode 100644 docs/examples/infra-features/low-precision.md create mode 100644 docs/examples/infra-features/p2p-weight-transfer.md create mode 100644 docs/examples/infra-features/random-async.md create mode 100644 docs/examples/infra-features/train-infer-mismatch-helper.md create mode 100644 docs/examples/infra-features/true-on-policy.md delete mode 100644 docs/examples/multi-agent.md create mode 100644 docs/examples/multi-lora.md create mode 100644 docs/examples/on-policy-distillation.md create mode 100644 docs/examples/on-policy-distillation/qwen3-5-35b-selfdistill.md delete mode 100644 docs/examples/openhermes-sft.md create mode 100644 docs/examples/ppo.md delete mode 100644 docs/examples/reproducibility.md create mode 100644 docs/examples/retool-v2.md delete mode 100644 docs/examples/search-r1.md create mode 100644 docs/examples/swe-agent-harbor-docker.md create mode 100644 scripts/tools/sync_example_docs.py create mode 100644 tests/fast/doc/test_sync_example_docs.py diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000..ab04112479 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +# Mirrored from examples/**/README.md by scripts/tools/sync_example_docs.py. +# linguist-generated collapses these files in PR diffs; review the README instead. +docs/examples/** linguist-generated=true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d1b42920a4..5f6fadfc94 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -51,6 +51,13 @@ repos: - repo: local hooks: + - id: sync-example-docs + name: Mirror examples/ READMEs into the docs site + language: system + entry: python3 scripts/tools/sync_example_docs.py + files: '^(examples/.*README\.md|docs/examples/.*|docs/docs\.json|scripts/tools/sync_example_docs\.py)$' + pass_filenames: false + - id: ban-mpu-get name: Ban direct mpu.get_* calls (use ParallelState) language: pygrep diff --git a/docs/README.md b/docs/README.md index 0683c9cb1a..1a5b52affd 100644 --- a/docs/README.md +++ b/docs/README.md @@ -35,7 +35,12 @@ Then open http://localhost:3000. the sitemap and out of search results entirely. 3. When linking between pages, use absolute paths: `[Quick Start](/getting-started/quick-start)`. Drop the `.md` extension. -4. Images and other assets go in `assets/` and are referenced the same way: +4. Do not edit anything under `examples/`. Those pages, and the Examples tab of + `docs.json`, are generated from the `README.md` files under the repository's + `examples/` directory, which is the single source of truth. Edit the README and run + `python scripts/tools/sync_example_docs.py` — pre-commit runs it for you and fails if + the two ever diverge. +5. Images and other assets go in `assets/` and are referenced the same way: `/assets/images/arch.png`. Group them into a subdirectory once a topic has more than one image, named after the page or area that uses them: `assets/images/dashboard/` for the dashboard screenshots, `assets/images/brand/` for the logo and favicon. A one-off diff --git a/docs/developer/debug.md b/docs/developer/debug.md index a3ac554d0b..acfd06b936 100644 --- a/docs/developer/debug.md +++ b/docs/developer/debug.md @@ -67,7 +67,7 @@ bitwise comparable, and it is how the Megatron-versus-FSDP alignment test can as equality at all. It is slow; never enable it in production. For end-to-end reproducibility of a whole run, including the sampling path, see the -[Reproducibility recipe](/examples/reproducibility). +[Reproducibility recipe](https://github.com/radixark/miles/tree/main/examples/experimental/reproducibility). ## The assertion harness CI uses @@ -191,7 +191,7 @@ dashboard. ## When all else fails * Drop to a tiny model on a known-good recipe (the - [Reproducibility](/examples/reproducibility) one) to separate framework from model. + [Reproducibility](https://github.com/radixark/miles/tree/main/examples/experimental/reproducibility) one) to separate framework from model. * `git bisect` between a known-good commit and HEAD, with the record-and-replay pattern above pinning the inputs and `--debug-deterministic-collective` pinning the reductions. * Open a GitHub issue with the launch script, `pip freeze`, the first 200 lines of trainer diff --git a/docs/developer/index.md b/docs/developer/index.md index 7bc7852573..574de66ca2 100644 --- a/docs/developer/index.md +++ b/docs/developer/index.md @@ -42,7 +42,7 @@ short tour for new contributors. ## TL;DR for first-time contributors 1. Pick something small from `good first issue` on [GitHub](https://github.com/radixark/miles/issues). -2. Run the [Reproducibility recipe](/examples/reproducibility) so you can be sure +2. Run the [Reproducibility recipe](https://github.com/radixark/miles/tree/main/examples/experimental/reproducibility) so you can be sure "I changed X and it broke" actually means that. 3. Use `--debug-train-only` or `--debug-rollout-only` to scope your changes, and `--list-only` to confirm your test is actually registered in CI. diff --git a/docs/docs.json b/docs/docs.json index 68a8cc2324..fc6a9b98b0 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -255,11 +255,33 @@ "group": "Examples", "root": "examples/index", "pages": [ - "examples/fully-async", - "examples/search-r1", - "examples/multi-agent", - "examples/reproducibility", - "examples/openhermes-sft" + { + "group": "Recipes", + "pages": [ + "examples/fully-async", + "examples/geo3k-vlm", + "examples/geo3k-vlm/multi-turn", + "examples/multi-lora", + "examples/on-policy-distillation", + "examples/on-policy-distillation/qwen3-5-35b-selfdistill", + "examples/ppo", + "examples/retool-v2", + "examples/swe-agent-harbor-docker" + ], + "expanded": true + }, + { + "group": "Infra Features", + "root": "examples/infra-features", + "pages": [ + "examples/infra-features/low-precision", + "examples/infra-features/p2p-weight-transfer", + "examples/infra-features/random-async", + "examples/infra-features/train-infer-mismatch-helper", + "examples/infra-features/true-on-policy" + ], + "expanded": false + } ] } ] @@ -347,6 +369,26 @@ "source": "/advanced/fp8-low-precision", "destination": "/advanced/low-precision", "permanent": true + }, + { + "source": "/examples/search-r1", + "destination": "/examples", + "permanent": true + }, + { + "source": "/examples/multi-agent", + "destination": "/examples", + "permanent": true + }, + { + "source": "/examples/reproducibility", + "destination": "/examples", + "permanent": true + }, + { + "source": "/examples/openhermes-sft", + "destination": "/models/qwen/qwen3", + "permanent": true } ], "background": { diff --git a/docs/examples/fully-async.md b/docs/examples/fully-async.md index 3a565f243d..8770bbe213 100644 --- a/docs/examples/fully-async.md +++ b/docs/examples/fully-async.md @@ -1,208 +1,54 @@ --- -title: Fully Async Rollout -description: Keep generation running continuously in the background so the trainer never waits. +title: "Fully Asynchronous Rollout Example" +description: "Demonstrates fully asynchronous rollout generation for higher efficiency." +# Generated from examples/fully_async/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. --- -**What you'll learn:** how to make rollout production and trainer consumption fully -parallel, with a queue in between, by using a custom rollout function. +This example shows a simple way to make rollout generation **fully asynchronous**: a single global worker is created once and then keeps running in the background, continuously pulling prompts and launching generation tasks. Training only needs to fetch already finished results. This removes the per‑step wait that happens in the normal synchronous style. -In the default training loop, every iteration looks like: - -```text -for it in range(N): - prompts = sample() # cheap - responses = generate() # 10-30s - rewards = score() # 1-3s - loss = train_step() # 5-20s - sync_weights() # 1-10s -``` - -`generate()` blocks `train_step()`. With Async Rollout the loop is split: a background -thread runs `generate` continuously, and the trainer drains a queue. The two run in -parallel and the wall-clock time per iteration drops to roughly `max(generate, train)` -instead of the sum. - -## Prerequisites - -* You completed the [Qwen3-4B](/models/qwen/qwen3) recipe (or have an - equivalent model + dataset). -* Comfortable with [Customization](/user-guide/customization) — async rollout uses - a custom rollout function. +The implementation lives in the core library at `miles/rollout/fully_async_rollout.py` (`FullyAsyncRolloutFn`, a class-based rollout function that owns a persistent background worker). It requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. ## Files +* `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. +* `run_qwen3_5_4b_fully_async_eval.py`: Qwen3.5‑4B with async checkpoint eval — `--eval-backend fleet` (dedicated eval fleet) or `--eval-backend external` (fn-launched sglang server). +* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it. -```text -miles/rollout/fully_async_rollout.py # FullyAsyncRolloutFn (worker + drain) -examples/fully_async/ -├── run-qwen3-4b-fully_async.sh # launch script (Qwen3-4B) -└── run_qwen3_30b_a3b_fully_async.py # MoE variant -``` - -## Quick start +## Prerequisite +First set up model & environment following the Qwen3-4B example. +## Quick Start ```bash -cd /root/miles +cd miles bash examples/fully_async/run-qwen3-4b-fully_async.sh ``` - -You should see: - -```text +You should see log lines like: +``` Started fully-async rollout worker -First rollout sample: ... ``` -## What changes vs. the default recipe +## How It Works (Very Short) +* First train call: the rollout fn starts a persistent worker task on the shared rollout event loop. +* The worker keeps up to `--rollout-batch-size` groups in flight using `generate_and_rm_group`. +* 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. -Just two changes: - -```diff -- python3 train.py ... -+ MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1 python3 train_async.py ... -+ --fully-async -``` - -Everything else — model args, optimizer, GRPO config — stays the same. - -## Walkthrough - -The interesting code is small. `FullyAsyncRolloutFn` is a class-based rollout -function: the constructor receives the args and data source once, and the worker is a -long-lived task on the shared rollout event loop, started lazily on the first train -call: - -```python miles/rollout/fully_async_rollout.py -class FullyAsyncRolloutFn: - async def __call__(self, input): - if input.evaluation: - raise ValueError(...) - if self._worker is None: - self._output = asyncio.Queue(maxsize=OUTPUT_QUEUE_MAX_GROUPS) - self._worker = asyncio.create_task(self._worker_loop()) - return await self._drain(input.rollout_id) -``` - -Key points: - -* **Instance state, no globals.** The worker, queue, and `GenerateState` live on the - rollout-fn instance that `RolloutManager` holds for the process lifetime. -* **One event loop.** Worker, drain, and eval coroutines all run on the shared rollout - loop — plain `asyncio.Queue`, no threads, no locks, no `atexit`. -* **Errors are loud.** A failed generation task kills the worker, and the next drain - raises instead of hanging. - -The worker keeps `--rollout-batch-size` groups in flight using -`generate_and_rm_group`: - -```python -async def _worker_loop(self): - active = set() - while True: - while len(active) < self._max_in_flight_groups(): - active.add(self._submit_one_group()) - done, active = await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED) - for task in done: - await self._output.put(task.result()) -``` - -And each training step simply drains, recycling aborted or stale groups back into the -data source: - -```python -async def _drain(self, rollout_id): - data = [] - while len(data) < self.args.rollout_batch_size: - group = await self._next_group() - if any(s.status == Sample.Status.ABORTED for s in _iter_samples(group)): - self._recycle(group) - continue - data.append(group) - return RolloutFnTrainOutput(samples=data, metrics=...) -``` - -## What's happening underneath - -```mermaid -sequenceDiagram - participant T as train_async.py - participant W as FullyAsyncRolloutFn worker - participant S as SGLang engines - - par Background - loop forever - W->>S: generate(prompt) - S-->>W: response - W->>W: enqueue - end - and Foreground - loop each iteration - T->>W: drain N samples - W-->>T: samples - T->>T: train_step - T->>S: weight_sync - end - end -``` - -The producer loop is decoupled from the consumer loop. As long as the queue stays -populated, the trainer never blocks on generation. - -## Tuning knobs - -| Knob | Effect | -|---|---| -| `--rollout-batch-size` | Worker target in-flight count | -| `--sglang-server-concurrency` | Per-engine concurrency cap | -| `--num-steps-per-rollout` | Increase to consume more per drain (off-policy) | - -If queue depth grows unbounded, training is slower than rollout — bump -`--num-steps-per-rollout` (you'll be slightly off-policy) or scale up trainer -parallelism. - -If queue depth stays at 0, rollout is the bottleneck — that's where async helps least -because there's nothing waiting to be consumed. - -## What to watch - -```text -async/queue_depth stable (50-200 typical) -async/producer_throughput_qps consistent -async/consumer_drain_seconds < producer cycle time -``` - -If `consumer_drain_seconds > producer_cycle_time`, your trainer is starving the queue — -check GPU utilization. +## Evaluation +Without extra GPUs, eval shares the rollout engines (producer pauses during the blocking +eval). For eval that never pauses training, `run_qwen3_5_4b_fully_async_eval.py` shows both +checkpoint-pinned backends behind the same contract: `--eval-backend fleet` (in-job eval +fleet via `--eval-num-gpus`) and `--eval-backend external` (`--eval-function-path` pointed +at `external_eval_fn.ExternalSglangEvalFn`, which launches or attaches its own sglang +server). See the fully-async docs for the posture trade-offs. ## Limitations +* Ordering is best effort (sorted at the end by index). -* **Shared-engine eval pauses production.** Without `--eval-num-gpus`, eval runs on the - rollout engines and the producer stops submitting for its duration; see the - [fully-async guide's Evaluation section](/user-guide/fully-async#evaluation) for the - dedicated-fleet and external-service postures that keep training unblocked. -* **Requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`.** Class-based rollout functions - only load on the new rollout API. -* **Best-effort ordering.** Samples are sorted by index at drain time, but exact-order - guarantees aren't provided. +## Config Differences (3 Key Points) +To enable the fully async pattern there are only three changes compared to a normal run: -## Variations +1. Use the async training driver: `train_async.py` (not `train.py`). +2. Enable the class-based rollout API: `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. +3. Pass `--fully-async`. -### Async on a 30 B MoE +Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? -`run_qwen3_30b_a3b_fully_async.py` shows the same pattern with `tp=4 ep=8` and -`--sglang-enable-ep-moe`. The only practical difference is increasing -`--rollout-batch-size` to 64+ to keep the larger engine pool fed. - -### Async + R3 - -Async rollout and R3 stack cleanly. Add: - -```bash -GRPO_ARGS+=( --use-rollout-routing-replay ) -``` - -The custom rollout function automatically passes `return_routed_experts=true` because -it uses `generate_and_rm_group` under the hood. - -### Async + partial rollout - -If you also use `--partial-rollout`, half-finished trajectories are saved to disk and -resumed — useful when the worker is killed mid-flight. +Because the real generation work is done by a **persistent background worker** owned by `FullyAsyncRolloutFn`. Each call from `train_async.py` only drains already completed samples from the worker's output queue; the worker has been continuously generating since the first call. Thus rollout production (model inference) and training consume happen in parallel with minimal waiting. diff --git a/docs/examples/geo3k-vlm.md b/docs/examples/geo3k-vlm.md new file mode 100644 index 0000000000..f3fa8b92fb --- /dev/null +++ b/docs/examples/geo3k-vlm.md @@ -0,0 +1,106 @@ +--- +title: "VLM Single-Turn RL (FSDP & Megatron)" +description: "Training VLMs with FSDP using GRPO on the GEO3K dataset." +# Generated from examples/geo3k_vlm/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Training VLMs with FSDP or Megatron on single-turn reasoning task using GRPO on the [GEO3K dataset](https://huggingface.co/datasets/hiyouga/geometry3k). We used processed version [here](https://huggingface.co/datasets/chenhegu/geo3k_imgurl). + +Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prevent severe performance regression in conv3d in torch 2.9 mentioned in https://github.com/pytorch/pytorch/issues/168167. Otherwise, you can reinstall cudnn with: +```bash +pip install nvidia-cudnn-cu12==9.16.0.29 +``` + +![FSDP vs Megatron Reward Plot](https://raw.githubusercontent.com/radixark/miles/main/examples/geo3k_vlm/fsdp_vs_megatron.png) + +## Data Preparation (For SFT Training) + +The [geo3k_imgurl](https://huggingface.co/datasets/chenhegu/geo3k_imgurl) dataset contains: +- `problem`: The math problem text (string) +- `answer`: The answer (string, e.g., "270") +- `images`: Image data (list) + +For SFT training, we need to format the `answer` field for `\boxed{}` format and the messages. You can use the following script to format the answer field: + +```python +from datasets import load_dataset +import pandas as pd + +ds = load_dataset("chenhegu/geo3k_imgurl", split="train") + +def format_answer(answer: str) -> str: + """Format answer to include \\boxed{} format.""" + return f"Answer: \\boxed{{{answer}}}" + +def process_sample(sample): + formatted_answer = f"Answer: \\boxed{{{sample['answer']}}}" + + sample["messages"] = [ + {"role": "user", "content": sample["problem"]}, + {"role": "assistant", "content": formatted_answer} + ] + return sample + +ds = ds.map(process_sample) +ds.to_parquet("/root/datasets/geo3k_imgurl/train_formatted.parquet") +``` + +## Reproduce + +```bash +export WANDB_API_KEY=your_wandb_api_key + +# Megatron backend (default -> Qwen3-VL-8B-Instruct + Megatron) +./examples/geo3k_vlm/run_geo3k_vlm.sh + +# FSDP backend +MILES_SCRIPT_TRAIN_BACKEND=fsdp ./examples/geo3k_vlm/run_geo3k_vlm.sh + +# With different model +MILES_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.sh + +# SFT +./examples/geo_3k_vlm/run_geo3k_vlm_sft.sh +``` + +### Configuration + +| Environment Variable | Default | Description | +|---------------------|---------|-------------| +| `MILES_SCRIPT_TRAIN_BACKEND` | `megatron` | Training backend (`megatron` or `fsdp`) | +| `MILES_SCRIPT_MODEL_NAME` | `Qwen3-VL-8B-Instruct` | Model name | +| `MILES_SCRIPT_DATASET_NAME` | `chenhegu/geo3k_imgurl` | HuggingFace dataset name | +| `MILES_SCRIPT_NUM_GPUS` | `8` | Number of GPUs | +| `MILES_SCRIPT_EXTERNAL_RAY` | `0` | Use external Ray cluster (`1` to enable) | + +### Supported Models + +- `Qwen3-VL-2B-Instruct` +- `Qwen3-VL-4B-Instruct` +- `Qwen3-VL-8B-Instruct` +- `Qwen3-VL-30B-A3B-Instruct` +- `Qwen3-VL-235B-A22B-Instruct` +- `Qwen3-VL-2B-Thinking` +- `Qwen3-VL-4B-Thinking` +- `Qwen3-VL-8B-Thinking` +- `Qwen3-VL-30B-A3B-Thinking` +- `Qwen3-VL-235B-A22B-Thinking` + +## Notes + +### Reward Model Configuration + +We experimented with three reward model configurations: +1. A geo3k-specific RM with tolerance=0.05 (to handle rounding in ground truth labels) +2. A geo3k-specific RM with tolerance=0.0 (strict matching) +3. The default math RM + +All three performed similarly, so we use the default math RM for simplicity. + +### Numerical Precision with Non-Binary Rewards + +Our initial geo3k-specific verifier produced "format scores" (**0 and 0.9**) instead of clean binary rewards. Under **fp32**, fractional values like 0.9 can't be exactly represented, so when all samples in a group have the same reward, `reward - mean` doesn't equal zero—creating spurious gradient signal. + +We fixed this by switching to the default math RM with clean **binary 0/1 rewards**. If you encounter similar precision issues with non-binary rewards, you can change the reward tensor dtype from `torch.float` to `torch.float16` in `miles/ray/rollout.py` (`_post_process_rewards` method) to truncate precision artifacts. + +## B200 +Blackwell currently does not support fa3, we need to use `--sglang-mm-attention-backend sdpa` and `--attn-implementation flash_attention_2` diff --git a/docs/examples/geo3k-vlm/multi-turn.md b/docs/examples/geo3k-vlm/multi-turn.md new file mode 100644 index 0000000000..fdddf12190 --- /dev/null +++ b/docs/examples/geo3k-vlm/multi-turn.md @@ -0,0 +1,53 @@ +--- +title: "VLM Multi-Turn (geo3k dataset)" +description: "The same dataset over multiple turns, with the model cropping images through an interactive environment." +# Generated from examples/geo3k_vlm/multi_turn/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Training VLM on [geo3k dataset](https://huggingface.co/datasets/hiyouga/geometry3k) with multi-turn reasoning with interactive environment feedback, using GRPO. For the dataset, we used the [processed version](https://huggingface.co/datasets/VeraIsHere/geo3k_imgurl_processed). + +Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prevent severe performance regression in conv3d in torch 2.9 mentioned in https://github.com/pytorch/pytorch/issues/168167. Otherwise, you can reinstall cudnn with: +```bash +pip install nvidia-cudnn-cu12==9.16.0.29 +``` + +The multi-turn rollout is implemented through a [custom generate function](https://github.com/radixark/miles/blob/main/examples/geo3k_vlm/multi_turn/rollout.py#L309), overriding the original generate function. + +In terms of the environment interaction, this example initializes a [custom interactive environment](https://github.com/radixark/miles/blob/main/examples/geo3k_vlm/multi_turn/env_geo3k.py) with the APIs below. + + + +- `build_env(sample: Sample | None = None, args: Any | None = None, **_) -> Geo3kEnv`: constructs the env. +- `reset() -> tuple[dict, dict]`: clears internal state. +- `step(response_text: str) -> tuple[dict, bool, dict]`: parses the actor's response text and update the state. Return new observation, a flag that marks whether the task is done, and step_info. +- `format_observation(observation: dict) -> dict`: converts an env observation into a chat message. + + + +
+ +The reward model is the default math RM. + +![VLM multi-turn geo3k reward](https://raw.githubusercontent.com/radixark/miles/main/examples/geo3k_vlm/multi_turn/geo3k_vlm_multi_turn_reward.png) +![Rollout megatron](https://raw.githubusercontent.com/radixark/miles/main/examples/geo3k_vlm/multi_turn/rollout_experiment_result_megatron.png) + +## Reproduce +```bash +# 1) Set environment variable +export WANDB_API_KEY=... +export MILES_SCRIPT_MODEL_NAME=Qwen3-VL-2B-Instruct +export MILES_SCRIPT_NUM_GPUS=4 +export MILES_SCRIPT_TRAIN_BACKEND=megatron + +# 2) Download the dataset +hf download --repo-type dataset VeraIsHere/geo3k_imgurl_processed --local-dir /root/datasets/geo3k_imgurl_processed + +# 3) Run the script: +cd /root/miles +python examples/geo3k_vlm/multi_turn/run_geo3k_vlm_multi_turn.py +``` + +## What each file does +- `examples/geo3k_vlm/multi_turn/run_geo3k_vlm_multi_turn.py`: downloads model, sets training/rollout args, and launches the run. +- `examples/geo3k_vlm/multi_turn/geo3k_vlm_multi_turn_config.yaml`: specifies `max_turns` and `rollout_interaction_env_path` for the multi-turn rollout. +- `examples/geo3k_vlm/multi_turn/rollout.py`: custom multi-turn rollout that calls SGLang for token generation, builds loss masks/log_probs, enforces max_turns, and early-stops on max_new_tokens. +- `examples/geo3k_vlm/multi_turn/env_geo3k.py`: geo3k tool-calling env that parses <tool_call>\{...\}</tool_call>, scores math answers, and returns tool feedback per turn. diff --git a/docs/examples/index.md b/docs/examples/index.md index 473bcd8ffb..cbdce2eb5c 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -1,66 +1,51 @@ --- -title: Examples -description: Annotated end-to-end walkthroughs for the workflows people actually want to build. +title: "Examples" +description: "These examples are runnable starting points for your own RL workflow." +# Generated from examples/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. --- -The model recipes show you how to train a model. The examples below show you how to -*build something useful* with Miles — tools, search, multi-agent, distillation, and -async rollout. +A few are purely demonstrative, but most are verifiable against a concrete performance score. -Each example follows the same template: +## Recipes -1. **What you'll learn** — the takeaway in one sentence. -2. **Prerequisites** — what you need installed/downloaded first. -3. **Files** — what's in the example directory. -4. **Quick start** — single command to run. -5. **Walkthrough** — annotated tour of the key code. -6. **What's happening underneath** — the moving parts you can't see. -7. **Tuning knobs** — the levers that matter. -8. **Troubleshooting** — the failure modes we've actually hit. -9. **Variations** — common adaptations. +End-to-end training workflows — the place to start. -## The catalog +- **[fully_async](/examples/fully-async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. +- **[geo3k_vlm](/examples/geo3k-vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset. + - **[multi_turn](/examples/geo3k-vlm/multi-turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment. +- **[lora](https://github.com/radixark/miles/tree/main/examples/lora)**: LoRA fine-tuning with the Megatron backend. +- **[multi_lora](/examples/multi-lora)**: Fully-async multi-adapter LoRA training with a slot-keyed adapter page table. +- **[on_policy_distillation](/examples/on-policy-distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop. + - **[qwen3_5_35b_selfdistill](/examples/on-policy-distillation/qwen3-5-35b-selfdistill)**: Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher. +- **[ppo](/examples/ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs. +- **[retool_v2](/examples/retool-v2)**: Tool-enabled language model generation with sandboxed Python code execution interleaved with thinking. +- **[swe-agent-harbor-docker](/examples/swe-agent-harbor-docker)**: Trains coding and terminal agents with Harbor-managed local Docker sandboxes and verifier rewards. - +## [Infra Features](/examples/infra-features) - +Runtime and infrastructure plumbing rather than training recipes — how miles moves +data and weights around. - Continuous background generation with a queue between rollout and training. - Up to 2× end-to-end speedup. +- **[low_precision](/examples/infra-features/low-precision)**: Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability. +- **[p2p_weight_transfer](/examples/infra-features/p2p-weight-transfer)**: Point-to-point weight transfer between training and rollout engines. +- **[random_async](/examples/infra-features/random-async)**: Dataset-free stress test of the async rollout ↔ trainer loop. +- **[train_infer_mismatch_helper](/examples/infra-features/train-infer-mismatch-helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS). +- **[true_on_policy](/examples/infra-features/true-on-policy)**: Ensures strictly equal log probabilities between inference (SGLang) and training engines. - +## [Experimental](https://github.com/radixark/miles/tree/main/examples/experimental) - +Not fully verified — for experimental and development use. - Multi-turn rollout where the model can issue `...` actions, get - observations from a retrieval server, and produce a final answer. - - - - - - Two specialized agents (e.g. doctor + patient) train together and improve - each other. - - - - - - Bit-stable training across reruns. Determinism flags, seeds, and what to - watch. - - - - - - Plain SFT (no RL) — sometimes you just need a quick fine-tune. - - - - - -## Where to start - -* **Never used Miles for anything beyond GRPO?** → [Fully Async Rollout](/examples/fully-async). -* **Want tool use / RAG?** → [Search-R1](/examples/search-r1). -* **VLM / multi-agent?** → [Multi-Agent Co-Evolution](/examples/multi-agent). -* **Replay an old result?** → [Reproducibility Recipe](/examples/reproducibility). +- **[agentenv](https://github.com/radixark/miles/tree/main/examples/experimental/agentenv)**: Rollouts against AgentENV, a self-hosted platform running agent sandboxes on Firecracker microVMs. +- **[DrGRPO](https://github.com/radixark/miles/tree/main/examples/experimental/DrGRPO)**: Custom reducer for Dr.GRPO algorithm. +- **[eval](https://github.com/radixark/miles/tree/main/examples/experimental/eval)**: Documentation and setup for evaluation environments using NeMo-Skills. +- **[eval_multi_task](https://github.com/radixark/miles/tree/main/examples/experimental/eval_multi_task)**: Example for supporting OOD evaluation tasks, e.g., GPQA, IFBench. +- **[formal_math](https://github.com/radixark/miles/tree/main/examples/experimental/formal_math)**: Examples related to formal math reasoning tasks, including a single round demo. +- **[multi_agent](https://github.com/radixark/miles/tree/main/examples/experimental/multi_agent)**: Example of running multi-agent RL with `miles`. +- **[nemo-gym](https://github.com/radixark/miles/tree/main/examples/experimental/nemo-gym)**: SWE-agent training with NVIDIA NeMo Gym as the environment ecosystem. +- **[openenv](https://github.com/radixark/miles/tree/main/examples/experimental/openenv)**: Rollouts against OpenEnv-hosted environments. +- **[reproducibility](https://github.com/radixark/miles/tree/main/examples/experimental/reproducibility)**: Guides on achieving bitwise experiment reproduction using deterministic modes. +- **[search-r1](https://github.com/radixark/miles/tree/main/examples/experimental/search-r1)**: A minimal reproduction of Search-R1, featuring multi-turn conversation and tool-calling. +- **[strands_sglang](https://github.com/radixark/miles/tree/main/examples/experimental/strands_sglang)**: Integration example with the Strands-Agents scaffolding framework. +- **[swe-agent-harbor-daytona](https://github.com/radixark/miles/tree/main/examples/experimental/swe-agent-harbor-daytona)**: The `swe-agent-harbor-docker` pipeline with task sandboxes hosted on Daytona instead of local Docker. +- **[tau-bench](https://github.com/radixark/miles/tree/main/examples/experimental/tau-bench)**: Training in an agentic multi-turn tool use environment (Tau-bench). +- **[verifiers](https://github.com/radixark/miles/tree/main/examples/experimental/verifiers)**: Training on a Prime Intellect Verifiers environment instead of a Miles prompt dataset. diff --git a/docs/examples/infra-features.md b/docs/examples/infra-features.md new file mode 100644 index 0000000000..bd793bb8c1 --- /dev/null +++ b/docs/examples/infra-features.md @@ -0,0 +1,8 @@ +--- +title: "Infra Features" +description: "These examples exercise runtime and infrastructure behaviour rather than a training recipe." +# Generated from examples/infra_features/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +They cover weight transfer, low precision, train/inference agreement, and the async rollout +loop. Reach for them when the question is how miles moves data and weights around, not how to +train a particular model. diff --git a/docs/examples/infra-features/low-precision.md b/docs/examples/infra-features/low-precision.md new file mode 100644 index 0000000000..205510aa2b --- /dev/null +++ b/docs/examples/infra-features/low-precision.md @@ -0,0 +1,132 @@ +--- +title: "FP8 training examples" +description: "Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability." +# Generated from examples/infra_features/low_precision/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +This is an example of FP8 training and FP8 inference. Under FP8 training and inference, it can achieve more efficient inference throughput and lower training-inference mismatch, resulting in more stable training. More details can be found in [this blog](https://lmsys.org/blog/2025-11-25-fp8-rl/). + +## Files + +* `run-qwen3-4b-fp8.sh`: example launch script with Qwen3‑4B in FP8. + +* `run-qwen3-30b-a3b-fp8-two-nodes.sh`: example launch script for running Qwen3‑30B‑A3B in FP8 across two nodes. + +## Quick Start + +1. Check if your training script is properly configured. + + For training tasks, we need to add these flags: + ```bash + --fp8-format e4m3 + --fp8-recipe blockwise + # --fp8-param-gather # [optional] Currently incompatible with CPU Adam + ``` + Then ensure the `NVTE_FP8_BLOCK_SCALING_FP32_SCALES` environment variable is enabled. + + Note that only `Linear` and `GroupLinear` layers in TransformerEngine use fp8 format. `embedding` and `lm_head` remain in their original precision. If `--fp8-param-gather` is not enabled, weights in TransformerEngine remain in bf16 format, only being cast to fp8 format during `GEMM` or `GroupGEMM` operations. + +2. Convert your HuggingFace model weights to FP8 format. + + You can use `tools/convert_hf_to_fp8.py` to convert bf16 weights to fp8 format. Ensure that the `--hf-checkpoint` parameter points to a directory where the `config.json` contains the correct `quantization_config`. miles will automatically use FP8 quantization during weight updates. + +3. Start FP8 training. + + ```bash + cd miles + + # Qwen3‑4B FP8 training (single node) + bash examples/infra_features/low_precision/run-qwen3-4b-fp8.sh + + # Qwen3‑30B‑A3B FP8 training (two nodes) + bash examples/infra_features/low_precision/run-qwen3-30b-a3b-fp8-two-nodes.sh + ``` + Following the above command will launch FP8 training. + +4. Use the saved checkpoint for evaluation. + + Note that TransformerEngine does not specifically save FP8 quantized weights; the saved torch dist remains in original precision (usually bf16). If you want to evaluate under FP8, you need to convert the checkpoint from `torch_dist` to HuggingFace format, then convert to FP8 HuggingFace format. + +## Quick Explanation + +Here's a quick explanation of how FP8 training is currently implemented in miles: + +1. Initialization: If FP8 recipe is enabled, layers will be built in FP8 context. + +2. Training: During training, weights and activations are quantized online to nvfp8 format, and cuBLAS FP8 GEMM is called for various GEMM computations in forward and backward passes. + +3. Weight updates: During RL weight updates, Megatron first dequantizes FP8 weights to bf16 format, then miles quantizes these bf16 weights to fp8 format and sends them to sglang. (This additional dequantization and quantization is not elegant, but we haven't modified the interface yet for framework compatibility.) + +4. Save checkpoint: Similar to weight updates, if checkpoints need to be saved from the training engine, they will also be dequantized back to bf16 and saved to `torch_dist` format checkpoints. + +## TODO + +Currently, FP8 is far from being a complete feature and still has the following bugs, for examples: + +- FP8 weights (`--fp8-param-gather`) can provide memory savings benefits, but currently FP8 weights must be used with TransformerEngine's FusedAdam, which conflicts with the commonly used Adam CPU offload technique in Megatron-LM. + +The miles team will continue to collaborate with the NVIDIA team to contribute more complete FP8 training infrastructure to the community. + +*** + +## INT4 Training Examples + +This guide provides examples for INT4 STE (Straight-Through Estimator) training and INT4 inference. Utilizing INT4 inference significantly improves throughput, thereby accelerating the training pipeline (specifically during the rollout generation phase). + +### Files + +* `run-qwen3‑30B‑A3B-int4.sh`: Launch script for **Qwen3‑30B‑A3B** (INT4) on 8x H200 GPUs. +* `run-qwen3-235B-A22B-int4.sh`: Launch script for **Qwen3-235B-A22B** (INT4) on 64x H200 GPUs. +* `run-kimi-k2-Thinking-int4.sh`: Launch script for **Kimi-k2-Thinking** (INT4) on 256x H200 GPUs. + +### Quick Start + +#### 1. Convert HuggingFace Weights to INT4 +First, download the PTQ (Post-Training Quantization) calibration dataset from HuggingFace: +[https://huggingface.co/datasets/Salesforce/wikitext/tree/main/wikitext-2-raw-v1](https://huggingface.co/datasets/Salesforce/wikitext/tree/main/wikitext-2-raw-v1) + +Next, use the `tools/convert_hf_to_hf_int4.py` script to convert BF16 weights to INT4 format. Ensure that the `--hf-checkpoint` parameter points to a directory where `config.json` contains the correct `quantization_config`. miles will automatically utilize INT4 quantization during weight updates. + +```bash +python tools/convert_hf_to_hf_int4.py \ + --input-dir /path/to/your/original/models \ + --output-dir /path/to/your/save/models \ + --data-dir /path/to/your/wikitext +``` + +#### 2. Start INT4 Training + +You need to configure the specific environment variables for quantization settings. + +**Environment Variables:** + +* **`OPEN_TRAINING_INT4_FAKE_QAT_FLAG`**: Enables fake quantization operations for INT4 training. +* **`OPEN_TRAINING_INT4_GROUP_SIZE`**: Specifies the block size (group size) for model quantization. + * Set to **128** for `qwen3-30B-A3B` and `qwen3-235B-A22B-int4`. + * Set to **32** for `kimi-k2-Thinking-int4`. + +**Configuration Example:** + +```json +RUNTIME_ENV_JSON="{ + \"env_vars\": { + ... + \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", + \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" + } +}" +``` + +**Launch Commands:** + +```bash +# Qwen3‑30B‑A3B Int4 training +bash examples/infra_features/low_precision/run-qwen3‑30B‑A3B-int4.sh + +# Qwen3-235B-A22B Int4 training (8 nodes) +bash examples/infra_features/low_precision/run-qwen3-235B-A22B-int4.sh + +# Kimi-k2-Thinking Int4 training (32 nodes) +bash examples/infra_features/low_precision/run-kimi-k2-Thinking-int4.sh +``` + +- For multi-node environments, please start the Ray service according to your cluster configuration. diff --git a/docs/examples/infra-features/p2p-weight-transfer.md b/docs/examples/infra-features/p2p-weight-transfer.md new file mode 100644 index 0000000000..d1a4cea1da --- /dev/null +++ b/docs/examples/infra-features/p2p-weight-transfer.md @@ -0,0 +1,90 @@ +--- +title: "P2P Weight Transfer Examples" +description: "Point-to-point weight transfer between training and rollout engines." +# Generated from examples/infra_features/p2p_weight_transfer/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Example scripts for running P2P (RDMA) and broadcast (NCCL) weight transfer between +Megatron training and sglang rollout engines. + +See [docs/advanced/p2p-weight-transfer.md](/advanced/p2p-weight-transfer) +for full documentation, architecture details, and profiling results. + +## Quick Start + +All examples use a single `run.py` script with two subcommands (`prepare` and `run`) and +per-model thin bash wrappers: + +``` +examples/infra_features/p2p_weight_transfer/ +├── run.py # Shared logic: prepare + run (model registry) +├── Qwen3-4B.sh # 1 node – bash Qwen3-4B.sh [MODE] +├── GLM-4.7-Flash.sh # 2 nodes +├── GLM-5.sh # 2/12/32 nodes – bash GLM-5.sh [MODE] [NODE_RANK] [HEAD_IP] +├── Qwen3-30B-A3B.sh # 4 nodes +├── GLM-4.5-Air.sh # 8 nodes +├── Qwen3-235B-A22B.sh # 16 nodes +└── Kimi-K2.sh # 64 nodes +``` + +## Usage + +### Single-node models (Qwen3-4B) + +Single-node models auto-detect networking and don't require node rank or head IP: + +```bash +bash examples/infra_features/p2p_weight_transfer/Qwen3-4B.sh p2p +bash examples/infra_features/p2p_weight_transfer/Qwen3-4B.sh broadcast +``` + +### Multi-node models + +```bash +# GLM-4.7-Flash (2 nodes, disaggregated) +bash examples/infra_features/p2p_weight_transfer/GLM-4.7-Flash.sh p2p 0 $HEAD_NODE_IP # head +bash examples/infra_features/p2p_weight_transfer/GLM-4.7-Flash.sh p2p 1 $HEAD_NODE_IP # worker + +# Qwen3-30B-A3B (4 nodes) +bash examples/infra_features/p2p_weight_transfer/Qwen3-30B-A3B.sh p2p 0 $HEAD_NODE_IP # head +bash examples/infra_features/p2p_weight_transfer/Qwen3-30B-A3B.sh p2p 1 $HEAD_NODE_IP # worker 1 +# ... workers 2, 3 +``` + +### GLM-5 variants (2/12/32 nodes) + +The `GLM-5.sh` wrapper accepts a `VARIANT` argument: + +```bash +# GLM-5_4layer (2 nodes) +bash examples/infra_features/p2p_weight_transfer/GLM-5.sh GLM-5_4layer p2p 0 $HEAD_NODE_IP + +# GLM-5 full (32 nodes) +bash examples/infra_features/p2p_weight_transfer/GLM-5.sh GLM-5 p2p 0 $HEAD_NODE_IP +``` + +### Using run.py directly + +```bash +# Prepare: download model, datasets, convert checkpoint +python examples/infra_features/p2p_weight_transfer/run.py prepare GLM-4.7-Flash + +# Run: launch training with P2P weight transfer +python examples/infra_features/p2p_weight_transfer/run.py run GLM-4.7-Flash \ + --mode p2p --node-rank 0 --head-ip $HEAD_NODE_IP + +# List available models +python examples/infra_features/p2p_weight_transfer/run.py list +``` + +### Environment variables + +| Variable | Default | Description | +|---|---|---| +| `CKPT_SAVE_DIR` | `/root/multinode` | Override checkpoint save directory | +| `SKIP_VALIDATION` | `0` | Set to `1` to skip `--check-weight-update-equal` | +| `BUCKET_SIZE_GB` | `1` | Transfer bucket size in GB | +| `MILES_LOG_DIR` | — | Directory for training logs | + +## Notes + +\* **Kimi-K2**: The profiling RDMA time for Kimi-K2 includes ~884 ms of GPU-side `post_load_weights` requantization on the rollout engines, since this model requires weight requantization after RDMA transfer. diff --git a/docs/examples/infra-features/random-async.md b/docs/examples/infra-features/random-async.md new file mode 100644 index 0000000000..b2c01b377d --- /dev/null +++ b/docs/examples/infra-features/random-async.md @@ -0,0 +1,36 @@ +--- +title: "Random fully-async example" +description: "Dataset-free stress test of the async rollout ↔ trainer loop." +# Generated from examples/infra_features/random_async/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Minimal sibling of `examples/fully_async/`. Exercises the entire async +rollout ↔ trainer loop **without any real dataset, real reward model, or +meaningful generation** — useful as an agent infrastructure stress test +for bigger agentic workloads. + +## Quick start + +```bash +# default (Qwen3.5-35B-A3B), in_place pause + broadcast weight transfer +python run_random_async_3node.py + +# retract pause + p2p weight transfer +python run_random_async_3node.py \ + --pause-generation-mode retract \ + --update-weight-transfer-mode p2p + +# swap in a different model +python run_random_async_3node.py \ + --model-name Qwen3.5-35B-A3B --megatron-model-type qwen3.5-35B-A3B +``` + +## Notes + +- `--disable-rollout-global-dataset` is on, so no `--prompt-data` file is + required. The rollout function ignores the data buffer and constructs + `Sample` objects from scratch. +- The rollout uses Qwen3.5-35B's vocab size (151643) for the random + `input_ids`; any model with vocab ≥ that works without changes. +- `ignore_eos=True` in the sampling params means SGLang generates until + it hits `max_new_tokens` (drawn from `MAX_TOKENS_RANGE`). `Sample.status` + is set to `COMPLETED`. diff --git a/docs/examples/infra-features/train-infer-mismatch-helper.md b/docs/examples/infra-features/train-infer-mismatch-helper.md new file mode 100644 index 0000000000..3b0d591273 --- /dev/null +++ b/docs/examples/infra-features/train-infer-mismatch-helper.md @@ -0,0 +1,221 @@ +--- +title: "Rollout Correction Methods" +description: "Algorithmic methods for rollout correction (e.g., TIS, MIS)." +# Generated from examples/infra_features/train_infer_mismatch_helper/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Rollout correction (e.g, TIS, MIS) through algorithmic methods. + +## Quick Takeaway + +This function is used to solve offline scenarios through algorithmic adaptations, e.g. TIS/MIS. + +We included 3 rollout correction algorithms: + +1. decoupled, 3-policies PPO with rollout importance sampling +2. direct rollout policy overwriting in the standard PPO +3. pure REINFORCE loss (without PPO clipping) with rollout importance sampling + +`--use-tis`: use this flag to **turn on TIS/MIS** for rollout correction (details in **Algorithms**). +You may specify the **IS/RS configs** with a config file using `--custom-config-path`. + +`--use-rollout-logprobs`: When use this flag, the logprobs will **not** be recomputed by training engine - rollout log probs will be directly used in PPO/GRPO loss. + +`--get-mismatch-metrics`: When you don't want to add TIS/MIS, but still want to monitor the mismatch-related metrics (e.g. rollout-training KL). It will **only return mismatch metrics** but not change the loss in any way. + +## Algorithms + +We give examples of the algorithms for solving the training-inference mismatch issue. + +### [Baseline: No Mismatch Correction] Standard PPO + +This is the basic PPO algorithm with potentially training-inference mismatch issue when the output of SGLang and Megatron does not exactly match. + +$$ +L_{\text{PPO}}(\theta) += - \mathbb{E}_{x \sim \mathcal{D},\, y \sim \pi_{\textcolor{red}{\text{SGLang}}}} \left[ + \min \left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{blue}{\text{Megatron}}}(y \mid x)} A_t, + \mathrm{clip}\left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{blue}{\text{Megatron}}}(y \mid x)}, + 1 - \epsilon, + 1 + \epsilon + \right) A_t + \right) +\right]. +$$ + +### Bypassing PPO importance sampling + +Like REINFORCE, we directly use the rollout engine's log probs as the old policy in offline PPO's importance sampling, rather than the recomputed log-probs from the training engine. + +$$ +L_{\text{PPO-bypass}}(\theta) += - \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_{\textcolor{red}{\text{SGLang}}}} \left[ + \min \left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{red}{\text{SGLang}}}(y \mid x)} A_t, + \mathrm{clip}\left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{red}{\text{SGLang}}}(y \mid x)}, + 1 - \epsilon, + 1 + \epsilon + \right) A_t + \right) +\right]. +$$ + +Advantages: + +- Efficiency: skip `log_prob` recomputation on training engine. Reduce one expensive forward pass on all the generated trajectories. + +### Decoupled, 3-policy PPO Importance Sampling + +[Decoupled PPO](https://arxiv.org/pdf/2110.00641) achieves batch-independent PPO by decoupling two roles: Proximal Policy (anchor policy for PPO clipping, control update size) and Behavior Policy (for off-policy correction in importance sampling). Therefore, there are totally 3 roles engaged in this mode, **target policy** $\pi_\theta$, **proximal policy** $\pi_{\textcolor{blue}{\text{old}}}$, and **behavior policy** $\pi_{\textcolor{red}{\text{SGLang}}}$. $\pi_{\textcolor{blue}{\text{old}}}$ is recomputed with Megatron at the beginning of each training step. + +$$ +L_{\text{PPO-decoupled}}(\theta) += - \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_{\textcolor{red}{\text{SGLang}}}} \left[ + \frac{\pi_{\textcolor{blue}{\text{old}}}(y \mid x)}{\pi_{\textcolor{red}{\text{SGLang}}}(y \mid x)} + \min \left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{blue}{\text{old}}}(y \mid x)} A_t, + \mathrm{clip}\left( + \frac{\pi_\theta(y \mid x)}{\pi_{\textcolor{blue}{\text{old}}}(y \mid x)}, + 1 - \epsilon, + 1 + \epsilon + \right) A_t + \right) +\right]. +$$ + +Advantages: + +- Achieves batch size invariance and efficient stale data utilization +- Enables accurate off-policy metrics monitoring + +## APIs of Algorithms + +You may choose from above algorithms by specifying arguments below: + +`--use-rollout-logprobs`: True if only use `rollout_log_probs` to compute the loss, bypassing old_log_probs calculated by training engine; + +`--use-rollout-correction`: True if apply importance sampling/rejection sampling to loss. + +| `use_rollout_logprobs` | `use_rollout_correction` | Algorithm | Policies |Compute old_log_probs | Batch Invariant | Recommended TIS Mode | +|-----------------|-------------|-----------|--------------|---------------|-----------------|----------------------| +| False | False | Standard PPO (Algorithm 0) | 2 ($\pi_\theta$, $\pi_{\textcolor{blue}{\text{old}}}$)|Yes | No | N/A | +| True | False | Bypassing PPO (Algorithm 3) | 2 ($\pi_\theta$, $\pi_{\textcolor{red}{\text{SGLang}}}$) |🚀 Skipped | No | N/A | +| False | True | Decoupled PPO (Algorithm 2) | 3 ($\pi_\theta$, $\pi_{\textcolor{blue}{\text{old}}}$, $\pi_{\textcolor{red}{\text{SGLang}}}$) |Yes | Yes | token/seq/geo | + +## Configs and Recommended Settings + +When choosing to use importance sampling or rejection sampling for mismatch correction (`use-rollout-correction` enabled, Algorithm 2 & 3), you may specify the IS modes and applied levels. + +### Arguments + +`use-tis`: Enable importance sampling. The IS weight will be multiplied by the policy gradient loss. + +- `--tis-mode`: Mode for IS. Allowed mode: **truncate**, **clip**. +- `--tis-lower-bound`, `--tis-upper-bound`: Bounds for IS weights. +- `--tis-level`: Allowed levels: **token**, **sequence**, **geometric**. See explanations below. +- `--tis-batch-normalize`: Normalize IS weights to mean=1.0 across batch + +`use-rs`: Enable rejection sampling. When choosing to use rejection sampling, the tokens/sequences with an IS weight out of threshold will be directly masked. Those rejected tokens/sequences will not be considered for loss averaging. + +- `--rs-lower-bound`, `--rs-upper-bound`: Bounds for RS +- `--rs-level`: Allowed levels: **token**, **sequence**, **geometric**. See explanations below. +- `--rs-veto-threshold`: Sequence-level rejection threshold for catastrophic mismatches + +### Importance Sampling + +For both IS and RS, we provided 3 levels: **token**, **sequence**, **geometric**. + +**Token Level (default)**: + +Computes importance weights independently for each token: +$w_i = \exp\left(\log \pi_{\text{train}}(x_i) - \log \pi_{\text{rollout}}(x_i)\right)$ + +Characteristics: Biased but computationally simple, suitable for most scenarios + +**Sequence Level**: + +Uses the product of all token weights as the sequence weight: +$w_{\text{seq}} = \exp\left( \sum_i \left( \log \pi_{\text{train}}(x_i) - \log \pi_{\text{rollout}}(x_i) \right) \right)$ + +Characteristics: Unbiased but high variance, suitable for sequence-level optimization + +**Geometric Level**: + +Uses geometric mean to compute sequence weight: +$w_{\text{seq}} = \exp\left( \frac{1}{n} \sum_{i=1}^{n} \left( \log \pi_{\text{train}}(x_i) - \log \pi_{\text{rollout}}(x_i) \right) \right)$ + +Characteristics: Biased but low variance, balances bias and variance + +### Rejection Sampling + +**Token Level**: Reject tokens with IS weight out of threshold + +**Sequence Level:** Reject sequences with mean IS weight out of threshold + +**Geometric Level:** Reject sequences with geometric mean IS weight out of threshold + +- Extremely selective: Requires near-perfect policy match +- High rejection rate: Only suitable for very slight distribution shifts + +**Veto Mechanism**: + +Veto mechanism rejects sequences with catastrophically low token probabilities. +Reject entire sequence if $\exists t \in T$ such that $\rho_t < C_{\text{veto}}$ + +- Prevents catastrophic updates from tokens with near-zero probability under $\pi_{\text{old}}$ +- Independent of IS/RS settings + +*Typical values: $10^{-4}$ to $10^{-6}$* + +## Mismatch Metrics + +When rollout log probabilities are available, MILES automatically tracks comprehensive metrics to monitor training-inference mismatch and importance sampling weights. These metrics help diagnose policy divergence and guide hyperparameter tuning. + +### Mismatch Monitoring Metrics + +These metrics quantify the difference between training and rollout policies. They are computed automatically when `rollout_log_probs` are provided, regardless of whether TIS/MIS correction is enabled. + +| Metric Name | Description | +|------------|-------------| +| `mismatch_training_log_ppl` | Negative mean log probability under training policy: $-\mathbb{E}[\log \pi_{\text{train}}]$ | +| `mismatch_training_ppl` | Perplexity of training policy: $\exp(-\mathbb{E}[\log \pi_{\text{train}}])$ | +| `mismatch_rollout_log_ppl` | Negative mean log probability under rollout policy: $-\mathbb{E}[\log \pi_{\text{rollout}}]$ | +| `mismatch_rollout_ppl` | Perplexity of rollout policy: $\exp(-\mathbb{E}[\log \pi_{\text{rollout}}])$ | +| `mismatch_kl` | Forward KL divergence estimator: $\mathbb{E}[\log \pi_{\text{rollout}} - \log \pi_{\text{train}}]$ | +| `mismatch_k3_kl` | K3 KL estimator: $\mathbb{E}[\exp(r) - r - 1]$ where $r = \log \pi_{\text{train}} - \log \pi_{\text{rollout}}$ | +| `mismatch_log_ppl_diff` | Log perplexity difference| +| `mismatch_log_ppl_abs_diff` | Absolute log perplexity difference | +| `mismatch_ppl_ratio` | Perplexity ratio | +| `train_rollout_logprob_abs_diff` | Token-level absolute log probability difference | + +**Usage**: These metrics help you monitor policy drift. Large values indicate a significant mismatch between the training and rollout engines. + +### IS/RS Correction Metrics + +These metrics track importance sampling weights and corrections. They are only computed when `--use-tis` is enabled. + +When using `--custom-tis-function-path` pointing to MIS implementation (e.g., `mis.py`), additional fine-grained metrics become available: + +| Metric Name | Description | Required Args | Optional Control Args | +|------------|-------------|---------------|----------------------| +| `ois` | On-policy importance sampling ratio: $\exp(\log \pi_{\text{train}} - \log \pi_{\text{old}})$ | `--use-tis` | Only for Algorithm 2 (Decoupled PPO) | +| `mis_mean_is_weight_before_clip` | Raw IS weights before any correction: $\exp(\text{log-ratio})$ | `--use-tis` | `--mis-level` (token/sequence/geometric) | +| `mis_ratio_mean_after_mis` | IS weights after correction (bounded or masked) | `--use-tis` | `--mis-mode`, bounds | +| `mis_truncate_fraction` | Fraction of weights truncated (mode-specific) | `--use-tis`, `--mis-mode=truncate` | `--mis-upper-bound` | +| `mis_clip_fraction_low` | Fraction of weights clipped below lower bound | `--use-tis`, `--mis-mode=clip` | `--mis-lower-bound`, `--mis-upper-bound` | +| `mis_clip_fraction_high` | Fraction of weights clipped above upper bound | `--use-tis`, `--mis-mode=clip` | `--mis-lower-bound`, `--mis-upper-bound` | +| `mis_mask_fraction_low` | Fraction of tokens rejected (below lower bound) | `--use-tis`, `--mis-mode=mask` | `--mis-lower-bound`, `--mis-upper-bound` | +| `mis_mask_fraction_high` | Fraction of tokens rejected (above upper bound) | `--use-tis`, `--mis-mode=mask` | `--mis-lower-bound`, `--mis-upper-bound` | +| `mis_catastrophic_token_fraction` | Fraction of catastrophic tokens (veto-specific) | `--use-tis`, `--mis-veto-threshold` set | Sequence-level rejection | +| `mis_catastrophic_seq_fraction` | Fraction of sequences with catastrophic tokens | `--use-tis`, `--mis-veto-threshold` set | Sequence-level rejection | +| `mis_batch_norm_factor` | Batch normalization factor applied to weights | `--use-tis`, `--mis-batch-normalize` | Normalizes mean to 1.0 | + +## Reference + +We thank the materials below for their excellent findings and theories. + +1. [Mathematical Formulations of Rollout Correction Methods in verl (Yingru Li)](https://github.com/szrlee/verl/blob/yingru/rollout_correction/docs/advance/rollout_corr_math.md). +2. [Your Efficient RL Framework Secretly Brings You Off-Policy RL Training](https://fengyao.notion.site/off-policy-rl) +3. [When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch](https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda) diff --git a/docs/examples/infra-features/true-on-policy.md b/docs/examples/infra-features/true-on-policy.md new file mode 100644 index 0000000000..0910a5ff29 --- /dev/null +++ b/docs/examples/infra-features/true-on-policy.md @@ -0,0 +1,77 @@ +--- +title: "True On-Policy between Training and Inference" +description: "Ensures strictly equal log probabilities between inference (SGLang) and training engines." +# Generated from examples/infra_features/true_on_policy/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +True on-policy ensures that the log probs generated by inference engine (SGLang) is strictly equal to the one generated by the training Engine. + +## Examples + +### Example 1 + +In this script, we provide a minimal example to use true-on-policy. + +```bash +python examples/infra_features/true_on_policy/run_simple.py +``` + +### Example 2 + +This script contains more features for various use cases, and one flag is about the true on policy feature. + +```bash +python scripts/run_qwen3_4b.py --train-backend fsdp --true-on-policy +``` + +In order to quickly see the curve, you may use `--mode debug_minimal`, which will skip evaluation and run generation with a very short output sequence length. Since true on policy is unrelated to OSL or answer correctness, this can be used for quick experiments. + +### Other Cases + +In order to support true on policy for other cases, please refer to the flags changed in the examples above. + +### What is Expected to Observe + +After running the training, you can see in wandb that the metric `train/train_rollout_logprob_abs_diff` should be exactly `0`. This indicates that there is no difference between the log probabilities from the training and the inference. Without the feature enabled, this value should be nonzero. + +### Setup & Results +We fine-tune Qwen3-4B-Base on dapo-math-17k dataset with max_new_tokens = 2048, and evaluate on aime-2024 dataset with max_new_tokens = 8192. +Global batch size is 64 × 16. Results are summarized below. + +![diff](https://raw.githubusercontent.com/radixark/miles/main/examples/infra_features/true_on_policy/src/train_rollout_abs_diff.png) + +![step_time](https://raw.githubusercontent.com/radixark/miles/main/examples/infra_features/true_on_policy/src/step_time.png) + +![rollout_time](https://raw.githubusercontent.com/radixark/miles/main/examples/infra_features/true_on_policy/src/rollout_time.png) + +![raw_rewards](https://raw.githubusercontent.com/radixark/miles/main/examples/infra_features/true_on_policy/src/raw_reward.png) + +![eval](https://raw.githubusercontent.com/radixark/miles/main/examples/infra_features/true_on_policy/src/aime.png) + +### Observations + +Train–inference-diff is strictly 0, verifying full numerical equivalence between training and inference forward passes. Raw rewards perfectly match the baseline, and rollout time shows an acceptable slowdown. + +### Reproduction + +Detailed reproduction refers to [this](https://gist.github.com/fzyzcjy/46f9fc096258cf6fb4516ad2ffcefa8c). + +## How it is Implemented + +The core idea is to make each and every operation in training and inference be bitwise equal. The main code is implemented in [#566](https://github.com/THUDM/slime/pull/566) and [SGLang#12058](https://github.com/sgl-project/sglang/pull/12058). + +Briefly speaking, we handled the following components to make them aligned: + +* Attention: We use the [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) backend for both training and inference, since it achieves bitwise equal between prefill and decode operations. +* GEMM: We use [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) for fast matrix multiplication while preserving true-on-policy, thanks to its algorithm to pick things like tensor core instructions ([SGLang#12142](https://github.com/sgl-project/sglang/pull/12142)). +* Batch invariant kernels: This is a prerequisite for true on-policy, and we use [the ones](https://github.com/thinking-machines-lab/batch_invariant_ops) from the Thinking Machines Lab. +* Torch compile: We also utilize [`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) to speed up by avoiding many tiny kernels. +* We align numeric operation details between the two systems for simplicity, such as op dtype, detailed kernels, etc. Some operations can also be compiled to speedup ([#603](https://github.com/THUDM/slime/pull/603), [SGLang#12161](https://github.com/sgl-project/sglang/pull/12161)). + +In order to more easily align the two parts, we use SGLang's [dumper](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/debug_utils/dumper.py) tool for quick comparisons. (Need [#12622](https://github.com/sgl-project/sglang/pull/12622) and [#12623](https://github.com/sgl-project/sglang/pull/12623) for most convenience.) + +## Future Works + +We will keep maintaining and enhancing this feature. More specifically, we will: + +* Further validate it with more experiments, and compare it with algorithmic fixes such as TIS. +* Potentially scale it to larger scale models if it is proven to be effective. diff --git a/docs/examples/multi-agent.md b/docs/examples/multi-agent.md deleted file mode 100644 index e19f2e8611..0000000000 --- a/docs/examples/multi-agent.md +++ /dev/null @@ -1,197 +0,0 @@ ---- -title: Multi-Agent Co-Evolution -description: Two specialized agents train together and improve each other. ---- -**What you'll learn:** how to wire up an asynchronous multi-agent system in Miles, where -two (or more) specialized agents take alternating turns and the joint outcome drives a -single shared reward. - -This example uses a dual-agent setup that interleaves a "thinker" and a "verifier", but -the same pattern scales to: - -* Doctor / patient simulations. -* Multi-step DeepResearch pipelines. -* Adversarial games (proposer / solver). - -The supporting framework for the production version of this is -[MrlX](https://github.com/AQ-MedAI/MrlX) — Miles ships the kernel of the same idea so -you can hack on it without pulling in MrlX's full dependency tree. - -## Prerequisites - -* You've completed the [Qwen3-30B-A3B](/models/qwen/qwen3-moe) recipe (the - example uses that model). -* Familiar with [Customization](/user-guide/customization). - -## Files - -```text -examples/experimental/multi_agent/ -├── agent_system.py # the agent state machine -├── prompts.py # role / system prompts -├── rollout_with_multi_agents.py # custom rollout (calls agent_system) -└── run-qwen3-30B-A3B-multi-agent.sh # launch script -``` - -## Quick start - -```bash -cd /root/miles -bash examples/experimental/multi_agent/run-qwen3-30B-A3B-multi-agent.sh -``` - -## Configuration - -```python -MULTI_AGENT_CONFIGS = { - "custom_multi_agent_function_path": - "examples.experimental.multi_agent.agent_system.run_agent_system", - "num_parallel": 5, # parallel agent runs per prompt - "incorrect_reward_weight": 0.8, # weight on agent A's reward when wrong - "correct_reward_weight": 1.2, # weight on agent A's reward when right -} -``` - -Asymmetric reward weighting (0.8 / 1.2) gives a small bias toward upweighting "correct" -trajectories, which empirically stabilizes early training when most attempts fail. - -## Launch script highlights - -```bash -ROLLOUT_ARGS=( - --custom-generate-function-path \ - examples.experimental.multi_agent.rollout_with_multi_agents.generate_with_multi_agents - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt --label-key label - --apply-chat-template --rollout-shuffle - --rm-type deepscaler - - --num-rollout 3000 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - - --rollout-max-context-len 16384 # entire conversation budget - --rollout-max-response-len 8192 # per-turn cap - - --global-batch-size 256 - --balance-data -) -``` - -Two flags matter most: - -* `--rollout-max-context-len` — total context budget across all turns. Larger than - `--rollout-max-response-len` because we accumulate. -* `--global-batch-size 256 = 32 × 8` — matches the rollout invariant. - -## Walkthrough — the agent loop - -The shipped pipeline is **solver → rewriter → selector**: `num_parallel` solver -attempts run in parallel, each rewriter rewrites the previous solutions, and a -`SelectorAgent` picks one. Sampling params are set on `args` upstream by the rollout -helper, so `run_agent_system` only takes `(args, sample)`. - -```python agent_system.py -async def run_agent_system(args, sample): - args = deepcopy(args) - args.sample = sample - args.results_dict = {"solver": [], "rewriter": [], "selector": []} - - problem_statement = sample.prompt - - # 1. Solver: num_parallel attempts in parallel. - tasks = [solver_worker(args, problem_statement, i) - for i in range(args.num_parallel)] - solver_solutions = await asyncio.gather(*tasks, return_exceptions=True) - rewards = await batched_async_rm(args, args.results_dict["solver"]) - for s, r in zip(args.results_dict["solver"], rewards): - s.reward = r - - previous = [r for r in solver_solutions if isinstance(r, str)] - - # 2. Rewriter: each worker rewrites the previous solutions. - tasks = [rewrite_worker(args, previous, problem_statement, i) - for i in range(args.num_parallel)] - rewritten = [r for r in await asyncio.gather(*tasks, return_exceptions=True) - if isinstance(r, str)] - - # 3. Selector: pick one of the rewritten solutions. - selector = SelectorAgent() - response = await selector.select(args, problem_statement, rewritten) - - # ... apply asymmetric reward weighting using - # args.incorrect_reward_weight / args.correct_reward_weight on the solver - # and rewriter samples, then return them all together. - return args.results_dict["solver"] + args.results_dict["rewriter"] + ... -``` - -Both roles share the same SGLang process — `solver_worker`, `rewrite_worker`, and -`SelectorAgent.select` all post to the same engine, just with different prompts. So -**both agents are the same model** updating in lockstep. For *architecturally distinct* -agents (separate models), see the MrlX repo. - -## Walkthrough — rollout integration - -`rollout_with_multi_agents.py` exposes `generate_with_multi_agents(args, sample, -sampling_params, evaluation=False)`. Internally it: - -1. Sets `args.sampling_params = sampling_params` and `args.tokenizer`, then loads the - custom multi-agent function from `args.custom_multi_agent_function_path`. -2. Calls `await custom_multi_agent_func(args, sample)` to get the list of samples - produced by the solver / rewriter / selector pipeline. -3. Returns the shuffled list of `Sample`s for the trainer to pack. - -The per-sample tokenization and reward already happen inside `solver_worker` / -`rewrite_worker` / `SelectorAgent.select` (which call `batched_async_rm`), so the -rollout integration itself is a thin wrapper. - -## Tuning knobs - -| Knob | Effect | -|---|---| -| `MAX_TURNS` | Conversation depth — longer = more context = slower | -| `incorrect_reward_weight` / `correct_reward_weight` | Asymmetric shaping | -| `num_parallel` | Rollouts per prompt running concurrently | -| `--rollout-max-context-len` | Stops the conversation when budget is hit | - -## What to watch - -```text -multi_agent/avg_turns 2.5 – 4.0 -multi_agent/early_termination_rate 0.4 – 0.6 (reaches ) -multi_agent/conversation_token_count 4096 – 12288 -loss_mask/role_split balanced (~50/50) -reward/avg trending up -``` - -If `loss_mask/role_split` is heavily skewed, one role is dominating — typically the -verifier becomes verbose. Tighten its system prompt or reduce its `max_tokens`. - -## Troubleshooting - -| Symptom | Fix | -|---|---| -| OOM mid-rollout | Reduce `MAX_TURNS` or `--rollout-max-context-len` | -| Both agents repeat each other | Verifier prompt is too permissive — make it adversarial | -| Reward never moves | Check that `` extraction matches the verifier output | -| Rollout much slower than baseline | Per-turn SGLang RTT × MAX_TURNS — consider async rollout | - -## Variations - -### VLM multi-turn - -Replace `call_role` with a VLM-aware caller that includes images in messages. Miles -supports VLM multi-turn natively — same pattern, just `multimodal_train_inputs` in the -sample dict (see [Customization #13](/user-guide/customization#training)). - -### True asymmetric agents - -Run two SGLang services — one per agent — and have your rollout function call the -appropriate URL per turn. The trainer can either train both jointly (one optimizer per -model) or train one and freeze the other (PvE). - -### Adversarial pairing - -Instead of a verifier, the second agent is an adversary that tries to find weaknesses -in the thinker's answer. Reward both: thinker for surviving, adversary for breaking. -This is the seed of self-play RLHF. diff --git a/docs/examples/multi-lora.md b/docs/examples/multi-lora.md new file mode 100644 index 0000000000..91778cb21c --- /dev/null +++ b/docs/examples/multi-lora.md @@ -0,0 +1,149 @@ +--- +title: "Multi-LoRA Training Example (fully-async)" +description: "Fully-async multi-adapter LoRA training with a slot-keyed adapter page table." +# Generated from examples/multi_lora/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +Train multiple LoRA adapters concurrently against a shared base model, using a +fully-async rollout (continuous producer) + a slot-keyed LoRA page table on the +SGLang engines (in-place upsert, no unload, no drain). + +This example trains two adapters on Qwen3-4B: + +- **gsm8k** — grade-school math, `rm_type: math` +- **dapo_math** — competition math (DAPO-Math-17k), `rm_type: deepscaler` + +## Layout + +``` +run_multi_lora.py # launcher: prepare / train / full-train / serve +service_smoke.py # register/deregister smoke test against the API +adapters/ + gsm8k.yaml + dapo_math.yaml +``` + +The implementation lives in the library: the driver is `train_multi_lora_async.py` +at the repo root (next to `train.py`/`train_async.py`), the rollout fn and data +source are `miles/rollout/multi_lora/`, and the controller is +`miles/ray/multi_lora/` (registry + backend + HTTP API, plus the named Ray +actor pinned to the head node). + +## Design (decoupled per-adapter optimizers) + +- **Controller** (Ray actor + control-plane HTTP API) is the source of truth: + `POST/GET/DELETE /adapter_runs` plus `GET /adapter_runs/state`. The data source + reads it; the trainer reads it. Generation traffic goes straight to the router; + on deregister the controller aborts the adapter's in-flight requests + engine-side by rid prefix (`rid = {adapter}::{uuid}`, set in `generate`). +- **Per-adapter gradient accumulation.** Each adapter has its own batch shape: + `rollout_batch_size` prompt groups per optimizer step, each group holding + `n_samples_per_prompt` responses (`adapter_global_batch_size = + rollout_batch_size x n_samples_per_prompt` samples per step). Completed + prompt groups flow into training continuously in multiples of the + adapter's `min_groups_per_dp_split` (the smallest group count whose samples + split evenly across data-parallel ranks), gradients + accumulate in the DDP buffers across train batches, and an adapter's + optimizer steps exactly when its adapter batch fills — independent of every other + adapter. The controller tracks adapter batch progress (`accumulated_groups`) and commits + it only after a successful train call. +- **Per-slot optimizers.** One Adam per adapter slot under Megatron's + `LayerWiseDistributedOptimizer` (whole-parameter ZeRO-1): per-slot state, + step counts, and gradient clipping; optimizer state sharded across DP ranks; + plain DDP all-reduce (no distributed optimizer) makes cross-batch gradient + retention idempotent. +- **Batch collection.** The collection loop (same shape as fully_async's) + pops groups from the per-adapter buffers round-robin, one + `min_groups_per_dp_split` at a time, capped at each adapter's remaining + batch, until the batch reaches `--global-batch-size` samples or a non-empty + batch makes no progress for `--multi-lora-max-coalesce-wait-s` (the target + can be permanently unreachable, so it trains on whatever is ready) — a + single adapter with a small batch trains alone without waiting for + anyone. Samples enter the gradient buffers with weight 1; at step time the + slot's accumulated gradient is scaled by `1/adapter_global_batch_size` + (a constant known in advance), so an adapter's update is identical to what + it would get training alone. +- **Selective weight sync.** Only adapters whose optimizer stepped are pushed + to the engines (upsert into the slot-keyed page table); only their slot + versions bump, keeping staleness filtering per-adapter accurate. +- Adapters deregister on committed optimizer-step count (`num_step`) in the + controller's train-commit path (`mark_batch_trained`), so stop checks happen + exactly when steps advance. `num_step` is relative to the adapter's + start/resume step. When an adapter doesn't set `num_step`, it is derived + from `num_epoch` (default 1) as `num_epoch x len(dataset) // + rollout_batch_size` once the data source loads the dataset (post-filter + length). The trainer's + `reconcile_adapters` (before each generate) retires it at the next sync + point and cleans up (save ckpt + clear Megatron slot + zero its optimizer + state and retained gradients). The adapter's untrained tail — buffered + groups and any partially accumulated gradients — is discarded. +- **Batch ⊆ loaded property:** `reconcile_adapters` runs before `generate`, so the + batch is fetched with loaded = active; active only shrinks during generate, so every + adapter in the batch is live on the trainer. + +## Provision (once) + +```bash +python examples/multi_lora/run_multi_lora.py prepare +``` + +Downloads `Qwen/Qwen3-4B` (to `/root/models`), `zhuzilin/dapo-math-17k`, and +`zhuzilin/gsm8k` (to `/root/datasets`). + +## Run + +```bash +python examples/multi_lora/run_multi_lora.py train # or: full-train (prepare + train) +``` + +Registers the two adapters from CLI flags and trains until each hits its `num_step`, +then exits. + +## Service mode + +```bash +python examples/multi_lora/run_multi_lora.py serve +``` + +Starts with no adapters and idles; register/deregister at runtime through the +control-plane API (port 8068): + +```bash +python examples/multi_lora/service_smoke.py --api-url http://127.0.0.1:8068 \ + --data /root/datasets/gsm8k/train.parquet --input-key messages --label-key label --rm-type math +``` + +## Multi-LoRA CLI flags + +| Flag | Purpose | +| --- | --- | +| `--multi-lora-n-adapters N` | Max concurrent adapter slots. `0` disables (default); `> 0` enables. | +| `--multi-lora-adapter NAME PATH` | Register an adapter at startup. Repeatable. `PATH` → an `adapter.yaml`. | + +Per-adapter `rank` in `adapter.yaml` must be `<= --lora-rank`. + +## adapter.yaml + +```yaml +rank: 16 +alpha: 16 +rollout_batch_size: 32 # prompt groups per optimizer step (defaults to --rollout-batch-size) +n_samples_per_prompt: 4 # group shape (defaults to --n-samples-per-prompt) +data: /root/datasets/gsm8k/train.parquet +input_key: messages +label_key: label +rm_type: math +num_step: 400 # stop adapter after N optimizer steps + # (default: derived from num_epoch, itself default 1) +# optional: save, num_epoch, custom_rm_path, ... +``` + +The derived `adapter_global_batch_size = rollout_batch_size x +n_samples_per_prompt` is the adapter's samples-per-optimizer-step (the +per-adapter analog of `--global-batch-size`). + +Batch-shape constraints (validated at registration, not at runtime): +`n_samples_per_prompt` must be a divisor or multiple of the trainer's +data-parallel size; `rollout_batch_size` must be a multiple of the adapter's +`min_groups_per_dp_split`; +`adapter_global_batch_size` is capped by +`--multi-lora-max-adapter-global-batch-size` (default 4x `--global-batch-size`). diff --git a/docs/examples/on-policy-distillation.md b/docs/examples/on-policy-distillation.md new file mode 100644 index 0000000000..28a8f99741 --- /dev/null +++ b/docs/examples/on-policy-distillation.md @@ -0,0 +1,22 @@ +--- +title: "On-Policy Distillation Examples" +description: "Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop." +# Generated from examples/on_policy_distillation/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +The canonical OPD documentation lives in +[`docs/advanced/on-policy-distillation.md`](/advanced/on-policy-distillation). +Keep the algorithm description, arguments, teacher-mode comparison, and +Rethinking OPD top-k recipe there so we do not maintain two copies. + +This directory contains runnable examples: + +- `run-qwen3-8B-opd.sh`: SGLang teacher server OPD. This script enables + Rethinking OPD with `--opd-log-prob-top-k 16`, `--opd-top-k-strategy only-student`, + and `--opd-reward-weight-mode student_p`. +- `run-qwen3-8B-opd-multi-teacher.sh`: Multi-teacher OPD with per-sample routing. + Math prompts are scored by a Qwen3-32B teacher and code prompts by a + Qwen3-Coder-30B-A3B teacher, selected via `--opd-teacher-urls` and a per-row + `{"metadata": {"opd_teacher": ...}}` tag in the dataset. +- `run-qwen3-8B-opd-megatron.sh`: Megatron-loaded teacher OPD. + +Use `--opd-log-prob-top-k 0` to run the original sampled-token OPD path. diff --git a/docs/examples/on-policy-distillation/qwen3-5-35b-selfdistill.md b/docs/examples/on-policy-distillation/qwen3-5-35b-selfdistill.md new file mode 100644 index 0000000000..ad80da5a01 --- /dev/null +++ b/docs/examples/on-policy-distillation/qwen3-5-35b-selfdistill.md @@ -0,0 +1,213 @@ +--- +title: "Qwen3.5-35B-A3B Self-Distillation on a Single Node (RLVR teacher → OPD)" +description: "Two-phase self-distillation of Qwen3.5-35B-A3B on one 8xH200 node, with an in-process Megatron teacher." +# Generated from examples/on_policy_distillation/qwen3_5_35b_selfdistill/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +A reproducible two-phase on-policy-distillation (OPD) example for the +**Qwen3.5-35B-A3B** MoE on a **single 8×H200 node**, using the **in-process +Megatron teacher** (`--opd-type megatron`, no separate teacher server). + +It differs from the sibling examples in three ways: + +1. **Real MoE at scale on one node.** The 2-node/16-GPU recipe is re-tiled to 8 GPUs. +2. **A genuinely diverged teacher.** `run-qwen3-8B-opd-megatron.sh` uses `teacher == base` + (a mechanism demo where the reverse-KL is ~0). Here Phase 1 *trains* the teacher + with RLVR so it is measurably better and more concise than the base — the + prerequisite for OPD to actually move the student. +3. **Self-distillation is the only valid option here.** Qwen3.5 has its own tokenizer + (vocab 248320); the smaller Qwen3 models (vocab 151936) are not token-compatible, + so a cross-model teacher would be invalid. Teacher and student are the same family. + +## Pipeline + +``` +Phase 1 (phase1_rlvr_teacher.sh) Phase 2 (phase2_opd_selfdistill.sh) +base 35B --RLVR (GRPO, lr 1e-5)--> teacher base 35B (student) + better + concise | <-- reverse-KL (--opd-type megatron) + (eval 0.83 -> 0.89) teacher (Phase-1 ckpt, in-process) +``` + +## Single-node parallelism (world = 8) + +The original recipe was 2 nodes × 8 GPUs (`TP2 PP1 CP2 EP8 ETP1`, DP4). On one node +we keep the same dims and only halve DP: + +| dim | value | check | +|-----|-------|-------| +| TP | 2 | decoder `TP*PP*CP = 2` ; `8 % 2 = 0` → DP = 4 | +| PP | 1 | | +| CP | 2 | shards the long (~17k) sequence so 24k context fits | +| EP | 8 | `num_experts 256 % 8 = 0` ; expert `ETP*EP*PP = 8` → expert_dp = 1 | +| ETP | 1 | expert_dp(1) ≠ dp(4) is allowed (miles rank order ends in `pp`) | + +`--colocate` time-shares the train and rollout phases (each fits 143 GB separately, +not summed); `--optimizer-cpu-offload` puts Adam state on host RAM; the model is a +hybrid linear-attention MoE so the KV cache is small. Peak ≈ 124 GB / 143 GB per GPU. + +## Reproduce + +**0. Prereqs** — model + torch_dist checkpoint, and the train/eval split: + +```bash +# model (and mcore conversion, see ../README.md for convert_hf_to_torch_dist usage) +# ${MODEL_DIR}/Qwen3.5-35B-A3B and ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist +# disjoint, seeded train/eval split (eval is held out from BOTH phases): +python make_split.py --src /path/to/dapo-math-17k.jsonl --out-dir ${DATA_DIR} +# -> ${DATA_DIR}/dapo_train.jsonl (16886) ${DATA_DIR}/dapo_eval.jsonl (512) +``` + +**1. Phase 1 — train the teacher** (watch `rollout/raw_reward` climb and +`eval/dapo_heldout` rise above the base ~0.83): + +```bash +MODEL_DIR=... DATA_DIR=... OUTPUT_DIR=/persistent/ckpt-teacher \ + bash phase1_rlvr_teacher.sh +``` + +**2. Phase 2 — distill the teacher into the base student**: + +```bash +# pure OPD (default): training reward = 0, only the teacher reverse-KL drives learning +TEACHER_LOAD=/persistent/ckpt-teacher DATA_DIR=... \ + bash phase2_opd_selfdistill.sh + +# grounded OPD: correctness reward (raw_reward == accuracy, climbs) + teacher reverse-KL +MODE=grounded TEACHER_LOAD=/persistent/ckpt-teacher DATA_DIR=... \ + bash phase2_opd_selfdistill.sh +``` + +`OUTPUT_DIR` / the teacher checkpoint must live on **persistent** storage. On a +KubeRay pod the head can be recreated and wipe the container overlay (`/root`); a +node-local disk (e.g. `/node_public`) survives and makes runs resumable. + +## Run on GB200 / GB300 (CUDA 13, Blackwell) — `phase2_gb200.sh` + +The recipe above targets a single **8×H200** node. Blackwell nodes (GB200/GB300) +have **4 GPUs/node**, so `world = 8` becomes **2 nodes × 4 GPUs** — same parallel +dims (`TP2 PP1 CP2 EP8 ETP1`, DP4), only the node tiling changes. `phase2_gb200.sh` +is the GB200 variant of `phase2_opd_selfdistill.sh`; the deltas (all validated on +2× GB200, reproducing the base eval `0.84` / `~14k`) are: + +- **Tiling** — `--actor-num-nodes 2 --num-gpus-per-node 4` (override via + `ACTOR_NUM_NODES` / `GPUS_PER_NODE`). Pin both nodes to one NVLink (MNNVL) domain + so the EP8 all-to-all stays on the NVLink fabric. +- **sglang backends** (cf. `scripts/run_qwen3_5_35b_a3b_mtp.py`) — + `--sglang-moe-runner-backend flashinfer_cutlass`, `--sglang-attention-backend + trtllm_mha`, and `--moe-token-dispatcher-type flex`. The default triton fused-MoE + mis-shards routed experts on the megatron→sglang weight sync + (`fused_moe_triton ... _load_w13`: `tensor a (64) vs b (2048)`), and FA3 is SM≤90 + only (Blackwell is SM 10.x). +- **NCCL** — `NCCL_NVLS_ENABLE=0` (multi-node Blackwell NVLS bind fails + `ncclCommInitRank`); keep `NCCL_MNNVL_ENABLE=1`. +- **k8s** — if a `prometheus` Service exists in the namespace, set + `PROMETHEUS_PORT=9090` (kube injects a `tcp://…:9090` URL that breaks miles' + `int(PROMETHEUS_PORT)`). + +`phase2_gb200.sh` already sets the sglang/MoE backends and folds +`NCCL_NVLS_ENABLE=0` + `PROMETHEUS_PORT=9090` into the Ray runtime env. Run it on the +Ray head in the CUDA-13 ARM64 miles image, with `MILES_DIR` pointing at the repo: + +```bash +ACTOR_NUM_NODES=2 GPUS_PER_NODE=4 MILES_DIR=/workspace/miles \ +MODEL_DIR=... DATA_DIR=... TEACHER_LOAD=/persistent/ckpt-teacher OUTPUT_DIR=/persistent/ckpt-opd-pure \ + bash phase2_gb200.sh # MODE=pure (default) | MODE=grounded +``` + +## Run Phase 2 only (skip Phase 1) + +If you already have a teacher checkpoint, skip Phase 1 and run Phase 2 directly — +point `--opd-teacher-load` (`TEACHER_LOAD`) at the teacher's **torch_dist parent +dir** (the one containing `latest_checkpointed_iteration.txt`). You still need the +base model (`--hf-checkpoint` + the `--ref-load` torch_dist) and the data split, but +no Phase-1 run. + +If your teacher is in **HuggingFace** format, convert it to torch_dist first with +`convert_gb200.sh` (a thin wrapper over `tools/convert_hf_to_torch_dist.py` carrying +the Qwen3.5 `MODEL_ARGS`): + +```bash +# teacher: HF safetensors -> Megatron torch_dist parent dir +bash convert_gb200.sh /path/to/teacher-hf /persistent/ckpt-teacher +# (and the base, if you don't have Qwen3.5-35B-A3B_torch_dist yet) +bash convert_gb200.sh ${MODEL_DIR}/Qwen3.5-35B-A3B ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist + +TEACHER_LOAD=/persistent/ckpt-teacher MODEL_DIR=... DATA_DIR=... \ + bash phase2_gb200.sh # or phase2_opd_selfdistill.sh on 8×H200 +``` + +> **Teacher expert layout.** The public `Qwen/Qwen3.5-35B-A3B` ships *fused* experts +> (`mlp.experts.gate_up_proj`); a teacher round-tripped through +> `convert_torch_dist_to_hf` may ship *unfused* per-expert weights +> (`mlp.experts.{i}.gate_proj.weight`). `miles_plugins/mbridge/qwen3_5.py` now +> autodetects both for the main layers (mirroring the existing MTP-expert +> autodetect), so either layout converts without manual re-fusing. + +## Results (DAPO-math, held-out 512, eval @ 24k cap, temp 0.6) + +**Phase 1 — RLVR teacher** (lr 1e-5): + +| step | eval/dapo_heldout | eval response length | +|------|-------------------|----------------------| +| 0 (base) | 0.828 | 14,070 | +| 5 | **0.887** | **6,248** | + +The teacher becomes both more accurate **and** ~2× more concise. This Phase-1 +teacher checkpoint is published at +[**cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher**](https://huggingface.co/cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher) +(weights only) and can be used directly as the Phase-2 teacher via +`--opd-teacher-load` after `convert_hf_to_torch_dist.py`. + +**Phase 2 — pure OPD** (student = base, teacher = Phase-1 step-5 ckpt; reward = 0): + +| step | eval/dapo_heldout | eval response length | opd_reverse_kl | +|------|-------------------|----------------------|----------------| +| 0 (base) | 0.840 | 14,070 | — | +| 5 | 0.852 | **6,132** | 0.045 → 0.013 | + +With **zero task reward**, pure reverse-KL distillation transfers the teacher's +concise behavior to the base student — eval length **−57%** with accuracy +preserved/slightly up (the +1.2 pt is within the ~1.6 pt eval SE; the robust, +headline effect is the efficiency transfer). A nonzero, shrinking `opd_reverse_kl` +confirms the teacher genuinely differs from the student and the student is +converging onto it. + +**Phase 2 — grounded OPD** (correctness reward + teacher reverse-KL): + +| step | rollout/raw_reward | train length | opd_reverse_kl | +|------|--------------------|--------------|----------------| +| 1 | 0.637 | 18,778 | 0.045 | +| 2 | **0.910** | **7,665** | 0.014 | + +With the correctness reward kept, `rollout/raw_reward` (== accuracy) climbs while +the student simultaneously adopts the teacher's concise responses (18.8k → 7.7k). +The shrinking `opd_reverse_kl` (0.045 → 0.014) shows the student converging onto +the teacher. (At lr 1e-5 the RLVR reward alone also drives accuracy up — Phase 1 +is the controlled view of that — so grounded OPD's `raw_reward` climb reflects +RLVR + the teacher pull combined; the pure-OPD run above isolates OPD's effect.) + +## Gotchas (each cost a wasted run to find) + +- **Reward grader.** `--rm-type deepscaler` requires a `` tag and returns 0 + otherwise; Qwen3.5 reasons inline (no tag) → every reward 0. `--rm-type math` + only reads `\boxed{}`; `--rm-type dapo` only `Answer:`. Use the format-agnostic + `rm.reward_func` (accepts either). Always pass `--label-key label` for the + `{prompt, label}` DAPO jsonl, or `Sample.label` is `None` and reward reads 0. +- **Context length.** The 35B's DAPO chain-of-thought is ~14–17k tokens. An 8k + response cap truncates ~95% of rollouts mid-reasoning → reward ~0. Use ≥24k + (CP2 makes 24–32k feasible). +- **`--opd-teacher-load` path.** Point at the checkpoint **parent** dir (contains + `latest_checkpointed_iteration.txt`), not an `iter_XXXXXXX` subdir. The subdir + has no metadata → silent fallback to base → teacher == student → `opd_reverse_kl ≈ 0`. + Sanity check: in the rollout log, `teacher_log_probs` should differ from + `rollout/log_probs`. +- **Teacher must diverge.** A few RLVR steps at lr 1e-6 barely move the weights, so + the teacher ≈ base and OPD is inert (`opd_reverse_kl ≈ 5e-4`). lr 1e-5 diverges it + fast (`opd_reverse_kl ≈ 5e-2`). `--opd-kl-coef` cannot amplify a ~0 KL. +- **Memory.** `with_ref = (--use-kl-loss or --kl-coef≠0)`. Dropping `--use-kl-loss` + keeps only student + teacher (2×35B) in memory; the teacher reverse-KL is the + regularizer. Adding it loads a 3rd model and risks OOM. + +## References +- Phase-1 teacher checkpoint: https://huggingface.co/cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher +- ../README.md (served-teacher OPD), ../run-qwen3-8B-opd-megatron.sh (in-process teacher) +- https://thinkingmachines.ai/blog/on-policy-distillation/ diff --git a/docs/examples/openhermes-sft.md b/docs/examples/openhermes-sft.md deleted file mode 100644 index 10a87acc88..0000000000 --- a/docs/examples/openhermes-sft.md +++ /dev/null @@ -1,163 +0,0 @@ ---- -title: SFT on OpenHermes -description: Plain supervised fine-tuning of Qwen3-4B-Base on the OpenHermes-2.5 dataset. ---- -**What you'll learn:** how to use Miles for plain supervised fine-tuning. No RL, no -rollout, no reward — just data → loss → optimizer. - -Why use Miles for SFT? Two reasons: - -1. **Same launch convention as your RL run** — one config, one Ray cluster. -2. **Async data prefetching** — the SFT loop reuses the rollout machinery to overlap - data loading with training. - -## Prerequisites - -* You completed the [Qwen3-4B](/models/qwen/qwen3) recipe (we reuse the - conversion). -* ~50 GB free disk for OpenHermes-2.5. - -## Quick start - -### 1. Convert Qwen3-4B-Base - -If you don't already have it: - -```bash -hf download Qwen/Qwen3-4B-Base --local-dir /root/models/Qwen3-4B-Base - -cd /root/miles -MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3-4B)" || exit 1 -read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}" -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/models/Qwen3-4B-Base \ - --save /root/models/Qwen3-4B-Base_torch_dist -``` - -### 2. Prepare the dataset - -OpenHermes ships in a custom shape. Convert to OpenAI messages format: - -```python -from datasets import load_dataset - -ds = load_dataset("teknium/OpenHermes-2.5")["train"] - -def convert(sample): - role_map = {"human": "user", "gpt": "assistant", "system": "system"} - return { - "messages": [ - {"role": role_map[turn["from"]], "content": turn["value"]} - for turn in sample["conversations"] - ] - } - -ds = ds.map(convert) -ds.to_parquet("/root/datasets/openhermes2_5.parquet") -``` - -### 3. Run - -```bash -python scripts/run_qwen3_sft.py --model-name Qwen3-4B-Base -``` - -The same launcher runs the 4-node Qwen3-235B-A22B SFT with `--model-name Qwen3-235B-A22B`. -The dataset is read from `--data-dir` (default `/root/datasets`), the checkpoints from -`--model-dir` (default `/root/models`), and the run writes to `--output-dir` -(default `/root/shared_data`). - -## What changes vs. the GRPO recipe - -Compare to [run_qwen3_dense.py](/models/qwen/qwen3). The deltas: - -```diff -- python3 train.py -+ python3 train_async.py # async for data prefetch - -- # GRPO knobs, n-samples-per-prompt, ... -+ --rollout-function-path miles.rollout.sft_rollout.generate_rollout -+ --prompt-data /openhermes2_5.parquet -+ --input-key messages -+ --rollout-shuffle -+ --num-epoch 3 -+ --rollout-batch-size 128 -+ --global-batch-size 128 -+ -+ --loss-type sft_loss -+ --calculate-per-token-loss -+ --disable-compute-advantages-and-returns -+ --debug-train-only - -- # the GRPO flags are gone entirely -- # so are the --sglang-* flags — no inference needed -``` - -## Why each flag - -| Flag | Why | -|---|---| -| `--rollout-function-path miles.rollout.sft_rollout.generate_rollout` | Read from disk instead of generating | -| `--rollout-batch-size = --global-batch-size` | One batch read = one optimizer step | -| **No** `--n-samples-per-prompt` | SFT has one target per input | -| `--loss-type sft_loss` | Cross-entropy instead of policy-gradient | -| `--calculate-per-token-loss` | Standard SFT averages over unmasked tokens | -| `--disable-compute-advantages-and-returns` | No advantage / return needed | -| `--debug-train-only` | Skip SGLang init (we don't need rollout) | -| `train_async.py` | Async data prefetch overlaps load with train | - -## What to watch - -```text -sft/loss decreasing -sft/per_token_loss decreasing (mirrors loss when using per-token) -sft/tokens_seen steadily increasing -sft/epoch 0 → num_epoch -data/prefetch_queue_depth > 0 (else loader is the bottleneck) -``` - -If `data/prefetch_queue_depth` stays at 0, your data loader is too slow — increase -worker count or use parquet (we already do). - -## Tuning knobs - -| Knob | Effect | -|---|---| -| `--num-epoch` | Total passes over dataset | -| `--rollout-batch-size` | Bigger = better GPU utilization, more memory | -| `--max-tokens-per-gpu` | As always — push it up until OOM | -| `--lr` | SFT typically `1e-5` to `5e-5` (10× higher than RL) | -| `--lr-decay-style cosine --lr-warmup-iters 100` | Standard SFT schedule | - -## Variations - -### Mix datasets - -Pass multiple `--prompt-data` entries: - -```bash ---extra-args "--prompt-data \ - hermes /root/datasets/openhermes2_5.parquet \ - slimorca /data/slimorca.parquet" -``` - -Per-source loss is logged separately. - -### Continue with RL - -After SFT, point the RL run at the SFT checkpoint: - -```bash ---hf-checkpoint /root/models/Qwen3-4B-Base ---ref-load /root/models/Qwen3-4B-Base_torch_dist # original (anchor) ---load /root/shared_data/checkpoints # SFT output (start point) ---save /root/shared_data_rl/checkpoints -``` - -Both runs derive `--load` / `--save` from `--output-dir`, so give the RL run its own. - -### LoRA SFT - -Use the LoRA hooks (`--lora-rank 16`) to keep VRAM low when fine-tuning a -larger model. See `examples/lora/` in the repo. diff --git a/docs/examples/ppo.md b/docs/examples/ppo.md new file mode 100644 index 0000000000..7ccc067107 --- /dev/null +++ b/docs/examples/ppo.md @@ -0,0 +1,95 @@ +--- +title: "PPO Example" +description: "Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs." +# Generated from examples/ppo/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +This example trains Qwen3-4B with **PPO** — the actor-critic algorithm, with a learned value +model and GAE advantages — on a single node with the Megatron backend. + +## PPO vs. GRPO in one paragraph + +To turn a reward into a learning signal you need a baseline: "was this response better or worse +than expected?" GRPO gets that baseline for free by sampling a *group* of responses per prompt and +comparing each against the group average. PPO instead trains a second network, the **critic**, +whose only job is to predict the expected reward of a partial response; the advantage is then how +much better the actual outcome was than the critic's prediction. The trade-off: PPO carries a +second model (more memory, more code paths), but its baseline is per-token rather than +per-group, and it does not need a large `--n-samples-per-prompt` to be well-behaved. + +In miles the critic is **colocated on the actor's train GPUs**, so PPO needs no extra GPUs over +the GRPO equivalent. It pays for that in memory, which is why `--offload-train` is turned on for +you — see [Constraints](#constraints-worth-knowing-before-you-debug). + +## Files + +* `run_qwen3_4b_ppo.py`: single-node launch script for Qwen3-4B. + +## Quick Start + +```bash +cd miles +python examples/ppo/run_qwen3_4b_ppo.py +``` + +The script's `prepare` step downloads Qwen3-4B and the DAPO-Math-17k dataset and converts the +checkpoint to Megatron `torch_dist` format, so there is nothing to set up by hand. Conversion is +skipped on reruns. + +## Turning PPO on + +The only flag that selects the algorithm is: + +```bash +--advantage-estimator ppo +``` + +Everything else is tuning. Passing it sets `use_critic`, which builds the critic and switches +advantage computation to GAE. + +## Critic flags + +| Flag | Default | Meaning | +|---|---|---| +| `--critic-lr` | falls back to `--lr` | Critic learning rate. Usually wants to be larger than the actor's — this example uses `1e-5` against an actor `1e-6`. | +| `--critic-load` | falls back to `--load` | Critic init checkpoint. | +| `--critic-save` | `--save` + `_critic` | Sibling directory, so the two models do not clobber each other's iteration tracker. | +| `--critic-lr-warmup-iters` | `0` | Linear warmup for the critic only. | +| `--num-critic-only-steps` | `0` | Value-function warmup: the actor stays frozen for this many initial rollout steps while the critic learns. A critic that starts from noise otherwise injects noisy advantages into the very first actor updates. | +| `--critic-num-nodes`, `--critic-num-gpus-per-node` | inherited from the actor | Set automatically — see the colocation constraint below. | + +## Constraints worth knowing before you debug + +These are enforced at argument validation, so you get an error rather than a silent wrong result: + +* **The critic is colocated with the actor, and inherits its parallelism.** The critic is placed + on exactly the same GPUs as the actor — `--critic-num-nodes` and `--critic-num-gpus-per-node` + are overwritten with the actor's values — and it currently reuses the actor's TP/PP/CP as well, + so there is no way to give the critic its own parallelism. Two consequences: **`--offload-train` + is forced on**, because both models resident on the same devices at once is usually too much + (`--no-offload-train` is accepted but warns, and is meant for offload debugging only); and when + you scale, you only ever change the actor's placement — the actor world size is + `--actor-num-nodes` × `--actor-num-gpus-per-node`, and `TP × PP × CP` must divide it. +* **Megatron only.** PPO raises with any other train backend, and is unsupported with + `--megatron-to-hf-mode bridge`. +* **`--kl-coef` must be 0.** Reward-level KL is rejected because the critic trains *before* the + actor and never sees ref log probs, so its value targets would silently exclude the KL penalty + applied to the actor's rewards. Use loss-level `--use-kl-loss` / `--kl-loss-coef` instead. +* **Not compatible with `MILES_EXPERIMENTAL_FT_TRAINER=1`.** The v2 fault-tolerant train group + cannot route critic values yet. + +## Which numbers here are verified + +The parallelism (`TP=1`, `PP=2`, `CP=2` over 4 GPUs), the GPU count, and the PPO flag set follow +`tests/e2e/megatron/test_qwen3_4B_ppo.py`, which runs in CI. + +Three values are deliberately **not** the CI ones, because the CI test is a 3-step smoke test +rather than a training recipe: + +* `--eps-clip 0.2` here vs. `4e-4` in CI. `4e-4` pins the actor almost in place, which is useful + for a fast deterministic test and wrong for actual training. `0.2` is the standard PPO value. +* `--num-rollout 300` here vs. `3` in CI. +* `--rollout-num-gpus-per-engine 1` here vs. `2` in CI. Qwen3-4B fits comfortably on one GPU, so + one engine per GPU avoids paying tensor-parallel communication for no capacity gain. + +Treat the rest — learning rates, `--kl-loss-coef`, `--entropy-coef` — as starting points to tune, +not as tuned values. diff --git a/docs/examples/reproducibility.md b/docs/examples/reproducibility.md deleted file mode 100644 index 28ee85a595..0000000000 --- a/docs/examples/reproducibility.md +++ /dev/null @@ -1,144 +0,0 @@ ---- -title: Reproducibility Recipe -description: Bit-stable training across reruns. Determinism flags, seeds, and what to watch. ---- -**What you'll learn:** how to configure Miles + SGLang + Megatron for **bit-wise -reproducible** RL training. Same inputs → identical outputs across reruns, machines, -and time. - -Reproducibility matters more than people admit: it's the only way to A/B test changes -in a system this complex. If you can't trust that "I changed flag X and reward went up", -you're flying blind. - -## How it works - -Bit-wise reproducibility requires **three** independent stacks to be deterministic: - -1. **Inference (SGLang)** — every kernel must be deterministic. -2. **Training (Megatron-LM)** — same. -3. **Communication (NCCL)** — algorithm choice and CUBLAS workspace can be - non-deterministic by default. - -Miles ships verified configurations that switch all three into deterministic mode. - -## Quick start - -We use the smallest Miles model (Qwen2.5-0.5B) on GSM8K so the loop fits in 5 minutes -and you can reproduce the bit-stability check yourself. - -### 1. Disable FA3 - -Flash-Attention 3 currently has non-deterministic backward kernels. Drop it: - -```bash -pip uninstall flash_attn_3 -y -``` - -### 2. Set the deterministic flags - -```bash -SGLANG_ARGS+=( - --sglang-enable-deterministic-inference - --sglang-attention-backend flashinfer -) - -PERF_ARGS+=( - --deterministic-mode -) -``` - -### 3. Set the env vars (Ray `env_vars`) - -```python -"env_vars": { - "NCCL_ALGO": "Ring", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", - "CUBLAS_WORKSPACE_CONFIG": ":4096:8" -} -``` - -| Variable | Why | -|---|---| -| `NCCL_ALGO=Ring` | Forces a deterministic collective algorithm | -| `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` | Disables non-deterministic Transformer-Engine kernels | -| `CUBLAS_WORKSPACE_CONFIG=:4096:8` | cuBLAS's deterministic workspace allocation | - -### 4. Download + convert + run - -```bash -hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/gsm8k -hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir /root/Qwen2.5-0.5B-Instruct - -cd /root/miles -MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen2.5-0.5B)" || exit 1 -read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}" -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen2.5-0.5B-Instruct \ - --save /root/Qwen2.5-0.5B-Instruct_torch_dist/ - -bash examples/experimental/reproducibility/run-qwen2.5-0.5B-gsm8k.sh -``` - -### 5. Verify - -Run twice, then: - -```bash -md5sum /root/Qwen2.5-0.5B-Instruct_torch_dist_run1/iter_0000020/*.pt -md5sum /root/Qwen2.5-0.5B-Instruct_torch_dist_run2/iter_0000020/*.pt -``` - -The hashes should match exactly. If they don't, something escaped the deterministic -configuration — see troubleshooting below. - -## What's deterministic and what isn't - -| Component | Default | Deterministic mode | -|---|---|---| -| Megatron forward | non-det | ✅ via `--deterministic-mode` | -| Megatron backward | non-det | ✅ | -| SGLang kernels | non-det | ✅ via `--sglang-enable-deterministic-inference` | -| Flash-Attn 3 | non-det | ❌ — uninstall | -| NCCL collectives | non-det | ✅ via `NCCL_ALGO=Ring` | -| cuBLAS GEMM | non-det | ✅ via `CUBLAS_WORKSPACE_CONFIG` | -| TE fused kernels | non-det | ✅ via `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` | -| Python dataloader shuffle | seeded | ✅ already | - -## Troubleshooting - -| Symptom | Likely cause | -|---|---| -| Hashes diverge after iter 1 | Flash-Attn 3 still installed | -| Hashes match for trainer but not SGLang | `--sglang-attention-backend flashinfer` not set | -| Hashes diverge across nodes | `NCCL_ALGO=Ring` not propagated to all workers | -| Hashes match locally but not on a different machine | cuDNN version mismatch | - -## Cost of determinism - -Roughly: - -| Component | Throughput cost | -|---|---| -| Megatron deterministic mode | -3% to -8% | -| SGLang deterministic | -10% to -15% | -| NCCL Ring | -2% (vs. Tree) | -| Drop FA3 | -10% to -25% | - -Total: ~25% slower for full bit-wise determinism. Worth it for debugging, science, and -A/B tests; usually disabled for production runs. - -## When to disable determinism - -* Production training runs where the cost is too high. -* When you've already nailed the result and want maximum throughput. -* On hardware that doesn't support deterministic kernels. - -For most other cases — **keep it on while developing**. The hours you save when "I -swear that worked yesterday" stops being a phrase you say will pay back the throughput -many times over. - -## References - -* [SGLang deterministic inference blog](https://lmsys.org/blog/2025-09-22-sglang-deterministic/) -* Megatron-LM deterministic-mode docs diff --git a/docs/examples/retool-v2.md b/docs/examples/retool-v2.md new file mode 100644 index 0000000000..bc687a78da --- /dev/null +++ b/docs/examples/retool-v2.md @@ -0,0 +1,36 @@ +--- +title: "Retool v2" +description: "Tool-enabled language model generation with sandboxed Python code execution interleaved with thinking." +# Generated from examples/retool_v2/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +This example is an upgraded version of the original retool example, using the updated interfaces provided by the miles framework to implement multi-turn RL training with tool calls in a cleaner way. + +## Key Differences from v1 + +**v1 (retool)** requires manually implementing the full multi-turn conversation loop in `generate_with_retool.py`, directly depending on low-level `GenerateState` and `sglang_rollout` interfaces — resulting in verbose code tightly coupled to the framework internals. + +**v2 (retool_v2)** uses the framework's standard plugin interfaces. Users only need to implement three functions and mount them via command-line arguments: + +| Argument | Description | +|----------|-------------| +| `--custom-generate-function-path` | Uses the built-in `miles.rollout.generate_hub.multi_turn.generate` — no need to implement the multi-turn loop yourself | +| `--generate-tool-specs-path` | Declare tool definitions (user-implemented) | +| `--generate-execute-tool-function-path` | Implement tool execution logic (user-implemented) | +| `--custom-rm-path` | Implement the reward function (user-implemented) | + +Users only need to focus on business logic (tool definitions, tool execution, reward calculation). Multi-turn scheduling, token concatenation, loss masking, etc. are all handled by the framework. + +## Files + +- `tool_sandbox.py`: Tool definitions (`tool_specs`), tool execution (`execute_tool`), reward function (`reward_func`), and sandboxed safe execution environment +- `run_retool_multi_turn.py`: Training launch script + +## Quick Start + +```bash +python examples/retool_v2/run_retool_multi_turn.py +``` + +The launch script prepares everything it needs on its own: it downloads the dapo-math-17k +training set and the aime-2024 eval set, downloads the checkpoint, and converts it to +`torch_dist` before training starts. diff --git a/docs/examples/search-r1.md b/docs/examples/search-r1.md deleted file mode 100644 index c66ea66472..0000000000 --- a/docs/examples/search-r1.md +++ /dev/null @@ -1,275 +0,0 @@ ---- -title: Search-R1 (Tool Use) -description: Train a model to issue search queries, integrate observations, and answer multi-turn QA. ---- -**What you'll learn:** how to wire up a tool (web search) into a Miles training loop — -custom multi-turn rollout, observation interleaving, reward function, and TIS to keep -training stable when train ≠ inference. - -This is a Miles-friendly reproduction of the original -[Search-R1](https://github.com/PeterGriffinJin/Search-R1). - -## Prerequisites - -* `radixark/miles:latest` container. -* Either a serper.dev API key (Google search backend) or ~135 GB free disk for the - local Wikipedia retriever (see [appendix](#appendix-local-wikipedia-retriever)). -* You completed [Customization](/user-guide/customization) — this example uses a - custom rollout function and reward. - -## Files - -```text -examples/experimental/search-r1/ -├── generate_with_search.py # custom rollout (multi-turn loop) -├── google_search_server.py # serper.dev wrapper -├── local_search_server.py # FastAPI server in front of FAISS index -├── local_dense_retriever/ # E5-base index/corpus downloader -├── qa_em_format.py # exact-match reward -└── run_qwen2.5_3B.sh # launch script -``` - -## Quick start - -### 1. Set up environment - -```bash -cd /root && git clone https://github.com/radixark/miles.git -cd miles && pip install -e . --no-deps && pip install chardet -``` - -### 2. Prepare data - -```bash -git clone https://github.com/PeterGriffinJin/Search-R1.git -cd Search-R1 && pip install -e . --no-deps && pip install tensordict - -WORK_DIR=/root/Search-R1 -LOCAL_DIR=$WORK_DIR/data/nq_hotpotqa_train -python $WORK_DIR/scripts/data_process/qa_search_train_merge.py \ - --local_dir $LOCAL_DIR \ - --data_sources nq,hotpotqa -``` - -### 3. Convert the model - -```bash -hf download Qwen/Qwen2.5-3B --local-dir /root/Qwen2.5-3B -cd /root/miles -MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen2.5-3B)" || exit 1 -read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}" -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen2.5-3B \ - --save /root/Qwen2.5-3B_torch_dist -``` - -### 4. Run - -```bash -bash examples/experimental/search-r1/run_qwen2.5_3B.sh -``` - -## Configuration - -Open `generate_with_search.py` and edit `SEARCH_R1_CONFIGS`: - -```python -SEARCH_R1_CONFIGS = { - "max_turns": 2, - "topk": 3, - "search_concurrency": 256, - "search_backend": "local", # or "google" - - "local": { - "search_url": "http://127.0.0.1:8000/retrieve", - "proxy": None, - }, - - "google": { - "api_key": "your_serper_key", - "snippet_only": True, - "proxy": None, - }, - - "return_logprob": True, # required for TIS - "format_score": 0.2, -} -``` - -## Walkthrough — multi-turn rollout - -The custom rollout lives in `generate_with_search.py:generate`. The loop is -straightforward but every step matters: - -```python -async def generate(args, sample: Sample, sampling_params) -> Sample: - prompt = build_prompt(sample) - full_response, loss_masks, tokens = "", [], [] - - for turn in range(SEARCH_R1_CONFIGS["max_turns"]): - # 1. Model generates an action - out = await call_sglang(prompt + full_response, sampling_params) - toks = tokenize(out.text) - full_response += out.text - tokens += toks - loss_masks += [1] * len(toks) # model tokens count toward loss - - # 2. Parse action - action, content = parse_action(out.text) - - # 3. Run the tool - if action == "search": - result = await search_backend(content, topk=SEARCH_R1_CONFIGS["topk"]) - obs_text = render_observation(result) - obs_toks = tokenize(obs_text) - full_response += obs_text - tokens += obs_toks - loss_masks += [0] * len(obs_toks) # observation tokens MASKED OUT - elif action == "answer": - break - - sample.response = full_response - sample.tokens = tokens - sample.loss_mask = loss_masks - sample.metadata["turns_used"] = turn + 1 - return sample -``` - -### The two crucial details - -1. **Loss masking.** Tool/observation tokens get `loss_mask=0`. Without this, the model - learns to *predict the search results*, which is both wrong and wildly unhelpful. -2. **Tokenization alignment.** The model must see and the trainer must score the - *exact same tokens*. Pre-tokenizing vs. re-tokenizing at training time can drift — - that's where the [chat template verifier](/user-guide/agentic-chat-template) - matters. - -## Walkthrough — reward - -```python -async def reward_func(args, sample: Sample, **kwargs) -> float: - answer = extract_final_answer(sample.response) - label = sample.label - em = exact_match(answer, label) - fmt = SEARCH_R1_CONFIGS["format_score"] if has_valid_format(sample.response) else 0 - return em + fmt -``` - -`format_score=0.2` gives partial credit for the correct `...` shape even if the -content is wrong — keeps gradient flowing during early exploration. - -## Enabling TIS - -The trajectory mixes model tokens (we want gradients) with tool tokens (we don't). -Without correction, the implicit policy ratio in the GRPO objective is *off-policy* — -the search results came from a stochastic environment, not the model. - -**Truncated Importance Sampling (TIS)** corrects for this. To enable: - -1. Set `"return_logprob": True` in `SEARCH_R1_CONFIGS`. -2. Uncomment the TIS flags in `run_qwen2.5_3B.sh`: - -```bash -GRPO_ARGS+=( --use-tis ) -CUSTOM_ARGS+=( - --custom-config-path examples/infra_features/train_infer_mismatch_helper/mis.yaml - --custom-tis-function-path examples.infra_features.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp -) -``` - -When `return_logprob=True`, response post-processing is automatically disabled to keep -token / logp alignment. - -## What to watch - -```text -search_r1/turns_per_sample ~1.5 (depends on max_turns) -search_r1/search_calls_per_sample ~1.0 -reward/exact_match trending up -reward/format ~0.18 (steady — most outputs are well-formed) -loss_mask/observation_fraction 0.4 – 0.7 (lots of obs tokens, all masked) -tis/effective_sample_size > 0.7 × batch_size -``` - -If `tis/effective_sample_size` collapses below 0.5, your inference distribution has -drifted too far. Lower `--lr` or shorten `max_turns`. - -## Tuning knobs - -| Knob | Effect | -|---|---| -| `max_turns` | More turns = more retrieval, more drift | -| `topk` | More retrieved snippets = longer context | -| `search_concurrency` | Cap on simultaneous tool calls (mind your QPS limit) | -| `format_score` | Partial credit for correct shape — higher = faster early shaping | - -## Troubleshooting - -| Problem | Fix | -|---|---| -| "Ray process stuck" | `rm -rf /root/.cache`, then `rm -rf /root/.*` if still stuck | -| Retriever 502 errors | `lsof -i :8000` — make sure your local server is alive | -| Conda activation collisions | Deactivate the `retriever` env before launching training | -| EM stays at 0 | Check the answer extractor — most often a regex mismatch | -| Loss masks shifted by one token | Tokenizer added a leading space; align with `add_special_tokens=False` | - -## Variations - -* **Use Google instead of local.** Set `"search_backend": "google"` and add an API key. -* **Different tool.** Replace `search_backend` with anything else — calculator, code - exec, internal API. The pattern is identical. -* **Group RM.** With multiple trajectories per prompt (GRPO), enable `--group-rm` so - rewards are computed in a batch. -* **Longer chains.** Bump `max_turns` to 8+ for deep-reasoning tasks. Watch - `loss_mask/observation_fraction` — if it dominates, the model is barely training. - -## Appendix — local Wikipedia retriever - -Heavy but completely offline. ~135 GB total disk and a separate conda env to avoid -conflicting with Miles. - -### One-time setup - -```bash -wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh -bash ~/miniconda.sh -b -p $HOME/miniconda3 -source ~/miniconda3/etc/profile.d/conda.sh - -conda create -n retriever python=3.10 -y && conda activate retriever -conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 \ - pytorch-cuda=12.1 -c pytorch -c nvidia -y -pip install transformers datasets pyserini huggingface_hub uvicorn fastapi -conda install faiss-gpu=1.8.0 -c pytorch -c nvidia -y - -# 2. Index + corpus (~135 GB) -save_path=/root/Index -python /root/miles/examples/experimental/search-r1/local_dense_retriever/download.py \ - --save_path $save_path -cat $save_path/part_* > $save_path/e5_Flat.index -gzip -d $save_path/wiki-18.jsonl.gz -``` - -### Run the server - -```bash -conda activate retriever -python /root/miles/examples/experimental/search-r1/local_dense_retriever/retrieval_server.py \ - --index_path /root/Index/e5_Flat.index \ - --corpus_path /root/Index/wiki-18.jsonl \ - --topk 3 \ - --retriever_name e5 \ - --retriever_model intfloat/e5-base-v2 \ - --faiss_gpu -``` - -5–7 GB of GPU memory per GPU. First startup is slow (model + index load); subsequent -restarts are 1–2 minutes. - -### Then launch training - -```bash -conda deactivate # don't train inside the retriever env! -cd /root/miles -bash examples/experimental/search-r1/run_qwen2.5_3B.sh -``` diff --git a/docs/examples/swe-agent-harbor-docker.md b/docs/examples/swe-agent-harbor-docker.md new file mode 100644 index 0000000000..9b37f71382 --- /dev/null +++ b/docs/examples/swe-agent-harbor-docker.md @@ -0,0 +1,142 @@ +--- +title: "SWE-Agent training with Harbor on Docker sandboxes" +description: "Trains coding and terminal agents with Harbor-managed local Docker sandboxes and verifier rewards." +# Generated from examples/swe-agent-harbor-docker/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +This example trains GLM-4.7-Flash on agentic coding and terminal tasks. Miles +runs synchronous GRPO and serves the policy through its session server; a +separate [Harbor](https://github.com/harbor-framework/harbor) agent server +creates the task sandboxes, runs the agents, and returns verifier rewards. + +The same pipeline supports Terminal-Bench, SWE-bench, and custom Harbor tasks. +Training records must contain a `prompt` and `metadata.instance_id` identifying +the Harbor task. + +## Files + +| File | Purpose | +| --- | --- | +| `run.py` | Validated synchronous GLM-4.7-Flash launcher. | +| `run-glm47-flash-agentic-async.py` | Disaggregated fully asynchronous launcher. | +| `run_glm52_lora_tb2_daytona.py` | Multi-node GLM-5.2 744B-A40B LoRA launcher (bf16 trainer, fp8 rollout). | +| `swe_agent_function.py` | Sends each rollout to the Harbor agent server. | +| `generate.py` | Builds rewards, metrics, and training samples. | +| `download_and_process_data.py` | Converts supported datasets to Miles JSONL. | + +## 1. Start the Harbor agent server + +Use the `harbor-miles-v0.20.0` branch of the `harbor-framework/harbor` +repository, which carries the Miles integration: + +```bash +git clone https://github.com/harbor-framework/harbor.git +cd harbor +git checkout harbor-miles-v0.20.0 +uv sync + +HARBOR_TASKS_DIR=/path/to/harbor_tasks uv run python miles_agent_server.py \ + --host 0.0.0.0 \ + --port 30000 \ + --dashboard-port 0 \ + --max-concurrent 32 \ + --agent-timeout 5400 \ + --trials-dir /path/to/trials +``` + +`HARBOR_TASKS_DIR` must contain one Harbor task directory for every +`metadata.instance_id` in the training data. The agent-server machine must have +Docker and enough capacity for the requested number of concurrent sandboxes; +set `--max-concurrent` to at least one sandbox per trajectory in a rollout step +(`--rollout-batch-size` times `--n-samples-per-prompt`). Keep `--agent-timeout` +generous — agentic trials routinely run past an hour, and a short timeout kills +them mid-episode. Verify `http://:30000/health` before launching +Miles. + +The two per-trial timeouts must be ordered. `--agent-timeout` is the authoritative +one: when it fires, the agent server ends the trial and frees its sandbox. The +rollout client applies a second ceiling, `AGENT_TRIAL_TIMEOUT` (default 7200 +seconds), which has to stay above `--agent-timeout`. If the client gives up first, +the trial is recorded as aborted while the agent server keeps running it, so the +sandbox and its `--max-concurrent` slot stay busy for the remaining difference, and +the aborted sample takes its whole GRPO group down with it. Raise it through the +launcher's generic env-var hook: + +```bash +python examples/swe-agent-harbor-docker/run.py ... --extra-env-vars 'AGENT_TRIAL_TIMEOUT=10800' +``` + +If the trainer reaches the agent server through a proxy or an in-cluster service +rather than directly, point `--agent-server-url` at that stable name rather than +an ephemeral pod address. The rollout client enables TCP keepalive probes so +long-running trials do not lose an idle connection while Harbor is working. + +## 2. Prepare Terminal-Bench data + +Convert a local JSONL whose rows include a task instruction and instance name: + +```bash +python examples/swe-agent-harbor-docker/download_and_process_data.py \ + --input /path/to/terminal-bench.jsonl \ + --output /path/to/tb2_train.jsonl \ + --agent-name mini-swe-agent \ + --prompt-key instruction +``` + +The resulting `metadata.instance_id` values must match task directories known to +the Harbor agent server. + +## 3. Launch synchronous GLM-4.7-Flash training + +The shape below is what a multi-day Terminal-Bench 2 run used on one node of 8 +H200 GPUs: 32 trajectories per GRPO step (4 prompts times 8 samples), each one a +full mini-swe-agent episode in its own Harbor sandbox. + +```bash +python examples/swe-agent-harbor-docker/run.py \ + --num-nodes 1 \ + --num-gpus-per-node 8 \ + --skip-prepare \ + --megatron-path /root/Megatron-LM \ + --hf-checkpoint /path/to/GLM-4.7-Flash \ + --ref-load /path/to/GLM-4.7-Flash_torch_dist \ + --save-dir /path/to/checkpoints \ + --prompt-data /path/to/tb2_train.jsonl \ + --max-seq-len 65536 \ + --rollout-batch-size 4 \ + --n-samples-per-prompt 8 \ + --global-batch-size 32 \ + --num-rollout 200 \ + --save-interval 20 \ + --agent-server-url http://:30000 \ + --router-external-host \ + --miles-host-ip 0.0.0.0 \ + --save-traces-dir /path/to/traces +``` + +For a smoke test, set `--num-rollout 1`. Expect roughly 10 minutes per step at +this shape; because synchronous rollout waits for the slowest trajectory in the +batch, a step that draws an unusually slow task can take several times that. + +`--router-external-host` is the address Harbor sandboxes use to call the Miles +session server and SGLang router. It must resolve and route from the agent-server +machine. `--miles-host-ip 0.0.0.0` is useful when those services must accept +connections forwarded from another host. Ensure ports 30000 and 31000 are +reachable end to end; Tailscale is one option when the machines are on different +networks. + +## 4. Verify progress + +Check both layers: + +1. Miles logs emit rollout metrics and write `rollout_data/*.pt` under the trace + directory. +2. Megatron logs emit `train/step` and the Ray job exits successfully. + +Confirm a suspected stall on disk before believing a dashboard. W&B uploads can +fail partway through a long run — dropping some metric rows while others keep +arriving — which looks exactly like a frozen reward curve. The per-step +`train_data/` and `rollout_data/.pt` dumps under `--save-traces-dir` +are written by the trainer itself and are the authoritative progress signal. + +The synchronous launcher uses GLM-4.7 tool-call and reasoning parsers, TITO, +the Miles session server, and the Megatron backend. diff --git a/docs/user-guide/rollout-endpoints.md b/docs/user-guide/rollout-endpoints.md index cd815e931f..1d09d4fe5b 100644 --- a/docs/user-guide/rollout-endpoints.md +++ b/docs/user-guide/rollout-endpoints.md @@ -266,4 +266,5 @@ inherited across turns. Each request is tokenized independently. - [Customization](/user-guide/customization): the full catalog of `--*-path` hooks. - [Agentic Rollout (TITO)](/user-guide/agentic-chat-template): verifying that a template is append-only across turns. -- [Multi-agent example](/examples/multi-agent): full agentic walkthrough. +- [Multi-agent example](https://github.com/radixark/miles/tree/main/examples/experimental/multi_agent): + full agentic walkthrough. diff --git a/examples/README.md b/examples/README.md index d79c63f257..b127b859ba 100644 --- a/examples/README.md +++ b/examples/README.md @@ -46,3 +46,19 @@ Not fully verified — for experimental and development use. - **[swe-agent-harbor-daytona](./experimental/swe-agent-harbor-daytona)**: The `swe-agent-harbor-docker` pipeline with task sandboxes hosted on Daytona instead of local Docker. - **[tau-bench](./experimental/tau-bench)**: Training in an agentic multi-turn tool use environment (Tau-bench). - **[verifiers](./experimental/verifiers)**: Training on a Prime Intellect Verifiers environment instead of a Miles prompt dataset. + + +## These READMEs are the documentation site + +Every README outside `experimental/` is mirrored onto +[miles.radixark.com/docs/examples](https://miles.radixark.com/docs/examples) by +`scripts/tools/sync_example_docs.py`, which pre-commit runs for you. The docs site is +generated from this directory and never edited directly, so a new example needs nothing +beyond its README and an entry in the list above — the sync fails if either is missing. +Three things that list controls: the +level-1 heading of each README becomes the page title, the one-line description becomes +the page's meta description (keep it under 160 characters), and the bullet order is the +sidebar order. Content between +`docs:exclude:start` / `docs:exclude:end` HTML comments (like this section) stays on +GitHub but is left out of the site. + diff --git a/scripts/tools/sync_example_docs.py b/scripts/tools/sync_example_docs.py new file mode 100644 index 0000000000..d5df9dced6 --- /dev/null +++ b/scripts/tools/sync_example_docs.py @@ -0,0 +1,514 @@ +#!/usr/bin/env python3 +"""Mirror the README files under examples/ into the Examples tab of the docs site. + +examples/ is the single source of truth. Every README.md outside examples/experimental/ +becomes one page under docs/examples/, and the Examples tab in docs/docs.json is +regenerated from the same tree, so the site cannot drift from the repository. + +Usage: + python scripts/tools/sync_example_docs.py # regenerate + python scripts/tools/sync_example_docs.py --check # fail if anything is stale + +The generated pages are Mintlify MDX, which is stricter than GitHub-flavored Markdown. +Two constructs in a README break the build and are rewritten here: unescaped braces +(parsed as JSX expressions) and non-self-closing void tags such as and
. +Relative links and images are rewritten to the site page when the target is mirrored, +and to GitHub otherwise. + +docs/docs.json is round-tripped through json.dumps on every run, so this script owns +that file's formatting (indent=1); hand-edits to other tabs keep their content but are +renormalized to that style. +""" + +import argparse +import json +import os +import re +import sys +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +EXAMPLES = REPO / "examples" +DOCS = REPO / "docs" +DOCS_JSON = DOCS / "docs.json" +OUT_DIR = DOCS / "examples" + +BRANCH = "main" +GITHUB_TREE = f"https://github.com/radixark/miles/tree/{BRANCH}" +GITHUB_BLOB = f"https://github.com/radixark/miles/blob/{BRANCH}" +GITHUB_RAW = f"https://raw.githubusercontent.com/radixark/miles/{BRANCH}" + +IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp"} +MAX_DESCRIPTION = 160 + +# Tags Mintlify renders as JSX. Anything else in prose (, , ...) is a +# model token, not markup, and is escaped so acorn never sees it. +HTML_TAGS = { + "a", + "b", + "blockquote", + "br", + "code", + "details", + "div", + "em", + "hr", + "i", + "img", + "kbd", + "li", + "ol", + "p", + "pre", + "span", + "strong", + "sub", + "summary", + "sup", + "table", + "tbody", + "td", + "th", + "thead", + "tr", + "ul", + "h1", + "h2", + "h3", + "h4", + "h5", + "h6", + "caption", + "center", + "dd", + "del", + "dl", + "dt", + "figcaption", + "figure", + "ins", + "mark", + "picture", + "s", + "small", + "source", + "u", + "video", + # Mintlify's own component, which
blocks are rewritten into. + "accordion", +} +VOID_TAGS = {"br", "hr", "img", "input"} + +# - **[fully_async](./fully_async)**: Demonstrates fully asynchronous rollout generation. +INDEX_BULLET = re.compile(r"^\s*[-*]\s+\*\*\[([^\]]+)\]\(([^)]+)\)\*\*:\s*(.+?)\s*$") +# ## [Infra Features](./infra_features) — a group root registered as a section heading. +INDEX_HEADING = re.compile(r"^#{2,}\s+\[([^\]]+)\]\(([^)]+)\)\s*$") +MD_LINK = re.compile(r"(!?)\[([^\]]*)\]\(\s*([^)\s]+)(\s+\"[^\"]*\")?\s*\)") +IMG_LINK = re.compile(r"!\[([^\]]*)\]\(\s*([^)\s]+)(\s+\"[^\"]*\")?\s*\)") +# Link text may carry one nested image ([![badge](img)](target)), which MD_LINK would +# mis-parse as text "![badge" with href "img". +OUTER_LINK = re.compile(r"(?\"'])*?)(/?)>") +IMG_TAG = re.compile(r"\"'])*)/?>", re.IGNORECASE) +ATTR = re.compile(r"([A-Za-z-]+)\s*=\s*\"([^\"]*)\"|([A-Za-z-]+)\s*=\s*'([^']*)'") +DETAILS = re.compile(r"]*>\s*(?:]*>(.*?))?(.*?)
", re.DOTALL | re.IGNORECASE) +# Content between these markers stays on GitHub but is left out of the docs site. +EXCLUDE_BLOCK = re.compile(r".*?", re.DOTALL) +HTML_COMMENT = re.compile(r"", re.DOTALL) +FENCE = re.compile(r"^\s*(```+|~~~+)") + + +class SyncError(Exception): + pass + + +def discover_pages(): + """Map each mirrored directory (relative to examples/, "" for the root) to its README.""" + pages = {} + for readme in sorted(EXAMPLES.rglob("README.md")): + rel = readme.parent.relative_to(EXAMPLES) + parts = [] if rel == Path(".") else list(rel.parts) + if "experimental" in parts or "__pycache__" in parts: + continue + pages["/".join(parts)] = readme + return pages + + +def slug_for(rel_dir): + if not rel_dir: + return "index" + return "/".join(part.replace("_", "-").lower() for part in rel_dir.split("/")) + + +def site_url(rel_dir): + return "/examples" if not rel_dir else f"/examples/{slug_for(rel_dir)}" + + +def repo_dir_of(rel_dir): + return "examples" if not rel_dir else f"examples/{rel_dir}" + + +def parse_index(index_readme): + """Read examples/README.md, the registry for every mirrored page. + + Returns (descriptions, registered): one-line descriptions keyed by directory in + bullet order, and the set of every directory the index mentions — as a bullet or, + for a group root like infra_features, as a section-heading link. + """ + descriptions, registered = {}, set() + for line in index_readme.read_text().splitlines(): + bullet = INDEX_BULLET.match(line) + heading = None if bullet else INDEX_HEADING.match(line) + m = bullet or heading + if not m: + continue + target = m.group(2).split("#")[0].strip().rstrip("/") + rel = os.path.normpath(os.path.join("examples", target)) + if not rel.startswith("examples/"): + continue + rel_dir = rel[len("examples/") :] + registered.add(rel_dir) + if bullet: + descriptions[rel_dir] = bullet.group(3) + return descriptions, registered + + +def first_sentence(text): + """First prose sentence of a README, used when the index has no bullet for a page.""" + body = [] + in_fence = False + for line in text.splitlines(): + if FENCE.match(line): + in_fence = not in_fence + continue + if in_fence or line.startswith("#"): + continue + if line.strip(): + body.append(line.strip()) + elif body: + break + para = " ".join(body) + para = MD_LINK.sub(lambda m: m.group(2) or m.group(3), para) + para = re.sub(r"[*`_]", "", para) + match = re.search(r"^(.+?[.!?])(\s|$)", para) + return (match.group(1) if match else para).strip() + + +def mask_code(text): + """Replace fenced blocks, inline code and display math with placeholders.""" + stash = [] + + def keep(chunk): + stash.append(chunk) + return f"\x00{len(stash) - 1}\x00" + + out_lines = [] + fence = None + buffer = [] + for line in text.split("\n"): + m = FENCE.match(line) + if fence is None and m: + fence = m.group(1) + buffer = [line] + elif fence is not None: + buffer.append(line) + if m and line.strip().startswith(fence): + out_lines.append(keep("\n".join(buffer))) + fence = None + else: + out_lines.append(line) + if fence is not None: + out_lines.append(keep("\n".join(buffer))) + text = "\n".join(out_lines) + + text = re.sub(r"\$\$.*?\$\$", lambda m: keep(m.group(0)), text, flags=re.DOTALL) + # A code span may wrap a line but not a paragraph, so an unpaired backtick cannot + # swallow the rest of the document into the stash. + text = re.sub(r"(`+)((?:(?!\n\n)[^`])+?)\1", lambda m: keep(m.group(0)), text) + # Inline math, bounded to one line so a stray "$5" in prose stays inert. The two + # alternatives are disjoint (an escape, or anything but a backslash), so an + # unterminated span cannot trigger exponential backtracking. + text = re.sub(r"(? {href} (escapes the repository)") + return None + suffix = "#" + fragment if sep else "" + + target = rel[: -len("/README.md")] if rel.endswith("/README.md") else rel + if target in mirrored: + return site_url(mirrored[target]) + suffix + + absolute = REPO / rel + # A README pointing at a docs source file should point at the published page instead. + if absolute.is_file() and rel.startswith("docs/") and absolute.suffix in {".md", ".mdx"}: + page = rel[len("docs/") : -len(absolute.suffix)] + return f"/{page[: -len('/index')] if page.endswith('/index') else page}{suffix}" + if absolute.is_dir(): + return f"{GITHUB_TREE}/{rel}{suffix}" + if absolute.is_file(): + if absolute.suffix.lower() in IMAGE_SUFFIXES: + return f"{GITHUB_RAW}/{rel}" + return f"{GITHUB_BLOB}/{rel}{suffix}" + broken.append(f"{repo_dir_of(cur_dir)}/README.md -> {href}") + return None + + +def convert_img_tags(text, cur_dir, mirrored, broken): + """... is not self-closing in most READMEs; markdown carries it fine.""" + + def replace(match): + attrs = dict() + for m in ATTR.finditer(match.group(1)): + key = (m.group(1) or m.group(3)).lower() + attrs[key] = m.group(2) if m.group(2) is not None else m.group(4) + src = attrs.get("src", "") + resolved = resolve_link(src, cur_dir, mirrored, broken) + # A row of side-by-side images relies on the width attributes we just dropped, so + # give each one its own paragraph rather than letting them run into each other. + return f"\n\n![{attrs.get('alt', '')}]({resolved or src})\n\n" + + text = IMG_TAG.sub(replace, text) + # The wrapper

around those images carries no meaning in markdown. + text = re.sub(r"]*>", "\n", text) + text = re.sub(r"^[ \t]+$", "", text, flags=re.MULTILINE) + return re.sub(r"\n{3,}", "\n\n", text) + + +def convert_details_blocks(text): + """A

block renders empty on the docs site; Mintlify's Accordion keeps it.""" + + def replace(match): + summary, inner = match.group(1) or "Details", match.group(2) + # Braces would be escaped to \{ later and render literally inside the attribute. + title = re.sub(r"[*`_{}]|<[^>]+>", "", summary).strip().replace('"', "'") + return f'\n\n\n\n{inner.strip()}\n\n\n\n' + + return DETAILS.sub(replace, text) + + +def escape_html(text): + """Self-close void tags; escape anything that is a model token rather than markup.""" + + def replace(match): + closing, name, attrs, self_closed = match.groups() + if name.lower() not in HTML_TAGS: + return "<" + match.group(0)[1:] + if name.lower() in VOID_TAGS and not self_closed and not closing: + return f"<{name}{attrs.rstrip()} />" + return match.group(0) + + return HTML_TAG.sub(replace, text) + + +def convert(readme_text, rel_dir, mirrored, broken): + readme_text = EXCLUDE_BLOCK.sub("", readme_text) + # Mask before anything else, including the title scan: a "# comment" inside a + # fenced block must not be mistaken for the page's level-1 heading. + masked, stash = mask_code(readme_text) + lines = masked.split("\n") + title = None + for i, line in enumerate(lines): + if line.startswith("# "): + title = unmask_code(line[2:].strip(), stash) + del lines[i] + break + if title is None: + raise SyncError(f"{repo_dir_of(rel_dir)}/README.md has no level-1 heading to use as the page title") + body = "\n".join(lines) + + # GitHub-only annotations; also raw comments are not valid MDX. + body = HTML_COMMENT.sub("", body) + body = convert_img_tags(body, rel_dir, mirrored, broken) + + def rewrite(bang): + def replace(match): + text, href, hint = match.groups() + resolved = resolve_link(href, rel_dir, mirrored, broken) + return f"{bang}[{text}]({resolved or href}{hint or ''})" + + return replace + + body = IMG_LINK.sub(rewrite("!"), body) + body = OUTER_LINK.sub(rewrite(""), body) + body = convert_details_blocks(body) + body = escape_html(body) + body = body.replace("{", "\\{").replace("}", "\\}") + body = re.sub(r"\n{3,}", "\n\n", body) + body = unmask_code(body, stash) + return title, body.strip("\n") + + +def render_page(title, description, rel_dir, body): + source = f"{repo_dir_of(rel_dir)}/README.md" + return ( + "---\n" + f"title: {json.dumps(title, ensure_ascii=False)}\n" + f"description: {json.dumps(description, ensure_ascii=False)}\n" + f"# Generated from {source} by scripts/tools/sync_example_docs.py. Edit that README, not this file.\n" + "---\n" + f"{body}\n" + ) + + +def build_pages(): + pages = discover_pages() + if "" not in pages: + raise SyncError("examples/README.md is missing; it is the source of the Examples index page") + mirrored = {repo_dir_of(rel): rel for rel in pages} + descriptions, registered = parse_index(pages[""]) + + unregistered = sorted(rel_dir for rel_dir in pages if rel_dir and rel_dir not in registered) + if unregistered: + raise SyncError( + "mirrored but not listed in examples/README.md — add a bullet for:\n " + + "\n ".join(repo_dir_of(d) for d in unregistered) + ) + + broken, rendered, slug_owner = [], {}, {} + for rel_dir, readme in sorted(pages.items()): + text = readme.read_text() + title, body = convert(text, rel_dir, mirrored, broken) + description = descriptions.get(rel_dir) + if description is None: + # Derived from the README's own first sentence, which Mintlify already renders + # under the title as the description — drop the duplicate from the body. The + # comparison strips markdown the same way first_sentence does, so a link or + # emphasis in the opening sentence does not defeat the dedup. + description = first_sentence(text) + lead = re.match(r"\s*(.+?[.!?])(\s|$)", body, re.DOTALL) + if lead: + normalized = MD_LINK.sub(lambda m: m.group(2) or m.group(3), lead.group(1)) + if re.sub(r"[*`_]", "", normalized).strip() == description: + body = body[lead.end(1) :].lstrip() + if not description: + raise SyncError(f"{repo_dir_of(rel_dir)}/README.md has no description; add a bullet in examples/README.md") + if len(description) > MAX_DESCRIPTION: + raise SyncError( + f"description for {repo_dir_of(rel_dir)} is {len(description)} characters, " + f"over the {MAX_DESCRIPTION} the docs site allows; shorten it at the source" + ) + out_path = OUT_DIR / f"{slug_for(rel_dir)}.md" + if out_path in rendered: + raise SyncError( + f"slug collision: {repo_dir_of(slug_owner[out_path])} and {repo_dir_of(rel_dir)} " + f"both map to {out_path.relative_to(REPO)}" + ) + slug_owner[out_path] = rel_dir + rendered[out_path] = render_page(title, description, rel_dir, body) + + if broken: + raise SyncError("READMEs link to paths that do not exist:\n " + "\n ".join(sorted(set(broken)))) + # dicts preserve insertion order, so this is the bullet order of examples/README.md. + return pages, rendered, list(descriptions) + + +def build_navigation(pages, bullet_order): + """Examples tab, mirroring the directory layout: top-level recipes, then infra_features. + + Sidebar order follows the bullet order in examples/README.md — the index README owns + ordering along with titles and descriptions. Directories without a bullet sort last, + alphabetically. + """ + rank = {rel_dir: i for i, rel_dir in enumerate(bullet_order)} + recipes, infra = [], [] + for rel_dir in sorted(pages, key=lambda d: (rank.get(d, len(rank)), d)): + if not rel_dir: + continue + page = f"examples/{slug_for(rel_dir)}" + (infra if rel_dir.startswith("infra_features") else recipes).append(page) + group = {"group": "Examples", "root": "examples/index", "pages": []} + group["pages"].append({"group": "Recipes", "pages": recipes, "expanded": True}) + if infra: + infra_root = "examples/infra-features" + children = [p for p in infra if p != infra_root] + group["pages"].append({"group": "Infra Features", "root": infra_root, "pages": children, "expanded": False}) + return group + + +def examples_tab(config): + for tab in config["navigation"]["tabs"]: + if tab.get("tab") == "Examples": + return tab + raise SyncError('docs.json has no "Examples" tab') + + +def render_docs_json(config): + return json.dumps(config, indent=1, ensure_ascii=False) + "\n" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--check", action="store_true", help="fail instead of writing when something is stale") + args = parser.parse_args() + + try: + pages, rendered, bullet_order = build_pages() + except SyncError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + config = json.loads(DOCS_JSON.read_text()) + try: + tab = examples_tab(config) + except SyncError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + tab["groups"] = [build_navigation(pages, bullet_order)] + docs_json_text = render_docs_json(config) + + existing = {p for p in OUT_DIR.rglob("*.md")} if OUT_DIR.exists() else set() + stale = sorted(existing - set(rendered)) + changed = sorted(p for p, text in rendered.items() if not p.exists() or p.read_text() != text) + json_changed = DOCS_JSON.read_text() != docs_json_text + + if args.check: + problems = [f"stale: {p.relative_to(REPO)}" for p in stale] + problems += [f"out of date: {p.relative_to(REPO)}" for p in changed] + if json_changed: + problems.append("out of date: docs/docs.json") + if problems: + print("error: docs/examples is out of sync with examples/:", file=sys.stderr) + for problem in problems: + print(f" {problem}", file=sys.stderr) + print("run: python scripts/tools/sync_example_docs.py", file=sys.stderr) + return 1 + return 0 + + for path in stale: + path.unlink() + for path, text in rendered.items(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + if json_changed: + DOCS_JSON.write_text(docs_json_text) + + for path in stale: + print(f"removed {path.relative_to(REPO)}") + for path in changed: + print(f"wrote {path.relative_to(REPO)}") + if json_changed: + print("wrote docs/docs.json") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fast/doc/test_sync_example_docs.py b/tests/fast/doc/test_sync_example_docs.py new file mode 100644 index 0000000000..317efef7c0 --- /dev/null +++ b/tests/fast/doc/test_sync_example_docs.py @@ -0,0 +1,165 @@ +"""Pin the behavior of scripts/tools/sync_example_docs.py, the examples -> docs mirror. + +The generator lives in scripts/tools because it writes files; these tests import it the +way tests/fast/examples/infra_features/p2p_weight_transfer/test_run.py imports run.py. +Each converter case here is a hazard the docs build either rejects (unescaped braces, +non-self-closed void tags) or silently drops (
content, corrupted inline math), +so the suite is what keeps README edits from breaking the published site. +""" + +import importlib.util +import subprocess +import sys +from pathlib import Path +from types import ModuleType + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +SCRIPT = REPO_ROOT / "scripts" / "tools" / "sync_example_docs.py" + + +@pytest.fixture(scope="module") +def sync() -> ModuleType: + spec = importlib.util.spec_from_file_location("sync_example_docs", SCRIPT) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def convert(sync, readme, rel_dir="fully_async", mirrored=None, broken=None): + return sync.convert(readme, rel_dir, mirrored or {}, [] if broken is None else broken) + + +class TestTitle: + def test_fenced_comment_is_not_the_title(self, sync): + readme = "```bash\n# Download the model\nhf download x\n```\n\n# Real Title\n\nBody.\n" + title, body = convert(sync, readme) + assert title == "Real Title" + assert "# Download the model" in body # and the fence keeps its line + + def test_missing_title_is_an_error(self, sync): + with pytest.raises(sync.SyncError, match="no level-1 heading"): + convert(sync, "just prose\n") + + +class TestMasking: + def test_inline_math_braces_survive(self, sync): + _, body = convert(sync, "# T\n\nRoles $\\pi_{\\text{old}}$ and $\\pi_{\\text{new}}$.\n") + assert "$\\pi_{\\text{old}}$" in body + assert "\\{" not in body + + def test_display_math_survives(self, sync): + _, body = convert(sync, "# T\n\n$$\nL_{\\text{PPO}}(\\theta)\n$$\n") + assert "L_{\\text{PPO}}" in body + + def test_prose_braces_are_escaped(self, sync): + _, body = convert(sync, "# T\n\nParses {...} tokens.\n") + assert "\\{" in body + assert "<tool_call>" in body + + def test_unterminated_math_span_is_linear(self, sync): + # Disjoint alternatives in the inline-math pattern; a pathological line must + # fail to match in linear time rather than backtrack exponentially (CodeQL 133). + readme = "# T\n\nPrice $" + "\\a" * 5000 + " end.\n" + _, body = convert(sync, readme) + assert "Price $" in body + + def test_unpaired_backtick_stays_within_its_paragraph(self, sync): + readme = "# T\n\nBroken `tick here.\n\nNext [run](./run.py) paragraph.\n" + broken = [] + _, body = convert(sync, readme, broken=broken) + # The link in the following paragraph is still seen by the rewriter. + assert any("run.py" in b for b in broken) + + +class TestLinks: + def test_mirrored_readme_resolves_to_site_page(self, sync): + mirrored = {"examples/fully_async": "fully_async"} + _, body = convert(sync, "# T\n\nSee [it](./README.md).\n", mirrored=mirrored) + assert "](/examples/fully-async)" in body + + def test_badge_in_link_rewrites_both_hrefs(self, sync): + mirrored = {"examples/fully_async": "fully_async"} + broken = [] + _, body = convert(sync, "# T\n\n[![badge](./pic.png)](./README.md) end.\n", mirrored=mirrored, broken=broken) + assert "](/examples/fully-async)" in body # outer link + assert any("pic.png" in b for b in broken) # inner image checked too + + def test_repo_escaping_link_is_reported(self, sync): + broken = [] + convert(sync, "# T\n\nSee [up](../../../outside.md).\n", broken=broken) + assert any("escapes the repository" in b for b in broken) + + +class TestHtml: + def test_common_tags_stay_markup(self, sync): + _, body = convert(sync, '# T\n\n

Hello

\n') + assert "' in body + assert "content" in body + + def test_void_img_tag_becomes_markdown_image(self, sync): + _, body = convert( + sync, '# T\n\n

\nA\n

\n' + ) + assert "![A](https://x.test/a.png)" in body + assert "