# Ray and Kubernetes Backend
Source: https://miles.radixark.com/docs/advanced/cluster-backend

Run the same training script on Ray, or let Kubernetes schedule every worker of the run.

Miles runs a job on one of two cluster backends. Ray is the default and needs no configuration.
Kubernetes installs the run as a helm release, so the cluster schedules every worker. The training
script is the same either way.

## Ray

Ray is the default backend, so `python scripts/run_*.py train` launches a run as is, with nothing to
configure. Nothing in the Kubernetes chapter below applies to it.

## Kubernetes

> [!WARNING]
>
> **Status.** Under active development: flags, chart values and failure semantics still change. Ray
> also runs *inside* Kubernetes, and that is the well-trodden path — take this one only if you want
> Miles to create the cluster's objects itself.

### Launch

Run everything from the repository root, with `kubectl` and `helm` on your PATH.

Install a workbench — the long-lived pod you launch from — once per namespace, from your
cluster's [`infra.yaml`](#for-cluster-administrator):

```bash theme={null}
export MILES_NS="miles-$USER"
python -m miles.utils.external_utils.miles_workbench install -n "$MILES_NS" -f infra.yaml
```

It creates the namespace if missing, checks your rights, installs, and waits for Ready. Runs
launched from the pod inherit those values.

Launch from inside it:

```bash theme={null}
python -m miles.utils.external_utils.miles_workbench exec -n "$MILES_NS" -- \
  bash -lc "cd /root/miles && python scripts/run_qwen3_4b.py train"
```

The recipe defaults to `/root/models`, `/root/datasets` and `/root/shared_data`; the example
[`infra.yaml`](#for-cluster-administrator) mounts all three from the shared volume. Fill them once
with `python scripts/run_qwen3_4b.py prepare`, or pass `--model-dir`, `--data-dir` and
`--output-dir` to point the recipe elsewhere.

### Observability

**Built in**

* The launcher follows every pod's logs, prints status changes and warning events, and ends with
  the run's exit code.
* On a failure, collect the namespace's logs, describes and events into one directory — before
  cleaning up, and the whole namespace, because the explanation is usually the pod next to it:
  `python -m miles.utils.external_utils.miles_workbench collect-diagnosis -n "$MILES_NS" --output-dir ~/artifacts/miles`

**External**

* A run's pods are ordinary pods, so whatever the cluster already runs — a metrics stack, a log
  collector, the platform's own dashboards — sees them with no wiring from Miles.
* Prefer it at scale. The built-in following is meant for watching one run, not hundreds of pods.

### Clean up

```bash theme={null}
python -m miles.utils.external_utils.miles_workbench stop -n "$MILES_NS" 260811-143000-042
python -m miles.utils.external_utils.miles_workbench uninstall -n "$MILES_NS"
```

`stop` removes the run and frees its GPUs; `uninstall` removes the workbench.

### Folder convention

A run is many pods on many machines, and they share nothing but the volumes `infra.yaml` mounts.

* Every path your script names — `/root/models`, `/root/datasets`, `infra.paths.runsRoot` — has to
  be under one of those mounts.
* To run your own branch, mount it at `/root/miles`, `/root/Megatron-LM` or `/sgl-workspace/sglang`;
  a copy anywhere else is not imported.

### For cluster administrator

Everything above assumes this was done once.

**Install LWS.** Miles deploys its worker pools as
[LeaderWorkerSets](https://github.com/kubernetes-sigs/lws). Install the CRDs and controller, and
grant users rights over them explicitly.

**Give each user a namespace.** The namespace is the real boundary, not the Role: anything that
may create workloads can name another ServiceAccount and read its token. Keep privileged accounts
out of it.

**Write one `infra.yaml`.** The same file drives every Miles chart:

```yaml theme={null}
infra:
  image:
    repository: radixark/miles
    tag: dev
  volumes:
    - name: cluster-storage
      hostPath: {path: /cluster-storage, type: Directory}
      mounts:
        - {mountPath: /cluster-storage}
        - {mountPath: /root/miles, subPath: alice/miles}
        - {mountPath: /root/models, subPath: models}
        - {mountPath: /root/datasets, subPath: datasets}
        - {mountPath: /root/shared_data, subPath: alice/shared_data}
  paths:
    runsRoot: /cluster-storage/${NAMESPACE}/miles_data
```

The `hostPath` above stands for a cluster-wide shared filesystem (NFS, Lustre, a CSI mount) already
mounted at `/cluster-storage` on every node: a per-node directory would lose the orchestrator state
file and the shared checkpoints. Where Pod Security forbids `hostPath`, replace that key with an RWX
`persistentVolumeClaim` and keep the mounts.

Any path in `infra.yaml` — a `hostPath`, a `mountPath`, a `subPath`, `infra.paths.runsRoot` — may
name `${NAMESPACE}`, which the chart replaces with the namespace it is installed into.

`charts/miles-run/values.yaml` shows the full shape, and each chart's `values.schema.json` is the
authoritative field list.


# Disaggregated RL Rollout
Source: https://miles.radixark.com/docs/advanced/disaggregated-rollout

Scale rollout independently across clusters and regions while miles preserves policy publication and version attribution.

Policy training and rollout inference have different resource and scaling
profiles. A trainer is a long-lived, tightly coupled cluster whose ranks step
together. Rollout is request-driven and bursty: long and agentic trajectories
finish at different times, demand changes over a run, and inference capacity
may be available in different clusters, regions, or compute providers.

Keeping both roles inside one job ties rollout capacity to the trainer's
placement, network, and lifecycle. Disaggregated RL rollout instead makes
inference an independent service. The trainer remains on its stable cluster,
while rollout replicas can scale with the live queue and draw on compute across
the globe without joining the trainer's Ray placement group or NCCL fabric.

For miles, this boundary brings:

* **Independent scaling.** Add or remove rollout capacity without resizing or
  restarting the trainer.
* **Flexible placement.** Run inference where GPUs are available, including
  other clusters, regions, and providers.
* **Separate failure domains.** Replace rollout replicas without making their
  process lifecycle part of the training job.
* **Better asynchronous utilization.** Keep a changing pool of rollout workers
  busy while training consumes completed trajectories.
* **Policy correctness.** Publish immutable policy versions, constrain requests
  by version, and record which version generated every returned trajectory.

The last property is what turns remote inference capacity into an RL rollout
service rather than an ordinary model-serving endpoint. miles must update the
policy without owning every engine and must preserve the behavior-policy
identity required for staleness control and off-policy training.

This page is about separating **RL training from rollout inference**. It is
different from [PD disaggregation](/docs/advanced/pd-disaggregation), which separates
the prefill and decode phases inside an SGLang deployment.

## Topologies

| Topology | Engine lifecycle | What miles talks to | Policy updates |
| - | - | - | - |
| miles-managed rollout | Part of the miles job | A miles-managed router and engine handles | miles converts, transfers, pauses, updates, and resumes the engines |
| Attached SGLang engines | Owned by the deployment | Fixed engine addresses supplied with `--rollout-external-engine-addrs` | miles still addresses and controls each engine through its engine handle |
| External rollout service | Independent and dynamically scaled | One stable rollout endpoint | miles publishes versions; the rollout service materializes and activates them across its replicas |

miles supports the first two topologies today. The single-endpoint external
rollout service is coming soon and extends that separation from GPU placement
to independent scaling, routing, and lifecycle.

`--rollout-external-engine-addrs` is the second row, not the third. It prevents
miles from launching SGLang, but miles still knows the individual engine addresses, checks
their configuration, registers them with its router, and calls their
weight-update lifecycle.

An external rollout service is a narrower interface. miles sends rollout
requests to one endpoint and publishes new policy versions without needing an
engine handle for every replica. The service behind that endpoint can be
implemented by any deployment or control-plane package that satisfies the
request and policy-version contracts below.

## Separate trainer and rollout GPUs

Disaggregated placement is the miles default. Give the trainer and rollout
engines separate GPU counts and do not pass `--colocate`:

```bash theme={null}
--actor-num-nodes 1 \
--actor-num-gpus-per-node 8 \
--rollout-num-gpus 8 \
--rollout-num-gpus-per-engine 2
```

The trainer and engines remain resident on their own GPUs, so they can run at
the same time under `train_async.py`. In a colocated job the two roles share
GPUs and take turns; fully async rollout therefore requires the disaggregated
layout. See [Training Backends: Choosing the GPU layout](/docs/user-guide/training-backend#3-choosing-the-gpu-layout)
for placement and offload behavior.

To attach SGLang engines launched outside the miles Ray job, provide their
addresses explicitly:

```bash theme={null}
--rollout-external-engine-addrs 10.0.1.10:30000 10.0.1.11:30000
```

Setting the addresses is what makes the rollout external; miles discovers each
engine's topology from its `/server_info`. If the attached fleet runs PD
disaggregation, add `--rollout-external-router-pd`, because the router starts
before the engines are discovered.

The engines must be reachable from the miles job and must have server settings
compatible with the rollout configuration. Because miles retains individual
engine handles, attaching external engines does not hand off weight-update ownership:
miles still runs the selected weight-update lifecycle.

## Weight synchronization

Once training and rollout use different GPUs, updated weights have to cross the
boundary between them. `--update-weight-transfer-mode` selects the current
miles-managed path:

| Mode | Data path | Use when |
| - | - | - |
| `broadcast` | Gather and convert trainer weights, then broadcast each tensor separately to the SGLang ranks over NCCL | Trainer and rollout ranks have NCCL connectivity; this remains the default |
| `broadcast_packed` | Pack each weight bucket into one byte buffer and issue one NCCL broadcast | Non-colocated Megatron with SGLang mixed-dtype flattened-bucket support; trades contiguous bucket storage for fewer collectives |
| `p2p` | Convert and re-shard weights, then write them directly to rollout-rank memory over RDMA | miles-managed, in-cluster jobs with direct rank-to-rank connectivity; see [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer) |
| `disk-delta` | Publish changed canonical checkpoint bytes to shared storage, let rollout hosts materialize them locally, then reload | Trainer and rollout cannot share an NCCL fabric, or model-sized full-weight transfer dominates the update |

These are weight synchronization choices, not different rollout APIs. In the
broadcast and P2P modes, miles transfers tensors into known engine ranks directly.
P2P RDMA is an in-cluster transfer optimization; it is not the external rollout
service mechanism described later in this page. Disk-delta instead establishes
a versioned publication boundary that can also be consumed by an external
rollout system.

### Disk-delta publication and activation

Enable disk-delta on a non-colocated Megatron run with a publication directory
visible to the trainer and rollout hosts, plus a host-local checkpoint
directory for each rollout host:

```bash theme={null}
--update-weight-transfer-mode disk-delta \
--update-weight-disk-dir /shared/miles/weight-updates \
--update-weight-local-checkpoint-dir /local-nvme/miles-rollout-checkpoint
```

The current lifecycle is:

1. On the first `update_weights()` call, miles captures a CPU snapshot from
   `--hf-checkpoint`. No delta version is published. Each rollout host also
   materializes the same base checkpoint in its local directory.
2. At the next update boundary, source trainer ranks gather Megatron tensors
   under their canonical Hugging Face names and compare their bytes with the
   previous snapshot.
3. miles publishes `weight_vNNNNNN/` with compressed changed bytes and an index
   containing the version, base version, delta encoding, and final-state
   checksums. Files are written atomically before the version is consumed.
4. Each rollout host pulls the version and patches its host-local checkpoint.
   Delta application verifies the checksum of the resulting tensor, not only
   the transferred delta.
5. miles pauses generation, reloads the materialized checkpoint into SGLang,
   advances the engine weight version, and resumes generation.

Pull and local materialization happen before the engine pause and can overlap
inference. Reloading the prepared weights is the activation boundary that
requires the pause.

The default delta encoding is XOR. It produces a compact byte-wise delta but
must be applied exactly once to the declared base version. `overwrite` stores
changed positions and their new values; it is larger but idempotent. Both
encodings are byte-oriented: the base and exported policy must agree on tensor
names, dtypes, shapes, and byte layout.

Each published tensor carries a checksum of its complete target state.
`xxh3-128` is the default; `blake3` and `adler32` are also accepted. A lineage,
layout, or checksum mismatch fails the update instead of activating a partially
updated policy.

On a POSIX shared filesystem, the published version becomes visible through
the normal filesystem contract. Object-store-backed mounts can use
`--custom-update-weight-post-write-path` on the miles side and
`--sglang-custom-pull-weights-pre-read-hook` on the SGLang side to make writes
visible before a rollout host reads them.

The maintained end-to-end coverage is
[`tests/e2e/megatron/test_qwen3_4B_disk_delta.py`](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/test_qwen3_4B_disk_delta.py).
It exercises a Qwen3-4B Megatron trainer (TP2 on two GPUs) and two TP1 SGLang
rollout engines on the other two GPUs of a single 4-GPU node. The same storage contract supports separate hosts, but the
registered test does not reproduce a cross-cluster deployment.

Current `main` rejects disk-delta with `--colocate`, LoRA, or PD
disaggregation. It also requires `--hf-checkpoint` to be a local checkpoint
directory. The implementation is selected by the Megatron actor; it is not a
general FSDP weight-update path.

## External rollout service contract

The coming single-endpoint integration builds on the current disk-delta
publication path and removes the need for miles to hold one handle per rollout
engine. The commands above cover miles-managed and attached-engine deployments;
this section defines the external-service boundary.

The intended boundary has two independent data paths:

```mermaid theme={null}
flowchart LR
    T[miles trainer] -->|publish policy version| S[(Version store)]
    W[miles rollout worker] -->|version-constrained request| G[External rollout endpoint]
    S -->|materialize and verify| E[Rollout replicas]
    G --> E
    E -->|trajectory and served version| W
```

miles does not need to know how many replicas are behind the endpoint, where
they run, or how the service replaces them. The integration is defined by
ownership:

| miles owns | The external rollout service owns |
| - | - |
| Training state and optimizer steps | Replica placement, scaling, and health |
| Conversion to canonical rollout-visible tensors | Materializing published versions on rollout hosts |
| Policy version numbering and publication | Verifying lineage and final-state checksums before activation |
| The policy-version requirement attached to each rollout request | Routing requests only to replicas that satisfy that requirement |
| Recording the served version on returned samples | Staging, admission control, activation, and reporting the version that served a request |

The rollout endpoint and the version store are separate interfaces. The
request path should not have to carry model-sized weights, and publishing a new
version should not require miles to enumerate the current replicas.

One open-source package implementing the rollout-service side is
[Stitch](https://github.com/modal-projects/stitch). It connects miles policy
publication and version-constrained requests to an independently managed
rollout fleet. The miles contract remains package-agnostic.

Stitch pins the miles revision it integrates against;
[Miles fork](https://github.com/modal-projects/stitch/blob/main/cookbook/miles_disagg/MILES_FORK.md)
records the trainer-side commits that pin depends on, how the same ownership
split lands in code, and how to move the pin forward.

### Policy-version requirements

An external rollout request needs more than an inference payload. It must be
able to express the policy version the caller can accept:

* A **minimum version** permits any replica at or above the requested version.
  This bounds staleness while allowing the fleet to converge gradually.
* An **exact version** requires the replica to serve one particular policy.
  It is useful for reproducibility, evaluation, and update validation.

If no compatible replica is ready, the service returns a retryable response
instead of silently serving an older policy. Retrying a stateful generation is
safe only when the rollout integration defines an idempotency and session
contract; the generic miles request path should not assume that every failed
generation can be replayed.

The response reports which version served the request. For long or agentic
trajectories, recording both the version at generation start and at generation
end also exposes whether an activation occurred while the request was in
flight. miles can then attach the observed version to the sample rather than
inferring it from the trainer's latest published version.

### Publication and replica convergence

A policy version becomes eligible for rollout only after its complete artifact
is visible. The publication contract therefore needs:

1. an immutable version identifier and its base lineage;
2. atomic visibility of the completed version;
3. final-state checksums for every changed tensor;
4. a way for a new or restarted replica to materialize a valid base and catch
   up through later versions; and
5. an activation acknowledgement before a replica advertises the new version.

Disk-delta already provides version and base identity, atomic file writes, and
final-state checksums. For an external consumer, the publication hook and
version store must also expose the completed version as one commit boundary.
Replica catch-up, admission, and activation belong to the external rollout
system. Other publication formats can implement the same boundary without
using disk-delta, as long as miles and the rollout system agree on version
identity and rollout-visible weights.

## Fully async rollout

Disaggregation and fully async rollout answer different questions:

* Disaggregation decides where training and rollout run and who owns the
  rollout engines.
* [Fully Async RL](/docs/user-guide/fully-async) decides how generation, buffering,
  and optimizer steps overlap.

With miles-managed engines, fully async rollout keeps generation in flight
while the trainer consumes completed groups and takes optimizer steps. At each
configured update boundary, the current schedule still synchronizes the active
generation call before invoking the weight updater. Disk-delta host
materialization occurs before the subsequent engine pause, but current `main`
does not overlap an opaque external publication with that active generation
call.

An external rollout service allows policy publication and replica preparation
to proceed without draining the entire fleet. Only activation needs a short
engine-local pause. Different replicas may converge at different times, so
request constraints and served-version attribution become part of the RL data
contract rather than optional serving metadata.

miles already uses weight versions to measure fully async sample staleness. A
service integration must preserve that information so
`--max-weight-staleness` and custom data-buffer policies make decisions from
the policy that generated each sample.

## What to measure

End-to-end `update_weights` time is not enough for a disaggregated service. It
mixes phases with different effects on rollout availability:

| Measurement | What it isolates |
| - | - |
| Trainer encode and publish | Megatron-to-HF conversion, delta construction, compression, and publication |
| Publication size and changed-byte density | Storage and network demand per version |
| Replica materialization | Fetch, lineage verification, delta application, and runtime-layout preparation while serving |
| Activation pause | Time the engine cannot admit or advance generation while prepared weights become active |
| Version convergence | Time from completed publication until the required rollout capacity advertises the version |
| Request version lag | Difference between the requested, served, and newest published policy versions |

For current disk-delta runs, miles records
`perf/update_weights_density` and `perf/update_weights_wire_bytes` in addition
to the normal weight-update timing.

The open-source Stitch integration reports the following reference rollout-side
weight update timings:

| Model | Preparation (incl. weight delta apply) | Engine pause |
| - | -: | -: |
| GLM-4.7-Flash | 15.8 s | 0.75 s |
| Kimi K2.6 NVFP4 | 72.5 s | 2.82 s |

Preparation does not pause the engine: rollout generation continues while the
weight delta is applied and the next weights are prepared. This makes staged
weight updates a natural fit for fully async training, because only activation
requires a brief engine pause.

## Related guides

* [Fully Async RL](/docs/user-guide/fully-async)
* [Training Backends](/docs/user-guide/training-backend)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)
* [PD Disaggregation](/docs/advanced/pd-disaggregation)
* [Stitch: open-source disaggregated rollout service](https://github.com/modal-projects/stitch)

## Acknowledgment

Thanks to [Nan Jiang](https://www.nanjiangwill.com/),
[Jason Mancuso](https://github.com/jvmncs), and the Modal team for
[Stitch](https://github.com/modal-projects/stitch) and for their support in
putting this guide together.


# Disk Offload
Source: https://miles.radixark.com/docs/advanced/disk-offload

Spill the paused training actor to node-local disk when host RAM cannot hold it.

Colocated RL keeps a training actor and a rollout engine on the same GPUs, so the actor
must get out of the way while the engine generates. miles offloads the actor during that
window, and by default the backup lives in pinned host memory. For large models that copy
does not fit, and disk offload is the alternative.

## Usage

```bash theme={null}
--offload-train --offload-train-target disk \
--offload-train-disk-dir /scratch/miles_offload \
--offload-train-disk-chunk-mb 256
```

Instead of a pinned host copy, the paused actor is streamed to per-rank files through a
fixed-size pinned staging buffer, so host memory stays bounded by
`--offload-train-disk-chunk-mb` regardless of how much is offloaded. Each rank writes to
its own directory under `--offload-train-disk-dir` (defaults to
`$SCRATCH/miles_train_offload_<uid>`), the files are overwritten in place every step, and
they are removed when the actor exits.

Point the directory at real node-local NVMe. A tmpfs mount (including `/tmp` on many
systems) keeps the backup in RAM and defeats the purpose.

## How it works

This runs on [torch\_memory\_saver](https://github.com/fzyzcjy/torch_memory_saver), which
hooks the allocator, so it does not care what the memory holds — weights, gradient
buffers and optimizer state all move as one block when the actor is paused, and come back
on resume. miles' part is choosing the per-rank directory, launching each actor with the
matching `TMS_DISK_BACKUP_*` environment, and reclaiming the files at startup and exit.

Because pause and resume happen at phase boundaries, everything is resident again by the
time the optimizer step runs. If the binding constraint is instead that the optimizer
state does not fit the GPU *during* the step, actor offload cannot help; that case is what
`--stream-optimizer-state-to-disk` addresses, and the two compose.

## Streaming the optimizer state

```bash theme={null}
--offload-train --offload-train-target disk \
--stream-optimizer-state-to-disk \
--offload-train-disk-dir /scratch/miles_offload
```

The two together are what this is for. Actor offload gets the paused actor out of HBM for
the whole rollout window; streaming keeps the largest part of it — 12 bytes per parameter
of fp32 main params and Adam moments — from being in HBM in the first place, including
while the step runs, which offload cannot do. Data parallelism divides the optimizer state,
so a run with GPUs to spare shards it small enough to stay resident and needs neither; the
runs that need both are the ones tight enough to sit at DP=1. GLM-5.2 744B on 8 GB300 nodes
is the shape of it: 279 GB of state per rank against 276.6 GB of HBM, on half the nodes the
run would otherwise need.

The fp32 main params and Adam moments live in per-bucket files instead of on the GPU.
Each step brings in one bucket, updates it, and writes it back, so peak residency is one
bucket rather than the whole state. Buckets are capped at 200M elements independently of
DDP's bucket sizes, which reach tens of GB at DP=1. Native-fp32 model params (a router's
`expert_bias`, a GDN/Mamba `A_log`) stay GPU-resident under a small separate Adam: they
are tiny, and unlike the bf16 path their optimizer shards alias the model params directly.
The fp32 gradients the step reads are bounded the same way: instead of Megatron copying every
reduced bf16 gradient into an fp32 `.grad` up front (4 more bytes per local parameter), each
bucket's gradients are converted, times the clip coefficient, right before its Adam step and
dropped after it. The grad norm reads the reduced gradients in the DDP grad buffer directly and
matches Megatron's up to float summation order.

For the Adam/DistributedOptimizer path, streaming also bounds initialization: Megatron releases
each fp32 main shard's storage after creating its tensor handle, then Miles fills one existing
runtime bucket at a time directly in the final files. There is no temporary state file. Peak
initialization HBM is the largest individual shard during construction and one main-only bucket
afterward (normally at most about 200M fp32 elements, except for an oversized entry). Each
initialized range is synchronized and evicted from page cache before continuing.

With `fp32` storage the state round-trips bit for bit and only the grad norm's summation
order differs from keeping it on GPU, so turning this on changes results by rounding at
most — it trades step time for memory. The step is I/O bound and the moments
tolerate less precision than the master copy, so they can be stored narrower:

```bash theme={null}
--stream-optimizer-state-moment-dtype bf16
```

That cuts streaming volume by a third (12 bytes per param to 8). A checkpoint records the
dtypes it was written with and a resume verifies them, so bytes written as bf16 can never
be read back as fp32. The fp8 options work but are not recommended: `exp_avg_sq` needs
per-block scaling to survive 8-bit storage, which this does not implement.

Four limits to know about. Resume is same-topology only — the on-disk layout follows this
rank's DP shard, so changing TP/PP/DP/EP fails the layout assert rather than resharding.
A checkpoint written before streaming was enabled cannot be resumed with it: the streamed
state is the only optimizer state read, so miles refuses rather than silently restarting
Adam from zero — pass `--no-load-optim` to accept a fresh optimizer state. And the
optimizer state is copied to the checkpoint directory synchronously, outside
`--async-save`, so expect checkpoint saves to take noticeably longer. Finally, the streamed
params have fp32 gradients only inside the step, so streaming refuses what reads them outside
it: fp16 loss scaling, `--log-num-zeros-in-grad`, and `--enable-mtp-training`, whose detached MTP
heads are clipped by their own grad norm.

The two also help each other. With the optimizer state already on disk there is that much
less to move when the actor is paused: on Qwen3-30B-A3B, sleep/wake went from 24s/8.9s to
5.2s/1.3s once the paused actor no longer carried the state.

## Choosing

`--offload-train` is the base mechanism, and it is not tied to colocation: the actor is
resident only during `train()` and sleeps for the whole rollout window either way. What
changes is who takes the freed HBM. Colocated, it is the engine on the same GPUs, so
offload is mandatory and defaults on. Disaggregated, it is whatever else you fit on the
training GPUs, which is the point when you are sizing engines against a fixed cluster.

* Paused actor fits in host RAM: keep `--offload-train-target=cpu` (default). Fastest, no
  scratch space.
* It does not fit: `--offload-train-target=disk`.
* On top of that, the optimizer state does not fit the GPU *during the step*, which offload
  cannot help with: add `--stream-optimizer-state-to-disk`, and consider
  `--stream-optimizer-state-moment-dtype bf16` to claw back some of the I/O cost.

Streaming on its own is the disaggregated case: nothing else wants the training GPUs during
rollout, so there is no actor to park and `--offload-train-target` is never read. Pass
`--stream-optimizer-state-to-disk` alone there. Under `--offload-train`, though, the two go
together — a run that cannot hold the optimizer state on the GPU for the step will not hold
a pinned host copy of the whole actor either — and miles asserts that pairing.

Both mechanisms share `--offload-train-disk-dir` and `--offload-train-disk-chunk-mb`. Point the
directory at real node-local NVMe: a tmpfs mount, which `/tmp` is on many systems, keeps
the data in RAM and defeats both. The chunk is a pinned host staging buffer and each
mechanism allocates its own, so enabling both costs 2x that per rank.


# Fault Tolerance
Source: https://miles.radixark.com/docs/advanced/fault-tolerance

Rollout-side health checks and engine recovery, gated by --use-fault-tolerance.

> **Outdated.** This page describes the pre-refactor fault-tolerance stack; the new fault-tolerance machinery is still being merged and will be documented once it lands.

The `--use-fault-tolerance` flag enables Miles's rollout-side
fault-tolerance machinery. It gates two code paths:

1. A `RolloutHealthMonitor` thread per server group, started in
   `miles/ray/rollout/rollout_manager.py`, which periodically heart-beats
   each SGLang engine.
2. A recovery hook in the trainer's weight-update step
   (`miles/backends/megatron_utils/actor.py`), which restarts engines
   that the health monitor has killed.

```bash theme={null}
--use-fault-tolerance
```

The flag is `action="store_true"`, default `False`
(`miles/utils/arguments.py`). The rollout paths below additionally require
`rollout` in `--ft-components`, which is what that flag selects when omitted.

Rollout fault tolerance requires a partial-target weight update:
`--update-weight-transfer-mode p2p` without `--colocate`.

## Health monitor

`RolloutHealthMonitor` (`miles/utils/health_monitor.py`) runs in a daemon
thread. Lifecycle: `start` (called once during init), `pause` and `resume`
(called when engines offload / onload), `stop` (called during dispose).
`pause` / `resume` are wired up in `miles/ray/rollout/rollout_manager.py`
and called around offload / onload events.

Each loop iteration does:

1. After a `resume`, wait `--rollout-health-check-first-wait` seconds before
   the first check (intended to cover model compilation and initialization).
2. For every active engine in the group, call `engine.health_generate.remote(timeout=self._check_timeout)`.
3. If the call raises, run `_kill_engine`: `engine.shutdown.remote()`,
   `ray.kill(engine)`, and the engine slot is set to `None`
   (`miles/utils/health_monitor.py`).
4. Sleep `--rollout-health-check-interval` seconds, then repeat.

### Flags

| Flag | Default | Source help text |
| - | - | - |
| `--rollout-health-check-interval` | `30.0` | "Interval in seconds between rollout engine `/health_generate` checks during generate/eval." |
| `--rollout-health-check-timeout` | `30.0` | "Timeout in seconds to wait for a rollout engine `/health_generate` response before killing it." |
| `--rollout-health-check-first-wait` | `0` | "Initial grace period (in seconds) before starting health checks. This allows time for model compilation and initialization. Increase this value significantly when using deepgemm." |

## Engine recovery

When `--use-fault-tolerance` is on, `MegatronActor.update_weights` calls
`inference_controller.recover_updatable_engines` before each weight
update (`miles/backends/megatron_utils/actor.py`).

`recover_updatable_engines` (`miles/ray/rollout/inference_controller.py`):

1. Pauses health monitoring.
2. Calls `srv.recover()` on the updatable server.

`srv.recover()` (`miles/ray/rollout/rollout_server.py`, which fans out to `ServerCell.init` in `miles/ray/rollout/server_cell.py`):

1. Finds engine slots set to `None` (killed by the health monitor).
2. Calls `start_engines` for each affected group.
3. Releases memory occupation on the new engines.

After `recover_updatable_engines` returns, the weight updater connects to
the new engines and the next weight transfer proceeds normally.

## P2P weight transfer timeouts

When `--update-weight-transfer-mode p2p` is on, each write is bounded by
Mooncake's own timeout (`MC_TRANSFER_TIMEOUT`, default 30 s, plus 1 ns per byte
written). At the end of the update the sender waits up to `--p2p-transfer-timeout`
(default `30.0`s) for writes still in flight; a write that failed or is still
running fails the update, and the error names each such rollout engine rank and
its session (`miles/backends/training_utils/weight_update/protocols/p2p.py`).
There is no automatic retry or broadcast-mode fallback.

## Dumper-mode interaction

In dumper mode (`miles/utils/arguments.py`), Miles forces
`use_fault_tolerance = False` and `rollout_health_check_interval = 1e18`
to keep heartbeats from firing.


# Advanced Features
Source: https://miles.radixark.com/docs/advanced/index

Systems-level features for large-scale and long-running RL.

This section covers the Miles features that the Core-features section of the
homepage points at: low-precision training (FP8 / MXFP8 / NVFP4 / INT4 QAT),
Rollout Routing Replay for MoE, fast weight updates over P2P RDMA,
disaggregated RL rollout through an external service, fault tolerance,
speculative decoding, sampling-support replay, and LoRA training and serving.

> **[Low Precision RL](/docs/advanced/low-precision)**
>
> Unified block-wise FP8, MXFP8, and NVFP4 recipes with matched training and
> rollout precision.

> **[INT4 QAT](/docs/advanced/int4-qat)**
>
> W4A16 quantization-aware training for fitting large models on a single
> 8-GPU node.

> **[Rollout Routing Replay (R3)](/docs/advanced/miles-router)**
>
> Capture expert routing during inference and replay during training. The
> mechanism that keeps MoE RL stable.

> **[Speculative Decoding](/docs/advanced/speculative-decoding)**
>
> Draft + target speculative rollout, with online MTP-SFT for the draft.

> **[Sampling-Support Replay](/docs/advanced/sampling-support-replay)**
>
> Preserve top-p and top-k rollout probabilities when the actor recomputes
> token log probabilities for training.

> **[On-Policy Distillation](/docs/advanced/on-policy-distillation)**
>
> Train a student on its own rollouts while matching teacher token
> probabilities through SGLang or Megatron teacher modes.

> **[Disaggregated RL Rollout](/docs/advanced/disaggregated-rollout)**
>
> Scale rollout across clusters and regions through an independent service,
> with versioned policy publication and request attribution.

> **[LoRA Training and Serving](/docs/advanced/lora)**
>
> Train LoRA adapters with SFT or RL and serve them through SGLang from the
> same checkpoint.


# INT4 Quantization-Aware Training
Source: https://miles.radixark.com/docs/advanced/int4-qat

Train MoE policies with fake-quantized expert weights in Megatron and packed W4A16 rollout weights in SGLang.

miles INT4 QAT keeps actor parameters in Megatron's training dtype (BF16 by
default) and applies symmetric INT4 fake quantization to routed MoE expert
weights during each Megatron forward. SGLang serves the routed expert
projections from a packed W4A16 checkpoint.

The purpose is to expose training forwards to the weight quantization used by
rollout. The hook does not store Megatron parameters, optimizer state,
gradients, or activations in INT4, so it does not reduce their storage.

INT4 QAT uses the Megatron backend and covers routed experts implemented by
Transformer Engine `GroupedLinear`. Attention, embeddings, routers, shared
experts, dense MLPs, and the LM head remain in their configured precision.

## Component roles

| Component | Role |
| - | - |
| Hugging Face checkpoint | `--hf-checkpoint` initializes SGLang and defines the rollout storage format. The generic miles exporter reads its quantization config when packing live updates. |
| Megatron | Keeps trainable weights in the configured training dtype. Before a grouped-expert GEMM, it quantizes each weight group to INT4 and dequantizes it back to the input weight's dtype. |
| Megatron Bridge or the miles raw-mode weight exporter | Maps updated Megatron weights to Hugging Face names. For an INT4 rollout checkpoint, the Kimi Bridge returns packed tensors itself; the generic miles exporter packs eligible `.weight` tensors. |
| SGLang | Loads the packed tensors with a compressed-tensors W4A16 MoE implementation and runs the expert GEMMs with 16-bit activations. |

The `W4` in W4A16 describes the stored SGLang expert weights. On the trainer,
INT4 is simulated in the forward pass; the underlying Megatron parameters stay
in the configured training dtype. `A16` means the expert GEMM uses BF16 or FP16
activations rather than quantized activations.

## Weight lifecycle

1. SGLang loads the compressed Hugging Face checkpoint from
   `--hf-checkpoint`.
2. Megatron initializes the actor in its training dtype. Kimi-K2.5 can load the
   packed checkpoint directly through Megatron Bridge; the maintained Qwen3
   raw-mode recipe initializes from a higher-precision Megatron checkpoint.
3. During training, Megatron fake-quantizes every routed-expert
   `GroupedLinear` weight immediately before the forward GEMM. Backward uses a
   straight-through estimator (STE), so the optimizer updates the original
   trainable weight.
4. At a weight-update boundary, Megatron Bridge or the miles raw-mode exporter
   exports Hugging Face-named tensors. The Kimi-K2.5 Bridge returns routed
   experts already packed as INT4; the generic miles path packs them from the
   checkpoint quantization config.
5. SGLang opens a weight-update session and loads the new packed tensors. On
   the CUDA/Marlin WNA16 path used by the registered H200 tests, the session
   restores the checkpoint-facing tensor shapes before loading and rebuilds
   the Marlin layout before rollout resumes.

For a Megatron weight group `w`, fake quantization is:

```text theme={null}
scale = max(max(abs(w)) / 7, 1e-5)
fake_w = clamp(round(w / scale), -7, 7) * scale
```

## Quantization settings

The rollout checkpoint and live exporter use the compressed-tensors storage
contract:

* Format: compressed-tensors `pack-quantized` with `num_bits: 4`,
  `strategy: group`, no activation quantization, and `symmetric: true`.
* Tensor scope: to match the current Megatron hook, pack the routed expert
  projections handled by `GroupedLinear` and keep other modules in their
  checkpoint dtype.
* Shape: the generic miles packer and direct converter require the final
  dimension of each packed matrix to be divisible by the group size.

Megatron's fake-QAT hook is configured separately through
`OPEN_TRAINING_INT4_GROUP_SIZE`. Set it to the checkpoint's `group_size` when
the training forward should use the same grouping as rollout. The two paths are
separate implementations, so matching this value does not by itself establish
bitwise-identical train and rollout weights.

Changing the environment variable does not convert a checkpoint. Changing only
the checkpoint config does not change the fake-quantization grid used by
Megatron.

## Kimi-K2.5

[Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) publishes routed
expert weights as symmetric group-size-32 INT4 tensors. The maintained Kimi
launcher exercises this Bridge path.

With `--megatron-to-hf-mode bridge`, the Kimi Bridge handles both directions:

* On Hugging Face to Megatron load, it unpacks the INT4 expert tensors and
  maps them to the Megatron parameter dtype, BF16 in the maintained recipe.
* On Megatron to Hugging Face export, it repacks the updated routed experts to
  group-size-32 INT4 and returns `weight_packed`, `weight_scale`, and
  `weight_shape` tensors for SGLang.

The Bridge can therefore initialize a fresh actor directly from the published
INT4 checkpoint:

```bash theme={null}
--hf-checkpoint /root/models/Kimi-K2.5
--megatron-to-hf-mode bridge
--model-name kimi_k25
```

When a fresh run has no usable `--load` checkpoint and no `--ref-load`, miles
falls back to `--hf-checkpoint` for Bridge initialization. The current
`scripts/run_kimi_k25.py` launcher also materializes a BF16 copy and passes it
through `--ref-load`; that is how the launcher is written today, not a
requirement of the Kimi Bridge.

Run the CI-sized two-layer recipe on one 4-GPU H200 node:

```bash theme={null}
python scripts/run_kimi_k25.py full-train \
  --model-name Kimi-K2.5-2layer \
  --num-nodes 1 \
  --num-gpus-per-node 4
```

For the full 32-node recipe, start Ray on the cluster and then run on the head
node:

```bash theme={null}
python scripts/run_kimi_k25.py prepare \
  --model-name Kimi-K2.5 \
  --num-nodes 32

MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_kimi_k25.py train \
  --model-name Kimi-K2.5 \
  --num-nodes 32
```

See the [Kimi-K2.5 model guide](/docs/models/kimi/kimi-k2.5) for its parallelism and
RL settings.

## Prepare another MoE checkpoint

The miles direct converter creates a symmetric, group-wise INT4 checkpoint
without a calibration dataset:

```bash theme={null}
python tools/convert_hf_to_int4_direct.py \
  --model-dir /root/models/MyMoE-BF16 \
  --save-dir /root/models/MyMoE-INT4 \
  --group-size 128
```

The ignore rules are name-based rather than architecture-aware. By default they
match `lm_head`, `norm`, `embed`, `self_attn`, `shared_experts`,
`mlp.(gate|up|gate_up|down)_proj`, and `mlp.gate`, plus names beginning with
`vision_tower` or `mm_projector`. Compare those patterns with a new model's
actual Hugging Face tensor names before conversion. The converter requires
CUDA and the `fake_int4_quant_cuda` extension installed by the miles image.

The converter defaults to group size 32. The current
`scripts/run_qwen3_30b_a3b.py --rollout-int4` path uses that default for the
rollout checkpoint and sets Megatron fake QAT to group size 128. The path is
therefore an example where the rollout storage group and training fake-QAT
group are configured independently; it does not simulate the same grouping on
both sides.

In raw mode, use the INT4 checkpoint for SGLang and a higher-precision
`torch_dist` checkpoint to initialize Megatron:

```bash theme={null}
--hf-checkpoint /root/models/MyMoE-INT4
--ref-load /root/models/MyMoE-BF16_torch_dist
```

Bridge mode can load the packed checkpoint directly only when that model's
Bridge implements the corresponding dequantization path, as Kimi-K2.5 does.

## Enable fake QAT

Set both variables in the Ray runtime environment used by every Megatron
worker:

```bash theme={null}
RUNTIME_ENV_JSON='{
  "env_vars": {
    "OPEN_TRAINING_INT4_FAKE_QAT_FLAG": "1",
    "OPEN_TRAINING_INT4_GROUP_SIZE": "128"
  }
}'

ray job submit --address="http://127.0.0.1:8265" \
  --runtime-env-json="${RUNTIME_ENV_JSON}" \
  -- python train.py ...
```

The Kimi launcher sets the same variables through `U.execute_train`, using
group size 32. If the selected Megatron model does not build its routed experts
with Transformer Engine `GroupedLinear`, the environment variables do not
enable QAT for that model.

## Validate the setup

Inspect the checkpoint config before launch. The expression handles both a
top-level quantization config and Kimi-K2.5's nested text config:

```bash theme={null}
jq '(.quantization_config // .text_config.quantization_config) | {
  format,
  quant_method,
  ignore,
  weights: .config_groups.group_0.weights
}' /root/models/MyMoE-INT4/config.json
```

Confirm that:

* `format` is `pack-quantized` and `quant_method` is `compressed-tensors`.
* `num_bits` is `4`, `strategy` is `group`, and `symmetric` is `true`.
* `group_size` matches `OPEN_TRAINING_INT4_GROUP_SIZE` if train and rollout
  should use the same grouping.
* The safetensors index has `weight_packed`, `weight_scale`, and
  `weight_shape` entries for the intended routed expert projections, and that
  other modules remain unpacked.

During a smoke test, confirm that SGLang selects a compressed-tensors W4A16 MoE
scheme and that a live weight update completes. For a new model, also compare
reward, KL, gradient norm, and train-versus-rollout log-probability differences
with a BF16 rollout baseline.

| Symptom | Check |
| - | - |
| SGLang fails at load or weight update | Check the packed tensor names, shapes, group size, and checkpoint ignore rules. |
| Training runs but QAT has no effect | Check that the environment reached every Megatron worker and that routed experts use Transformer Engine `GroupedLinear`. |
| Train and rollout log-probabilities diverge | Check the quantized tensor scope and group size first; MoE routing can be a separate source of mismatch. |
| INT4 QAT does not reduce trainer parameter or optimizer storage | Expected. Those states are not stored in INT4. |
| The converter cannot import `fake_int4_quant_cuda` | Use the miles CUDA image or build the repository's INT4 QAT extension. |

## Hardware and test coverage

The direct converter and generic miles INT4 packer use the CUDA
`fake_int4_quant_cuda` extension. SGLang's CUDA WNA16 MoE path requires NVIDIA
compute capability 8.0 or newer. SGLang also contains a ROCm WNA16 MoE
implementation. The repository registers two INT4 end-to-end tests, the
Kimi-K2.5 two-layer Bridge recipe and the Qwen3-30B-A3B raw-mode recipe, both on
4-GPU H200. It does not register an INT4 QAT end-to-end test on ROCm.

INT4 reduces storage and weight-update traffic for the packed routed experts.
The whole-model reduction is smaller because scales, metadata, and unquantized
tensors remain. Rollout throughput depends on the model, batch shape, expert
parallelism, and the W4A16 backend SGLang selects.

## Related guides

* [Low Precision RL](/docs/advanced/low-precision)
* [Kimi-K2.5 model guide](/docs/models/kimi/kimi-k2.5)
* [P2P weight transfer](/docs/advanced/p2p-weight-transfer)
* [slime low-precision training and rollout](https://thudm.github.io/slime/advanced/low-precision.html)


# LoRA Training and Serving
Source: https://miles.radixark.com/docs/advanced/lora

Train LoRA adapters with miles, synchronize them to SGLang, and run dense, MoE, quantized-rollout, multi-LoRA, and agentic recipes.

miles trains the adapter matrices in Megatron and serves the same live adapter
from SGLang. The base model stays frozen; at each configured weight-update
boundary, miles exports and synchronizes the updated LoRA weights. The default
boundary is once per rollout/train iteration. No merge-to-base or checkpoint
conversion is required inside the training-rollout loop.

## Training and rollout lifecycle

For a single adapter, miles exports the updated LoRA tensors at each configured
publish boundary. Colocated jobs load them through the local tensor/IPC path;
disaggregated Bridge jobs broadcast them to remote SGLang engines over NCCL.
Rollout requests then select the live named adapter with `lora_path`.

Multi-LoRA currently supports disaggregated rollout only and rejects
`--colocate` at launch. Each SGLang engine keeps the base checkpoint resident;
adapter weights reach the engines through explicit pushes: each push registers
an immutable version under its own engine-side name, so in-flight sampling
against an older version is never disturbed.

Model support is therefore a three-way contract rather than a hard-coded
allowlist:

1. Megatron-Bridge (or a model-specific provider) must build and wrap the
   training model.
2. miles must map and export the selected module names correctly.
3. SGLang must be able to allocate and apply the same adapter modules.

A model can support attention-only LoRA while still having caveats for expert,
linear-attention, or model-specific projections. Start from a validated recipe
instead of assuming that any checkpoint with familiar module names will work.

## LoRA training implementations

Both implementations use `--train-backend megatron`. The difference is how
LoRA is attached to the Megatron model and converted for serving; this is
separate from SGLang's serving-side `--sglang-lora-backend` choice.

| Implementation | How the adapter is built | Model coverage on current `main` | Status |
| - | - | - | - |
| **Megatron-Bridge PEFT** | `AutoBridge` builds the provider, applies Bridge LoRA before DDP, and exports HF-named adapter tensors. Select it with `--megatron-to-hf-mode bridge`. | Qwen2.5, Qwen3, GPT-OSS, Kimi K2.5, GLM-5/5.1/5.2/5.3, and Qwen3.5/3.6, subject to the evidence and module caveats below. | General production path. Current multi-LoRA also requires this path. |
| **Native / raw-mode LoRA** | Under `--megatron-to-hf-mode raw`, miles builds the model with its own Megatron provider and the `miles.backends.megatron_utils.lora.native` plugin attaches adapter modules directly before DDP. | Architectures registered in `miles/backends/megatron_utils/lora/native/registry.py`: fused-QKV GQA (Llama, Qwen2/3, GLM-4), Qwen3.5/3.6/Next hybrids (GQA layers), MLA (DeepSeek-V3/V3.2, Kimi K2/K2.5, GLM-4.7-Flash, GLM-5.x), Inkling, and Kimi K3. | Single-LoRA, colocated. Unregistered models fail at startup. |

> [!NOTE]
>
> The planned maintenance direction is native-first: after the generalized native plugin
> lands, new model enablement and ongoing LoRA maintenance will primarily target
> the native path. Bridge remains the compatibility path for existing validated
> recipes during the transition; this is not an immediate Bridge deprecation.

### Native LoRA

Native LoRA attaches adapter modules directly to raw-mode Megatron models
instead of relying on Bridge PEFT conversion. An architecture spec declares, once,
which HF projections exist on which physical linear and how each is sharded; the
same declaration drives attachment, HF export/import, and SGLang serving targets.
The spec is selected from the checkpoint's `model_type`.

Adapters cover attention projections (GQA Q/K/V/O, MLA Q-A/Q-B/KV-A/KV-B/O), dense
MLPs, and shared experts. Routed experts carry shared-outer adapters only for Inkling
and Kimi K3 (which require `--experts-shared-outer-loras`); Inkling's spec also covers
its output head. Inkling and Kimi K3 accept only their verified adapter layout (Kimi
K3 may omit the routed-expert down projection), so omit `--target-modules` for them.
Selecting a projection the spec does not implement
(for example routed experts or GDN input projections) fails at startup; use
`--megatron-to-hf-mode bridge` for those. A partial fused selection (for example
only `q_proj`) trains only the selected projection; SGLang serves its fused-buffer
siblings as zero adapters. A shared expert's `up_proj` needs `gate_proj` selected too.

## Validated models and recipes

The table distinguishes CI-tested configurations from maintained recipes or
full-scale experiment evidence. It is not an exhaustive model whitelist.

| Implementation | Model family | Architecture exercised | Evidence | Notes |
| - | - | - | - | - |
| Bridge | Qwen2.5 0.5B / 3B | Dense | [0.5B CUDA and ROCm E2E](https://github.com/radixark/miles/blob/main/tests/e2e/lora/test_lora_qwen2.5_0.5B.py), [3B disaggregated recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-qwen2.5-3B-megatron-lora-disaggregated.sh) | The simplest starting point; `all-linear` works. |
| Bridge | Qwen3 4B | Dense | [Single-LoRA recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-qwen3-4B-megatron-lora.sh), [AMD launcher](https://github.com/radixark/miles/blob/main/scripts/amd/run_qwen3_4b_lora.py) | The AMD launcher uses Triton attention and LoRA backends. |
| Bridge | Qwen3-30B-A3B | MoE | [Multi-LoRA gateway example](https://github.com/radixark/miles/tree/main/examples/multi_lora) | Disaggregated Tinker training and versioned sampling. |
| Bridge | GPT-OSS 20B | MoE | [Recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-gpt-oss-20B-megatron-moe-lora.sh), [AMD launcher](https://github.com/radixark/miles/blob/main/scripts/amd/run_gpt_oss_20b_lora.py), [MoE LoRA E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_gpt_oss_20b_moe_lora_ci.py) | Uses the SGLang `triton` LoRA backend; the AMD launcher also uses Triton attention and MoE backends. |
| Bridge | Kimi K2.5 | Multimodal MoE + MLA | [16-node recipe](https://github.com/radixark/miles/blob/main/examples/lora/run-kimi-k25-megatron-lora.sh) | Demonstrates shared-outer expert LoRA and an INT4 rollout / fake-QAT setup. |
| Bridge | GLM-5 / 5.1 / 5.2 744B-A40B | MoE + MLA + DSA | [GLM-5.1 launcher](https://github.com/radixark/miles/blob/main/scripts/run_glm5_1_744b_a40b_lora.py), [GLM-5.2 launcher](https://github.com/radixark/miles/blob/main/scripts/run_glm5_2_744b_a40b_lora.py) | CI covers reduced 6-layer / 5-layer checkpoints; historical full-744B results are described below. |
| Bridge | Qwen3.5 / Qwen3.6 35B-A3B | Hybrid GDN + MoE | [Launcher](https://github.com/radixark/miles/blob/main/scripts/run_qwen3_5_35b_a3b_lora.py), [Qwen3.5 E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py) | Uses explicit wildcard targets to exclude MTP and vision modules. |
| Native / raw | Inkling / Inkling-Small | Native multimodal MoE | [Launcher](https://github.com/radixark/miles/blob/main/scripts/run_inkling.py), [Inkling-Small 6-layer E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_inkling_small_6layer_lora_ci.py) | Larger profiles are launcher/experiment evidence rather than LoRA CI. |
| Native / raw | Qwen3 0.6B, Qwen3.5-35B-A3B | Dense GQA; gated-GQA + GDN hybrid MoE | [Launcher](https://github.com/radixark/miles/blob/main/examples/lora/run_qwen3_lora_native.py), [Qwen3 0.6B E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_qwen3_0_6b_lora_native_ci.py), [Qwen3.5 E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_qwen3_5_35b_a3b_lora_native_ci.py) | Attention (+ dense/shared MLP) adapters; GDN layers carry none. |
| Native / raw | GLM-4.7-Flash, GLM-5.2, Kimi K2.5 | MLA MoE | [Launcher](https://github.com/radixark/miles/blob/main/examples/lora/run_lora_native.py), [GLM-5.2 5-layer E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_glm5_2_5layer_lora_native_ci.py), [Kimi K2.5 2-layer E2E](https://github.com/radixark/miles/blob/main/tests/e2e/megatron/model_scripts/test_kimi_k25_2layer_lora_native_ci.py) | MLA attention adapters; routed experts require Bridge. |
| Native / raw | GLM-5.3-Flash | Hybrid KDA + DSA MoE + mHC | [Launcher](https://github.com/radixark/miles/blob/main/scripts/run_glm5_3_flash_lora.py), [model page](/docs/models/glm/glm5-3-flash) | KDA, DSA (except `kv_b_proj`) and dense/shared MLP adapters; routed experts stay frozen. Serving needs SGLang LoRA on the KDA projections. |

## Quick start

Append the following to a Megatron dense-model recipe:

```bash theme={null}
LORA_ARGS=(
  --lora-rank 32
  --lora-alpha 32
  --lora-dropout 0.0
  --target-modules all-linear
  --megatron-to-hf-mode bridge
)
```

`--lora-rank 0` is the default and disables LoRA when no adapter path is
provided. For RL, dropout is normally set to zero. Alpha is recipe-dependent:
maintained recipes use both `alpha = rank` and `alpha = 2 * rank`.

> [!WARNING]
>
> This quick start is the Bridge path. Native LoRA uses `--megatron-to-hf-mode raw`
> with a torch\_dist base checkpoint and a registered architecture (see Native LoRA
> above). FSDP does not currently implement LoRA training.

Omitting `--target-modules` or passing `all-linear` uses the model defaults from
`miles/utils/lora/hf_lora_targets.py`: attention + MLP, with model-specific exclusions
and output-head defaults. Multi-LoRA without explicit targets selects all three
training groups. Use `--target-modules attn,mlp,unembed` to select groups explicitly.
Ordinary LoRA also accepts specific HF targets mixed with group names, such as
`--target-modules attn,lm_head`; Tinker accepts only group names.

### Core arguments

| Flag | Default | Purpose |
| - | -: | - |
| `--lora-rank` | `0` | Adapter rank; a positive value enables LoRA. Supplying an adapter path also marks LoRA enabled, but still requires matching positive-rank configuration. |
| `--lora-alpha` | `16` | Adapter scaling factor. |
| `--lora-dropout` | `0.0` | Dropout on the adapter path. |
| `--lora-type` | `lora` | `lora` uses fused Megatron projections; `canonical_lora` uses split Q/K/V and gate/up projections. The canonical path is implemented and covered by fast name-mapping tests, but has no maintained recipe or E2E validation. |
| `--target-modules` | none | Uses HF model defaults when omitted. Accepts `all-linear`, `attn/mlp/unembed` groups, HF leaf names or scoped HF paths; Bridge also accepts Megatron selectors. |
| `--exclude-modules` | none | Comma-separated HF leaf names or scoped HF paths removed after selection; Bridge also accepts Megatron selectors. |
| `--lora-adapter-path` | none | Warm-start/resume path. Also provide the matching positive rank, alpha, and target modules. Bridge training resume currently requires miles' per-rank adapter shards and the same parallel topology; an HF PEFT-only adapter cannot yet be loaded directly into the Bridge model. Inkling native has its own HF adapter loader. |
| `--lora-base-cpu-backup` | off | Colocated mode only: keep a CPU mirror of the frozen SGLang base and avoid re-sending base weights. This trades host RAM for faster and more reliable pause/resume. |
| `--lora-train-only` | off | Train the adapter while keeping ordinary rollout engines on the frozen base policy. |
| `--experts-shared-outer-loras` | off | Use shared outer factors for grouped MoE experts. This layout is not checkpoint-compatible with per-expert LoRA. |
| `--check-lora-weight-equal` | off | On the colocated path, verify each synchronized adapter tensor with SHA-256. |
| `--update-weights-interval` | `1` | Publish new weights every N rollout/train iterations. This is not LoRA-specific, but it controls when the live adapter is synchronized. |

### HF target source of truth

`miles/utils/lora/hf_lora_targets.py` owns **HF target groups and defaults** for all
backends. It derives attention, MLP, and output-head paths from the model config,
including nested text models, optional MLA projections, expert layouts, and
hybrid attention. Vision towers, routers, norms, and GDN convolutions are excluded.

| Model layout | Attention | MLP | Ordinary LoRA default |
| - | - | - | - |
| Llama, Qwen2/2.5, Qwen3 | Q/K/V/O | Dense gate/up/down | Attention + MLP |
| Qwen3 MoE | Q/K/V/O | Routed experts; dense layers when configured | Attention + MLP |
| Qwen3-Next | Q/K/V/O + GDN qkvz/ba/out | Routed + shared experts; configured dense layers | Attention + MLP |
| Qwen3.5/3.6, text and multimodal | Q/K/V/O + GDN qkv/z/b/a/out | Dense or packed routed + shared experts | Attention + MLP |
| GPT-OSS | Q/K/V/O | Packed gate\_up/down experts | Attention + MLP |
| DeepSeek V2/V3, Kimi K2/K2.5 | MLA | Dense + routed + shared experts as configured | Attention + MLP |
| DeepSeek V3.2, GLM-5/5.1/5.2 | MLA + DSA indexer | Dense + routed + shared experts as configured | Attention + MLP, excluding indexer |
| GLM-4 MoE | Q/K/V/O | Dense + routed + shared experts as configured | Attention + MLP |
| Inkling | Q/K/V/R/O | Dense + packed routed + shared expert projections | Attention + MLP + output head |

`--target-modules` accepts comma-separated `attn`, `mlp`, and `unembed` groups,
explicit HF targets, or a mix of both. Group names expand through the model's HF
definition, so `mlp` includes dense, shared, and routed expert projections, including
packed `gate_up_proj` weights. Only exact group names expand; full paths remain
literal target patterns. Duplicates are removed, then `--exclude-modules` applies.
Explicit module selections must not overlap exclusions; defaults and group
selections may still be narrowed with exclusions.

A matching pair of `gate_proj` and `up_proj` selectors also selects
`gate_up_proj` when the HF model has packed projections, with a warning.
Scoped pairs retain their path prefix; selecting only gate or only up does not
expand to the packed projection. Split selectors remain where the model has
split modules and are replaced where it has only packed modules.

Without explicit targets, ordinary LoRA uses its model defaults and multi-LoRA
selects all three groups. `all-linear` selects the ordinary model defaults and
must be used alone. Explicit lists replace the defaults.

Packed expert entries identify HF parameters rather than `nn.Linear` modules.
Inkling entries follow the native HF model namespace. Its native Megatron LoRA
implementation requires the complete fixed layout. SGLang discovers its adapter
modules with `all`, and online weight-sync config uses `all-linear`; adapter
checkpoints derive explicit names from exported tensors. Adapter factors, tensor
packing, and checkpoint names are unchanged. The pinned Transformers version does
not yet include native Inkling, so its HF structure is not covered by the native
meta-model tests. A layout entry is not a backend support claim.

Ordinary LoRA and Tinker share these HF target groups. HF targets retain their
meaning throughout training and serving. `miles/utils/hf_utils/weight_mapping.py`
uses Transformers conversion rules and a meta model's parameter names to relate
checkpoint keys to the current HF model namespace, without loading base weights.
This resolves renamed and packed parameter names. Custom models absent from
native Transformers keep their existing checkpoint namespace.

Bridge resolves mappings against parameters that actually exist across PP/EP
ranks before checking fused selections. Explicit Megatron selectors retain a
compatibility conversion at startup. Adapter factories receive only Megatron
targets; they do not overwrite the HF selection. Arbitrary layer/expert subsets
within a registry template are not implemented and are rejected.
Standard LoRA requires all projections of a fused weight together;
`canonical_lora` supports individual Q/K/V and dense gate/up selections.
Grouped-expert FC1 and GDN input projections remain fused and require all of
their HF projections even in canonical mode.

Tinker uses `--target-modules attn,mlp,unembed` by default; for example,
`--target-modules attn,mlp` disables output-head training. Client SDK flags must
match the selected server groups. Tinker rejects explicit module names,
`all-linear`, and `--exclude-modules` because the SDK describes only whole groups.

For Bridge, SGLang receives the selected HF paths and normalizes them into buffer types
(for example, Q/K/V become `qkv_proj`); it does not own the selection policy.
Adapter checkpoints derive their concrete `target_modules` from exported tensor
keys rather than a separate Megatron-to-HF name table. Online adapter registration
uses the resolved HF selection for Bridge and native LoRA, and `all-linear` for
native Inkling, without an extra
tensor export.
This requires [SGLang's HF-path normalization support](https://github.com/sgl-project/sglang/pull/40242),
including GDN split names. FSDP LoRA injection remains
unsupported.

Native LoRA validates the resolved HF selection against the architecture spec
before injection and matches each adapter by its full HF path, so layer-scoped
selections work. Native Inkling uses a fixed adapter schema: defaults select that
complete schema, and partial/custom layouts are rejected. `canonical_lora` is
Bridge-only.

### Rollout topology

| Topology | Adapter transport | Requirements |
| - | - | - |
| Colocated | Local tensor serialization / CUDA IPC | Add `--colocate`; large-model recipes generally also use `--lora-base-cpu-backup`. Pipeline parallel adapter shards are assembled before the load. |
| Disaggregated, including multi-node | NCCL broadcast to remote SGLang engines | Use `--update-weight-transfer-mode broadcast`, Bridge mode, and `--pipeline-model-parallel-size 1`. |

Native LoRA is colocated only. Remote/disaggregated LoRA sync requires Bridge.

On SGLang attention DP (`--sglang-attn-dp-size N`, or the deprecated
`--sglang-enable-dp-attention --sglang-dp-size N`), SGLang serves LoRA only from
pinned adapters and requires the attention DP size to equal the engine's TP size
(sgl-project/sglang#36389). miles pins the rollout adapter there and gives the engine
a second LoRA slot, which SGLang keeps for the base model. Multi-LoRA cannot run on
attention DP: pinned adapters are never evicted.

LoRA is not supported by the P2P/RDMA or disk-delta weight-transfer modes. A
hybrid job with both colocated and additional remote rollout engines also cannot
send LoRA weights to the remote engines today.

The legacy `--lora-sync-from-tensor` flag is still accepted by the parser but is
not consumed by the implementation. The topology selects the synchronization
path automatically; do not rely on this flag.

## MoE adapters

This section describes the Bridge path; native Inkling uses its fixed custom
expert layout. For grouped MoE experts, the training and serving layouts must
agree. A typical Bridge configuration is:

```bash theme={null}
LORA_ARGS=(
  --lora-rank 32
  --lora-alpha 32
  --lora-dropout 0.0
  --target-modules "gate_up_proj,down_proj"
  --sglang-lora-backend triton
  --megatron-to-hf-mode bridge
)
```

The SGLang `triton` backend is required by the maintained MoE recipes; the
default backend can skip MoE layers. Per-expert adapters are the default.
`--experts-shared-outer-loras` selects the shared-outer layout, and miles enables
SGLang virtual-expert serving for expert adapters by default. Use
`--no-sglang-lora-use-virtual-experts` only when intentionally selecting the
alternative aligned-expert path.

## Supported features

* **Losses and algorithms.** LoRA wraps the actor model independently of the
  policy loss. The maintained RL recipes use critic-free estimators such as
  GRPO. The SFT loss path can use the same adapter hooks, although the repository
  does not yet have a dedicated LoRA-SFT E2E test. Shared actor/critic PPO with
  Bridge LoRA is untested: the critic is always built as a full model without
  adapters.
* **Checkpoints.** miles saves native per-rank adapter shards and
  optimizer/scheduler state. Exact resume expects the same TP/PP topology. It
  also attempts a best-effort HF PEFT `adapter_model.safetensors` plus
  `adapter_config.json` export in Bridge mode, through the same snapshot publisher
  used by Tinker. HF export errors are logged while native checkpoint saving
  continues. Raw mode saves native shards and a rank-sharded adapter config without
  HF export. `--save-hf` exports a merged model and an HF adapter without native
  training shards; native LoRA merges each adapter into its host weights before the
  ordinary base-weight conversion, so the merged model uses the same converters as
  full-model exports. Direct HF PEFT-to-Bridge resume is not implemented yet;
  native LoRA imports HF adapters through `--lora-adapter-path`.
* **Weight synchronization.** Colocated IPC and remote NCCL broadcast both ship
  adapter tensors at each configured update boundary without merging them into
  the base. A checksum checker is available for the colocated path.
* **Model structures.** Dense attention/MLP, MLA, GDN, multimodal models, models
  containing DSA, and both per-expert and shared-outer MoE adapters have
  validated recipes or tests. DSA-indexer adapters themselves are not validated;
  see the model-specific caveats in the table above.
* **Agentic sessions.** Session-server requests carry the active single-LoRA
  name, so multi-turn trajectories are generated by the updated policy rather
  than the frozen base. See the experimental example below.

### Quantized rollout

Precision support is asymmetric: the validated FP8 configuration trains the
LoRA actor in BF16 while SGLang serves a separately quantized FP8 base
checkpoint. Adapter tensors remain unquantized and are synchronized at each
configured weight-update boundary. The GLM-5.2 launcher exposes this as its
`--fp8-rollout` option and writes an SGLang config with `update_weights: true`.
See [Low Precision RL](/docs/advanced/low-precision) and
[INT4 QAT](/docs/advanced/int4-qat) for the underlying precision features.

Historical GLM-5.2 validation measured:

| Configuration | Train/rollout abs diff | KL |
| - | -: | -: |
| 5-layer GLM-5.2, BF16 rollout | `0.0104` | `1.2e-4` |
| 5-layer GLM-5.2, FP8 rollout | `0.0347` | `9.8e-4` |
| Full GLM-5.2 744B, FP8 rollout, two 16-GPU engines | `0.0196` | `2.7e-3` |

These values are configuration-specific rather than universal FP8 thresholds.
The full-744B row is historical multi-node evidence and cannot be reproduced
directly by the current single-node `main` launcher; multi-node launcher work is
tracked in [PR #2033](https://github.com/radixark/miles/pull/2033).

The same work also produced the following curves from a different, longer
validation run. The exact configuration for these curves was not published, so
their plotted levels should not be mapped to any one of the final point
measurements in the table:

![GLM-5.2 LoRA FP8 rollout train-rollout log-prob difference](https://mintcdn.com/radixark/1Ix2r8uv4-Uc24OW/assets/images/lora-glm5-2-fp8-logprob.png?fit=max\&auto=format\&n=1Ix2r8uv4-Uc24OW\&q=85\&s=c43e2ce2ca7940207739b1bfed383be3)

*Train/rollout log-prob difference in the longer FP8 rollout validation.*

![GLM-5.2 LoRA FP8 rollout reward](https://mintcdn.com/radixark/1Ix2r8uv4-Uc24OW/assets/images/lora-glm5-2-fp8-reward.png?fit=max\&auto=format\&n=1Ix2r8uv4-Uc24OW\&q=85\&s=b3aed42e101b9f90ef3df4c608ce8c37)

*Raw reward in the longer FP8 rollout validation.*

Kimi K2.5 provides a separate model-specific INT4 example: SGLang serves the
INT4 checkpoint while the trainer uses BF16 fake-QAT and TIS correction. Do not
generalize these two recipes into blanket support for FP8, MXFP8, or INT4 LoRA
training on every model. In particular, multi-LoRA on MoE expert leaves rejects
FP8/FP4 experts.

## GLM-5.2 744B BF16 validation

Historical full-scale runs validated the GLM-5/5.1/5.2 Bridge LoRA path across
MoE and MLA on models containing DSA; the DSA indexer itself was excluded from
the adapter targets. The full GLM-5.2 744B run used 64 GPUs and completed more
than 50 rollout -> train -> save steps. The validation also compared expert
target selections, including a run that excluded `down_proj`.

Treat this as historical implementation and scale evidence. The reported
train/rollout log-prob gaps are configuration-specific and are not established
as a portable acceptance threshold, so the corresponding curves are not
reproduced here.

> [!WARNING]
>
> This historical validation is not a claim that the current `main` launcher
> reproduces the 744B run on one node. The checked-in launcher is single-node; the
> repository's copy-paste validation path uses a reduced checkpoint. Multi-node
> full-744B launcher work is tracked in
> [PR #2033](https://github.com/radixark/miles/pull/2033).

## Agentic RL with LoRA (experimental)

Agentic rollout has one extra correctness requirement: every turn sent through
the session server must select the newly synchronized adapter. The session
server therefore attaches `lora_path=miles_lora` to its requests; otherwise the
trainer could update LoRA while the agent continues collecting trajectories
from the frozen base policy. The current session integration selects the fixed
single-adapter name; it is not multi-LoRA slot routing, and it should not be
combined with `--lora-train-only`.

[Draft PR #2280](https://github.com/radixark/miles/pull/2280) adds the first
LoRA-specific agentic recipe: GLM-5.2 744B-A40B with synchronous GRPO on
Terminal-Bench-2-style tasks, a Harbor agent server, and Daytona sandboxes. Its
reference configuration uses:

* 4 nodes x 8 H200, TP8 / EP32 / PP1 / CP1;
* BF16 training with an FP8 rollout checkpoint;
* rank 16, alpha 32, attention/MLA-only targets;
* 64K sessions, TileLang DSA, full recompute, and disk actor offload; and
* two 16-GPU rollout engines with `--lora-base-cpu-backup`.

The author reports 30+ 64K-context GRPO rollouts with stable rewards and no
NaN/OOM, but the PR does not publish a reward value, solve rate, or learning
curve. It remains a draft and depends on an open memory-headroom companion
change ([PR #2199](https://github.com/radixark/miles/pull/2199)), so treat it as
stability evidence rather than a released benchmark.

## Multi-LoRA training

Multi-LoRA training is served through the Tinker protocol: `serve_tinker.py`
turns a miles deployment into a training service where each client drives its
own adapter slot with explicit forward\_backward / optim\_step operations (the
server owns no datasets or schedules). The gateway lives in `miles/tinker/`,
the slot mechanics in `miles/backends/megatron_utils/lora/`. The dataset-driven
multi-LoRA v1 backend (fully-async driver, adapter controller, per-adapter data
sources) has been removed.

Clients use the official `tinker` SDK with the gateway URL and a distinct API key
per tenant. The key identifies ownership of models, futures, and checkpoints.
Each adapter accumulates gradients until its client submits an optimizer step;
saving weights for sampling publishes an immutable adapter version.

`save_state` saves adapter parameters and optimizer state, including FP32 masters.
It waits for preceding commands but neither steps nor saves pending gradients;
call it after `optim_step` to save the effect of the accumulated training work.
Sampler saves commit an immutable adapter directory on shared storage, without
contacting inference engines. Engines load the same frozen base at startup and
load adapter snapshots from disk on demand. A failed engine load fails sampling
without invalidating the saved snapshot or model training.

Futures, model leases, and sequence deduplication are in memory and are lost on
gateway restart. Saved checkpoints retain their ownership and adapter metadata
and can be used to create a new training or sampling client.

See the [gateway example](/docs/examples/multi-lora) for the launcher and smoke client.


# Low Precision RL
Source: https://miles.radixark.com/docs/advanced/low-precision

Unified low-precision pipelines for RL — block-wise FP8, MXFP8, and NVFP4 across rollout and training.

A common failure mode in MoE RL is precision drift between training and
inference. Pipelines that train in BF16 and serve in FP8 accumulate per-layer
numerical disagreement, which compounds into divergent log-probabilities and
gradients pointing in unintended directions.

Miles supports a unified low-precision path where rollout and training share
the same quantization logic on the forward pass. The same path is wired up for
three formats today — **block-wise FP8**, **MXFP8**, and **NVFP4** — plus the
lower-friction "BF16 train + FP8 inference" mode that's useful when standing
up a new model architecture.

## Choose a precision

| Format | Block layout | Hardware | Models tested | Maturity |
| - | - | - | - | - |
| **BF16** | — | All NVIDIA + AMD MI300X / MI325X / MI350X / MI355X | All | Baseline |
| **FP8 block-wise** (DeepSeek-style) | 128×128, FP32 scales | Hopper (H100 / H200), Blackwell (B200+), AMD MI350X / MI355X | Qwen3-4B, Qwen3-30B-A3B, DeepSeek-V4-Flash | Generally available |
| **MXFP8** | 1×32, UE8M0 scales | Blackwell only (B200, B300, GB200, GB300) | Qwen3-30B-A3B, DeepSeek-V3.2 | Beta |
| **NVFP4** (E2M1) | 1×16, two-level (FP8 + FP32) scales | Blackwell only (B200, B300, GB200, GB300) | Qwen3-30B-A3B | Beta |

## Rollout × training compatibility

Each row is a rollout (inference) precision; each column is the trainer's
forward precision. ✅ = supported; ✗ = not supported.

| Rollout \ Train | BF16 | FP8 block-wise | MXFP8 | NVFP4 |
| - | - | - | - | - |
| **BF16** | ✅ baseline | ✗ | ✗ | ✗ |
| **FP8 block-wise** | ✅ | ✅ Hopper + Blackwell + MI350X / MI355X | ✗ | ✗ |
| **MXFP8** | ✅ | ✗ | ✅ Blackwell | ✗ |
| **NVFP4** | ✗ | ✗ | ✗ | ✅ Blackwell |

The reference script (`scripts/run_qwen3_30b_a3b.py`) allows only one rollout
precision and one training precision at a time. Use paired rollout and training
flags for the end-to-end MXFP8 and NVFP4 recipes below.

## Unified training recipe

| Stage | Typical pipeline | Miles unified low-precision |
| - | - | - |
| Rollout (forward) | FP8 / MXFP8 / NVFP4 GEMM | Matching low-precision GEMM |
| Trainer (forward) | BF16 GEMM | Matching low-precision GEMM |
| Trainer (backward) | BF16 grads | Recipe-specific precision |
| Optimizer | BF16 master | Higher-precision master weights |

The precision contract covers checkpoint conversion, the trainer forward pass,
SGLang rollout, and live weight export. High-precision tensor exceptions must
match across all four stages.

## Modes

### 1. BF16 train + FP8 inference

The lowest-friction path. SGLang loads FP8 weights while the trainer keeps a
BF16 `torch_dist` checkpoint. There is precision drift between the two paths;
on MoE workloads, pair this with R3 (and optionally TIS).

```bash theme={null}
hf download Qwen/Qwen3-30B-A3B-FP8 --local-dir /root/Qwen3-30B-A3B-FP8

CKPT_ARGS=(
   --hf-checkpoint /root/Qwen3-30B-A3B-FP8        # FP8 weights for SGLang
   --ref-load      /root/Qwen3-30B-A3B_torch_dist  # BF16 torch_dist for trainer
)
```

Reference recipe:
[`examples/infra_features/low_precision/run-qwen3-4b-fp8.sh`](https://github.com/radixark/miles/blob/main/examples/infra_features/low_precision/run-qwen3-4b-fp8.sh)
— single-node Qwen3-4B. It serves an FP8 checkpoint to SGLang and trains from a
BF16 `torch_dist` checkpoint; it sets no `--fp8-recipe`, so the trainer forward
stays BF16.

### 2. Unified block-wise FP8 (DeepSeek-style)

Rollout and training share the same block-wise FP8 quantization. This is the
recipe to use on Hopper, and the layout DeepSeek ships its FP8 checkpoints in.
Block layout is 128×128 with FP32 scales.

```bash theme={null}
--transformer-impl transformer_engine
--bf16
--fp8-format e4m3
--fp8-recipe blockwise

# Optional, for MoE numerical stability
--use-tis
```

| Flag | Effect |
| - | - |
| `--transformer-impl transformer_engine` | Routes Megatron's forward through TransformerEngine so FP8 GEMM is engaged. |
| `--fp8-format e4m3` | Forward FP8 format used by TransformerEngine. |
| `--fp8-recipe blockwise` | 128×128 block-wise quantization; sglang must serve weights in the matching layout. |
| `--use-tis` | Truncated Importance Sampling for residual precision drift. |

`NVTE_FP8_BLOCK_SCALING_FP32_SCALES` is set for you in the actor env
(`miles/ray/train/actor_factory.py`), defaulting by hardware: `1` on Hopper, and
`0` on Blackwell, where TransformerEngine emulates the block-wise recipe with
MXFP8 and needs power-of-two scales. Override it only if you know you want the
non-default for your GPU.

For models that already ship 128×128 block-wise FP8 weights (DeepSeek-V3.2,
`Qwen/Qwen3-30B-A3B-FP8`), point `--hf-checkpoint` at the
block-wise FP8 directory and let SGLang autodetect. Otherwise convert with
`tools/convert_hf_to_fp8.py`.

For MoE workloads, also consider `--use-rollout-routing-replay` (R3). The
canonical recipe leaves it commented out by default but the flag is available.

Reference recipe:
[`examples/infra_features/low_precision/run-qwen3-30b-a3b-fp8-two-nodes.sh`](https://github.com/radixark/miles/blob/main/examples/infra_features/low_precision/run-qwen3-30b-a3b-fp8-two-nodes.sh)
— two-node Qwen3-30B-A3B.

### 3. Unified MXFP8 (Blackwell)

MXFP8 uses one-dimensional microscaling blocks: 32 consecutive E4M3 values
share one UE8M0 scale. The end-to-end recipe uses MXFP8 for rollout, forward
propagation, weight-gradient GEMMs, and data-gradient GEMMs while preserving
configured tensors in BF16.

![End-to-end MXFP8 RL recipe](https://mintcdn.com/radixark/Cy4ImdPUZ2E-qOic/assets/images/low-precision/mxfp8-e2e.png?fit=max\&auto=format\&n=Cy4ImdPUZ2E-qOic\&q=85\&s=d60de262ad9dec1d60d192d2e719ff3a)

**Hardware:** B200, B300, GB200, or GB300.

The Qwen3-30B-A3B launcher prepares the MXFP8 Hugging Face checkpoint and
configures both sides of the RL loop:

```bash theme={null}
python scripts/run_qwen3_30b_a3b.py prepare \
  --hardware B200 \
  --rollout-mxfp8 \
  --train-mxfp8

python scripts/run_qwen3_30b_a3b.py execute \
  --hardware B200 \
  --rollout-mxfp8 \
  --train-mxfp8
```

The training path uses:

```bash theme={null}
--transformer-impl transformer_engine
--bf16
--fp8-format e4m3
--fp8-recipe mxfp8
```

SGLang selects the MoE and dense linear GEMM backends independently. Common
MXFP8 choices are:

| Workload | Miles flag | Backends |
| - | - | - |
| MoE | `--sglang-moe-runner-backend` | `flashinfer_trtllm_routed`, `flashinfer_trtllm`, `cutlass` |
| Dense linear GEMM | `--sglang-fp8-gemm-backend` | `flashinfer_trtllm`, `flashinfer_cutlass`, `triton` |

Choose the pair that matches the model, parallel layout, and installed SGLang
stack. Convert a checkpoint outside the launcher with:

```bash theme={null}
python tools/convert_hf_to_mxfp8.py \
  --model-dir /root/models/Qwen3-30B-A3B \
  --save-dir /root/models/Qwen3-30B-A3B-MXFP8
```

The converter records a 1x32 weight block and UE8M0 scale layout. It excludes
norms, embeddings, routers, and configured high-precision tensors.

Low-precision parameter gather is not enabled in the reference recipe, so a
higher-precision master weight copy remains present during training.

### 4. NVFP4 (Blackwell)

NVFP4 stores E2M1 values in 16-value blocks with an E4M3 scale for each block
and an outer FP32 scale. The Miles reference recipe combines NVFP4 and BF16
through tensor-level precision configuration.

Activation scaling is computed per token. Gate and up projections are
quantized together so the fused rollout GEMM uses the same outer weight scale.
The trainer, checkpoint converter, and rollout kernels must use the same
scaling contract.

**Hardware:** B200, B300, GB200, or GB300.

Run the Qwen3-30B-A3B reference recipe with paired rollout and training flags:

```bash theme={null}
python scripts/run_qwen3_30b_a3b.py prepare \
  --hardware B200 \
  --rollout-nvfp4 \
  --train-nvfp4

python scripts/run_qwen3_30b_a3b.py execute \
  --hardware B200 \
  --rollout-nvfp4 \
  --train-nvfp4
```

The launcher selects per-token activation scaling, disables the incompatible
2D quantization, RHT, and stochastic-rounding paths, and loads the matching
tensor-level precision configuration. NVFP4 applies only to the routed-expert
FC1 and FC2 GEMMs; shared experts and all other unmatched tensors stay in BF16
across checkpoint conversion, training, and live weight export. Rollout uses a
BF16 KV cache.

The training path uses:

```bash theme={null}
--transformer-impl transformer_engine
--bf16
--fp4-format e2m1
--fp4-recipe nvfp4
```

Note the `--fp4-` prefix: NVFP4 has its own format and recipe flags rather than
reusing the `--fp8-` pair. The launcher also adds `--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`.

The base NVFP4 recipe uses **high-precision backward**: the forward pass uses
NVFP4 while the BF16 backward GEMMs consume the original BF16 operands.

![NVFP4 with high-precision backward](https://mintcdn.com/radixark/Cy4ImdPUZ2E-qOic/assets/images/low-precision/nvfp4-high-precision-backward.png?fit=max\&auto=format\&n=Cy4ImdPUZ2E-qOic\&q=85\&s=4dee9789d03b2efe660b29d114490cc3)

The base recipe settings used by the launcher are:

```bash theme={null}
NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1
NVTE_NVFP4_DISABLE_RHT=1
NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1
NVTE_NVFP4_ROW_SCALED_ACTIVATION=1
NVTE_BACKWARD_OVERRIDE=high_precision
NVTE_USE_FAST_MATH=0
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1
TRTLLM_DISABLE_FP4_QUANT_FAST_MATH=1
FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH=1
```

Convert a checkpoint outside the launcher with the same `NVTE_*` recipe
environment:

```bash theme={null}
python tools/convert_hf_to_nvfp4.py \
  --model-dir /root/models/Qwen3-30B-A3B \
  --save-dir /root/models/Qwen3-30B-A3B-NVFP4
```

Checkpoint conversion and live weight export quantize only routed experts in
the main language decoder. MTP weights (including appended decoder layers),
vision/audio components, shared experts, and dense layers keep their source
precision. The HF converter uses the language `num_hidden_layers` from
`text_config` when present, otherwise the top-level config, and recognizes
`model.layers`, `language_model.model.layers`, `model.language_model.layers`,
and `language_model.layers`. First/last-layer BF16 controls apply to those
language layers. Unsupported namespaces remain unquantized. Training-time
precision is still controlled separately by `--te-precision-config-file`.

#### Advanced: dequantized backward

Dequantized backward keeps the backward GEMMs in BF16 but uses BF16
dequantizations of the NVFP4 operands produced during the forward pass.
See the humans& discussion of
[gradient stability](https://humansand.ai/blog/nvfp4-rl#improving-gradient-stability)
for the motivation and ablations.

![NVFP4 with dequantized backward](https://mintcdn.com/radixark/Cy4ImdPUZ2E-qOic/assets/images/low-precision/nvfp4-dequantized-backward.png?fit=max\&auto=format\&n=Cy4ImdPUZ2E-qOic\&q=85\&s=5a03566ad1d59884b672c745faf712b7)

Select this mode in the environment used to launch the job:

```bash theme={null}
export NVTE_BACKWARD_OVERRIDE=dequantized
```

#### Advanced: four-over-six

[Four Over Six: More Accurate NVFP4 Quantization with Adaptive Block Scaling](https://arxiv.org/abs/2512.02010)
optionally chooses, for each NVFP4 block, whether mapping the largest FP4
magnitude to 4 or 6 produces less quantization error.
See the humans&
[four-over-six analysis](https://humansand.ai/blog/nvfp4-rl#four-over-six-for-rl-weights-and-activations)
for the RL recipe discussion.

An example setting is:

```bash theme={null}
export NVTE_NVFP4_4OVER6=all
export FLASHINFER_NVFP4_4OVER6=1
export NVTE_NVFP4_4OVER6_E4M3_USE_256=all
export FLASHINFER_NVFP4_4OVER6_E4M3_USE_256=1
export NVTE_NVFP4_4OVER6_ERR_MODE=MSE
export FLASHINFER_NVFP4_4OVER6_ERR_MODE=MSE
export NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=0
export FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH=0
```

## Fine-grained BF16 exceptions

Miles supports per-layer precision configuration across checkpoint
conversion, Megatron training, SGLang rollout, and live weight export.

| Control | Purpose |
| - | - |
| `--num-layers-at-start-in-bf16` | Keep an initial layer range in BF16. |
| `--num-layers-at-end-in-bf16` | Keep a final layer range in BF16. |
| `--first-last-layers-bf16` | Apply the corresponding Megatron training rule. |
| `--extra-high-precision-layers-hf` | Exclude matching Hugging Face tensor names during conversion and rollout export. |
| `--extra-high-precision-layers-megatron` | Exclude matching Megatron tensor names during training and live export. |
| `--te-precision-config-file` | Select Transformer Engine recipes by Megatron tensor name. |

Use equivalent Hugging Face and Megatron name matchers for exceptions beyond a
recipe's default scope. NVFP4 checkpoint conversion and live export exclude
shared experts, MTP, and non-language components automatically. Common
additional exceptions include final transformer layers and MLA projections
whose contraction axis does not match a one-dimensional scaling layout.

## Hardware support

| GPU | BF16 | FP8 block-wise | MXFP8 | NVFP4 |
| - | - | - | - | - |
| NVIDIA H100 / H200 | ✅ | ✅ | ✗ | ✗ |
| NVIDIA B200 / B300 / GB200 / GB300 | ✅ | ✅ | ✅ | ✅ |
| NVIDIA A100 | ✅ | ✗ | ✗ | ✗ |
| AMD MI350X / MI355X | ✅ | ✅ | ✗ | ✗ |
| AMD MI300X / MI325X | ✅ | ✗ | ✗ | ✗ |

## When BF16 is enough

* Dense models below \~30 B.
* A100 hardware (no FP8 GEMM).
* AMD MI300X / MI325X (no FP8 block-wise path).
* Bring-up of a new model architecture, where clean BF16 numerics simplify
  debugging.

## Further reading

* [Blackwell MXFP8 and NVFP4 RL roadmap](https://github.com/radixark/miles/issues/615)
* [Towards Blackwell-Native 8-bit and 4-bit RL: End-to-End MXFP8 and NVFP4 RL in Miles](https://www.lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl)
* [The 4-bitter Lesson](https://humansand.ai/blog/nvfp4-rl)
* [Faster-than-light NVFP4 decode—with precision](https://humansand.ai/blog/faster-than-light-nvfp4-decode)


# Rollout Routing Replay (R3)
Source: https://miles.radixark.com/docs/advanced/miles-router

Capture expert routing during inference and replay it during training to stabilize RL.

Rollout Routing Replay (R3) records the expert routing decisions made during
inference and replays them during training, producing bit-identical expert
allocation between rollout and training.

## Why MoE RL is unstable without R3

For each token, an MoE router picks `top-k` experts. The choice depends on the
input through a soft router and a top-k operation. In production the router is a
learned `nn.Linear` with non-deterministic kernels and FP8 quantization, so tiny
numerical differences flip routes at the per-layer, per-token level.

An example without R3:

* Rollout selects experts `{2, 7}` for token 314.
* Training (with the same weights but slightly different precision and kernels)
  selects experts `{2, 8}` for token 314.
* The gradient is computed against the wrong expert. Multiplied by hundreds of
  layers, tens of thousands of tokens, and thousands of training steps, the
  policy diverges.

With R3, the trainer replays the rollout router's expert assignments verbatim,
so numerical noise no longer flips routes.

## How R3 wires up

**SGLang side.** Miles enables `enable_return_routed_experts` automatically when
`--use-rollout-routing-replay` is on. SGLang then includes `routed_experts`
in `meta_info` of each response, with shape
`(seq_len - 1, num_layers, top_k)` and dtype `int32`.

**Miles side.** Enable R3 with:

```bash theme={null}
--use-rollout-routing-replay
```

Rollout sends `return_routed_experts=true` in each request and stores the
results in `sample.rollout_routed_experts` (`miles/utils/types.py`). The
trainer pushes the arrays through `RoutingReplayManager`
(`miles/utils/replay_base.py`), and `replay_utils.py` plugs them into the
forward pass so recorded routes are used instead of recomputed ones.

## Memory cost

`(num_tokens - 1) × num_layers × top_k × 4 bytes` (int32 per element, see
`miles/utils/types.py`). For a 32K-token sequence, 60 layers, and
`top_k = 8`, that is roughly 60 MB per sample of routing metadata.

## When R3 is not required

* The model is dense.
* The `--advantage-estimator` is `reinforce_plus_plus` and `--use-tis` already
  masks the off-policy term.

For MoE training with `--advantage-estimator grpo`, current recipes turn R3
on (for example `scripts/run_glm47_flash.py`).

## References

* R3 paper: [arXiv 2510.11370](https://arxiv.org/pdf/2510.11370).
* Routing replay: [arXiv 2507.18071](https://arxiv.org/abs/2507.18071).


# Mooncake Rollout Data Transfer
Source: https://miles.radixark.com/docs/advanced/mooncake-rollout-transfer

Configure Mooncake Store as the rollout object-store backend in Miles.

Miles uses Ray's object store by default to pass rollout data from rollout workers to
trainers. Mooncake Store can replace Ray for this handoff. Both `train.py` and
`train_async.py` support the Mooncake backend, and the trainer receives the same
rollout dictionary with either backend.

Rollout-data transfer and model-weight transfer are separate settings:

* `--object-store-backend` selects the rollout object store.
* `--update-weight-transfer-mode` selects the model-weight transfer path.

## Requirements

Before starting a Miles job:

* Run the same Miles revision and Mooncake version on every Ray node and on the host
  that runs `mooncake_master`.
* Start `mooncake_master`, or provide a Mooncake HA endpoint, and make the endpoint
  reachable from every node. Miles connects as a client and does not manage the
  endpoint lifecycle.
* Use routable data-network addresses for Ray and Mooncake clients.
* Reserve enough host memory for `global_segment_size` and `local_buffer_size`.
* For RDMA, expose the RDMA device to the Miles environment, allow memory locking,
  and set the local device name on each node. Leaving it to automatic selection can
  spread host-memory transfers across every HCA and run far slower
  ([kvcache-ai/Mooncake#4322](https://github.com/kvcache-ai/Mooncake/issues/4322)).

This backend imports the structured-object API from `mooncake.structured_object_store`
(`FieldSchema`, `MooncakeBundleTransfer`, `export_ref`, `import_ref`), which Mooncake
provides from 0.3.12.post1 on; earlier releases ship the module without these names. The
Miles CUDA images use the Mooncake package that ships in the SGLang base image, and the
ROCm images add Mooncake's Python files from `mooncake-transfer-engine-non-cuda` (see
`docker/Dockerfile.rocm`). For a custom environment, install a Mooncake package that
matches the CUDA runtime and is compatible with the Miles revision. Follow the
[Mooncake installation guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html)
for current package names and supported platforms. See
[Versions and Images](/docs/developer/versions) for how the images pick their Mooncake.

## Configure the backend

Choose the transfer protocol before starting Ray. TCP works on any routable data
network. RDMA also requires a local RDMA device on every node; device names may
differ between nodes.

If Ray and the Mooncake endpoint are already running, set the protocol and endpoint:

```bash theme={null}
export MOONCAKE_PROTOCOL="<tcp-or-rdma>"
export MOONCAKE_MASTER_ADDR="<mooncake-master-host>:50051"

# RDMA only. Set this to the device on the current node before starting Ray.
# export MOONCAKE_DEVICE="<local-rdma-device>"
```

Then add these options to an existing Miles training command:

```bash theme={null}
--object-store-backend mooncake \
--mooncake-store-init-kwargs \
  "{\"protocol\":\"${MOONCAKE_PROTOCOL}\",\"master_server_address\":\"${MOONCAKE_MASTER_ADDR}\"}"
```

Use a data-network address that every client can reach. The JSON value must remain
one shell argument. `MOONCAKE_MASTER_ADDR` above is only a shell variable for these
examples; the environment variable Miles itself falls back to is `MOONCAKE_MASTER`.

Launch scripts and the Mooncake end-to-end tests assemble both pieces with
`start_mooncake_master()` and `get_mooncake_object_store_args()` from
`miles/utils/external_utils/command_utils.py`. The walkthrough below spells them out
instead, so every setting stays visible.

## Two-node walkthrough

The example below runs three synchronous rollout and training iterations with FSDP.
It uses one eight-GPU node for rollout and one eight-GPU node for training. Set the
variables to match your cluster, then complete the following steps in order.

| Node | Address variable | Role |
| - | - | - |
| Head | `HEAD_IP` | Ray head, Mooncake master, rollout manager, rollout engines |
| Worker | `WORKER_IP` | Training actors and Mooncake Store segment |

Both nodes must use the same Python environment, Miles revision, and Mooncake
version. The model and dataset paths referenced by the job must be available to the
processes that use them.

### 1. Choose the protocol and set node addresses

On the head node, select TCP or RDMA and set its data-network address:

```bash theme={null}
export HEAD_IP="<head-data-network-ip>"
export MOONCAKE_PROTOCOL="<tcp-or-rdma>"

# RDMA only. Use the device attached to HEAD_IP.
# export MOONCAKE_DEVICE="<head-rdma-device>"
```

On the worker node, select the same protocol and set both node addresses:

```bash theme={null}
export HEAD_IP="<head-data-network-ip>"
export WORKER_IP="<worker-data-network-ip>"
export MOONCAKE_PROTOCOL="<tcp-or-rdma>"

# RDMA only. Use the device attached to WORKER_IP.
# export MOONCAKE_DEVICE="<worker-rdma-device>"
```

For RDMA, set `MOONCAKE_DEVICE` before starting Ray so that Ray workers inherit the
node-local device setting.

### 2. Start Mooncake and Ray on the head

Activate the Miles environment first, then run:

```bash theme={null}
export MOONCAKE_LOG="${MOONCAKE_LOG:-mooncake_master.log}"

mooncake_master --rpc_port 50051 --metrics_port 50052 \
  >"${MOONCAKE_LOG}" 2>&1 &

ray stop --force
ray start --head \
  --node-ip-address "${HEAD_IP}" \
  --num-gpus 8 \
  --disable-usage-stats \
  --dashboard-host 0.0.0.0 \
  --dashboard-port 8265
```

### 3. Join the worker

Activate the same Miles environment on the worker, then run:

```bash theme={null}
ray stop --force
ray start \
  --address="${HEAD_IP}:6379" \
  --node-ip-address "${WORKER_IP}" \
  --num-gpus 8 \
  --disable-usage-stats
```

Mooncake uses each Ray node IP as `local_hostname` by default. Starting Ray with the
data-network addresses above therefore keeps both systems on the same network.

### 4. Submit training from the head

Set the Store configuration and the repository, model, and dataset paths before
submitting the job. The 2 GiB values are suitable for this small example; production
jobs should size them for live rollout data, replicas, and concurrent transfers.

```bash theme={null}
export MOONCAKE_MASTER_ADDR="${HEAD_IP}:50051"
export MOONCAKE_STORE_INIT_KWARGS="{\"protocol\":\"${MOONCAKE_PROTOCOL}\",\"master_server_address\":\"${MOONCAKE_MASTER_ADDR}\",\"global_segment_size\":\"2gb\",\"local_buffer_size\":\"2gb\",\"chunk_bytes\":67108864}"

export MILES_HOME="<path-to-miles>"
export MODEL_PATH="<path-to-Qwen2.5-0.5B-Instruct>"
export DATASET_PATH="<path-to-gsm8k-train.parquet>"
export RAY_DASHBOARD_ADDR="${RAY_DASHBOARD_ADDR:-http://127.0.0.1:8265}"

RUNTIME_ENV_JSON='{
  "env_vars": {
    "CUDA_DEVICE_MAX_CONNECTIONS": "1"
  }
}'

ray job submit --address="${RAY_DASHBOARD_ADDR}" \
  --runtime-env-json="${RUNTIME_ENV_JSON}" \
  -- python3 "${MILES_HOME}/train.py" \
  --train-backend fsdp \
  --hf-checkpoint "${MODEL_PATH}" \
  --prompt-data "${DATASET_PATH}" \
  --input-key messages \
  --label-key label \
  --apply-chat-template \
  --rollout-shuffle \
  --rm-type math \
  --num-rollout 3 \
  --rollout-batch-size 8 \
  --n-samples-per-prompt 2 \
  --rollout-max-response-len 64 \
  --rollout-temperature 0.8 \
  --global-batch-size 16 \
  --advantage-estimator grpo \
  --eps-clip 0.2 \
  --optimizer adam \
  --lr 1e-6 \
  --lr-decay-style constant \
  --weight-decay 0.1 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  --rollout-num-gpus 8 \
  --rollout-num-gpus-per-engine 1 \
  --num-gpus-per-node 8 \
  --sglang-mem-fraction-static 0.7 \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 8 \
  --pin-rollout-manager-to-head \
  --update-weight-transfer-mode broadcast \
  --object-store-backend mooncake \
  --mooncake-store-init-kwargs "${MOONCAKE_STORE_INIT_KWARGS}" \
  --wandb-mode disabled
```

For fully asynchronous training, keep the same Mooncake options and use the normal
Miles async entrypoint and flags:

```diff theme={null}
- python3 "${MILES_HOME}/train.py" ...
+ python3 "${MILES_HOME}/train_async.py" ...
+   --fully-async
```

See [Fully Async Rollout](/docs/user-guide/fully-async) for the remaining async settings.

## Configuration reference

`--mooncake-store-init-kwargs` accepts the following fields:

| Field | Default source | Description |
| - | - | - |
| `master_server_address` | `MOONCAKE_MASTER` | Mooncake master address or HA endpoint. |
| `local_hostname` | `MOONCAKE_LOCAL_HOSTNAME`, then Ray node IP | Data-network address advertised by this client. |
| `metadata_server` | `MOONCAKE_TE_META_DATA_SERVER`, then `P2PHANDSHAKE` | Transfer Engine metadata endpoint. |
| `protocol` | `MOONCAKE_PROTOCOL`, then `rdma` | Transfer protocol, normally `tcp` or `rdma`. |
| `device_name` | `MOONCAKE_DEVICE` | Local RDMA device. Prefer the environment variable when device names differ by node. |
| `global_segment_size` | `MOONCAKE_GLOBAL_SEGMENT_SIZE`, then `8gb` | Store capacity contributed by an eligible client. |
| `local_buffer_size` | `MOONCAKE_LOCAL_BUFFER_SIZE`, then `32gb` | Local transfer and staging capacity. |
| `chunk_bytes` | Mooncake default | Structured-object PUT chunk size. |

Values supplied in `--mooncake-store-init-kwargs` take precedence over environment
variables. Use data-network addresses in multi-node jobs; loopback addresses work
only when all clients run on one node.

Not every client contributes a segment: the rollout manager never does, and among
training actors one per node does. Size `global_segment_size` for that per-node
contribution rather than for the total client count.

Mooncake stores one replica by default. Request additional replicas with:

```bash theme={null}
--mooncake-replica-num 2
```

Additional replicas require additional Store capacity.


# On-Policy Distillation
Source: https://miles.radixark.com/docs/advanced/on-policy-distillation

Train a student on its own rollouts with a teacher's token-level probabilities as a reverse-KL signal, composable with GRPO, PPO, and other estimators.

On-policy distillation (OPD) trains a student model on its own rollouts while using a teacher model's token-level probabilities as the distillation signal. In Miles, the teacher signal is converted into a per-token reverse-KL penalty and applied after the selected RL advantage estimator has produced token advantages. This lets the same OPD recipe compose with GRPO, PPO, REINFORCE++, GSPO, and other estimators.

## Key Arguments

| Argument | Description |
| - | - |
| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. |
| `--opd-type` | Type of OPD: `sglang` or `megatron`. Required when `--use-opd` is set. |
| `--opd-kl-coef` | OPD KL penalty coefficient (default: 1.0). Controls the weight of the distillation signal relative to the RL advantage. |
| `--opd-log-prob-top-k` | Number of top-k tokens retained for the Rethinking OPD token reward. `0` uses sampled-token OPD; `16` matches the paper recipe default. |
| `--opd-top-k-strategy` | Top-k token set strategy: `only-student`, `only-teacher`, `intersection`, `union`, or `xor`. |
| `--opd-reward-weight-mode` | Weighting scheme for top-k rewards: `student_p`, `teacher_p`, or `none`. |
| `--opd-teacher-urls` | Optional multi-teacher routing map (`NAME=URL` pairs, SGLang mode only). Routes each sample to a teacher by `sample.metadata[--opd-teacher-key]`; reserved name `default` is the fallback. Unset = single teacher at `--rm-url`. |
| `--opd-teacher-key` | Metadata key holding the teacher name for routing (default: `opd_teacher`). |
| `--opd-teacher-load` | Path to teacher Megatron checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=sglang`. |
| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. |

## How It Works

OPD modifies the advantage computation by subtracting a KL penalty term that encourages the student to match the teacher's output distribution:

$$
\hat{A}_t = A_t - \lambda_{\text{opd}} \cdot D_{\text{KL}}(P_{\text{student}} \| P_{\text{teacher}})_t
$$

Where $A_t$ is the original advantage from the base estimator (e.g., GRPO), $\lambda_{\text{opd}}$ is `--opd-kl-coef`, and $D_{\text{KL}}$ is the token-level reverse KL divergence.

The implementation follows the additive OPD training recipe described in the [Thinking Machines OPD blog](https://thinkingmachines.ai/blog/on-policy-distillation/), with an additional SGLang top-k reward mode from [Rethinking On-Policy Distillation](https://arxiv.org/abs/2604.13016).

## Rethinking OPD Top-K Reward

SGLang OPD supports the top-k token reward recipe from [Rethinking On-Policy Distillation](https://arxiv.org/abs/2604.13016). Set `--opd-log-prob-top-k` above zero to request student rollout top-logprobs, score the same sequence with the teacher, and aggregate a weighted reverse-KL estimate over a selected token set at each response position.

The token set is controlled by `--opd-top-k-strategy`:

| Strategy | Token set |
| - | - |
| `only-student` | Student top-k tokens, with teacher logprobs queried for those IDs. |
| `only-teacher` | Teacher top-k tokens, with student logprobs queried for those IDs. |
| `intersection` | Tokens appearing in both top-k sets. |
| `union` | Tokens appearing in either top-k set, with duplicates removed. |
| `xor` | Tokens appearing in exactly one top-k set. |

`--opd-reward-weight-mode` controls whether each selected token is weighted by student probability, teacher probability, or uniformly. For compatibility, `--opd-log-prob-top-k=0` keeps the original sampled-token OPD path.

## Two Teacher Modes

### SGLang Mode (`--opd-type sglang`)

The teacher runs on an external SGLang server. Teacher log-probs are obtained during the rollout phase.

**When to use**: The teacher has a different architecture from the student, or the teacher is too large to load alongside the training model.

**How it works**:

1. An external SGLang server runs the teacher model.
2. During rollout, the custom reward function (`miles.rollout.on_policy_distillation.reward_func`) sends each sample to the teacher server to obtain token-level log-probs.
3. With `--opd-log-prob-top-k=0`, the custom post-processing function trims sampled-token teacher log-probs to the response span and stores them in `sample.teacher_log_probs`.
4. With `--opd-log-prob-top-k>0`, it computes the Rethinking OPD weighted top-k reverse-KL estimate and stores it in `sample.opd_reverse_kl`.
5. During training, the stored OPD penalty is subtracted from the selected estimator's advantages.

**Configuration**:

```bash theme={null}
--use-opd
--opd-type sglang
--opd-kl-coef 1.0
--opd-log-prob-top-k 16
--opd-top-k-strategy only-student
--opd-reward-weight-mode student_p
--custom-rm-path miles.rollout.on_policy_distillation.reward_func
--custom-reward-post-process-path miles.rollout.on_policy_distillation.post_process_rewards
--rm-url http://<TEACHER_IP>:<TEACHER_PORT>/generate
```

### Multi-Teacher Routing (SGLang mode only)

`--opd-teacher-urls` routes each sample to a task-specific teacher, e.g. a math
specialist for math prompts and a code specialist for code prompts. Each sample
is still scored by exactly one teacher, so scoring cost is identical to
single-teacher OPD and the loss is unchanged — `teacher_log_probs` /
`opd_reverse_kl` are per-sample and do not care which teacher produced them.

**How it works**:

1. Tag each prompt with a teacher name in its dataset metadata column
   (read via `--metadata-key`, default `metadata`):
   ```json theme={null}
   {"prompt": "...", "metadata": {"opd_teacher": "math"}}
   {"prompt": "...", "metadata": {"opd_teacher": "code"}}
   ```
2. Map names to teacher endpoints with `--opd-teacher-urls NAME=URL ...`.
   The reserved name `default` is the fallback for samples whose name is
   missing or unknown; without a `default` entry such samples raise an error
   (failing loudly beats silently distilling from the wrong teacher).
3. `reward_func` resolves the URL per sample; everything downstream is
   unchanged.

**Configuration** (on top of the SGLang-mode flags above; `--rm-url` is ignored
when the routing map is set):

```bash theme={null}
--opd-teacher-urls math=http://<H1>:<P1>/generate code=http://<H2>:<P2>/generate default=http://<H1>:<P1>/generate
--opd-teacher-key opd_teacher   # metadata key holding the teacher name (default)
```

> **Notes**: All teachers must share the student's tokenizer — scoring sends
> `input_ids` and gathers per-token-id log-probs. Works with both the
> sampled-token path and the top-k path (student-side scoring still goes to the
> student router). For throughput, point multiple names (or one name backed by
> an sglang router) at replicas; `--opd-teacher-urls` is for *different*
> teachers, not load balancing.

### Megatron Mode (`--opd-type megatron`)

The teacher model is loaded directly into Megatron via `--opd-teacher-load`. Teacher log-probs are computed during the training forward pass.

**When to use**: The teacher has the same architecture as the student/reference model and fits in GPU memory.

**How it works**:

1. The teacher model is loaded as an additional Megatron model during initialization.
2. During the training forward pass, the teacher model computes log-probs for each sample.
3. The KL penalty is computed inline and applied to advantages.

**Configuration**:

```bash theme={null}
--use-opd
--opd-type megatron
--opd-kl-coef 1.0
--opd-teacher-load /path/to/teacher_torch_dist
```

> **Note**: The teacher checkpoint must be in Megatron format (`torch_dist` or `torch`). You can convert from HuggingFace format using `tools/convert_hf_to_torch_dist.py`.

## Running the Examples

Complete example scripts are provided in `examples/on_policy_distillation/`:

### SGLang Teacher

```bash theme={null}
# 1. Download models and data
hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B
hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k

# 2. Convert student model
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3-8B)" || 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/Qwen3-8B \
    --save /root/Qwen3-8B_torch_dist

# 3. Run
bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
```

### Megatron Teacher

```bash theme={null}
# 1. Convert both student and teacher models to Megatron format
# 2. Run
bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
```

## Preliminary Results

Using Qwen3-8B-Base model SFT-ed on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, on-policy distillation with a Qwen3-32B teacher on the remaining data yields:

| | Pass\@1 |
| - | - |
| Qwen3-8B-Base + SFT | 76% |
| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% |

## References

* [Thinking Machines: On-Policy Distillation](https://thinkingmachines.ai/blog/on-policy-distillation/)
* [Rethinking On-Policy Distillation](https://arxiv.org/abs/2604.13016)


# P2P Weight Transfer
Source: https://miles.radixark.com/docs/advanced/p2p-weight-transfer

Direct rank-to-rank weight sync from actor to rollout via RDMA writes.

miles supports P2P (point-to-point) weight transfer between training and rollout engines. By using `--update-weight-transfer-mode p2p`, miles enables more efficient weight transfer from training ranks to rollout engine ranks. More details on the design and implementation can be found in [this issue](https://github.com/radixark/miles/issues/755).

## Usage

To enable P2P weight transfer, add the following flag to your training command:

```
--update-weight-transfer-mode p2p
```

## How It Works

The default weight transfer mode in miles is `broadcast`: after training, updated weights are broadcast via NCCL to all rollout engine ranks. This works but does not fully utilize the available bandwidth, as redundant copies of the same weights are transferred to multiple ranks.

P2P mode addresses this by having each training rank transfer only the specific weight shards required by its target rollout engine rank(s), writing them directly to remote memory without redundant copies. The key steps are:

1. **Initialization**: Training ranks establish point-to-point connections (via RDMA) to their target rollout engine ranks. Including:
   * Create a transfer plan that maps each training rank to its target rollout rank(s) based on GPU counts and parallelism configuration.
   * Query remote rollout engines for their weight memory registration info (addresses and sizes for RDMA writes).
   * Query each target rank's parallelism config and server args, and build a model replica of its sharding layout without parameter storage, which converts HF weights into the exact bytes that rank's loader would write.

2. **Weight gather**: Megatron TP/EP shards are all-gathered and converted to HF format, same as the broadcast path.

3. **P2P transfer**: Instead of a collective broadcast, each source rank writes bucketed weight tensors directly to the destination rollout rank's memory, in a write-only fashion.

4. **Synchronization**: Once all RDMA writes are confirmed complete, rollout engines increment their weight version and resume generation for the next training step.

## Architecture

Both broadcast and P2P modes share the same bucketed weight-update pipeline in `miles/backends/megatron_utils/update_weight/`. The diagram below shows which components are shared and which are P2P-specific.

### Shared components (broadcast & P2P)

| Component | Description |
| - | - |
| **TP/EP all-gather** | Megatron TP shards are all-gathered within each PP stage; EP shards are gathered per-bucket when the accumulated expert data exceeds `buffer_size * ep_size`. Both modes perform this identically via `common.py`. |
| **Bucketed update** | Weights are not transferred one parameter at a time. Instead, converted tensors are accumulated into a fixed-size buffer (`--update-weight-buffer-size`, default 512 MB). When the buffer is full, the entire bucket is flushed — via NCCL broadcast or RDMA write depending on the mode. This amortizes per-transfer overhead. Non-expert and expert weights use separate buckets. |
| **PP independence** | Each pipeline-parallel stage updates its own weights independently. In broadcast mode, each PP rank has its own NCCL group (`miles-pp_{pp_rank}`). In P2P mode, each PP stage writes its own params to every engine rank (with Megatron Bridge, which gathers PP, each engine rank is written once). No cross-PP synchronization is needed during weight transfer, which is key to scaling. |
| **HF format conversion** | After all-gather, Megatron-format tensors (with custom naming and sharding) are converted to HuggingFace-format names expected by the sglang rollout engine. |

### P2P-specific components

| Component | File | Description |
| - | - | - |
| **Rollout engine rank assignment** | `protocols/utils/rollout_engine_rank_assignment.py` | Decides which rollout engine ranks each training rank sends weights to, for the rollout engines the rollout manager passes in (`engine_gpu_counts`). Senders are the training ranks that hold the same weights: those of one PP stage, or all ranks when PP is gathered (Megatron Bridge). Every (rollout engine, rollout engine rank) gets exactly one sender. First every sender gets one, so all of them send in parallel; the rest go to a sender that already sends the same rollout engine rank, so it prepares those weights once and sends them to several rollout engines. |
| **Model replica** | `protocols/utils/model_replica.py` | An sglang model of the target rank's layout, built on `meta` with no parameter storage and brought, one module at a time on the GPU, to the state the engine loads updates into: its startup postprocess, then the restore every update session begins with. Its `weight_loader` functions re-shard all-gathered HF weights into the exact bytes the target rank expects, loaded into a transfer buffer. `ModelReplicas` keeps one replica per shard layout (the rank's parallelism and the sglang arguments that change its bytes and may differ between server groups, such as quantization and MoE and GEMM backends) for the whole trainer process, so a reconnect rebuilds only the connections to the new engines. Rollout engines whose expert placement a replica cannot reproduce (redundant experts, a non-trivial `init_expert_location`, EPLB, elastic EP, DWDP, KTransformers) are rejected at connect. |
| **Transfer buffers** | `protocols/utils/transfer_buffers.py` | Two pinned host buffers per sender, sized to the larger of `--update-weight-buffer-size` and the largest complete parameter group, including alignment. Parameters sharing HF inputs stay in one group. Buffers are registered once per trainer process and reused only after every write reading them has returned; a write that failed or is still running after `--p2p-transfer-timeout` fails the update there. |
| **Speculative draft** | `p2p.py`, `protocols/utils/model_replica.py` | With MTP training and same-model EAGLE, senders update both `target` and `draft` roles. Each runner receives every weight its own loader loads, as in the engine's own weight updates; a weight both loaders load into storage the draft shares with the target is written twice with the same bytes. Draft replicas use the engine's speculative MoE settings. Multi-layer EAGLE, separate draft checkpoints and a draft MoE runner or a2a backend other than the target's are rejected at connect, since sglang reloads a draft under the target's MoE backends. |
| **Pipelined transfer** | `p2p.py`, `protocols/transports/mooncake.py` | While the writes reading one transfer buffer are in flight, the sender loads the next group of parameters into the other. Each rollout engine has its own write thread, so a stuck engine delays only its own writes. The end of the update waits for the remaining writes and fails it if any failed or is still running after `--p2p-transfer-timeout`. |

The model replica runs the rollout engine's own sglang code in the trainer process: its model construction, startup postprocess and weight loaders. The trainer and the rollout engines therefore run the same sglang build, and their GPUs share an architecture: postprocess chooses kernels and layouts by architecture (TRT-LLM MoE layouts and UE8M0 scales on Blackwell, for example), so a replica on a Hopper trainer would lay out bytes for Hopper while Blackwell engines expect their own.

## Supported Model Architectures

P2P derives HF-to-parameter mappings by probing each replica's `load_weights` at the first update (`protocols/utils/loader_probe.py`). The probe uses real GPU inputs and meta parameters, then caches the mapping for later updates. Models need no separate mapping implementation.

P2P writes **parameters only**. The engine recomputes derived buffers, such as Gemma norm's `weight + 1`, after the update; `--check-weight-update-equal` checks those buffers too.

The following sglang model classes are tested:

| sglang Model Class | Model Family | Example Models |
| - | - | - |
| `Qwen2ForCausalLM` | Qwen2 (dense) | Qwen2.5-0.5B, Qwen2.5-7B |
| `Qwen3ForCausalLM` | Qwen3 (dense) | Qwen3-4B, Qwen3-8B |
| `Qwen3MoeForCausalLM` | Qwen3-MoE | Qwen3-30B-A3B, Qwen3-235B-A22B |
| `Glm4MoeForCausalLM` | GLM4-MoE | GLM-4.5-Air |
| `Glm4MoeLiteForCausalLM` | GLM4-MoE | GLM-4.7-9B-Flash |
| `DeepseekV3ForCausalLM` | DeepSeek V3p2 | GLM-5 (744B-A40B) |
| `DeepseekV3ForCausalLM` | DeepSeek V3 | Kimi-K2 (1T) \*, Kimi-K2.6 (1T) |

> **Note:** All the above models are tested on H100-80GB clusters.
>
> For Kimi-K2, we set `training nodes = inference nodes = 32` to ensure sufficient memory.
> The original Kimi-K2 checkpoint uses a block-quant size of `[128, 128]`, which will trigger errors if `sglang-tp-size = 32`.
> To work around this, we re-quantize to `[64, 64]` and update all affected scale tensors accordingly.

## Validated Models

All models below have been validated with `--check-weight-update-equal` in P2P mode.

| Model | sglang Model Class | Nodes |
| - | - | - |
| Qwen3-4B | `Qwen3ForCausalLM` | 1 |
| GLM-4.7-9B-Flash | `Glm4MoeLiteForCausalLM` | 2 |
| GLM-5\_4layer | `DeepseekV3ForCausalLM` | 2 |
| Qwen3-30B-A3B | `Qwen3MoeForCausalLM` | 4 |
| GLM-4.5-Air | `Glm4MoeForCausalLM` | 8 |

> Enabling `--check-weight-update-equal` for Kimi-K2 is non-trivial due to several issues:
>
> * The user must first dequantize the Kimi-K2 model to BF16 for training, then re-quantize the parameters for weight updating. Meanwhile, the rollout side snapshots the original Kimi-K2 tensors as the reference when enabling `--check-weight-update-equal`. This means the rollout SGLang engine should load the checkpoints processed by the `dequant-requant` pipeline.
> * We use modified checkpoints with block-quant size `[64, 64]`, so any `quant/dequant` code that hard-codes the block-quant size as `[128, 128]` will break.
> * Certain tensors that are only initialized on the rollout side (e.g., `k_scale` / `v_scale`) must be skipped during the weight-check process.
>
> We verified P2P correctness for Kimi-K2 by enabling `--check-weight-update-equal` with hard-coded workarounds for the issues above. The transferred weights were confirmed correct. These hard-coded workarounds are verification-only and will not be merged into the main branch to keep the codebase maintainable.

## Profiling Results

For M source ranks and N target ranks, with source pp size `src_pp` and target ep size `sgl_ep`, the benefit of P2P transfer is approximately:

1. utilizes `M // src_pp` times more source transmission bandwidth.
2. each target rank receives `sgl_ep` times less data.

Thus we expect our solution to scale better, especially on the MoE models.

All profiling is run on H100-80GB clusters with 1GB transfer bucket.
Timing measures after `pause_generation` call returns and before `update_weight` call exits to exclude request queue abortion time.
Table includes steady-state steps 3–12 average. For Kimi-K2, the RDMA (ms) column includes GPU-side post-processing time (`post_load_weights` \~884ms) since this model requires GPU-side weight requantization after RDMA transfer.

Models marked with ★ are MoE architectures, where P2P benefits are most pronounced due to expert-parallel sharding reducing per-target transfer volume.

| Model Family | Model Name | Total Param | sglang Model Class | Train Config | Inference Config | NCCL (ms) | RDMA (ms) | Delta |
| - | - | - | - | - | - | - | - | - |
| GLM4-MoE ★ | GLM-4.7-9B-Flash | 30B(3B) | `Glm4MoeLiteForCausalLM` | TP=4, PP=1, CP=1, EP=8, ETP=1, 1 node | TP=4, EP=4, 1 node | 2,508.6 | 4,229.0 | +68.6% |
| DeepSeek-V3 ★ | GLM-5\_4layer | 4-layer | `DeepseekV3ForCausalLM` | TP=4, PP=1, CP=1, EP=8, ETP=1, 1 node | TP=8, EP=8, 1 node | 732.2 | 1,260.8 | +72.2% |
| Qwen3-MoE ★ | Qwen3-30B-A3B | 30B(3B) | `Qwen3MoeForCausalLM` | TP=4, PP=1, CP=1, EP=8, ETP=1, 2 nodes | TP=8, EP=8, 2 nodes | 2,670.0 | 2,160.2 | **−19.1%** |
| GLM4-MoE ★ | GLM-4.5-Air | 106B(12B) | `Glm4MoeForCausalLM` | TP=1, PP=4, CP=1, EP=8, ETP=1, 4 nodes | TP=8, EP=8, 4 nodes | 5,001.1 | 2,637.2 | **−47.3%** |
| Qwen3-MoE ★ | Qwen3-235B-A22B | 235B(22B) | `Qwen3MoeForCausalLM` | TP=4, PP=4, CP=2, EP=16, ETP=1, 8 nodes | TP=32, EP=32, 8 nodes | 10,753.6 | 3,162.0 | **−70.6%** |
| DeepSeek-V3p2 ★ | GLM-5 | 744B(40B) | `DeepseekV3ForCausalLM` | TP=4, PP=8, CP=2, EP=16, ETP=1, 16 nodes | TP=64, EP=64, 16 nodes | 58,301.5 | 8,479.7 | **−85.5%** |
| DeepSeek-V3 ★ | Kimi-K2 | 1T(64B) | `DeepseekV3ForCausalLM` | TP=8, PP=8, CP=4, EP=32, ETP=1, 32 nodes | TP=32, EP=32, 32 nodes | 53,279.1 | 7,227.3 | **−86.4%** |

![P2P vs NCCL Broadcast Scaling](https://mintcdn.com/radixark/hlz3_cm2V1iMNGwr/assets/images/p2p_vs_nccl_scaling.png?fit=max\&auto=format\&n=hlz3_cm2V1iMNGwr\&q=85\&s=09e21aa6a89d0f421afce2c6f7dfc45a)

\* Kimi-K2 RDMA time includes \~884 ms GPU-side `post_load_weights` requantization on rollout engines.

## Examples

### CI Test (single-node, Qwen3-4B)

The P2P weight transfer E2E test validates correctness on a single node using `Qwen3-4B`:

```python theme={null}
#
# Train: 4 GPUs (TP=2, CP=2)
# Rollout: 4 GPUs (sglang, 2 engines × 2 GPUs each)
# Flags: --update-weight-transfer-mode p2p --check-weight-update-equal
```


# PD Disaggregation
Source: https://miles.radixark.com/docs/advanced/pd-disaggregation

Separate prefill and decode pools so each is sized for its workload.

In a typical SGLang deployment, every engine handles both **prefill** (the
one-shot forward over the prompt) and **decode** (the per-token autoregressive
loop). The two phases have different compute profiles:

This page covers separation inside the rollout deployment. To separate policy
training from rollout inference, see [Disaggregated RL Rollout](/docs/advanced/disaggregated-rollout).

| Phase | Compute pattern | Bottleneck |
| - | - | - |
| Prefill | Long sequence × full batch | FLOPs |
| Decode | One token × batch | Memory bandwidth |

Mixing them in one engine means each engine is sized for the worse case.
PD disaggregation splits them into two pools, each sized for its own workload.

## Enable

```bash theme={null}
--prefill-num-servers 2
```

`--prefill-num-servers` is a Miles-native flag added by
`add_prefill_decode_disaggregation_arguments` in `miles/utils/arguments.py`.
When set, `resolve_sglang_config` in
`miles/backends/sglang_utils/sglang_config.py` dedicates that many SGLang
servers to prefill, with the rest used for decode.

`--prefill-num-servers` is mutually exclusive with the `sglang_config`
attribute (the YAML `server_groups` config), and also cannot be combined
with external rollout engines (`--rollout-external-engine-addrs`); a PD
external fleet declares itself with `--rollout-external-router-pd` instead
(`arguments.py`).

## When PD is worth it

* Long prompts (≥ 4K). Prefill dominates total latency.
* High decode batch sizes. Decode is memory-bandwidth bound.
* MoE models. Decode benefits disproportionately from EP scaling that prefill
  does not need.

For typical post-training with 1–2K prompts, the routing and KV-transfer
overhead can outweigh the speedup. Measure first.

## How requests flow

```mermaid theme={null}
flowchart LR
    P[Prompt] --> Router
    Router -->|prefill request| Prefill[Prefill pool]
    Prefill -->|KV cache + first token| Router
    Router -->|decode request + KV| Decode[Decode pool]
    Decode -->|stream tokens| Router
    Router --> Out[Response]
```

The KV cache produced by the prefill pool is migrated to the decode pool.
SGLang handles the cache transfer when the
[SGLang Model Gateway](https://docs.sglang.io/advanced_features/sgl_model_gateway.html)
fronts the engines (PD support is a feature of the SGLang router).

## Sizing the pools

A starting heuristic:

```
prefill_servers = ceil(rollout_qps × avg_prompt_tokens / single_engine_prefill_tps)
decode_servers  = N - prefill_servers
```

Without measurements, start at `prefill_num_servers = N / 4` and adjust based
on observed queueing:

| Symptom | Action |
| - | - |
| Prefill queue backing up | Increase `prefill_num_servers` |
| Decode latency creeping up | Decrease `prefill_num_servers` |
| Both queues growing | Scale up both pools, then revisit the ratio |

## Pairs with

* [DeepSeek-V3.2 recipe](/docs/models/deepseek/deepseek-v3-2). PD is a clear win at
  671B scale.
* [Speculative decoding](/docs/advanced/speculative-decoding). Both are SGLang-side
  features; pool sizing should account for the verify-batch size when
  speculative is on.

## When PD is not useful

* Short prompts (under \~1K tokens). Prefill is already cheap.
* Single-node setups. Pool boundaries do not help.
* Highly variable workloads. Fixed pool sizes are wasteful.


# Sampling-Support Replay
Source: https://miles.radixark.com/docs/advanced/sampling-support-replay

Train against the same bounded sampling distribution used during rollout.

Top-p and top-k change the distribution that generates a rollout token. They
remove tokens from the vocabulary and renormalize the remaining probability
mass. Recomputing a token's log probability over the full vocabulary therefore
does not reproduce its behavior-policy probability, even when the rollout and
trainer weights are identical.

Sampling-support replay preserves that distribution. For every generated token,
SGLang returns the realized set of token IDs that remained after sampling
filters. Miles transports this ragged support with the sample and applies it to
the actor logits before the softmax:

$$
q_\theta(a_t \mid S_t) =
\frac{\exp(z_\theta(a_t) / T)}
{\sum_{j \in S_t} \exp(z_\theta(j) / T)}.
$$

The support `S_t` is fixed rollout data; `z_θ` comes from the actor being
trained. This gives PPO/GRPO a denominator and numerator in the same sampled
probability space while retaining gradients through the current actor logits.

## Enable replay

```bash theme={null}
ROLLOUT_ARGS+=(
  --rollout-temperature 1.0
  --rollout-top-p 0.95
  --rollout-top-k 64
)
```

Replay is enabled whenever `--rollout-top-p` is below `1` or
`--rollout-top-k` is positive. A finite top-p run must also set a positive
top-k so that the captured support is bounded. Miles asks SGLang to capture its
native sampling support with `return_sampling_mask` and uses the returned
support-normalized log probability as the rollout log probability.

`--rollout-top-k` is the default for rollout requests, not a global upper bound
on request-specific top-k values. SGLang owns support capacity through
`--sglang-sampling-mask-max-tokens` (or the matching server-group override) and
rejects requests whose realized support cannot be represented.

## Correctness requirements

Miles rejects configurations that it cannot replay faithfully:

* Every resolved training request must have top-p, top-k, and temperature. Agentic
  sessions fill omitted values from the sampling defaults registered when the
  session is created.
* Request temperature must match the registered training temperature
  (`--rollout-temperature` for the standard rollout paths).
* Frequency, presence, and repetition penalties, `logit_bias`, and custom logit
  processors are not supported because the trainer does not replay those logit
  transformations.
* `--recompute-logprobs-via-prefill` is incompatible because that path does not
  preserve the per-token support.
* The selected router must preserve `return_sampling_mask`. The Miles router
  forwards raw request bodies. The native SGLang router requires typed chat and
  generate request schemas that declare the field.

The SGLang backend must return one complete support and its normalized log
probability for every sampled token. Miles validates that response contract but
does not add separate restrictions for speculative decoding.

Tool and environment tokens are recorded with singleton support. Evaluation
requests do not capture this training-only metadata.

## Current objective limitation

Replay currently exposes only the support-normalized actor score to the loss.
Reference KL and on-policy distillation also need a full-vocabulary actor score,
so Miles rejects `--use-kl-loss`, nonzero `--kl-coef`, and `--use-opd` with
replay until the loss interface carries both scores.

## Monitor replay

On the first update from identical weights,
`train/train_rollout_logprob_abs_diff` should be near numerical tolerance.
Afterward, interpret it with version-lag and clipping metrics because it also
reflects policy updates and staleness.


# Speculative Decoding
Source: https://miles.radixark.com/docs/advanced/speculative-decoding

Draft + target speculative rollout, with online SFT for MTP-style drafts.

Speculative decoding accelerates rollout by letting a lightweight draft model
generate ahead a few tokens and then verifying them with a single batched
forward of the target model. When the draft is correct the target produces N
tokens for the cost of one forward pass.

## Enabling speculative decoding

For models with built-in MTP (Multi-Token Prediction) layers (GLM-4.7,
DeepSeek-V3, DeepSeek-R1):

```bash theme={null}
SGLANG_ARGS+=(
   --sglang-speculative-algorithm EAGLE
   --sglang-speculative-num-steps 3
   --sglang-speculative-eagle-topk 1
   --sglang-speculative-num-draft-tokens 4
)
```

These are passthrough flags forwarded to SGLang. Miles auto-enables
`enable_draft_weights_cpu_backup` so SGLang can run training without
MTP weights resident on GPU
(`miles/backends/sglang_utils/sglang_engine.py`).

For an externally trained draft model (for example, trained with
[SpecForge](https://docs.sglang.ai/SpecForge/)):

```bash theme={null}
SGLANG_ARGS+=(
   --sglang-speculative-draft-model-path /data/draft_model/
)
```

Full reference: [SGLang speculative decoding docs](https://docs.sglang.ai/advanced_features/speculative_decoding.html).

## Drift over a long RL run

As RL training progresses, the target model's distribution shifts away from the
draft. Fewer draft tokens pass verification, and over many steps speculative
decoding can become a net negative because the wasted draft compute outweighs
the verified speedup.

Miles supports training the draft alongside the target through online MTP-SFT.

## Online SFT for MTP-style draft models

```bash theme={null}
PERF_ARGS+=(
   --mtp-num-layers 1
   --enable-mtp-training
   --mtp-loss-scaling-factor 0.2
)
```

| Flag | Notes |
| - | - |
| `--mtp-num-layers` | MTP layers the model has (1 matches GLM/DeepSeek release defaults). Built only with `--enable-mtp-training`. |
| `--enable-mtp-training` | Build and train the MTP layers alongside the policy (Megatron only). Without it no MTP layer is built, so exported checkpoints carry no MTP weights. |
| `--mtp-loss-scaling-factor` | Weight of the MTP loss in the combined gradient (default `0.2`). |

> [!NOTE]
>
> **Checkpoint must contain MTP weights.** Pass `--mtp-num-layers 1` when running `convert_hf_to_torch_dist.py`.
> Without it the resulting `torch_dist` checkpoint will not contain the MTP
> layer to train.

## External draft model SFT

Training an external (non-MTP) draft model online is not yet supported in
Miles. The current path is to retrain the external draft offline every N
rollouts and reload it.

## Pairs with

* [Low Precision RL](/docs/advanced/low-precision). Draft and target both quantized the
  same way.
* [INT4 QAT](/docs/advanced/int4-qat). A quantized draft is cheaper to verify.
* [R3](/docs/advanced/miles-router). R3 captures routing for the verified tokens emitted
  by the target.

## When to skip

* Rollout-bound on dense models below \~13 B. The verification overhead can
  outweigh the benefit.
* Already at high draft acceptance and the bottleneck is verification compute,
  not generation.

## Reading

* SpecForge: [SGLang docs](https://docs.sglang.ai/SpecForge/).


# Blog
Source: https://miles.radixark.com/docs/blog/index

Miles engineering write-ups on RL systems, post-training recipes, and performance results, published on the LMSYS blog.

Check out Miles blog posts on the [LMSYS blog](https://www.lmsys.org/blog).


# Adding a Sandbox Provider
Source: https://miles.radixark.com/docs/developer/adding-a-sandbox-provider

How a cloud sandbox provider plugs into Miles, and what it takes to add one.

Connectors that run a task in a container (Harbor, OpenEnv) get that container
from a sandbox provider. A provider is **one entry in `PROVIDER_CREDENTIALS`**
(`miles/rollout/agentic/credentials.py`) — no connector code branches on it.
The in-process Harbor path passes `HARBOR_ENV_TYPE` straight to Harbor; the
golden test runs one case per registry entry, and the rollout e2e takes its
provider from `HARBOR_ENV_TYPE`. Setting a provider up as a user is
[Sandbox Providers](/docs/user-guide/sandbox-providers).

## Adding one

1. **The registry entry.** What credential a worker needs, the path-valued
   variable the launcher forwards in place of the secret, the address variables
   that are safe to forward by value, and the SDK. Field semantics are
   documented on `PROVIDER_CREDENTIALS` itself.
2. **A golden episode** on the real platform:
   `python scripts/sandbox_smoke/run.py --connector harbor --backend <name>`.
   This is the only step that proves the platform round trip — image
   resolution, sandbox create, exec, verifier, teardown — and nothing below it
   is claimable without one.
3. **A GPU rollout**: `HARBOR_ENV_TYPE=<name>` with
   `tests/e2e/agentic/test_harbor_rollout.py`. This is the bar for the
   [provider table](/docs/user-guide/sandbox-providers).
4. **A cell in that table**, plus a setup section on the same page if the
   provider needs anything beyond a credential.

Nothing in step 1 needs a test change: the golden episode runs as one case per
registry entry, and the rollout e2e reads the provider from the environment.

## The two tests

| | `tests/e2e/agentic/test_sandbox_golden.py` | `tests/e2e/agentic/test_harbor_rollout.py` |
| - | - | - |
| Proves | the platform round trip | harness, session server, TITO recording, reward |
| Needs | a provider credential | that, plus 2 GPUs and the model |
| Provider axis | one case per registry entry | `HARBOR_ENV_TYPE` |
| In CI | registered, skipped without a credential | registered, `disabled` |

The golden test drives `scripts/sandbox_smoke/run.py`, which stays the entry
point for the combinations CI does not fix — another agent harness, another
task, another connector. Its README covers those axes.


# Architecture Overview
Source: https://miles.radixark.com/docs/developer/architecture

The 30-minute tour of how Miles is organized internally.

A reading guide before you start patching.

## The processes

A Miles run is three kinds of processes wrapped in a Ray cluster:

```mermaid theme={null}
flowchart TB
    subgraph Ray cluster
        subgraph "Trainer (1+ Megatron group)"
            T1[Actor rank 0]
            T2[Actor rank 1]
            T3[Actor rank ...]
        end
        subgraph "Rollout (N SGLang servers + Miles Router)"
            R1[SGLang server 1]
            R2[SGLang server 2]
            MR[Miles Router]
            R1 -. health, route .- MR
            R2 -. health, route .- MR
        end
        D[Data Source<br/>RolloutDataSource]
        T1 <-- weight sync --> R1
        T1 <-- weight sync --> R2
        D --> MR
        MR --> D
    end
```

* **Trainer ranks** — Megatron processes that load `torch_dist` checkpoints and run the
  RL loop.
* **SGLang servers** — independent HTTP services that produce rollouts.
* **Miles Router** — FastAPI proxy that distributes rollout requests, preserves
  metadata (R3), and enforces health checks.
* **Data Source** — Python object owned by the trainer; reads prompt JSONL and acts as
  a buffer between rollout and training.

## The package layout

```text theme={null}
miles/
├── backends/             # one directory per backend, plus what they share
│   ├── megatron_utils/   # Megatron actor, update_weight/ (HF iterator), checkpointing, fp32 markers
│   ├── fsdp_utils/       # FSDP2 actor, adaptations/ per architecture, hf_weight_iterator.py,
│   │                     # checkpoint.py (DCP save/resume)
│   ├── torchtitan_utils/ # torchtitan actor, Trainer config tree, hf_weight_iterator.py
│   ├── sglang_utils/     # SGLang engine wrapper + argument glue
│   └── training_utils/   # what the backends share: loss.py / loss_hub/, ParallelState,
│                         # weight_update/ (engine session, transport protocols, WeightUpdater),
│                         # torch_native/ (the RL step FSDP and torchtitan share, the
│                         # StepRunner seam, the routing-replay surface), log + CI checkers
├── ray/                  # Ray actors, placement groups, train/ and rollout/ groups
├── rollout/
│   ├── sglang_rollout.py # legacy v1 rollout function
│   ├── data_source.py    # buffer + JSONL loader
│   ├── filter_hub/       # built-in filters
│   ├── rm_hub/           # built-in reward types (`--rm-type` dispatch)
│   ├── fully_async_*.py  # queue-backed producer for train_async.py
│   └── inference_rollout/# default class-based rollout
├── router/               # FastAPI proxy + worker load-balancer (router.py)
├── dashboard/            # run dashboard: collector, backend, dump reader
├── true_on_policy/       # true-on-policy contracts and per-model profiles
└── utils/                # arguments.py, async / IO / distributed helpers, audit_utils/
```

The `miles_plugins/` tree sits beside it. Nothing in `miles/` imports it directly: a plugin
is loaded only when a run names its import path in a flag (`--spec`, or one of the
`--custom-*-path` flags). `models/` holds Megatron specs and HF module wrappers, `mbridge/`
per-architecture weight bridges, `megatron_bridge/` the `megatron.bridge` shims, and
`optimizers/` optimizer plugins.

`train.py` and `train_async.py` are the entry points. They are
thin; most logic lives in the modules above.

## A request's life

For a single GRPO iteration:

```mermaid theme={null}
sequenceDiagram
    participant T as Trainer
    participant DS as DataSource
    participant MR as MilesRouter
    participant SG as SGLang
    participant RM as RewardFn

    T->>DS: get_samples(N)
    DS-->>T: prompts
    T->>MR: generate(prompts)
    MR->>SG: dispatch
    SG-->>MR: responses + meta_info
    MR-->>T: samples
    T->>RM: score(samples)
    RM-->>T: rewards
    T->>T: GRPO loss / step
    T->>SG: weight_sync(p2p)
```

This is the sync path. Fully async (`train_async.py --fully-async`) breaks the request
from the trainer loop and uses a continuously-running worker.

## Where common changes go

| You want to … | Edit |
| - | - |
| Add a new RL algorithm | `miles/backends/training_utils/loss/objective.py` and `loss/hub/`, plus the enum in `miles/utils/arguments.py` |
| Add a new built-in reward type | `miles/rollout/rm_hub/` (the `rm_type` dispatch lives in its `__init__.py`) |
| Add a new built-in filter | `miles/rollout/filter_hub/` |
| Support a new architecture on Megatron | `miles_plugins/models/<model>.py` + a bridge in `miles_plugins/mbridge/` |
| Support a new architecture on FSDP | `miles/backends/fsdp_utils/adaptations/specs/<arch>.py` |
| Add a new flag | `miles/utils/arguments.py` |
| Change the engine handshake (pause / flush / announce / resume) | `miles/backends/training_utils/weight_update/session.py` — shared by all backends |
| Change how weights are produced or transported | `miles/backends/training_utils/weight_update/` (transport protocols, shared), each backend's HF weight iterator beside it: `miles/backends/megatron_utils/update_weight/`, `miles/backends/fsdp_utils/hf_weight_iterator.py`, `miles/backends/torchtitan_utils/weight_bridge.py` |
| Add a training backend | a new `miles/backends/<name>_utils/` beside the existing three, a loader + validator in `miles/utils/arguments.py`, and an entry in `_TRAINER_ACTOR_CLASSES` in `miles/ray/specs/train.py` |
| Change rollout buffer | `miles/rollout/data_source.py` |

## Extension points (the right way)

The trainer is plug-in-friendly. Most extensions don't need a code change inside Miles —
just pass a `--something-path my_pkg.thing`. See [Customization](/docs/user-guide/customization)
for the full list.

If you find yourself patching the trainer to make something work, that's a sign we're
missing a hook. Open an issue.

## Tests

```text theme={null}
tests/
├── fast/             # CPU CI only — each test_*.py auto-registers as stage-a-cpu (register_cuda_ci is rejected here)
├── fast-gpu/         # GPU or CPU CI, registered explicitly (register_cuda_ci / register_cpu_ci)
├── ci/               # the suite runner + registry, with their own CPU CI
├── e2e/              # end-to-end (spins up Ray + SGLang); GPU or CPU CI, registered explicitly
├── manual/           # run on request, not discovered by the CI runner
└── snapshots/        # recorded fixtures the launch-script and other snapshot tests assert against
```

CI discovery is location-based. The `tests/fast/` folder may hold **only CPU CI**: every `test_*.py`
there auto-registers as `stage-a-cpu`, so no boilerplate is needed — write a literal `register_cpu_ci(...)`
only to override the defaults, and a `register_cuda_ci` under `tests/fast/` is an error (move the file
to `tests/fast-gpu/`). Every other folder may hold **GPU or CPU CI** and must register each test
explicitly with `register_cpu_ci` / `register_cuda_ci`. The runner collects `tests/fast/`,
`tests/fast-gpu/`, `tests/e2e/`, and `tests/ci/`.

Run `pytest tests/fast` for a quick CPU check (`pytest tests/fast-gpu` if you have a GPU);
run `tests/e2e` before landing anything that touches the train loop.

## Where to look first when reading the code

If you have 30 minutes and want to understand Miles end-to-end:

1. `train.py` — the loop, top-to-bottom.
2. `miles/rollout/sglang_rollout.py:generate_rollout` — how prompts become samples.
3. `miles/backends/training_utils/loss/objective.py` — the loss and advantage computation.
4. `miles/router/router.py` — the FastAPI proxy.
5. `miles/backends/training_utils/weight_update/` — how trained weights reach the engines:
   the session handshake, the transport protocols, and the per-backend HF weight iterators.

That's the spine. Everything else hangs off it.


# Wheels rebuild
Source: https://miles.radixark.com/docs/developer/ci/06-build-wheels

Rebuild and publish router, int4_qat, and Transformer Engine wheels with the independent build-wheels workflow.

`build-wheels.yml` builds prebuilt wheel releases independently of the [Docker image workflow](/docs/developer/ci/02-docker-build). The image workflow detects release changes and rebuilds the consuming image.

## Wheel sets

Three `radixark/miles-wheels` asset sets follow a moving source. Each build writes a `<dist>-source.json` release asset recording the commit it came from, and `build-wheels.yml` rebuilds a set only when its source has moved past that commit:

| set | source | x86\_64 | aarch64 |
| - | - | - | - |
| `sgl-router` (wheel + `sgl-model-gateway` binary) | `radixark/sgl-router-for-miles` `main` | GitHub-hosted `ubuntu-24.04` | GitHub-hosted `ubuntu-24.04-arm` |
| `int4_qat` (`fake_int4_quant_cuda`) | last `main` commit touching `miles/kernels/quant/int4_fake` | `docker-build` runner, in the SGLang base image | by hand |
| `te` (Transformer Engine triplet, `<release>+miles`) | `radixark/TransformerEngine` `miles-main`: NVIDIA's release commit plus the fixes Miles carries | `docker-build` runner | by hand |

## Publish setup

Use the shared `radixark-miles-ci` GitHub App installed on `radixark/miles` and `radixark/miles-wheels`. Grant `Contents: write` for wheel publishing and `Issues: write` plus `Pull requests: write` for [label commands](/docs/developer/ci/01-label). Store its client ID in the `radixark/miles` repository variable `CI_APP_CLIENT_ID` and its private key in the repository secret `CI_APP_PRIVATE_KEY`; both workflows use the same credentials.

The publish job requests only `Contents: write` on `radixark/miles-wheels`. `CI_APP_ENABLED` controls label commands and does not gate wheel publishing. When migrating existing configuration, add the new names before deploying the workflows and retain the old entries until workflows on the default branch and active runs no longer reference them.

## Schedule and manual runs

It runs hourly and by `workflow_dispatch` (`wheels`, `force`, `te_ref`, `keep_superseded`). `docker/Dockerfile` installs the `+miles` TE triplet, so scheduled runs rebuild TE like the other sets, and an upload replaces superseded wheels unless `keep_superseded` is true. There is no self-hosted ARM64 runner, so aarch64 `int4_qat` and `te` are rebuilt by hand with the miles-wheels README commands.

```bash theme={null}
# Rebuild stale TE wheels.
gh workflow run build-wheels.yml --repo radixark/miles -f wheels=te

# Rebuild the router even when its release manifest already matches.
gh workflow run build-wheels.yml --repo radixark/miles -f wheels=sgl-router -f force=true
```

The hourly schedule runs at minute 17. Manual dispatch defaults to `wheels=all`, `force=false`, `te_ref=miles-main`, and `keep_superseded=false`. `wheels` also accepts `sgl-router`, `int4_qat`, or `te` to select one set. `te_ref` selects the Transformer Engine branch or commit.

The `check` job summary lists each selected source commit and which architectures need rebuilding. A matching release manifest leaves the build unselected unless `force=true`.

## Build and publish flow

* **`check`** — resolves only the selected sources and reads the commit their manifests record in the `WHEELS_TAG_X86` / `WHEELS_TAG_ARM64` release. A lookup failure for a selected source or release fails the job. It takes the repository, tags and `SGLANG_IMAGE_TAG` from `docker/Dockerfile`, so assets land in the releases the image installs and compile against its torch.

* **`build-router`**, **`build-int4-x86`**, **`build-te-x86`** — build only the stale sets. The router job also checks that the wheel and the binary forward a `/generate` request. `te` runs NVIDIA's manylinux recipe on the host, then compiles `transformer_engine_torch` inside the SGLang base image (`lmsysorg/sglang:<SGLANG_IMAGE_TAG>`).

* **`publish`** (GitHub-hosted) — when no selected build failed, syncs each arch's assets and manifests into its release with a CI App token (`vars.CI_APP_CLIENT_ID` / `secrets.CI_APP_PRIVATE_KEY`) minted with only `Contents: write` on `radixark/miles-wheels`. The release fingerprint in `check-upstream` then rebuilds the image.

* **`notify-build-failure`** (GitHub-hosted) — after the check, build, and publish jobs finish, posts one failure card to the Miles CI Lark group if any failed. Scheduled and manual runs in `radixark/miles` use the existing repository secret `LARK_WEBHOOK`; the card names the failed jobs and steps and links to the run. Successful, skipped, or cancelled jobs alone do not trigger a notification.

Runs are serialized across the whole workflow to prevent overlapping source snapshots from replacing newer assets. A long manual TE build can therefore delay the hourly router/int4\_qat runs. A failed selected build blocks publishing for both architectures.


# Contributing
Source: https://miles.radixark.com/docs/developer/contributor-guide

Repository layout, the local loop, what enforces code style, what lives in .claude, and how to drive CI from a PR.

Miles is open source under the LICENSE file in the repo. Contributions of every size are
welcome: bug reports, doc fixes, new model recipes, full features.

## Repository layout

```text theme={null}
miles/
├── train.py                  # synchronous entry point
├── train_async.py            # fully-async entry point
├── miles/                    # the package
│   ├── backends/
│   │   ├── megatron_utils/   # Megatron actor, weight sync, checkpointing, fp32 markers
│   │   ├── fsdp_utils/       # FSDP2 actor, per-arch adaptations, MoE kernels
│   │   ├── sglang_utils/     # SGLang engine + argument glue
│   │   └── training_utils/   # loss / GRPO / PPO / GSPO / REINFORCE++, shared ParallelState
│   ├── ray/                  # Ray actors, placement groups, the train and rollout groups
│   ├── rollout/              # rollout functions, data source, filters, fully-async buffer
│   ├── router/               # Miles Router (FastAPI proxy in router.py)
│   ├── dashboard/            # the run dashboard (collector + backend)
│   ├── true_on_policy/       # true-on-policy contracts and model profiles
│   └── utils/                # arguments.py, async / IO / distributed helpers, audit utils
├── miles_plugins/            # opt-in plugins, imported by name from flags
│   ├── models/               # per-architecture Megatron specs and HF wrappers
│   ├── mbridge/              # per-architecture weight bridges
│   ├── megatron_bridge/      # megatron.bridge shims
│   └── optimizers/           # optimizer plugins (NVMe streaming store)
├── scripts/                  # launchers, one per recipe; scripts/models/ holds the architecture flags
├── tools/                    # checkpoint converters, quantizers, profilers
├── tests/                    # fast / fast-gpu / e2e / ci / manual (see Running CI)
├── docker/                   # Dockerfile, Dockerfile.rocm, build.py, patches
├── docs/                     # the source of this site, plus docs/developer/ci internals
└── .claude/                  # rules and skills (see What lives in .claude)
```

## The local loop

```bash theme={null}
git remote add me git@github.com:<your_user>/miles.git
git checkout -b feat/awesome

pip install -e . --no-deps       # editable install, deps come from the image
pre-commit install               # optional, runs the hooks on every commit

pytest tests/fast                # CPU suite, no GPU needed
pre-commit run --all-files       # what the pre-commit CI job runs

git commit -m "feat(rollout): add partial-rollout buffer"
git push me feat/awesome && gh pr create
```

## Code style

Formatting is not a matter of taste here, it is a hook. `.pre-commit-config.yaml` is the
enforcement, the `pre-commit` workflow runs `pre-commit run --all-files` on every PR, and
you can reproduce it exactly with the same command locally.

| Hook | What it enforces |
| - | - |
| `ruff-check --fix` | Pycodestyle errors, Pyflakes, bugbear, pyupgrade. `E402` and `E501` are ignored on purpose |
| `autoflake` | Removes unused imports in place |
| `isort` | Import order, `--profile=black`, first-party is `miles` and `miles_plugins` |
| `black` | Formatting at **line length 119** (`[tool.black]` in `pyproject.toml`) |
| `check-yaml`, `check-case-conflict`, `detect-private-key`, `check-added-large-files` | Hygiene; files cap at 1000 KB |
| `requirements-txt-fixer` | Keeps `requirements.txt` sorted |

Three hooks are Miles-specific bans, each pointing at the API you should use instead:

| Ban | Use instead | Why |
| - | - | - |
| `mpu.get_*` | `get_parallel_state()` from `miles.backends.training_utils.parallel` | The two backends share one `ParallelState`; reading Megatron's `mpu` directly does not work under FSDP |
| `AutoConfig.from_pretrained` / `AutoTokenizer.from_pretrained` | `load_hf_config` / `load_tokenizer` from `miles.utils` | The wrappers centralize trust, caching, and the multi-node file-system race |
| `huggingface-cli` | `hf` | The old CLI is deprecated upstream |

If a commit legitimately needs an exception, the hooks carry `exclude` patterns; extend
those in the same PR rather than disabling a hook.

Beyond formatting, the conventions a reviewer will hold you to live in
[`.claude/rules/general-code-style.md`](https://github.com/radixark/miles/blob/main/.claude/rules/general-code-style.md):
prefer stateless and immutable, keep functions under roughly 100 lines and files under
roughly 1000, initialize derived values once, keep imports at the top, use absolute
imports, prefer keyword arguments where they add clarity. It applies to `miles/**/*.py`,
`scripts/**/*.py`, `tools/**/*.py`, `train.py` and `train_async.py`.

## What lives in `.claude`

The `.claude` directory is how the repo hands its conventions to coding agents, and it is
worth reading even if you never run one, because it is where several rules are written
down exactly once.

**`.claude/rules/`** holds conventions. A file with a `paths:` front matter list applies
to any file matching it; a file without one applies everywhere. `general-code-style.md`
is the one described above, and `AGENTS.md` at the repo root points Codex at it too, so
both agents and humans review against one document. `ci-test-registration.md` holds the
rules for adding a CI test (see [Registering a test](#registering-a-test)), and
`ci-failure-triage.md` what to do when a CI check goes red.

**`.claude/skills/`** holds procedures, one directory per skill with a `SKILL.md`. They
are workflows rather than style rules:

| Skill | What it is for |
| - | - |
| `doc-dev` | Keeping a file and its governing document in sync (see below) |
| `ci-fetch-log` | Pulling complete GitHub Actions logs and diagnosing a failed run from saved evidence |
| `ci-e2e-time-tune` | Recalibrating `register_cuda_ci(est_time=...)` from real run times |
| `pr-ci-targeted` | Running only the GPU test files a PR's diff reaches, with per-file `/rerun-test`, and iterating fixes until they pass |
| `mechanical-refactor-verify` | Reviewing a file split or move by requiring a reproducible transform script |
| `setup-ci-host`, `manage-gh-runners` | Provisioning a CI host and its self-hosted runners |

### The `doc-dev` sentinel

Some files are bound to a document. A `# doc-dev:` line in a file's own comment syntax is
the opt-in: bare, it binds the file's own header block; with a repo-relative path, it also
binds that central document. Editing such a file means updating its documentation in the
same change, and editing the document means finding the files that name it.

```python theme={null}
# doc-dev: docs/developer/ci/02-docker-build.md
```

Current sentinels, so you know when you have walked into one:

| File | Governing document |
| - | - |
| `.github/workflows/pr-test.yml`, `pr-test-rocm.yml` | `docs/developer/ci/00-stage.md`, `docs/developer/ci/01-label.md` |
| `.github/workflows/bot-bump-miles-version.yml`, `bot-cherry-pick.yml`, `release-*.yml` | `docs/developer/ci/04-release.md` |
| `docker/build.py` | `docs/developer/ci/02-docker-build.md` |
| `tests/ci/metric_history/**` | `docs/developer/ci/03-metric-history-gate.md` |

Grep for `doc-dev:` before editing anything under `.github/workflows/` or `docker/`. A
change that lands the code and leaves the document stale is the failure mode this
convention exists to prevent.

## Running CI

### What a PR runs

Two things start automatically: the `pre-commit` workflow on every PR, and `PR Test`
(`.github/workflows/pr-test.yml`) on every PR based on `main`. A PR based on another branch,
such as a stacked PR, runs `PR Test` only with a `run-ci*` label (see
[Labels](/docs/developer/ci/01-label)). `PR Test` resolves a policy and an image, runs the two
CPU stages, and then the GPU stages, which are gated on `stage-a-cpu` succeeding so a
formatting or import error does not burn GPU time. A PR that touches `docker/Dockerfile`,
`docker/build.py`, `docker/verify_transformer_engine.py`, `docker/patch/**` or
`requirements.txt` additionally builds the image first and runs every GPU suite inside it.

### Registering a test

Selection is declared in the test file, never in the workflow YAML. Every `test_*.py`
under `tests/fast/` is auto-registered as a CPU test in `stage-a-cpu`; a file under
`tests/fast-gpu`, `tests/e2e` or `tests/ci` declares a top-level `register_*_ci(...)` call,
or collection fails with `No CI registry found`.

Where the file goes, which stage runs it, how it is declared, which PRs run it, how
`est_time` is measured with `/rerun-test`, and when `disabled=` is allowed are rules,
written once in
[`.claude/rules/ci-test-registration.md`](https://github.com/radixark/miles/blob/main/.claude/rules/ci-test-registration.md).
Telling your failure from an infra one, and reporting either, is in
[`.claude/rules/ci-failure-triage.md`](https://github.com/radixark/miles/blob/main/.claude/rules/ci-failure-triage.md).

### Labels

Labels are how a PR opts into the expensive matrix. A test's `labels=["megatron"]` is
triggered by the GitHub label `run-ci-megatron`: the workflow forwards the labels, Python
strips the `run-ci-` prefix and intersects with each test's list. The canonical set lives
in `tests/ci/labels.py`, and a value outside it is a collection-time error.

| Label | Effect |
| - | - |
| `run-ci-<x>` | Selects tests declaring `<x>` |
| `run-ci-all` | Every enabled tag |
| `nightly` | Nightly cadence: admits `nightly=True` tests, every tag except `long` and `ft-long`, fast-fail off |
| `run-ci-image` | Every tag except `long`, `ft-short`, `ft-long`; for validating an image bump |
| `bypass-fastfail` | Run GPU stages even if `stage-a-cpu` failed, and let each suite continue past the first failure |

If your fork PR sits waiting for approval, that is GitHub holding first-time contributor
runs; any maintainer-applied `run-ci-*` label doubles as the approval.

## PR-description CI tags

Three directives are read out of the PR description itself, one per line, matched at the
start of a line:

| Line in the PR description | Effect |
| - | - |
| `ci-image-tag: <tag>` | Run the GPU suites on `radixark/miles:<tag>` instead of `dev`. Must be a bare tag, not a full image reference |
| `ci-sglang-pr: <ref>` | Check the SGLang tree out at `<ref>` instead of the default `sglang-miles` branch |
| `ci-megatron-pr: <ref>` | Check Megatron-LM out at `<ref>` instead of the default `miles-main` branch |

For the two ref directives, `<ref>` is a branch or commit, and the shorthand `#1234`
resolves to `refs/pull/1234/head`, which is how you test against an unmerged SGLang or
Megatron-LM PR:

```text theme={null}
ci-sglang-pr: #12345
ci-megatron-pr: my-fix-branch
ci-image-tag: dev-202608100600
```

Precedence, when several sources disagree: a `workflow_dispatch` input wins, then the
PR-description line, then the default. `ci-image-tag:` has one more rule: on a PR that
built its own image, the fresh `pr-<number>` tag outranks the directive, so a docker PR
always tests what it just built.

## PR conventions

Commit subjects follow conventional commits, under 70 characters, and the body explains
*why*:

```text theme={null}
feat(rollout): add partial-rollout buffer
fix(megatron): correct fp32 marker on Qwen3.5 A_log
docs: clarify FP8 rationale for MoE
```

Before marking a PR ready for review:

* [ ] `pre-commit run --all-files` is clean.
* [ ] `pytest tests/fast` passes, plus `tests/fast-gpu` if you have a GPU.
* [ ] New behavior has a test.
* [ ] Each new or moved CUDA test has its `CI timing:` line in the PR description, from a
  `/rerun-test` run (see `.claude/rules/ci-test-registration.md`).
* [ ] A new flag appears in [CLI Reference](/docs/user-guide/cli-reference), and
  `python3 train.py --help` still parses.
* [ ] A change to a `doc-dev:` governed file updates its document in the same PR.
* [ ] The PR description carries the CI directives your change needs.

Commenting `/run-lint` on a PR runs the hooks in CI and pushes the autofixes back to your
branch, which is the quick way out of a red `pre-commit` job.

## Issue triage

| Label | Meaning |
| - | - |
| `good first issue` | Self-contained, no system knowledge needed |
| `help wanted` | Community PRs welcome |
| `bug` | Reproducible breakage |
| `enhancement` | Feature request |
| `discussion` | Design conversation, not yet a task |
| `needs-repro` | Not reproducible yet, please add a minimal example |

Comment to claim an issue before you start. For an infra failure or a flake, file the
issue with the fields `.claude/rules/ci-failure-triage.md` lists (job URL, runner name,
suite, log line), so a maintainer can map it to a host.

## Where to ask

* **Quick questions:** the `#miles-rl` channel of the [SGLang Slack](https://slack.sglang.ai).
* **Design discussions:** a GitHub Discussion, or an Issue labeled `discussion`.
* **CI internals:** [Stage](/docs/developer/ci/00-stage) (stages), [Labels](/docs/developer/ci/01-label) (label
  semantics), [Docker build](/docs/developer/ci/02-docker-build) (images), and
  [Metric history & regression gate](/docs/developer/ci/03-metric-history-gate) (metric gate).
  The CI rules for contributors are `.claude/rules/ci-test-registration.md` and
  `.claude/rules/ci-failure-triage.md`.


# Debugging
Source: https://miles.radixark.com/docs/developer/debug

Split a misbehaving run into rollout and training halves, then debug the side that is actually wrong.

When a Miles run misbehaves, the first question is always: **rollout or training?** The
flags on this page exist to answer it by cutting the loop into pieces you can run one at a
time. Once you know which side is wrong, it becomes an ordinary debugging session.

## Cut the loop in half

| Flag | What it does |
| - | - |
| `--debug-rollout-only` | Run generation only. The training backend (Megatron or FSDP) is never initialized. |
| `--debug-train-only` | Run training only. No rollout engines are started; a snapshot-eval fleet (`--eval-num-gpus`) still is when configured. |
| `--save-debug-rollout-data <path>` | Pickle every rollout to `path.format(rollout_id)`. The template must contain `{rollout_id}`, which also carries the eval and policy prefixes. |
| `--load-debug-rollout-data <path>` | Train from those recordings instead of generating. Implies `--debug-train-only`, since it does not start engines. |

The two `--debug-*-only` flags are mutually exclusive and argument validation rejects
setting both.

The pattern worth internalizing:

```bash theme={null}
# 1. capture a few well-formed rollouts, no trainer involved
--debug-rollout-only --save-debug-rollout-data /tmp/rollout_{rollout_id}.pt

# 2. iterate on training with those exact inputs, no engines involved
--load-debug-rollout-data /tmp/rollout_{rollout_id}.pt
```

Step 2 removes sampling randomness from the loop, which is what makes a bisect or an A/B
of a loss change trustworthy.

### Replaying without killing the engines

`--load-debug-rollout-data` gives up the engines, so it cannot tell you whether the
*weights that reached the engines* are right. `--ci-inject-rollout-data-path` is the
variant that keeps them:

```bash theme={null}
--ci-inject-rollout-data-path /tmp/rollout_{rollout_id}.pt \
--ci-inject-rollout-data-start-rollout-id 2 \
--ci-inject-rollout-data-min-match-ratio 0.9
```

From the start rollout id onwards, generation still runs and its output is still compared,
then discarded, and training consumes the recording instead. The comparison is the point:
if the mean response-token match ratio between the fresh generation and the recording
falls below `--ci-inject-rollout-data-min-match-ratio` (default `0.9`), the engines are
holding the wrong weights. Legitimate ulp-level drift only flips the occasional sampled
token, so the ratio stays high when the sync is correct.

## Cut a single step into pieces

| Flag | What it removes |
| - | - |
| `--debug-skip-weight-update` | The actor-to-rollout weight update itself, while keeping the offload and onload schedule around it. Separates "the sync is wrong" from "the schedule around the sync is wrong". |
| `--debug-disable-optimizer` | Optimizer and LR-scheduler construction, and the optimizer step. Rollout, log-prob forward and actor forward/backward still run, so this isolates optimizer-state memory and update behavior from the rest. |
| `--debug-exit-after-rollout <n>` | Everything after rollout `n`. Built for exercising checkpoint resume with consistent scheduler state. |

## Make two runs comparable

`--debug-deterministic-collective` runs the training world on the `det_nccl` backend from
`miles.utils.test_utils.det_process_group`, which folds order-sensitive SUM and AVG
reductions in a fixed tree order. That is what makes two different reduction topologies
bitwise comparable, and it is how the Megatron-versus-FSDP alignment test can assert
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](https://github.com/radixark/miles/tree/main/examples/experimental/reproducibility).

## The assertion harness CI uses

An e2e test in Miles is mostly a normal training run with `--ci-test` added. That flag
turns on a set of in-process assertions, so a violated invariant fails the run where it
happens instead of showing up as a wrong number hours later. You can use the same flags on
your own runs, and each checker has an escape hatch for the case where your change makes it
legitimately not hold.

| Checker | What it asserts | Turn off with |
| - | - | - |
| KL | At step 0, `abs(train/ppo_kl) < 1e-9` and `train/pg_clipfrac < 1e-10`, and `train/kl_loss` is zero. MLA relaxes to `1e-8`; LoRA relaxes because the Megatron-to-HF adapter conversion is not bit-exact | `--ci-disable-kl-checker` |
| Log-probs | At rollout 0, trainer `rollout/log_probs` matches `rollout/ref_log_probs` within `1e-8` (`5e-3` under R3, whose reference does not replay routing); trainer versus engine log-probs within `0.03`; rollout entropy in `(0, 0.7)`. Under `--true-on-policy-mode` the two must be exactly equal | `--ci-disable-logprobs-checker` |
| Weight update | Sets `check_weight_update_equal`, comparing trainer and engine weights after a sync. Skipped automatically under either `--debug-*-only` | `--ci-disable-weight-update-checker` |

Three more take values rather than switching off:

**Accuracy gate.** `--ci-metric-checker-key <key> --ci-metric-checker-threshold <x>`
checks eval metrics when the checker is disposed at the end of the run. By default, at
least one eval must report `key >= x`. Set `--ci-metric-checker-expect-num <n>` to
require exactly `n` eval checks and require all of them to succeed. A run with no eval
always fails.

**Gradient-norm comparison.** `--ci-save-grad-norm <path>` writes the grad norm, and
`--ci-load-grad-norm <path>` asserts a later run matches within `rel_tol=abs_tol=0.03`.
Both accept `{role}`, `{rollout_id}` and `{step_id}` in the path, and only rank 0
participates. This is how the Megatron-versus-FSDP alignment test compares two backends
that cannot share a process.

**Per-layer weight hashes.** `--ci-save-model-hash` and `--ci-check-model-hash` compute
SHA256 over each layer's parameter bytes, including name, shape and dtype, and write one
JSON per rank under `iter_<iteration>/model_hash_tp*_pp*_dp*_cp*.json`. Layer granularity
is deliberate: a mismatch names the layer instead of just saying the model differs.

**Fault injection.** `--ci-fault-hooks` takes a JSON array of fault hook requests, such as
`[{"request_id": "stop", "hook_name": "trainer_controller_step_end", "action": {"kind": "stop_cell", "cell_id": "trainer-engine-actor-00000"}, "rollout_id": 3}]`. A request waits at a named hook
(`trainer_step_before_allreduce`, `trainer_controller_step_start`, `trainer_controller_step_end`, `orchestrator_step_end`, the weight-update
hooks) and runs its action there (`kill_process`, `exit_process`, `segfault_process`, `freeze_process`, `stop_process`, `deadlock_thread`, `stop_cell`, `start_cell`,
`sleep`, `sleep_forever`, `observe`, `api_server_fault`, `restart_rollout_cell`). The rollout
restart runs immediately at `orchestrator_step_end` for its declared rollout, requires rollout fault tolerance
and synchronous single-policy training, and waits for a fresh Running replacement before the
next weight transaction. It is how the fault-tolerance suite kills things on purpose. See
[Fault Tolerance](/docs/advanced/fault-tolerance).

## Aligning precision

The most common class of bug. Walk these checks before anything else.

### Is the first rollout coherent?

If the very first rollout is gibberish:

* **Parameters didn't load.** Megatron logs a clear load line; if it is absent, fix
  `--load` / `--ref-load`.
* **Parameter mapping is wrong.** With `pp_size > 1`, second-stage layer IDs are a common
  offset bug. Dump the parameters in the SGLang model's `load_weights` and compare against
  the checkpoint.
* **SGLang dropped buffers.** Some buffers can be released during the parameter-release
  path; check they are re-loaded after weight sync.
* **Pretrained versus instruct.** If the instruct model of the same architecture works,
  your base model plus chat template combination is wrong.

### Are `log_probs` and `ref_log_probs` equal at step 1?

They should be, because the actor and the reference are the same weights. This is exactly
what the log-probs checker above asserts, so running with `--ci-test` turns the question
into a hard failure. If KL is non-zero:

* **Non-deterministic kernels.** Some Transformer Engine versions need
  `--attention-backend flash` to force deterministic Flash Attention under context
  parallelism.
* **KL below `1e-4`.** Kernel-level jitter, acceptable.
* **KL above 1.** A configuration error; re-check parallelism and precision.
* **About 0.8 per token on an instruct model.** Almost always a chat-template mismatch. Run
  the [chat template verifier](/docs/user-guide/agentic-rollout).

### Is `grad_norm` reasonable?

Step 1 with `num_steps_per_rollout=1` should produce a tiny gradient. If it does not, look
at MoE fusion first (`--moe-permute-fusion`), then at whether the reward is being computed
from the right key (a `--label-key` typo is the usual cause).

### Does step 2 OOM under colocate?

The trainer's offload and reload cycle is colliding with the engine's static memory. Lower
`--sglang-mem-fraction-static` to 0.7 or 0.6. See
[Training Backends](/docs/user-guide/training-backend) for the layout and offload knobs.

## Common kernel pitfalls

| Symptom | Likely culprit |
| - | - |
| Garbled rollout, parameters loaded fine | Chat-template mismatch, or a dropped buffer in SGLang |
| KL non-zero at step 1 | Non-deterministic fused attention; force `--attention-backend flash` |
| MoE training collapses after tens of steps | Routing not preserved; check `--use-rollout-routing-replay` |
| Gradient NaN or Inf | Bad chat template, or activation overflow in FP8 |
| `illegal memory access` in SGLang | OOM in disguise; lower `--sglang-mem-fraction-static` |
| `JSONDecodeError` from inductor | Corrupt compile cache; set `TORCHINDUCTOR_FORCE_DISABLE_CACHES=1` |
| NCCL hang during weight sync | Raise the NCCL timeout and re-run with `NCCL_DEBUG=INFO` |

## Reading logs

| Component | Where |
| - | - |
| Trainer stdout | Wherever you redirected `ray job submit` |
| Ray workers | `~/.ray/session_latest/logs/worker-*.{out,err}` |
| SGLang | Inside the Ray worker logs; verbosity via `--sglang-log-level` |
| NCCL | `NCCL_DEBUG=INFO NCCL_DEBUG_FILE=/tmp/nccl_%h_%p.log` |

Useful environment variables while debugging:

```bash theme={null}
RAY_DEDUP_LOGS=0                       # every rank's line, not one deduplicated line
NCCL_DEBUG=INFO
NCCL_DEBUG_SUBSYS=COLL,P2P
TORCHINDUCTOR_FORCE_DISABLE_CACHES=1
```

The per-step metrics are the real signal, and they are the same numbers the checkers above
assert on: `train/ppo_kl`, `train/pg_clipfrac`, `train/kl_loss`, `rollout/log_probs`,
`rollout/ref_log_probs`, `rollout/rollout_log_probs`, `rollout/entropy`. Trainer steps log
as `step <n>: {...}`, and the rollout side logs one reduced dict per rollout. See
[Monitoring and Logging](/docs/user-guide/monitoring) for the full metric surface.

## When all else fails

* Drop to a tiny model on a known-good recipe (the
  [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
  stdout, and what you have already ruled out.


# Developer Guide
Source: https://miles.radixark.com/docs/developer/index

Contribution conventions, internal architecture, dependency versions, and debugging.

You're here because you want to change Miles, not just use it. This section is the
short tour for new contributors.

> **[Contributing](/docs/developer/contributor-guide)**
>
> Repo layout, what enforces code style, what lives in .claude, and how to drive CI from a PR.

> **[Architecture Overview](/docs/developer/architecture)**
>
> The 30-minute tour of how Miles is organized internally.

> **[Versions and Images](/docs/developer/versions)**
>
> How the miles, SGLang and Megatron-LM trees fit together, and how to bump one.

> **[Debugging](/docs/developer/debug)**
>
> Isolating rollout from training, the debug and CI assertion flags, aligning precision.

> **[Training Backends](/docs/user-guide/training-backend)**
>
> Megatron-LM and FSDP: what each backend owns and where its code lives.

## 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](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.
4. Open a PR. We'll review within \~48h.


# Reconcile Loop
Source: https://miles.radixark.com/docs/developer/reconcile-loop

A minimal, level-triggered controller runtime for Miles, and how it maps onto the Kubernetes Go stack.

## What this is

* A deliberately small port of the [client-go informer stack](https://github.com/kubernetes/client-go/tree/master/tools/cache) plus [controller-runtime](https://github.com/kubernetes-sigs/controller-runtime).
* Tracks pools of processes — SGLang engine cells, trainer cells — that appear, disappear and restart mid-run.
* The alignment record: **every deviation from Go must appear below with a reason.**

## Modules

All under `miles/utils/workers/reconcile/`.

| Ours | Does | Go | Not 1:1 |
| - | - | - | - |
| `k8s_types.py` | The pod fields the loop reads, as pydantic models that validate straight off whatever the client yields | [core/v1 `Pod` types](https://github.com/kubernetes/api/blob/master/core/v1/types.go) | A partial view: only the fields miles reads, and every other one is ignored |
| `k8s_api.py` | LIST/WATCH calls; the only module importing `kubernetes_asyncio`, lazily so the loop stays importable without a Kubernetes backend | [typed client `List` / `Watch`](https://github.com/kubernetes/client-go/blob/master/kubernetes/typed/core/v1/pod.go) | |
| `k8s_reflector.py` | Cursor bookkeeping, relist on cursor rejection | [`cache.Reflector`](https://github.com/kubernetes/client-go/blob/master/tools/cache/reflector.go) | |
| `source_event.py` | Reflector-to-loop wire format: `UpsertEvent`, `DeleteEvent`, and a whole-world `ReplaceEvent` | [`watch.Event`](https://github.com/kubernetes/apimachinery/blob/master/pkg/watch/watch.go) + [`Store`'s write methods](https://github.com/kubernetes/client-go/blob/master/tools/cache/store.go) | Go writes to a store by method call; we send the same operations as values |
| `object_store.py` | Cache, parent index, replace with deletion synthesis | [`cache.Store`](https://github.com/kubernetes/client-go/blob/master/tools/cache/store.go) + [`DeltaFIFO.Replace()`](https://github.com/kubernetes/client-go/blob/master/tools/cache/delta_fifo.go) | Absorbs `Replace()`, the only part of `DeltaFIFO` we keep |
| `work_queue.py` | Insertion-ordered dedup with a wakeup, generic over the key type | [`workqueue`](https://github.com/kubernetes/client-go/blob/master/util/workqueue/queue.go) | |
| `retry_scheduler.py` | Per-key exponential backoff, latest-wins deadlines swept by one poll loop, generic over the key type | [rate limiter](https://github.com/kubernetes/client-go/blob/master/util/workqueue/default_rate_limiters.go) + [delaying queue](https://github.com/kubernetes/client-go/blob/master/util/workqueue/delaying_queue.go) | Absorbs both, because latest-wins is one mechanism, not two |
| `source_stream_driver.py` | Open, sync, reopen the stream; pump events into the store | [informer `Run` / `processLoop`](https://github.com/kubernetes/client-go/blob/master/tools/cache/controller.go) | |
| `loop.py` | Lifecycle, the single worker, resync | [controller-runtime `Controller`](https://github.com/kubernetes-sigs/controller-runtime/blob/main/pkg/internal/controller/controller.go) | |

An empty last column means 1:1. Each entry is the shadow of a **Dropped** / **Replaced** row below — the Go class died with its feature and the remainder landed on a neighbor — so the reason lives with that row.

## Decisions per module

### `k8s_types.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Generated core/v1 types | Name the fields of an object the apiserver owns | **Replaced** by hand-written pydantic models covering only the fields miles reads, with `extra="ignore"` | The store, the reflector and the pod views all hold `Pod`, so a typo is a validation error at the edge instead of an `AttributeError` deep in a view. Modelling the whole of `V1Pod` would only re-type what `kubernetes_asyncio` already ships |
| Two wire spellings of the same field | One parser for both | **Kept** as `from_attributes=True` plus `AliasChoices("pod_ip", "podIP")` on every compound field | `kubernetes_asyncio` deserializes into a snake\_case model unless it has none, and only then leaves the raw JSON, which is camelCase. One model reading both is what keeps the conversion in pydantic rather than in a hand-written field walker |
| An unset field | Tell "absent" from "empty" | **Dropped**; `None` becomes `{}` / `[]` before validation | The client spells an absent `labels` block as `None` and an empty one as `{}`, and no consumer here distinguishes them |

### `k8s_reflector.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Reflector | Move remote changes into a local cache reliably | **Kept**, Kubernetes only | Ray / external-URL backends emit in-process: no cursor, no replay window |
| `watchHandler` per-event metadata failure | One malformed frame must not stop the watch | **Kept**: log, skip, advance past it | Tearing down reconnects at the same cursor, replays the same frame, wedges the watch until expiry |
| Bookmarks | Keep an idle cursor fresh | **Kept** | Free server-side, avoids relists |
| `BackoffManager` around reconnects | Survive a watch that keeps dropping | **Kept**, ours only | Always sending `timeout_seconds` is load-bearing: `kubernetes_asyncio` reads its absence as `watch_forever` and then reconnects and retries a 410 behind our back. Sending it keeps every reopen in `watch()` |
| `IsTooLargeResourceVersion` | A cursor from the future is never satisfied | **Kept** as the code 504, **dropped** as a reason string | A rolled-back backend would otherwise freeze the store forever, and a plain gateway timeout costs one LIST. Go finds `ResourceVersionTooLarge` in `Status.Details.Causes[].Type` ([`isTooLargeResourceVersionError`](https://github.com/kubernetes/client-go/blob/master/tools/cache/reflector.go)), never in `Status.Reason`, so matching it as a reason string never fires |
| Watch frame shape | Read a cursor out of whatever the client yields | **Kept**: an attribute, or a camelCase key when the frame is a dict, both through the `WatchFrame` model | `kubernetes_asyncio` deserializes into a model unless it has none, and only then leaves the raw JSON, which is camelCase. A dict spelling `resource_version` is one of the accepted aliases, since one model reads both shapes |
| Frame payload type | Only some frames carry a pod | **Kept**: `PodWatchEvent.pod` is validated for `ADDED` / `MODIFIED` / `DELETED` only, and a pod frame that does not validate raises | `BOOKMARK` carries a bare `resourceVersion` and `ERROR` carries a `Status`, so validating either as a pod would fail on every frame. A pod frame miles cannot read is the apiserver breaking the contract, not an object to skip: skipping it would leave the store silently short of a pod. The envelope the reflector needs — the cursor, and a `Status`'s `code` and `reason` — is read from every frame by `WatchFrame` instead |
| `ApiException` shape | Tell a dead cursor from a transient failure | **Kept**: `status` against 410 / 504 only | `ApiException.__init__` sets `status` from `http_resp.status` or the int the watch layer passes, so it is always an int and never carries a separate `code` |
| An `ERROR` frame reaching the consumer | A watch that reports failure in-band | **Kept**, though `kubernetes_asyncio` cannot produce it | `Watch.unmarshal_event` raises `ApiException` on `type == "error"` rather than yielding, so the live path is the exception one. Kept because losing a cursor-death signal freezes the store forever and the dependency is unpinned |
| `ListAndWatch` → `Run` → LIST again | Refresh after every watch ends | **Dropped**; reopen WATCH from the cursor | A LIST per timeout dominates an idle reflector's cost, and the cursor is still valid |
| `BackoffUntil` | Keep a relist storm off the apiserver | **Replaced** by one flat `retry_delay` | One reflector, small label-scoped LIST |
| LIST pagination | Huge collections | **Dropped** | Thousands of pods at most |

### `object_store.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Store | Read without hitting the apiserver | **Kept**, a plain `dict` | Single-threaded asyncio: no locks |
| `Replace()` on relist | Deletions missed while disconnected | **Kept**, store-side | Ghost cells are forever. Store-side also survives a whole stream reopening, which a reflector-side diff cannot remember across. Costs one event type (`ReplaceEvent`) |
| Indexer | Large-scale reverse lookup | **Dropped**; `dict[ObjectKey, ParentKey]` scanned | The parent map is already the index |
| `EnqueueRequestForOwner` | Child event to parent key | **Kept** as `key_map`; an object it cannot map is evicted with an error, not stored | Cells are not Kubernetes objects, so the parent comes from labels, and one bad pod must not stall the pool. Go caches an object it cannot attribute and merely enqueues nothing, but this store is indexed by parent, so an object whose labels stop mapping leaves and re-drives the parent it left. Ignoring the update instead would serve it under that parent forever |
| DeltaFIFO | Delta coalescing | **Dropped** | Reconcile reads a snapshot, so the queue needs parent-key dedup, never a delta chain |

### `work_queue.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| workqueue | The scheduling core | **Kept** as a dedup set; delayed retry lives in `retry_scheduler.py` | With one worker, the dirty/processing protocol collapses into the set |
| `ShutDown` vs `ShutDownWithDrain` | Finish in-flight work first | **Dropped**; `stop()` cancels everything, then waits. Awaiting it inside reconcile asserts (use `asyncio.create_task`). It may only follow a `start()` call, but a `start()` that raised leaves no tasks and `stop()` is then a no-op that returns, so an aborted start unwinds carrying its own exception instead of an assertion | Drain exists for many Go workers. One worker means one in-flight key, and reconcile is idempotent, so abandoning it costs a re-derivation. A caller that wraps `start()` in `except: await stop()` is the normal shape, and it must not lose the failure it is unwinding |

### `retry_scheduler.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Delaying queue on earliest `readyAt` | Rate-limited retry | **Replaced** by latest-wins deadlines swept by one `POLL_INTERVAL` loop | A new failure overwrites a deadline instead of cancelling a task, so the delay always matches current state. Go sweeps a heap on an exact timer inside one `waitingLoop` goroutine; polling keeps the single loop and costs a retry up to one interval late, which a backoff heuristic absorbs. The poller is created in `__init__` and cancelled in `shutdown()`, so its lifetime is a straight line and it ticks even when idle, which costs one wakeup a second and buys having no state that decides when to run. Creating it there is why a `ReconcileLoop` must be constructed inside a running event loop |
| Bucket rate limiter | Cap retry pressure | **Dropped** | In-memory accounting behind one worker: no backend to protect |
| Delaying queue *is* a queue (`delayingType` embeds `Interface` and re-`Add`s to itself) | Retry without a second component | **Replaced** by an `on_retry` callback | Nothing here reads the queue, so a scheduler that only pushes needs no queue at all. `WorkQueue` and `RetryScheduler` then know no domain type, and `ReconcileLoop` names the key by instantiating them as `[ParentKey]` |
| `Result{RequeueAfter}` | Timed re-check | **Dropped** | Backoff plus resync covers it |

### `source_stream_driver.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Cache-before-notify ordering | Handlers never read a stale state | **Kept** | Correctness, not volume |
| WaitForCacheSync | Do not decide on a half-filled cache | **Kept**: `run()` + `wait_for_sync()`, the [`SyncingSource`](https://github.com/kubernetes-sigs/controller-runtime/blob/main/pkg/source/source.go) shape; `start()` awaits it | A partial engine list at step 0 would silently shrink the pool |
| `source.Channel` | External event injection | **Dropped** | Miles-internal events are method calls |
| Retry inside vs outside the stream | Recover without losing the cursor | **Kept** at both levels: `KubernetesReflector.retry_delay`, then `ReconcileLoop.source_retry_delay` | The inner one recovers a dropped watch, a failed LIST or an expired cursor without ending the stream. The outer one is the net for a stream that dies for good; in-process registries reach it |
| Closing a stream during teardown | Do not hide a worse error | **Replaced**: every close is logged, none propagates | `run()`'s `finally` closes the stream and clears it before `stop()` gathers the driver, so `aclose()` only ever finds nothing left to close. That leaves `_aclose_logging_failure` as the single close path, and it must swallow: a close failure would otherwise mask an unwinding cancellation, or re-raise the very failure the driver exists to recover from |
| Streaming a listing into the store | Sync a huge collection without holding it in memory | **Dropped**; a listing is one `ReplaceEvent`, and a stream that does not open with one is reopened | client-go buffers too: its watch-list path fills a `temporaryStore` until the `k8s.io/initial-events-end` bookmark, then calls `Replace` once. Ours materializes the page before yielding anyway, so bracketing it would only add a segment FSM and a half-applied state to police |

### `loop.py`

| Upstream | Solves | Decision | Reason |
| - | - | - | - |
| Resync period | Level-triggered backstop | **Kept**, on by default at Go's 10 h, re-enqueues parents that still have members | Go defaults it on for the same reason we need it: insurance against this controller, or something it depends on, failing to requeue. Go's own limit is kept too, in that a parent which lost its last member is not re-driven. The 10 percent jitter is **dropped**: it spreads the LIST storm of many controllers, and there is one loop here whose resync only adds keys to a dedup set |
| `MaxConcurrentReconciles` | Overlap reconcile I/O | **Dropped**; one worker | controller-runtime's default. No I/O to overlap |
| Predicates | Save reconciles at scale | **Dropped** | Reconcile is cheap |
| Lifecycle | Never leak a running loop | **`async with loop:`** | `start()` and `stop()` stay public, because `stop()` sometimes has to be scheduled from inside reconcile. Every caller that owns both ends writes the `async with` instead, and a caller that hands teardown back to someone else keeps the pairing with `AsyncExitStack.enter_async_context(loop)` |

### Dropped wholesale

| Upstream | Reason |
| - | - |
| SharedInformer | One or two consumers per object type |
| Manager, leader election, metrics, webhooks, multi-GVK cache, cached client | In-process objects, not a deployed operator |
| kubebuilder / Operator SDK | One resource type, hand-written |
| Thread-safe store / DeltaFIFO / workqueue | Under asyncio state mutates only between `await` points. Removes the data-race class, not the interleaving class: "check, await, mutate" still needs FSM discipline |

## Test layers

The fake-based tests remain always-on. The real-apiserver and kind tests run in Stage B CPU when selected by `run-ci-k8s` or a broad CI scope such as `run-ci-all`, `run-ci-image`, or nightly. `pytest tests/e2e/k8s_apiserver tests/e2e/k8s_kind` self-provisions everything below from a Docker daemon, and the modules that provision it depend on nothing from the reconcile package, so an environment can be verified before anything uses it.

| Layer | Where | Proves | Provisions |
| - | - | - | - |
| Fakes + fake clock | `tests/fast/utils/workers/reconcile/` | Control flow, timing, shutdown | Nothing |
| Real apiserver, no kubelet | `tests/e2e/k8s_apiserver/` | API semantics: cursors, watch timeouts, real 410, relist | etcd and the apiserver as containers, plus a second apiserver with `--watch-cache=false` — with the cache on, a compacted `resourceVersion` is still served, so nothing can invalidate a live cursor |
| kind cluster | `tests/e2e/k8s_kind/` | What only a kubelet produces: Running, restarts, graceful deletion, bookmarks | A pinned kind binary: a `kind` on `PATH` is used only when its version matches the pin, otherwise the pinned build is downloaded on demand and checked against its pinned SHA-256 |

| Env var | Effect |
| - | - |
| `MILES_K8S_KEEP=1` | Leaves the environment up |
| `MILES_K8S_KUBECONFIG=<path>` | Reuses an existing cluster |
| `MILES_K8S_REQUIRE=1` | Default in CI; missing Docker becomes a failure |


# Versions and Images
Source: https://miles.radixark.com/docs/developer/versions

How the Miles, SGLang, and Megatron-LM trees fit together, which artifact owns each release identity, and what the Docker images pin.

A Miles run is three Python source trees on one `PYTHONPATH`. Understanding which tree owns
a given behavior, and which file pins its version, is most of what you need to debug a
version problem or land a bump.

## The three trees

| Tree | Comes from | Installed as | Lives at |
| - | - | - | - |
| Miles | this repository | `pip install -e . --no-deps` | `/root/miles` |
| SGLang | the `sglang-miles` branch of [`sgl-project/sglang`](https://github.com/sgl-project/sglang), on top of the `lmsysorg/sglang:<tag>` base image | `pip install -e "python[all]" --no-deps` | `/sgl-workspace/sglang` |
| Megatron-LM | the `miles-main` branch of [`radixark/Megatron-LM`](https://github.com/radixark/Megatron-LM) | `pip install -e .` | `/root/Megatron-LM` |

Two things follow from that table.

**The dependency branches are moving development sources.** Both carry patches Miles needs before they land upstream, which is why installing stock `sglang` or upstream Megatron-LM next to Miles does not work. Ordinary `dev` images follow those branches; a versioned release records their exact commits in `release-lock.json`.

**All three are editable installs.** Nothing is copied into `site-packages`, so changing a
file in any of the three trees changes the next run. That is also what lets CI move the
SGLang and Megatron-LM checkouts to a different ref without reinstalling anything.

The `version` in `setup.py` owns the base Miles release version `X.Y.Z`. Release candidates and post releases keep that base in `setup.py` and add `rcN` or `.postN` only to the exact Git and Docker tags.

## Where each pin is written

```text theme={null}
docker/Dockerfile     the image recipe and its default build-args
docker/build.py       the variant table: which build-args each variant overrides
requirements.txt      Miles' own Python dependencies
release-lock.json     generated on a release branch: dependency commits, CI image tag, and audit fingerprints
```

The default build-args are the version surface:

| Build-arg | Default | What it selects |
| - | - | - |
| `SGLANG_IMAGE_TAG` | `v0.5.22` | The `lmsysorg/sglang` base image, which brings torch, CUDA, cuDNN and Mooncake |
| `SGLANG_BRANCH` | `sglang-miles` | The branch fetched into the base image's SGLang checkout |
| `SGLANG_COMMIT` | empty | Empty means the branch HEAD at build time; set it to freeze one commit |
| `MEGATRON_REPO` / `MEGATRON_BRANCH` / `MEGATRON_COMMIT` | `radixark/Megatron-LM` / `miles-main` / empty | The Megatron-LM checkout; an empty commit follows branch HEAD, while a release build supplies the locked commit |
| `MILES_COMMIT` | `main` | The Miles checkout baked into the image |
| `WHEELS_REPO` | `radixark/miles-wheels` | The prebuilt-wheels repository |
| `WHEELS_TAG_X86` / `WHEELS_TAG_ARM64` | `cu130-torch214-x86_64` / `cu130-torch214-aarch64` | Two complete wheels releases, selected by `TARGETARCH` and installed verbatim |

Two design choices are worth naming. The Dockerfile holds the defaults and `build.py` owns
the per-variant deltas, so there is exactly one place to look for what a variant changes.
And a wheels release is installed as a whole release, with no tag assembled at build time,
so the set of kernels in an image is one auditable name rather than a computed string.

Everything else is pinned inline where it is installed: `mbridge` and `Megatron-Bridge` and
`torch_memory_saver` at explicit commits, `flash-linear-attention`, `tilelang` and friends at
explicit versions. Transformer Engine is special: the wheels release carries the `+miles`
triplet built from `radixark/TransformerEngine` `miles-main` (NVIDIA's release commit plus
the fixes Miles carries), and `docker/verify_transformer_engine.py` asserts the installed
triplet is `2.17.0+miles`. Mooncake comes from the base image unchanged; Miles'
Mooncake object-store backend needs 0.3.12.post1 or newer, the first release whose
`mooncake.structured_object_store` has the API Miles imports.

`requirements.txt` is Miles' own dependency list, and the convention there is that **a pin
carries its reason inline**: `transformers==5.19.0` names the HF-native weight conversion and
the SGLang release it matches, `blake3` / `xxhash` / `zstandard` name the disk-delta weight sync,
`onnxscript` records that TE 2.17 imports it at import time. Follow that when you add one.

## Release identities and source ownership

Release values represent different facts. Each row below names the source used by the supported automatic path.

| Fact | Authoritative source | Derived or consumed values |
| - | - | - |
| Base Miles version | `setup.py` | Release branch `release/vX.Y.Z`; base-version check before tagging |
| Exact published version | `version` input to `release-tag.yml`, persisted as annotated Git tag `v<exact-version>` | CUDA 13 image tag `v<exact-version>` |
| Frozen dependency selection | `release-lock.json` committed on the release branch | SGLang commit, Megatron-LM commit, and the CUDA image tag used by release CI |

For example, base version `0.3.0` owns branch `release/v0.3.0`; that branch can produce exact tags `v0.3.0rc0`, `v0.3.0`, and `v0.3.0.post1` without changing `setup.py` between tags.

`release-lock.json` freezes the two dependency source commits and names the prune-exempt CUDA image used by release CI. Miles itself is frozen by the release commit and final Git tag. Wheels asset fingerprints are audit records, not content-addressed inputs.

Use [Release a Version](/docs/developer/ci/04-release) for the maintainer procedure. That runbook owns workflow order, inputs, success signals, and recovery; this page owns the identity and pinning model.

## The images

CUDA variants are published to [`radixark/miles`](https://hub.docker.com/r/radixark/miles)
on Docker Hub, ROCm variants to `rocm/sgl-dev`. The
[`dev` tag](https://hub.docker.com/r/radixark/miles?tag=dev) is what CI runs on unless a job
says otherwise, so its last-pushed timestamp is the fastest way to see how current the
fleet's image is.

| Variant | Tag written by `--image-tag dev` | Platforms |
| - | - | - |
| `cu13` | `radixark/miles:dev` | `linux/amd64` + `linux/arm64`, one manifest. This is the daily image |
| `cu13-x86` / `cu13-aarch64` | `radixark/miles:dev` | Single-arch rebuilds of the same image |
| `rocm724-mi35x` / `rocm10-mi35x` | `rocm/sgl-dev:miles-rocm*-mi35x` | Native |
| `rocm10-mi30x` | `rocm/sgl-dev:miles-rocm10-mi30x` | Native, MI300X / MI325X |

`--image-tag dev` also publishes a timestamped sibling. Scheduled retention and manual tag behavior are documented in [Docker build](/docs/developer/ci/02-docker-build).

Build one yourself with `docker/build.py`:

```bash theme={null}
python docker/build.py --variant cu13-x86 --image-tag custom --custom-tag my-experiment --push
```

[Docker build](/docs/developer/ci/02-docker-build) is the full reference for the build script, the
workflow and the tag rules.

Official versioned releases add `radixark/miles:v<exact-version>` for the CUDA 13 multi-arch image. CUDA 12 image builds are retired, including rolling builds; previously published CUDA 12 tags remain available. Publishing a release does not move the rolling `dev` or `latest` families.

## What CI moves, and what it does not

This is the part that decides whether your change needs a new image. A CUDA CI job starts
from `radixark/miles:<tag>` and then:

1. Installs `examples/multi_lora/requirements.txt` before `requirements.txt` through `tests/ci/reconcile_dependencies.py`. It preserves installed CUDA 12/13 cuDNN versions with uv overrides so a resolve cannot move cuDNN off the image's version; images without cuDNN use pip. Before each install, a dry-run gate rejects replacement of installed GPU runtimes or packages with at least 100 MiB of recorded installed files, including same-version reinstalls. Update these packages in the image.
2. Resets both dependency checkouts and fetches the selected refs. Explicit dispatch or PR-body overrides win first, `release-lock.json` commits win when no override exists, and the moving `sglang-miles` / `miles-main` heads are the final defaults.
3. Sets `PYTHONPATH` to the Miles workspace plus both source roots.

It never reinstalls the three source trees, because they are editable installs. So:

| Your change | Needs a new image? |
| - | - |
| `requirements.txt` | Only when replacing an installed GPU runtime or a package with at least 100 MiB of recorded installed files; other changes install in the next CI run. |
| A Dockerfile layer: a pinned wheel, an inline commit, a TE patch, the base image | Yes |
| SGLang or Megatron-LM code | No. Point CI at a ref instead |
| Miles code | No |

The ROCm stage is the exception: it takes SGLang and Megatron-LM from `rocm/sgl-dev` unless the run names a ref for one, and never reads `release-lock.json`. ROCm CI runs independently and is not part of the versioned release gate.

## Bumping principle

**Bump where the pin lives.** Python dependency requirements live in `requirements.txt`; image layers live in `docker/Dockerfile`; variant-only differences live in `docker/build.py`. A requirement change that replaces an installed GPU runtime or a package with at least 100 MiB of recorded installed files also needs an image rebuild so the dependency gate can retain the new image version.

**Prefer moving the branch to pinning a commit during rolling development.** `SGLANG_COMMIT` and `MEGATRON_COMMIT` are empty by default, so ordinary images follow `sglang-miles` and `miles-main` together. A versioned release is the deliberate exception: its lockfile supplies both exact commits to CI and the final image build.

**Validate an image-affecting bump on the PR that makes it.** A PR touching
`docker/Dockerfile`, `docker/build.py`, `docker/verify_transformer_engine.py`,
`docker/patch/**` or `requirements.txt` builds a `pr-<number>` image first and runs every
GPU suite inside it, and that fresh tag outranks a `ci-image-tag:` directive so the PR
cannot accidentally test the old image. Add the `run-ci-image` label to widen the selection
to every tag except the long-running and fault-tolerance ones.

**Validate a dependency-ref bump without building anything.** Put the ref in the PR
description and CI checks out that instead:

```text theme={null}
ci-sglang-pr: #12345
ci-megatron-pr: my-fix-branch
```

`#12345` resolves to `refs/pull/12345/head`, so an unmerged upstream fix can be tested
before it lands on `sglang-miles`. See
[Contributing](/docs/developer/contributor-guide#pr-description-ci-tags) for all three
directives.

**Expect `dev` to move on its own, within a bound.** The scheduled check (every 10 minutes)
polls the SGLang and Megatron-LM branch heads plus a fingerprint of the wheels release,
and starts a rebuild within about 30 minutes of any of them moving. It deliberately does not
poll Miles, which would rebuild constantly, and instead forces a build once the last one is
12 hours old. So `dev` follows its dependencies closely and trails Miles `main` by at most
half a day. When you need that to stop moving underneath you, pin `ci-image-tag:` to a
timestamped tag; the scheduled prune keeps every timestamped tag for at least 14 days.

**The ROCm images move daily too.** The sgl-project/sglang nightlies rebuild the undated
`rocm/sgl-dev:miles-rocm*-mi35x` and `miles-rocm10-mi30x` tags from Miles `main` every day and publish a dated
`-YYYYMMDD` sibling; an out-of-band rebuild is a `workflow_dispatch` on the variant you want.

## After a bump, the usual suspects

| Symptom | Look at |
| - | - |
| `CUDNN_STATUS_BAD_PARAM` in a fused-attention backward | Something re-resolved cuDNN below the image's version |
| Build fails with "expected exactly one TransformerEngine 2.17.0+miles wheel" | The wheels release lacks the TE triplet for that arch, or `miles-main` moved off `2.17.0`; rebuild it (`build-wheels.yml` on x86\_64, by hand on aarch64) and update the Dockerfile globs together with `docker/verify_transformer_engine.py` |
| `#error C++20 or later compatible compiler is required` while building a wheel | The new torch requires C++20 and the extension hard-codes `-std=c++17`; raise it in the extension's build flags, patching `setup.py` for a PyPI sdist |
| `cannot import name 'FieldSchema' from 'mooncake.structured_object_store'` | The base image's Mooncake predates 0.3.12.post1, typically because `SGLANG_IMAGE_TAG` points at an older base |
| A test passes locally but fails in CI, or the reverse | Compare the image tag and the two dependency refs or commits CI resolved. Every job logs all three |

## Related

* [Installation](/docs/getting-started/installation) for pulling and running the image.
* [Release a Version](/docs/developer/ci/04-release) for the maintainer release runbook.
* [Contributing](/docs/developer/contributor-guide) for the CI labels and PR-description
  directives.
* [Docker build](/docs/developer/ci/02-docker-build) for the build script, workflow triggers and tag
  mechanics.


# Deterministic Training
Source: https://miles.radixark.com/docs/diffusion/advanced/deterministic

What --deterministic-mode covers, which attention backends it accepts, and what it deliberately does not fix.

`--deterministic-mode` configures the **training actor's** forward and backward for repeatable execution with the same
hardware, topology, software stack, and inputs. However, some argument gates are still incomplete, so configurations
that do not support deterministic execution may still pass validation.

## What it turns on

### At actor spawn (`miles/ray/actor_group.py`)

```bash theme={null}
NCCL_DETERMINISTIC=1
CUBLAS_WORKSPACE_CONFIG=:4096:8
```

NCCL reads its variable at `init_process_group`, which runs in the base actor's `init` before the
deterministic setup below, and cuBLAS at the first matmul — so both go into the spawn environment rather
than being set from Python. Both are `setdefault`, so an
explicit value in `--train-env-vars` wins.

`:4096:8` rather than `:16:8`: both are deterministic, but the larger workspace keeps cuBLASLt
from being workspace-limited. Costs roughly 32 MiB per handle.

### In the actor (`miles/backends/fsdp_utils/actor.py`)

```python theme={null}
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.use_deterministic_algorithms(True, warn_only=False)
```

`warn_only=False` is load-bearing: SDPA's deterministic backward is gated on `!warnOnly` inside
`aten`'s `attention_backward.cu`, so `warn_only=True` would silently be a no-op on the native
backend — the most common one.

## Attention is the hard part

`torch.use_deterministic_algorithms` only governs torch-native ops. A custom attention kernel is
opaque to it and would run nondeterministically **without any warning**, so backends are
classified explicitly:

| `--fsdp-attention-backend` | How determinism is obtained |
| - | - |
| unset, `*native*`, `*math*` (SDPA) | torch's global flag covers it. `math` backends are deterministic by construction. |
| `*flash*` (flash-attn, FA3) | torch's flag cannot reach it — miles patches `deterministic=True` onto the kernel entry points. |
| `sage`, `xformers`, `flex`, `aiter`, … | No hook exists. **Rejected.** |

The check runs **driver-side, before any actor launches** (`validate_attention_args`), so a
misconfiguration fails in seconds instead of after a multi-node startup.

A flash backend is also rejected when no diffusers flash entry point exposes a `deterministic` parameter (driver-side,
same check). Native model packages (LTX) additionally check their own kernel at actor init and fail with a message
listing which kernels were patched.

### How the flash patch works

For diffusers-backed families, miles wraps the dispatch functions diffusers routes flash through:

```
flash_attn_func         flash_attn_varlen_func
flash_attn_3_func       flash_attn_3_varlen_func
```

Each entry point with a `deterministic` parameter is replaced by
`functools.partial(fn, deterministic=True)`; the rest are skipped.

## What it does not cover

* **Rollout determinism.** Guaranteed by sglang-d and its post-training support.
* **Train/rollout agreement.** Determinism removes run-to-run variance; it does not make the two
  forwards equal. That is a precision problem — see [Dtype Control](/docs/diffusion/advanced/dtype-control).

## Cost

Real but usually acceptable: `cudnn.benchmark=False` gives up kernel autotuning, deterministic
kernels are chosen over faster nondeterministic ones, and cuBLAS reserves \~32 MiB per handle.
Recipes that need reproducibility more than throughput keep it on permanently.

## In CI

The e2e suite is built on this flag. Each test under `tests/e2e/short/` runs a real recipe with
`--deterministic-mode` for its registered short run (currently two or four rollouts) and compares every registered
metric series — reward statistics, old/new log-probs, grad norm — **bit for bit** against a standard committed under
`tests/ci/fixtures/e2e_standards/` (strict unless the test registers a per-metric tolerance).
Determinism is what makes strict comparison viable: a tolerance wide enough to absorb run-to-run
noise would also absorb small regressions. Standards are re-recorded by the PR author
(`python tests/ci/e2e_metrics_registry.py record --test <test-file>`, or the `record-e2e-standards` workflow, which
uploads the result as an artifact) and never committed by CI.

Because reward metrics are compared too, the CI recipes also pin the engine side — e.g.
`--sglang-dit-precision fp16` — matching the
"what it does not cover" list: bit-stable curves need both halves.

## Related

* [Dtype Control](/docs/diffusion/advanced/dtype-control) — the other half of train/rollout numeric agreement.


# Dtype Control
Source: https://miles.radixark.com/docs/diffusion/advanced/dtype-control

The global dtype flags, the model-boundary input cast policy, and per-parameter dtype overrides under FSDP2.

Diffusion RL is unusually sensitive to precision. The PPO ratio is
`exp(log_prob_new − log_prob_old)`, where `log_prob_old` came from the **rollout engine** (without recompute flag) and `log_prob_new` from the **trainer**. Any numeric difference between the two forwards shows up as a ratio drifting off 1.0 — an update signal made of nothing but rounding.

## 1. Global dtype control flags

| Flag | Default | Applies to |
| - | - | - |
| `--fsdp-master-dtype` | `fp32` | The FSDP-sharded master copy: load dtype, shard dtype, optimizer state. |
| `--fsdp-reduce-dtype` | `fp32` | `MixedPrecisionPolicy.reduce_dtype` — the gradient reduce-scatter. |
| `--diffusion-forward-dtype` | `bf16` | The DiT forward compute on the training side. Must match `--sglang-dit-precision`. |
| `--sglang-dit-precision` | `bf16` | The DiT forward compute in the sglang-d rollout engine. |

It is **strongly not recommended** to modify any of these dtype settings unless you have a
thorough understanding of what the change does. The defaults are what keeps `log_prob_new`
comparable to `log_prob_old`; in particular, `--diffusion-forward-dtype` must be set to the same
value as `--sglang-dit-precision`, and `--fsdp-reduce-dtype bf16` trades multi-rank gradient
stability for nothing you want.

## 2. Input dtype control

Autocast is not enough on its own. It governs op *interiors* — matmuls, convolutions — but
element-wise ops run at whatever dtype their inputs already carry. So the dtype of the tensors
handed to the DiT is itself a semantic choice, and it has to match what the family's sglang-d
pipeline feeds its DiT.

Each family therefore declares an `input_dtype_policy` over three boundary inputs:

```python theme={null}
input_dtype_policy = {"latents": ..., "cond": ..., "timestep": ...}
```

| Value | Meaning |
| - | - |
| `"default"` | Cast to the run's `--diffusion-forward-dtype`. |
| `"fp32"` / `"bf16"` / `"fp16"` | Cast to that dtype specifically. |
| `None` | Passthrough — keep whatever dtype rollout handed us. |

The base class defaults to passthrough on all three, but the model family dtypes in the
`TrainPipelineConfig` are aligned with sglang-d by default. LTX-2 and Cosmos3 declare their own
policies; every other family inherits passthrough:

```python theme={null}
# miles/backends/fsdp_utils/configs/ltx.py
input_dtype_policy = {"latents": "default", "cond": "default", "timestep": None}

# miles/backends/fsdp_utils/configs/cosmos3.py
input_dtype_policy = {"latents": "default", "cond": None, "timestep": "fp32"}
```

**Division of labour:** `input_dtype_policy` owns the boundary; autocast owns the interior. That
split also keeps gradient-checkpointing recompute consistent, since the recomputed forward sees
the same ambient autocast as the original. Autocast is off when `--diffusion-forward-dtype` is
`fp32` and for families that set `enable_autocast = False` (MiniMax H3).

## 3. Per-parameter dtype overrides

Some parameters must stay fp32 even in a bf16 forward — timestep embedders,
`scale_shift_table`-style modulation parameters. The patterns match parameters only; buffers
(such as RoPE frequencies) are outside this mechanism. FSDP2's stock `MixedPrecisionPolicy`
is per-wrap, not per-parameter, so miles-diffusion adds a targeted patch.

### Declaring the rule

A model's `FSDPParallelPlan` carries FQN glob patterns, matched against **root-relative** names:

```python theme={null}
# miles/backends/fsdp_utils/models/diffusers/wan2_2/parallel_plan.py
FSDP_PARALLEL_PLAN = FSDPParallelPlan(
    param_dtype_patterns={
        "*.norm2.*": "fp32",
    },
)
```

Semantics:

* Patterns apply **in declaration order**, and a later pattern overrides an earlier one — a narrow
  rule can carve a parameter back out of a broad one.
* A pattern matching **nothing** is an error, not a no-op. Rules do not rot silently when a model
  is renamed.
* An assignment equal to the group default compiles to nothing.

### How it is compiled

`compile_param_dtype_maps` (`miles/backends/fsdp_utils/mixed_precision.py`) does two passes:

```mermaid theme={null}
flowchart LR
    A["param_dtype_patterns<br/>root-relative globs"]
    B["one dtype per<br/>matched parameter"]
    C["per-wrap map<br/>wrap-local FQNs"]
    D["root_map<br/>root FQNs"]
    E["mp_policy of that<br/>fully_shard(block)"]
    F["mp_policy of the root<br/>fully_shard(model)"]
    A -- "pass 1: expand" --> B
    B -- "pass 2: claimed<br/>by a child wrap" --> C
    B -- "claimed by<br/>no wrap" --> D
    C --> E
    D --> F
```

Put simply: the rules name parameters from the model root, but at runtime each `fully_shard`
call looks its parameters up by wrap-local name — the compile step bridges the two. A parameter
belongs to the first `fully_shard` call that reaches it: the block wraps first, the final root
call takes whatever is left, exactly the order in which FSDP2 itself claims parameters. Each
call then gets its own small map in its own namespace; at cast time a hit uses the override
dtype, a miss falls back to the group `param_dtype`. The one thing this cannot express — two
parameters in one wrap group sharing a local name but wanting different dtypes — is rejected at
compile time.

### Expand FSDP2 with param-level dtype control

When any override exists, the actor swaps in `ParamDtypeMixedPrecisionPolicy` — a
`MixedPrecisionPolicy` extended with a per-parameter `param_dtype_map` — and patches FSDP2's
casting path to honour it. Because the patch reaches into FSDP2 internals, it is version-gated on
`torch==2.11.0` and raises on any other version; `requirements.txt` pins the matching torch.

## 4. Verifying the result

To measure any of this, run with:

```bash theme={null}
--diffusion-debug-mode --debug-skip-optimizer-step
```

The engine then returns per-step debug tensors and the trainer never updates weights, so any
divergence is pure forward-path difference. Watch these metrics:

| Metric | What it tells you |
| - | - |
| `train/model_output_mean_abs_diff` | Average trainer-vs-rollout `noise_pred` gap. |
| `train/model_output_max_abs_diff` | Worst element. |
| `train/model_output_rel_max` | The worst element relative to the rollout output's scale. |
| `train/log_prob_mean_abs_diff` | The gap that actually feeds the PPO ratio. |
| `train/ratio_abs_minus_1` | How far the ratio sits from 1.0 in a live run. |

For a sense of scale: the Qwen-Image RoPE cache bug — CPU-built vs CUDA-built frequency tables
differing by fp32 ULPs — produced a frozen-weight `noise_pred` mean |Δ| of about **2e-2**. Small
absolute numbers here are not automatically fine; compare against a known-good run. With CFG
enabled, `noise_pred` is the guided combination of two forwards weighted by the CFG scale (`true_cfg_scale` when
set, else `--diffusion-guidance-scale`; Qwen-Image may also rescale it to the conditional branch's norm, see its
[model page](/docs/diffusion/models/qwen-image/qwen-image)), so these
metrics depend on the CFG settings — compare runs only at the same ones.


# LoRA Training and Weight Sync
Source: https://miles.radixark.com/docs/diffusion/advanced/lora

PEFT LoRA on FSDP diffusion actors and weight sync to sglang-diffusion rollout engines over CUDA IPC or NCCL.

Miles-diffusion trains LoRA adapters on the FSDP actor and syncs them to
sglang-diffusion rollout engines each iteration. The rollout engine has no PEFT
layers — weights arrive either as merged base tensors or as raw `lora_A` /
`lora_B` pairs for local merge.

## 1. Recommended flags

For colocated LoRA training, prefer **IPC merge** — push only `lora_A` /
`lora_B` and let the rollout engine merge locally:

```bash theme={null}
--use-lora \
--lora-ipc-weight-sync \
--lora-rank 32 \
--lora-alpha 64 \
--colocate
```

`--lora-rank` / `--lora-alpha` vary by recipe (e.g. SD3 uses 32/64; some others
use 64/128). Without `--lora-ipc-weight-sync`, LoRA still trains but merges on
the train side and pushes full merged weights (§3). IPC merge requires
`--colocate`.

Without `--colocate` (separate train and rollout GPU pools, e.g. async
training), leave out `--lora-ipc-weight-sync`: LoRA always pushes only
`lora_A` / `lora_B` over NCCL, and the rollout engine merges them the same way.

## 2. Key flags

| Flag | Purpose |
| - | - |
| `--use-lora` | Enable PEFT LoRA on the FSDP actor |
| `--lora-ipc-weight-sync` | Push only `lora_A`/`lora_B` over CUDA IPC; rollout merges locally (requires `--colocate`) |
| `--lora-rank` | LoRA rank (recipe-specific; often 32 or 64) |
| `--lora-alpha` | LoRA alpha (typically 2× rank) |
| `--lora-target-modules` | Override family defaults (optional) |
| `--lora-init-weights` | Init scheme, e.g. `gaussian` |
| `--update-weight-buffer-size` | Sync bucket size in bytes (most recipes use 2 GB) |
| `--update-weight-target-module` | Component to sync (SD3 default: `transformer`) |
| `--colocate` | Train and rollout share GPU visibility and sync over CUDA IPC; without it, weights sync over NCCL |

> [!WARNING]
>
> `--lora-ipc-weight-sync` requires `--use-lora`, and `--colocate` whenever weights
> are synced (any run without `--train-only` or `--debug-rollout-only`). Without
> colocation, CUDA IPC handles cannot cross the train/rollout process boundary;
> LoRA then always syncs `lora_A` / `lora_B` over NCCL, and the flag is rejected.

LoRA target modules default from the model family's
`TrainPipelineConfig.lora_target_modules`. For SD3, see
[SD3 model guide](/docs/diffusion/models/sd3/sd3).

## 3. Weight-sync strategies

Selection logic in `miles/backends/fsdp_utils/actor.py`:

| Condition | Updater class | Behavior |
| - | - | - |
| No LoRA, `--colocate` | `DiffusionUpdateWeightFromTensor` | Full base-weight IPC |
| No LoRA, no `--colocate` | `DiffusionUpdateWeightFromDistributed` | Full base weights over NCCL |
| `--use-lora --colocate`, no IPC | `DiffusionUpdateWeightFromTensorLoRA` | Merge `W + αBA/r` on train side, push merged weights |
| `--use-lora --colocate --lora-ipc-weight-sync` | `DiffusionUpdateWeightFromTensorLoRAIPC` | Push only `lora_A`/`lora_B` over IPC (`weight_update_mode=lora_merge`); the rollout merges them or keeps them unmerged per `--sglang-lora-merge-mode` |
| `--use-lora`, no `--colocate` | `DiffusionUpdateWeightLoRADistributed` | Push only `lora_A`/`lora_B` over NCCL (`weight_update_mode=lora_merge`); the rollout merges them or keeps them unmerged per `--sglang-lora-merge-mode` |

`rollout_merges_lora(args)` in `miles/utils/arguments.py` is the single check
for the last two rows; the engine's LoRA target modules, the merge precision and
the adapter checks (§5) all follow it.

Implementation: `miles/backends/fsdp_utils/diffusion_update_weight_utils.py`.

### Full-weight sync (no LoRA)

FSDP shards are all-gathered into `FlattenedTensorBucket` objects, serialized
via CUDA IPC, and sent to the rollout engine's
`update_weights_from_tensor(load_format="flattened_bucket")`.

### Train-side merge (LoRA, no IPC)

For each base layer with adapters, the updater computes:

```
W_merged = W_base + (B @ A) * scaling
```

on the fly, strips PEFT key prefixes, and pushes standard weight names that
sglang-d expects (e.g. `transformer_blocks.0.attn.to_q.weight`).

### NCCL sync (no `--colocate`)

Trainer rank 0 and every rollout engine GPU join one NCCL group (trainer rank
0, engine *i* from rank `1 + i × --rollout-num-gpus-per-engine`). For each
bucket, rank 0 broadcasts the tensors while the engines receive them through
`update_weights_from_distributed`. This needs an sglang build with distributed
weight updates for diffusion engines.

### Rollout-side LoRA merge (IPC or NCCL)

`DiffusionUpdateWeightFromTensorLoRAIPC` (IPC) and
`DiffusionUpdateWeightLoRADistributed` (NCCL) share this path:

1. `FSDPHfWeightIterator.iter_hf_adapter_weights()` turns each LoRA module into one unit,
   its **lora\_A and lora\_B**, named by `to_hf_name()`
   (e.g. `transformer_blocks.0.attn.to_q.lora_A`). Fused families such as H3
   keep these diffusers names; sglang-d `lora_merge` applies
   `param_names_mapping` and the disk-load FFN swap.
2. FSDP shard all-gather → `pack_units_by_size()` packs units into buckets of up to
   **`--update-weight-buffer-size`** (most recipes use 2 GB; a unit is never split, so one larger than that ships
   as its own bucket) → CUDA IPC or NCCL
   broadcast.
3. Rollout engine receives `weight_update_mode="lora_merge"` with
   `lora_alpha` and `lora_rank`.

**Bucket packing:** every updater packs **units**, not individual tensors. When
adding the next unit would reach or exceed `--update-weight-buffer-size`, the current
bucket is flushed first; the whole unit then starts the next bucket. A unit is
never split across buckets, so a `lora_A` / `lora_B` pair always arrives together.
For a family whose rollout fuses projections into one layer (H3 `to_q/k/v`, Qwen-Image and
Cosmos3 `add_q/k/v_proj`), the adapters of those projections form one **atomic update group**
and share a bucket too: sglang-d's `lora_merge` replaces a fused layer's LoRA with the sections
one request carries.

Constant: `LORA_WEIGHT_UPDATE_MODE = "lora_merge"`.

Rollout-side merge precision is controlled by environment variable
`SGLANG_DIFFUSION_LORA_MERGE_FP32`:

* `"1"` when `--diffusion-forward-dtype fp32`
* `"0"` otherwise (no fp32 upcast: the engine merges in its own weight dtype, e.g. bf16)

Set automatically in `RolloutManager` when spawning engines that merge LoRA.

On the first few syncs, rank 0 logs lines like:

```text theme={null}
LoRA weight sync v1 [transformer]: pushed N lora tensors in K buckets
```

After FSDP all-gather, IPC buckets are collected on the **gather-src rank**
only, and that rank calls the rollout engine; over NCCL, trainer rank 0
broadcasts. If sync stalls or VRAM grows across rollouts, check trainer logs
for `LoRA weight sync` lines and Ray worker stderr under
`~/.ray/session_latest/logs/`.

## 4. Internals

| File | Role |
| - | - |
| `miles/backends/fsdp_utils/diffusion_update_weight_utils.py` | Updater classes: send buckets over CUDA IPC or NCCL |
| `miles/backends/fsdp_utils/adaptations/weight_bridge.py` | `ParamTransform` registry (train -> rollout name/shape), as in miles |
| `miles/backends/fsdp_utils/hf_weight_iterator.py` | `FSDPHfWeightIterator`: FSDP shards -> HF-named units (full, LoRA-merged, LoRA adapters) |
| `miles/backends/training_utils/weight_update/hf_weight_iterator/` | `HfWeightIteratorBase` and bucketing (atomic groups, size-bounded packing), as in miles |
| `miles/backends/training_utils/weight_update/hf_weight_iterator/atomic_groups.py` | Per-family atomic groups for weights and adapters (today: LoRA adapters the rollout fuses into one layer) |
| `miles/backends/fsdp_utils/actor.py` | Updater selection |
| `miles/backends/fsdp_utils/model_loader.py` | LoRA apply via PEFT (`apply_lora`, `load_lora_adapter`) |
| `miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py` | HTTP `update_weights_from_tensor` (IPC) / `update_weights_from_distributed` (NCCL) to rollout |
| `miles/utils/arguments.py` | `rollout_merges_lora(args)` and LoRA weight-sync validation |
| `miles/ray/rollout.py` | Engine env vars (`SGLANG_DIFFUSION_LORA_MERGE_FP32`) |

## 5. Limitations

* **Rollout-side merge expresses only uniform standard LoRA** — adapters from
  `--lora-adapter-path` with rsLoRA or per-layer rank/alpha are rejected for it.
  They need `--colocate` without `--lora-ipc-weight-sync` (train-side merge),
  which has no NCCL path yet.
* **One adapter per trained component** — every component in `--update-weight-target-module` gets its own adapter,
  all from one set of `--lora-*` flags per job.
* **FSDP backend** — LoRA weight sync is implemented for the FSDP diffusion
  actor; Megatron LLM LoRA (in upstream Miles) uses a separate path.


# SDE Step Backend
Source: https://miles.radixark.com/docs/diffusion/advanced/sde-backend

Train-side SDE dynamics for Flow-GRPO — when it applies, which flags to set, and how to check train/rollout alignment.

Flow-matching RL algorithms fall into two paradigms:

* **Coupled (Flow-GRPO)** — training re-scores the same `(x_t → x_{t+1})`
  transitions that rollout produced, so it needs tractable log-probs via
  `SdeStepBackend`.
* **Decoupled (DiffusionNFT, …)** — training re-noises the final image at the
  rollout's sampling sigmas; rollout dynamics are irrelevant, so this backend is unused.

## 1. When it applies

Flow-GRPO (`--loss-type policy_loss`, the default) calls
`SdeStepBackend.sde_step_logprob` on each recorded transition. The backend
mirrors sglang-diffusion's rollout stepping so train-side log-probs match
rollout-side values.

| Objective | Miles flag | Uses `SdeStepBackend`? | Paper |
| - | - | - | - |
| Flow-GRPO | `--loss-type policy_loss` | **Yes** | [Flow-GRPO](https://arxiv.org/abs/2505.05470) |
| DiffusionNFT | `--loss-type nft` | **No** | [DiffusionNFT](https://arxiv.org/abs/2509.16117) |

Canonical Flow-GRPO recipe: `scripts/run_diffusion_grpo_sd3_ocr_sglang.py`.

Implementation: `miles/backends/fsdp_utils/sde_step_backend.py`. Override with
`--sde-step-backend-path` if you need a custom kernel.

## 2. Dynamics types

| Formulation | `--diffusion-sde-type` | Backend | Notes |
| - | - | - | - |
| Flow-SDE | `sde` (default) | `FlowSdeStepBackend` | Transition noise std η·√(σ/(1−σ))·√(−dt) with the rollout's recorded σ (the grid's second σ replaces the denominator's σ at σ = 1); [Flow-GRPO](https://arxiv.org/abs/2505.05470) |
| CPS | `cps` | `CpsSdeStepBackend` | Transition noise std σ\_next·sin(η·π/2); log-prob is −(x−μ)², without the Gaussian constants or the 1/(2σ²) scale (forces `--diffusion-log-prob-no-const`); [FlowCPS](https://arxiv.org/abs/2509.05952) |
| ODE | `ode` | `FlowSdeStepBackend` (placeholder; NFT never calls it) | Deterministic; NFT (or a custom loss formula) only — with the Flow-GRPO loss it is rejected at parse time |

`--diffusion-noise-level` (η) sets the SDE noise scale. DanceGRPO-style constant-η
schedules use the same `sde` backend with a tuned η
([DanceGRPO](https://arxiv.org/abs/2505.07818)).

## 3. Key flags

These CLI flags configure sglang-diffusion rollout stepping and must stay
consistent with the train-side backend:

| CLI flag | Default | Role |
| - | - | - |
| `--diffusion-sde-type` | `sde` | Dynamics family (auto-selects backend) |
| `--diffusion-noise-level` | `0.7` | η in the SDE kernel |
| `--diffusion-num-steps` | `10` | Denoise schedule length |
| `--diffusion-num-sde-steps` | `0` | Steps scored by the step strategy |
| `--diffusion-sde-window-range` | None | e.g. `0,10` for a full window |
| `--diffusion-step-strategy-path` | None | Which steps enter training |

Example (SD3 Flow-GRPO OCR — full-window SDE):

```bash theme={null}
--diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window \
--diffusion-num-sde-steps 10 \
--diffusion-sde-window-range 0,10 \
--diffusion-noise-level 0.7
```

Step strategies live in `miles/rollout/step_strategy_hub.py`. `sde_window` picks a
contiguous window; `epoch_global_random_choice` picks a subset of
`--diffusion-sde-candidate-steps`, shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups). Partial windows follow the MixGRPO /
TempFlow-GRPO idea ([MixGRPO](https://arxiv.org/abs/2507.21802),
[TempFlow-GRPO](https://arxiv.org/abs/2508.04324)).

## 4. Train / rollout alignment

The train-side kernel must reproduce rollout dynamics. Mismatched σ resolution
or noise level shows up as a rising `train/log_prob_mean_abs_diff`.

Checklist:

1. Both sides read `--diffusion-sde-type` and `--diffusion-noise-level`, and the default backend is derived from the
   former, so only a custom `--sde-step-backend-path` can drift from the engine's dynamics.
2. Train-side scoring takes σ from the trajectory the rollout recorded (`sigma`, `next_sigma`, `sigma_max` per
   transition), not from a scheduler, so `--diffusion-num-steps` needs no train-side copy.
3. Match input dtypes between FSDP forward and the rollout engine for fp32 runs.
4. Watch `train/log_prob_mean_abs_diff` — near zero before the first optimizer step.

## 5. Pairs well with

* [Customization](/docs/diffusion/user-guide/customization) — `--diffusion-step-strategy-path` and `--sde-step-backend-path`.
* [SD3 model guide](/docs/diffusion/models/sd3/sd3) — Flow-GRPO vs NFT recipe flags.
* [Quick Start](/docs/diffusion/getting-started/quick-start) — SD3.5 Flow-GRPO OCR walkthrough.


# Single-Prompt Multi-Generation
Source: https://miles.radixark.com/docs/diffusion/advanced/single-prompt-multi-gen

One rollout request, N outputs — engine-side conditioning expansion, microgroup mechanics, and seed layout.

GRPO needs `n_samples_per_prompt` outputs per prompt to normalize advantages within a group. The naive way is N separate
requests, which encodes the same prompt N times and denoises N batches of size 1. Miles-diffusion instead sends **one
request per microgroup** of up to M of those N samples: sglang-diffusion encodes the conditioning once, expands it
engine-side, and denoises the microgroup's outputs as a batch. Batched expansion is sglang-diffusion's default; Krea-2
and FLUX 3 Image Klein validate a microgroup of 1 instead.

## 1. Microgroups

`--rollout-microgroup-size M` splits each prompt group of `n_samples_per_prompt` samples into requests of at most M
(`generate_and_rm_group` in `miles/rollout/sglang_diffusion_rollout.py`). Each request carries
`num_outputs_per_prompt` = its microgroup's size: M, or less for the last microgroup when M does not divide
`n_samples_per_prompt`:

```
group (1 prompt × n_samples_per_prompt samples)
  ├─ microgroup 0  → POST /rollout/generate  num_outputs_per_prompt=M
  ├─ microgroup 1  → POST /rollout/generate  num_outputs_per_prompt=M
  └─ ...           (the last one may be smaller)
```

Microgroups of one group run as concurrent asyncio tasks, each routed to the least-loaded engine; in-flight requests
across all engines are capped at `--sglang-server-concurrency` × the number of engines.

Canonical values on `main`:

| Recipe | `n_samples_per_prompt` | microgroup size |
| - | - | - |
| SD3.5 GRPO + OCR | 16 | 8 |
| Qwen-Image GRPO + PickScore | 16 | 8 |
| Wan2.2 GRPO + PickScore (5-GPU) | 16 | 8 |
| Wan2.2 GRPO + PickScore (17-GPU) | 16 | 4 |
| LTX-2.3 GRPO + PickScore | 8 | 1 — multi-gen works but single-gen is faster |
| Cosmos3 GRPO | 16 | 1 — packed forward is single-sample |

LTX-2.3 supports engine-side expansion, but the recipe generates one output per request: a multi-output video response
is large enough that deserializing and scoring it stops overlapping with other in-flight requests, so single-output
requests pipeline better end to end.

## 2. Seed layout

Rollout stays deterministic and collision-free per sample:

* sgl-d expands one request's `seed` into `seed, seed+1, …, seed+M−1` — one RNG stream per output.
* The rollout (`generate_and_rm_group`) spaces request seeds so streams never overlap:

```python theme={null}
seed_base = (rollout_seed + group_index * n_samples_per_prompt) % 2**31
microgroup_seed = seed_base + idx   # idx = index of the microgroup's first sample: 0, M, 2M, …
```

`group_index` is monotonic across the run, so every `(rollout, prompt-group, sample)` triple gets a distinct seed
(until `seed_base` wraps at 2^31). This
is what makes rollout replay and [deterministic mode](/docs/diffusion/advanced/deterministic) possible at microgroup granularity.

## 3. Pairs well with

* [Streaming Reward and Deserialization](/docs/diffusion/advanced/streaming-reward) — what happens to a microgroup response after
  generation.
* [Deterministic Training](/docs/diffusion/advanced/deterministic) — seed layout is half of run reproducibility.
* [Core Concepts](/docs/diffusion/user-guide/concepts) — where microgroup size sits among the batch knobs.


# Streaming Reward and Deserialization
Source: https://miles.radixark.com/docs/diffusion/advanced/streaming-reward

How rollout responses stream through msgpack transport, a Ray parser-actor pool, and per-microgroup async reward scoring.

Rollout in miles-diffusion is not "generate everything, then score everything". Each microgroup is generated,
deserialized, and scored as an independent asyncio task, so reward computation and tensor decoding overlap with
generation that is still in flight. This page covers the three stages of that pipeline and the knobs that keep each one
off the critical path.

## 1. The per-microgroup pipeline

`miles/rollout/sglang_diffusion_rollout.py` runs, per microgroup:

```mermaid theme={null}
flowchart LR
    G[generate<br/>POST /rollout/generate] --> D[deserialize<br/>parser actor pool]
    D --> R[reward<br/>rm_hub actor pool]
```

The rollout pipeline decomposes into `generate` / `deserialize` / `reward` components. Because every microgroup is its
own `asyncio` task, the engine starts the next request while previous responses are still being unpacked and scored.
Nothing waits for the full batch.

## 2. Deserialization: msgpack + parser-actor pool

Trajectory tensors are large — for video models the response for one microgroup can be gigabytes, and encoding tensors
into base64 further increases the size. In a naive implementation, tensors were base64-encoded inside a JSON body and
parsed on the main asyncio event loop, one sample at a time.

The current path:

* **msgpack raw-bytes transport.** The engine responds with `application/msgpack`; `post(..., raw=True)` returns the
  body untouched, and tensors decode directly from safetensors raw bytes — no base64.
* **Unpacking runs inside Ray actors, not the event loop.** `RolloutImageResponseParserActor.apply_raw(samples, raw)`
  does `msgpack.unpackb` + tensor decode in a separate process (`miles/utils/diffusion_rollout_response.py`), so the
  rollout event loop never blocks on a multi-GB unpack.
* **One call per microgroup.** A whole microgroup is parsed in a single `apply_raw` call — fewer Ray RPCs, one unpack
  per response.
* **A pool, round-robin dispatched.** `--rollout-parser-num-workers N` spins up N parser actors so multiple microgroups
  deserialize in parallel.
* **Optionally, the parser fetches too.** With `--rollout-fetch-in-parser` (requires `--use-miles-router`; six shipped
  recipes set it, the LTX-2.3 one included) the parser actor also makes the request: it asks the router for the
  least-loaded engine, posts `/rollout/generate` itself, and parses the response in place (`fetch_and_apply`), so the
  response never passes through the rollout process. A parser then stays busy for the whole generation, so
  `perf/parser_max_queue_depth` counts queued generations, not just queued unpacks.

Measured on the LTX-2.3 recipe (H200, same config, averaged over two consecutive RL steps, base64/JSON vs the current
path):

| Metric | base64/JSON | msgpack + pool | Speedup |
| - | - | - | - |
| `perf/rollout_time` | 157.4 s | 87.6 s | **\~1.8× (−44 %)** |
| `perf/step_time` | 321.9 s | 252.1 s | **\~1.28× (−22 %)** |

## 3. Reward: scored as soon as a microgroup lands

`generate_and_rm_microgroup` calls `batched_async_rm` immediately after parsing, per microgroup — not once at the end of
the rollout:

```python theme={null}
microgroup = await generate_microgroup(...)   # generate + deserialize
rewards    = await batched_async_rm(args, microgroup)   # score right away
```

Reward workers are Ray actor pools (see [Rewards](/docs/diffusion/user-guide/rewards) for per-reward flags), so scoring one microgroup
overlaps with generation and deserialization of the others.

## 4. Diagnosing the pipeline

Both actor pools report how deep their backlog was, which is what tells you whether a stage is keeping up. When a
microgroup is dispatched, the pool records how many calls that worker already had in flight and keeps the largest value
seen over the rollout. A reward pool keeps counting a scoring call's batches until the last of them returns:

| Metric | Meaning |
| - | - |
| `perf/parser_max_queue_depth` | Deepest backlog on a parser actor when a response was handed to it |
| `perf/reward_max_queue_depth_<pool>` | Deepest backlog on that reward pool's actors when a batch was handed to them, one metric per pool (`hps`, `pickscore`, `ocr`, `openai_api`, `custom_api`) |

`0` means every dispatch found its worker idle — that stage never made anything wait. A number that climbs with the
number of concurrent microgroups means the stage is saturated and requests are lining up behind it. The API pools
(`openai_api`, `custom_api`) are one actor that runs `max_concurrency` requests at once, so for them a depth below
`max_concurrency` means nothing waited; at or above it, requests may have waited.

Read them together with `perf/rollout_time`:

| Symptom | Look at | Likely fix |
| - | - | - |
| `perf/rollout_time` high | `perf/parser_max_queue_depth` > 0 | Responses are queuing on the parsers (deserialization, or with `--rollout-fetch-in-parser` whole generations); raise `--rollout-parser-num-workers` |
| `perf/rollout_time` high | `perf/reward_max_queue_depth_<pool>` > 0 (API pools: ≥ `max_concurrency`) | That pool's scoring is the bottleneck; add its workers (API pools: raise `max_concurrency`), or give it a dedicated GPU |
| `perf/rollout_time` high, both depths `0` | Nothing queued at dispatch | Usually the engines — check `--rollout-microgroup-size` and `--sglang-server-concurrency`. Parsing also runs inside a `--sglang-server-concurrency` slot, so a slow parse can hold back generation without a parser backlog |

Both metrics are emitted only when the corresponding pool ran, so a reward with no actor pool leaves
its `perf/reward_max_queue_depth_<pool>` absent rather than zero.


# Contributing
Source: https://miles.radixark.com/docs/diffusion/developer/contributor-guide

Repository layout, the test suites, CI labels, and PR conventions for miles-diffusion.

miles-diffusion is the diffusion-model sibling of [miles](https://github.com/radixark/miles). It
shares miles' conventions — conventional commits, `pre-commit`, English-only PRs — but has its own
test topology, because almost everything meaningful needs GPUs.

## Repository layout

```text theme={null}
miles_diffusion/
├── train_diffusion.py             # the entry point — the whole train loop, ~90 lines
├── train_sft.py                   # train-only entry point: generates rollout i+1 while i trains
├── miles/
│   ├── backends/
│   │   ├── fsdp_utils/            # the training half
│   │   │   ├── actor.py           # FSDPTrainRayActor: wrap, forward, loss, step
│   │   │   ├── configs/           # TrainPipelineConfig per model family
│   │   │   ├── models/            # native model packages (ltx) + diffusers parallel plans
│   │   │   ├── loss_hub/          # flow_grpo / nft / sft objectives
│   │   │   ├── sequence_parallel/ # USP = Ulysses × Ring
│   │   │   ├── monkey_patches/    # FSDP2 param-dtype patch (torch 2.11 gated)
│   │   │   ├── mixed_precision.py # per-parameter dtype map compilation
│   │   │   └── sde_step_backend.py# SDE dynamics + log-prob scoring
│   │   └── sglang_diffusion_utils/# the rollout half: engine wrapper + parity patches
│   ├── ray/                       # RolloutManager, train actor group, placement groups
│   │   └── data_conversion_hub/   # samples → train pairs
│   ├── rollout/                   # rollout fn, data source, rm_hub, filters, step strategies
│   ├── router/                    # miles router
│   ├── dashboard/                 # offline telemetry
│   └── utils/                     # arguments.py, types, metrics, debug tooling
├── scripts/                       # run_*.py launchers — Typer CLIs, also the e2e entry points
├── tests/                         # fast / fast-gpu / e2e / ci
├── docker/                        # Dockerfile + pinned version tags
└── docs/                          # this site
```

If you are adding a **model family**, you will touch:
`miles/backends/fsdp_utils/configs/<family>.py` (register with
`@register_train_pipeline_config`), optionally
`miles/backends/fsdp_utils/models/...` for a native package or FSDP plan, a launcher in
`scripts/`, and a page under `docs/models/`.

## Local dev loop

```bash theme={null}
cd /root/miles_diffusion
git remote add me git@github.com:<your_user>/miles_diffusion.git
git checkout -b feat/awesome

pip install -e . --no-deps          # editable install picks up changes
pytest tests/fast -x -q             # CPU suite: no GPU, seconds
pytest tests/fast/backends/fsdp_utils/test_mixed_precision.py -xvs

git add -p && git commit -m "feat(fsdp): short imperative description"
git push me feat/awesome
gh pr create
```

## Tests

Four tiers, by what they need:

| Directory | Needs | What it is |
| - | - | - |
| `tests/fast/` | CPU only | Argument validation, config registry, dtype-map compilation, loss math, data splitting. The bulk of the suite. |
| `tests/fast-gpu/` | 1-4 GPUs | FSDP behaviour that only exists on device: param-dtype maps, hybrid shard, SP attention parity, LoRA weight sync. Includes ported upstream PyTorch FSDP tests. |
| `tests/e2e/short/` | 2-5 GPUs | Runs an actual launch script for a few rollouts and compares its **metric series** against a recorded standard. |
| `tests/ci/` | — | The harness itself: registry, label filter, suite runner, e2e standards. |

Run the CPU tier before every push; it is fast and catches most regressions. On a CPU-only box
without sglang's GPU kernels, install the stubs first the way CI does:
`uv pip install tests/ci/cpu_stubs`.

CI does not call `pytest` directly — it goes through the suite runner, which is also the way to
reproduce a CI failure locally:

```bash theme={null}
python tests/ci/run_suite.py --hw cpu  --suite stage-a-cpu --labels run-ci-fsdp
python tests/ci/run_suite.py --hw cuda --suite stage-b-3-gpu-h200 --match-all-labels
```

Stage A (CPU) gates every GPU stage: a broken import should not burn GPU runner slots.

### Registering a test with CI

CI discovers tests by **AST-parsing** a module-level `register_*_ci(...)` call — scheduling never
imports the file, and `register_cpu_ci` / `register_cuda_ci` are runtime no-ops (`register_e2e_ci`,
below, also runs the recipe when the file executes).

```python theme={null}
from tests.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=["fsdp"])
```

| Parameter | Meaning |
| - | - |
| `est_time` | Seconds. Used to pack the suite. |
| `suite` | Which runner it lands on (below). |
| `labels` | Domain gate. Empty/omitted = runs on **every** PR. |
| `nightly` | Nightly only. |
| `disabled` | A string reason; the test is skipped. |

### Suites

| Suite | Runner |
| - | - |
| `stage-a-cpu`, `stage-b-cpu` | CPU |
| `stage-b-3-gpu-h200`, `stage-b-5-gpu-h200` | H200, split into 3-GPU and 5-GPU runners via per-runner `CUDA_VISIBLE_DEVICES` |
| `stage-c-3-gpu-h200`, `stage-c-5-gpu-h200` | Same hardware; stage C holds the long e2e runs, plus the 4-GPU sequence-parallel attention parity test |

### Labels

A test with a non-empty `labels` list runs only when the PR carries `run-ci-<label>` for one of
them. The canonical registry is `tests/ci/labels.py`:

| Label | Domain |
| - | - |
| `sglang-diffusion` | Engine wrapper and monkey-patch tests |
| `fsdp` | FSDP backend and config tests |
| `torch` | Ported PyTorch regression tests |
| `rollout` | Sampling, filters, step strategies |
| `ray` | Actors and placement groups |
| `router` | Routing decisions |
| `arguments` | Top-level argparse / `validate_args` |
| `e2e` | Opt-in full-recipe GPU e2e runs (heavy) |
| `model-scripts` | `train_diffusion.py` + `scripts/*.py` smoke tests |

Two meta-labels bypass the filter and run everything: `run-ci-all` and `run-ci-image`. The filter is also bypassed for
docs-only PRs (which stop after the CPU stages), manual dispatches, and the nightly schedule.

Adding a label means adding an entry to `tests/ci/labels.py` **and** creating the matching
`run-ci-<key>` label in GitHub repo settings. The workflow needs no edit — the stage job filters
at runtime.

### E2E metric standards

An e2e test declares the recipe to run, the arguments to run it with, and the metrics to check:

```python theme={null}
register_e2e_ci(
    est_time=1200,
    suite="stage-c-3-gpu-h200",
    script="scripts/run_diffusion_grpo_sd3_ocr_sglang.py",
    args=["--num-rollout", "2"],
    metrics=["rollout/reward/raw_mean", "train/grad_norm", ...],
)
```

`script` is a Python recipe under `scripts/`, run with the same interpreter. Because the recipes
are Typer CLIs, CI configures them through `args` rather than through environment variables —
`env` still exists for anything the recipe reads from the environment instead.

This is also why launch scripts must stay runnable with no arguments and configurable through
their `ScriptArgs` dataclass: CI drives them the same way you do.

The recorded series live in `tests/ci/fixtures/e2e_standards/`. Because these runs use
`--deterministic-mode`, comparison is **strict, bit for bit** — which is exactly what makes them
useful, and also why any intentional numeric change requires re-recording the standard via the
`record-e2e-standards` workflow.

If your PR legitimately changes numerics, say so in the PR body and re-record. Do not loosen a
tolerance to make a test pass.

## Style

`pre-commit` is the enforcement point:

```bash theme={null}
pip install pre-commit
pre-commit install
pre-commit run --all-files --show-diff-on-failure --color=always
```

Hooks: `ruff` (with `--fix`), `autoflake`, `isort` (black profile), `black`, plus the standard
YAML / case-conflict / large-file / private-key checks and `requirements-txt-fixer`.

* **Line length 119** (`black` and `isort` in `pyproject.toml`).
* **Python ≥ 3.12** — modern syntax is fine (`X | None`, `match`, PEP 695 where it reads well).
* **Type hints** on new code.
* **Comments explain why, not what.** The codebase leans on this heavily: most of the tricky code
  here exists because of a specific numeric mismatch, and the comment recording *which* mismatch
  is the most valuable line in the file.

## Docker changes

A PR touching `docker/Dockerfile` or `requirements.txt` triggers an image build; every GPU suite
then runs inside `radixark/miles_diffusion:pr-<num>` instead of `latest`, and a failed build stops
the matrix. The fresh build outranks a `ci-image-tag:` directive in the PR body. Fork PRs skip
the build and run on the `ci-image-tag:` image if their body names one, else on `latest`. The tag is deleted when the PR closes.

## PR conventions

**English only** — title, body, commit messages, code comments.

Conventional commits, first line under \~70 characters:

```
feat(rollout): add epoch-global SDE step strategy
fix(fsdp): keep RoPE freq caches on CUDA for Qwen-Image
docs(models): document the LTX-2 velocity conversion
test(ci): cover the label filter
```

The body explains **why**; the diff already shows what.

### Before requesting review

* [ ] `pre-commit run --all-files` passes.
* [ ] `pytest tests/fast -x -q` is green.
* [ ] `python3 train_diffusion.py --help` still parses (any argparse change).
* [ ] A new public flag is documented in [CLI Reference](/docs/diffusion/user-guide/cli-reference).
* [ ] A new model family has a page under `docs/models/`.
* [ ] Numeric changes are called out, and e2e standards re-recorded if they moved.
* [ ] New behaviour has a test, registered with the right suite and label.

## Where to ask

* Design discussion: open a GitHub Issue or Discussion on
  [radixark/miles\_diffusion](https://github.com/radixark/miles_diffusion).


# FSDP Model Roles and Sleep/Wake
Source: https://miles.radixark.com/docs/diffusion/fsdp_model_roles

The FSDP trainer's actor, frozen reference and teacher models, and EMA — where each comes from, how they are checkpointed, and how sleep/wake moves them.

The FSDP trainer owns one trainable actor, optional independent frozen reference
and teacher models, and an optional EMA optimizer. Each model contains the same
configured components; Wan2.2 can therefore have both `transformer` and
`transformer_2` in every role.

```text theme={null}
train worker
  actor components -------- AdamW + scheduler
         |
         +----------------- EMAOptimizer.state[parameter]["ema"]
  reference components ---- frozen, eval, independent checkpoint
  teacher components ------ frozen, eval, independent checkpoint
```

Only actor parameters belong to AdamW. EMA stores averages for trainable actor
parameters in a separate `torch.optim.Optimizer`; frozen base parameters and all
buffers are excluded from averaging.

## Model sources

Use an independent reference with:

```bash theme={null}
--hf-checkpoint /models/student \
--ref-mode ref --ref-load /models/reference
```

Both sources use the same family backend, component names, input preparation, and
device mesh. This provides a separate model instance, not cross-family distillation.
The existing loss receives the reference prediction, and custom prepare/loss hooks can
access `ctx.reference_models` by component name.

A third source, the teacher (`--teacher-load`, with `--teacher-lora-adapter-path` for its
adapters), is reserved for a DMD2-style loss that miles does not ship yet. Until then it
has no built-in consumer: `--teacher-load` requires a custom prepare/loss hook, which reads
it from `ctx.teacher_models`, and teacher inference should run inside `torch.no_grad()`.

Each source may load pretrained PEFT LoRA adapters, one plain PEFT directory (what
`save_pretrained` writes) per `--update-weight-target-module` entry, in the same order:

```bash theme={null}
--use-lora --lora-adapter-path /adapters/student --lora-rank 64 --lora-alpha 128 \
--ref-mode ref --ref-load /models/reference \
--ref-lora-adapter-path /adapters/reference
```

Wan2.2, for example, takes two directories, one per expert; any other number is rejected:

```bash theme={null}
--update-weight-target-module transformer,transformer_2 \
--use-lora --lora-adapter-path /adapters/high_noise /adapters/low_noise
```

Without a role's adapter path, that frozen role loads ordinary full weights, even
when the actor uses LoRA. Full weights and base-only weights use the same loader.
`--lora-rank` and `--lora-alpha` must match the actor adapter's `r` and `lora_alpha`. Actor adapter weights are loaded
before a training checkpoint is restored. Independent reference/teacher weights
are always reconstructed from their configured sources, so retain those sources
and their adapter checkpoints when resuming training.

`--ref-mode lora_base` keeps the existing low-memory alternative: temporarily
disable the actor adapter for a gradient-free base forward. It does not construct
another base model. FSDP's lazy init runs at load time, while the adapters are
trainable, so the gradient dtype does not depend on which forward runs first. FSDP is
resharded before the adapters are re-enabled, so PEFT restores adapter gradients on
the shards rather than on gathered copies.

## EMA and checkpointing

`--use-ema` creates the EMA optimizer. Under `--loss-type nft` it is π\_old: the rollout
samples with it (`--rollout-weights ema`) and the loss trains against its prediction.
`--ref-mode` only selects the KL reference (`none`, `lora_base`, or `ref`; it defaults to
`lora_base` when `--diffusion-kl-beta > 0`, else `none`), so NFT can use EMA and a reference together.
EMA updates once after a rollout's training completes, and its decay follows the actor's
optimizer step count, as in DiffusionNFT. Publishing or retrying a weight update never advances
EMA. EMA evaluation and publication swap the EMA values into the actor's shards and
swap them back afterward, dropping FSDP's gathered copies on both sides so no forward
reads stale weights; the actor's backward re-gathers its restored shards.

EMA state lives on the same device as its actor parameter. Sleep moves it to
pinned CPU memory together with AdamW state, and wake moves both back.

Rollout publication pushes the current actor weights unless
`--rollout-weights ema` asks for the EMA weights.

```text theme={null}
iter_0000001/
  model/         actor model state (adapter-only for LoRA)
  optimizer/     AdamW state
  lr_scheduler/  scheduler state
  ema/           EMA optimizer state and decay schedule
  rng.pt         torch and CUDA RNG state
  meta.json      step counters and next rollout id
```

EMA uses the existing distributed optimizer-state checkpoint wrapper, including
parameter names and DTensor sharding. `--no-save-optim` skips saving both AdamW and
LR-scheduler state; `--no-load-optim` skips loading AdamW state only, so a saved
LR-scheduler state is still restored. EMA remains independently saved and restored.
Loading a legacy checkpoint without EMA seeds EMA from the loaded actor and logs that
its averaging history starts over. Disabling EMA simply omits this state from loading.

## Sleep/wake and native CPU offload

| Setting | Parameter residence during training | Transfer boundary |
| - | - | - |
| No offload | GPU | None |
| `--offload-train` | GPU while awake | Training phase sleep/wake |
| `--fsdp-cpu-offload` | Actor shards on CPU | FSDP unshard |
| `--ref-cpu-offload` | Reference shards on CPU | FSDP unshard |
| `--teacher-cpu-offload` | Teacher shards on CPU | FSDP unshard |

Reference and teacher CPU offload are opt-in. Sleep covers actor, reference,
teacher, their buffers, AdamW state, and EMA state. Sleep clears completed-step
gradients instead of transferring them. Transfers use pinned CPU
storage and enqueue asynchronous copies before synchronization; Parameter objects
and optimizer bindings are preserved. A natively CPU-offloaded model never sleeps,
as in miles: its parameter shards stay on CPU and its buffers stay on GPU.

Sleep retains the model's own CPU tensors, without a second actor backup. Reference and
teacher weights never change, so waking keeps their pinned host copies and each
later sleep points the parameters back at them instead of copying them off the GPU.
The helper lets FSDP rebuild its shard views and padding through `Module._apply`;
uneven shards remain pinned. This requires a small FSDP-internal adaptation of its
pinning flag.

## Focused validation

CPU CI covers EMA averaging, DCP resume, model-source selection, loaded LoRA,
and argument validation.

```bash theme={null}
python -m pytest tests/fast/backends/fsdp_utils/test_ema_optimizer.py \
  tests/fast/backends/fsdp_utils/test_reference_models.py \
  tests/fast/utils/test_reference_arguments.py
python -m pytest tests/fast-gpu/backends/fsdp_utils/test_device_move.py \
  tests/fast-gpu/backends/fsdp_utils/test_reference_models.py
```


# Installation
Source: https://miles.radixark.com/docs/diffusion/getting-started/installation

Get a working miles-diffusion environment — Docker (recommended) or from source.

miles-diffusion trains a diffusion DiT with FSDP2 while
[sglang-diffusion](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen) serves the rollout.
The two halves must agree numerically, so the pinned versions matter more than usual: the rollout engine tracks
sglang's **`sglang-miles`** branch, not a release tag, and the training side pins **torch 2.11.0** (an FSDP monkey patch is version-gated
on it). A future update will upgrade the training stack to **torch 2.13**.

Use Docker unless you have a reason not to.

## Method 1: Docker (recommended)

```bash theme={null}
docker pull radixark/miles_diffusion:latest
```

`latest` carries the `sglang-miles` tip as of its last build, and is rebuilt when `docker/Dockerfile` or
`requirements.txt` changes on main (or on a manual dispatch); dated tags
(`dev-cu129-sglang-main-YYYYMMDD`) pin a specific build. CUDA 12.9 is the only supported
build — the Dockerfile carries a CUDA 13 recipe in a comment, but its `sglang-kernel` pin is
hardcoded to cu129 and no CI covers it.

### Build it yourself

```bash theme={null}
git clone https://github.com/radixark/miles_diffusion.git
cd miles_diffusion
docker build -f docker/Dockerfile -t miles-diffusion:$(cat docker/version.txt) .
```

Useful build args:

| Arg | Default | What it does |
| - | - | - |
| `SGLANG_IMAGE_TAG` | `v0.5.12-cu129` | Base `lmsysorg/sglang` image. |
| `SGLANG_DIFFUSION_BRANCH` | `sglang-miles` | sglang branch the rollout engine is built from; carries upstream `main` plus the diffusion rollouts this repo needs ahead of it. |
| `SGLANG_DIFFUSION_COMMIT` | `none` | Pin a sglang sha; `none` follows the branch tip. |
| `FA3_WHEELS_TAG` | `cu129-x86_64` | Which prebuilt FlashAttention-3 wheel to pull. |
| `MILES_DIFFUSION_COMMIT` | `main` | Ref of miles\_diffusion baked into the image. |

### Run

```bash theme={null}
docker run --rm \
  --gpus all --ipc=host --shm-size=32g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --network=host \
  -v /your/datasets:/root/datasets \
  -e HF_TOKEN=$HF_TOKEN \
  -it radixark/miles_diffusion:latest /bin/bash
```

The image ships with:

* PyTorch 2.11.0 and the sglang base image's CUDA stack
* sglang built from the `sglang-miles` branch (`/sgl-workspace/sglang`, editable) with `sglang.multimodal_gen`
* FlashAttention-3 (`flash_attn_interface`), `sglang-kernel==0.4.5`, `torch_memory_saver==0.0.9`
* `diffusers`, `peft`, `transformers`, `ray`, `wandb`, and `ltx-core` from `requirements.txt`
* miles\_diffusion installed editable at `/root/miles_diffusion`
* PaddleOCR's English det/rec/cls weights pre-downloaded (the OCR reward would otherwise
  race-download them at runtime)
* `nccl-tests` binaries on `PATH` for link diagnostics

To run your own working tree instead of the baked copy, bind-mount it and reinstall:

```bash theme={null}
docker run ... -v $PWD:/root/miles_diffusion -it radixark/miles_diffusion:latest /bin/bash
cd /root/miles_diffusion && pip install -e . --no-deps
```

## Method 2: Update an existing container

If you already run the image and want the latest code:

```bash theme={null}
cd /root/miles_diffusion
git pull --rebase
pip install -e . --no-deps
```

No Ray restart is needed — the launch scripts stop and restart the cluster themselves.

## Method 3: Bare metal, from source

> [!WARNING]
>
> **Strongly discouraged unless the image genuinely cannot run on your machine.** Unlike
> Docker, this mutates the host: it installs apt packages and replays the image's pinned
> package set into the **system** Python with `--no-deps`, overwriting whatever versions are
> already there. Do not run it on a machine you use for anything else.

```bash theme={null}
apt-get update && apt-get install -y --no-install-recommends git ca-certificates
git clone https://github.com/radixark/miles_diffusion.git
cd miles_diffusion
bash .claude/skills/install-miles-diffusion/install.sh
```

Written against a CUDA 12.9 / Ubuntu 24.04 host, to match the image. Nothing enforces that,
but the apt versions are unpinned, so another release hands out different versions of the
apt-sourced dists, which the verify step counts under `expected-absent` instead of flagging.

The installer has six idempotent steps, and `--from STEP` resumes a failed run. The last step diffs every installed
package against the snapshot of one image tag (`OFFICIAL_IMAGE` in `snapshot/pins.env`, refreshed by hand with
`refresh.sh`, so it can lag `latest`) — a clean run reports `mismatched 0   missing 0` and ends with
`environment matches the official image`; the `matched` and `expected-absent` counts depend on the snapshot
and on which apt-sourced dists the host provides.

To check an existing machine against the image without installing anything:
`python3 .claude/skills/install-miles-diffusion/verify_env.py`.

Plain `pip install -r requirements.txt && pip install -e . --no-deps` is not enough on its own:
`requirements.txt` has no sglang, and `miles.utils.arguments` imports `sglang.multimodal_gen` at load. Adding
sglang the way CI's CPU job does (`sglang-miles` installed `--no-deps`, plus `tests/ci/requirements-ci-cpu.txt`
and the CPU stubs) leaves a **trainer-only** environment: CPU tests run, but RL does not, and neither does MiniMax H3 SFT, whose
encoder calls sglang's FA3 kernel — the engine's GPU kernels and FlashAttention-3 are not part of that install. Resolving the full set through
pip is not possible (it is only consistent under `--no-deps` replay), which is what the
installer handles.

## Verify

```bash theme={null}
# package imports
python -c "import miles; print('miles-diffusion import OK')"

# rollout engine is present
python -c "from sglang.multimodal_gen.runtime.server_args import ServerArgs; print('sglang-d OK')"

# GPUs visible
nvidia-smi

# arguments parse
python3 train_diffusion.py --help | head
```

Then run the CPU test suite — seconds, and it catches most environment breakage:

```bash theme={null}
pytest tests/fast -x -q
```

On a machine with no sglang GPU kernels installed, install the CPU stubs first, the way CI does:
`uv pip install tests/ci/cpu_stubs`.

## Hardware

| | |
| - | - |
| Validated GPUs | H200 — what the CI runners are, and what every shipped recipe was tuned on. |
| Minimum for a real run | 2 GPUs — SD3.5-medium LoRA GRPO colocated (`scripts/run_diffusion_grpo_sd3_ocr_sglang.py`) |
| Typical | 4 train GPUs + 1 dedicated reward GPU (the Qwen-Image / Wan2.2 / LTX recipes) |

Under `--colocate` the training actor and the rollout engines time-share the same GPUs,
so the floor is set by whichever of the two needs more memory, not by their sum.

## Environment variables

| Variable | When you need it |
| - | - |
| `HF_TOKEN` | Gated checkpoints (SD3.5): required to download them. Keep it set on cached runs too — the engine re-fetches `model_index.json` from the Hub at startup and, without access, falls back to the cached copy with a warning that it may be out of date. |
| `WANDB_API_KEY` | Without it the launch scripts drop all wandb flags (`--use-wandb` included) and only print a one-line notice. |
| `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` | Set by most RL launch scripts (the Wan2.2 and Cosmos3 recipes set `expandable_segments:False`; the SFT scripts leave it unset); reduces fragmentation OOMs. |
| `MILES_DIFFUSION_MODEL_FAMILY` | Escape hatch when `--diffusion-model-family` is unset. It overrides checkpoint-name auto-detection, but an explicit CLI family takes precedence. |

## Next

* [Launch Scripts](/docs/diffusion/user-guide/launch-script) — what a launch script does and how to override a
  recipe.
* [CLI Reference](/docs/diffusion/user-guide/cli-reference) — every flag.


# Quick Start
Source: https://miles.radixark.com/docs/diffusion/getting-started/quick-start

A working Flow-GRPO training job on SD3.5 + OCR — default script uses 2 GPUs, using the Flow-GRPO OCR dataset (`flowgrpo_ocr`).

This page takes you from `docker pull` to a running **Flow-GRPO** job on
**Stable Diffusion 3.5 Medium** with **OCR** reward. The launch script below
targets **2 GPUs** (colocate train + rollout). A single-GPU topology is not
provided or verified. You need Hugging Face access to the gated SD3.5 checkpoint.

Installation and environment setup are documented separately — this page starts
inside a ready container or machine.

Deeper SD3 recipe config, reference curves, and the DiffusionNFT + PickScore
recipe are in the [SD3 model guide](/docs/diffusion/models/sd3/sd3).

## 1. Start the container

On the **host** (experimental image — see
[`docker/README.md`](https://github.com/radixark/miles_diffusion/blob/main/docker/README.md)):

```bash theme={null}
docker pull radixark/miles_diffusion:latest

docker run --rm \
  --gpus all --ipc=host --shm-size=32g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --network=host \
  -it radixark/miles_diffusion:latest /bin/bash
```

Refresh the editable install inside the container:

```bash theme={null}
cd /root/miles_diffusion && git pull && pip install -e . --no-deps
```

Steps 2–4 below run inside the container (or on any machine with the same deps).

> [!NOTE]
>
> Launch recipes live under `scripts/` as **Python modules** (not bash). Each
> script sets recipe-specific env vars (e.g. `PYTORCH_CUDA_ALLOC_CONF`) and submits `train_diffusion.py` through Ray via
> `miles.utils.external_utils.command_utils`.

## 2. Download model and data

SD3.5 is a **gated** Hugging Face model. Export your token before any download or
training:

```bash theme={null}
export HF_TOKEN=<your_hf_token>
```

The launch script loads the checkpoint by its Hub id and downloads it into the Hugging Face cache on
first run. To prefetch into that cache:

```bash theme={null}
hf download stabilityai/stable-diffusion-3.5-medium
```

Training prompts come from the **`flowgrpo_ocr`** subset of
[`rockdu/miles-diffusion-datasets`](https://huggingface.co/datasets/rockdu/miles-diffusion-datasets).
Each prompt embeds the target string in double quotes (e.g. a logo saying
`"Miles"`) — OCR reward compares PaddleOCR output against that target.

The script downloads the subset under `--data-dir` (default `/root/datasets`) on every launch;
`hf download` skips files that are already there. To prefetch into the default location:

```bash theme={null}
hf download --repo-type dataset rockdu/miles-diffusion-datasets \
  --include "flowgrpo_ocr/**" \
  --local-dir /root/datasets/miles-diffusion-datasets
```

| Split | Path |
| - | - |
| Train | `<data-dir>/miles-diffusion-datasets/flowgrpo_ocr/train.jsonl` |

See [Rewards](/docs/diffusion/user-guide/rewards) for OCR scoring and prompt format.

## 3. Launch training

```bash theme={null}
export HF_TOKEN=<your_hf_token>
python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py \
  --cuda-visible-devices 6,7
```

The script starts Ray, launches sglang-diffusion rollout engines, loads the
FSDP actor with LoRA, syncs weights via CUDA IPC (`--lora-ipc-weight-sync`),
and runs the Flow-GRPO rollout / train loop. With `WANDB_API_KEY` set, images
and metrics are logged to project `miles-diffusion-grpo`.

Once rollout begins, stdout prints a `[reward stats] raw mean=…` line per rollout and
`[train step N]` lines with metrics such as `train/loss` and `train/log_prob_mean_abs_diff`.
WandB (when configured) also gets the reward as `rollout/reward/raw_mean`.

To finish faster while debugging, override rollout count (default **600**):

```bash theme={null}
python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py \
  --cuda-visible-devices 6,7 \
  --num-rollout 2
```

Any `ScriptArgs` field also accepts `MILES_SCRIPT_<FIELD>` (e.g.
`MILES_SCRIPT_NUM_ROLLOUT=2`). For train/rollout alignment debugging (skips the
optimizer step; does not run IPC weight-sync checksums):

```bash theme={null}
MILES_SCRIPT_DEBUG_ALIGNMENT=1 python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py
```

## 4. What's happening

Each rollout iteration runs four steps:

```mermaid theme={null}
flowchart LR
    P[flowgrpo_ocr] --> R[sglang-diffusion rollout]
    R --> RM[OCR rm_hub]
    RM --> A[FSDP actor + LoRA]
    A == CUDA IPC weight sync ==> R
    A -. KL .-> Ref[(LoRA base)]
```

1. Sample prompts and generate images via sglang-diffusion (SDE window stepping).
2. Score images with OCR (`--rm-type ocr`) on CPU Ray actors — no reward GPU.
3. Compute the **Flow-GRPO** objective (GRPO advantage normalization + SDE
   log-prob ratios) and step the LoRA optimizer.
4. Push updated LoRA weights to rollout engines via `--lora-ipc-weight-sync`.

This recipe uses **Flow-GRPO** (`--loss-type policy_loss`, the default) with
LoRA-base KL (`--diffusion-kl-beta 0.04`), noise level 0.7, CFG 4.5, and batch-wide reward-standard-deviation
normalization (`--globalize-reward-std`). `--deterministic-mode` and `--global-batch-size 64` match the CI e2e recipe.

Which denoising steps enter the loss is selected by a **step strategy**
(`--diffusion-step-strategy-path` in `miles/rollout/step_strategy_hub.py`):

| Strategy | Behavior | Typical flags |
| - | - | - |
| `sde_window` | Random contiguous window inside the range (this recipe sets a 10-step window over all 10 steps, so every step is SDE) | `--diffusion-num-sde-steps 10`, `--diffusion-sde-window-range 0,10` |
| `epoch_global_random_choice` | Random subset of candidate steps, shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups) | `--diffusion-sde-candidate-steps …`, `--diffusion-num-sde-steps` |

Train-side dynamics and more on these strategies: [SDE step backend](/docs/diffusion/advanced/sde-backend).

**DiffusionNFT + PickScore** (3 GPUs, ODE rollouts from the EMA policy) is a separate recipe
covered in the SD3 model guide.

## Inspecting a run

| Question | Where to look |
| - | - |
| Is reward improving? | `rollout/reward/raw_mean` in WandB (`[reward stats] raw mean` in stdout) |
| Train/rollout SDE aligned? | `train/log_prob_mean_abs_diff` (should stay small) |
| Policy loss stable? | `train/loss`, `train/kl_loss` |
| Rollout or train bottleneck? | `perf/rollout_time`, `perf/train_time` |
| Generated images? | WandB `rollout_media/sample_images` |


# Miles-Diffusion Documentation
Source: https://miles.radixark.com/docs/diffusion/index



[Miles-diffusion](https://github.com/radixark/miles_diffusion) is currently a standalone repository built on
[Miles](https://github.com/radixark/miles)' design philosophy, focused on RL post-training for image and video diffusion
models. [sglang-diffusion](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen) serves the
rollout, and the DiT trains under **FSDP2** on a backend that co-evolves with Miles' own. Models load from a diffusers
pipeline, or from a native package when a family brings its own modeling. Shipped recipes carry explicit
[verification levels](/docs/diffusion/user-guide/recipe-verification). Custom rewards, losses, and rollout functions plug in through flags.

## Core features

* **Verified Recipes for Latest Diffusion Models.** Launchers for Wan2.2-T2V-A14B, Qwen-Image,
  LTX-2.3, Cosmos3-Nano, SD3.5, Krea-2, FLUX 3 Image Klein, and MiniMax H3. `TrainPipelineConfig` allows for easy
  model support.
* **Quality control on three fronts.** Deterministic mode supports bit-for-bit comparisons for recipes covered by
  committed E2E standards; sglang-side monkey patches reduce train/rollout mismatches; and an FSDP2 param-dtype patch
  provides per-parameter fp32 control under the mixed-precision policy. See [Deterministic
  Training](/docs/diffusion/advanced/deterministic) and [Dtype Control](/docs/diffusion/advanced/dtype-control).
* **SFT, DiffusionNFT, and Flow-GRPO under one trainer.** The loss type, training-batch preparation, rollout function,
  and reward function are all **replaceable components**, so integrating a new algorithm — or swapping in your own
  customized component — is easy.
* **Sglang native.** Rollout runs **on the inference engine itself** — the sglang-diffusion serving stack — with RL
  support and optimizations living engine-side. An optional curated set of monkey patches aligns selected engine
  operations with the training-side forward.
* **Multiple parallelisms.** The rollout engines scale with **tensor and sequence parallelism** to support large models
  and very long contexts; training scales with **USP (Ulysses × Ring)**, built from each family's diffusers `_cp_plan` —
  or a self-written one — for agile model integration.
* **LoRA training support.** With `--lora-ipc-weight-sync`, PEFT LoRA on the FSDP2 actor ships only
  `lora_A`/`lora_B` pairs to colocated rollout engines over CUDA IPC; the engine merges them into its weights or
  evaluates them unmerged, per `--sglang-lora-merge-mode`. See
  [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora).

## Supported models

Each model name links to its recipe page. Every documented recipe is labeled with a
[recipe verification level](/docs/diffusion/user-guide/recipe-verification). SD3.5, Qwen-Image, Wan2.2,
LTX-2.3, Cosmos3-Nano, and MiniMax H3 also appear in the
[Miles model list](https://miles.radixark.com/docs#supported-models).

| Model | Task | Canonical recipes |
| - | - | - |
| [Stable Diffusion 3.5](/docs/diffusion/models/sd3/sd3) | T2I | Flow-GRPO + OCR, Flow-GRPO + HPSv2.1, Flow-GRPO + OCR & PickScore, DiffusionNFT + PickScore |
| [Qwen-Image](/docs/diffusion/models/qwen-image/qwen-image) | T2I | Flow-GRPO + PickScore (flow\_grpo-aligned) |
| [Qwen-Image 2.1](/docs/diffusion/models/qwen-image/qwen-image-21) | T2I | Flow-GRPO + PickScore |
| [Wan2.2-T2V-A14B](/docs/diffusion/models/wan/wan2-2) | T2V | Flow-GRPO + PickScore, LoRA SFT |
| [LTX-2.3](/docs/diffusion/models/ltx/ltx2) | T2V | Flow-GRPO + PickScore |
| [Cosmos3-Nano](/docs/diffusion/models/cosmos/cosmos3) | T2I | Flow-GRPO + PickScore |
| [Krea-2](/docs/diffusion/models/krea/krea2) | T2I | DiffusionNFT + OCR (2- and 8-GPU sync, 6 + 2 async) |
| [FLUX 3 Image Klein](/docs/diffusion/models/flux3/flux3-image) | T2I | DiffusionNFT + PickScore (LoRA), Flow-GRPO + HPSv2.1 (full finetune) |
| [MiniMax H3](/docs/diffusion/models/h3/h3) | T2VA | Flow-GRPO + PickScore (t2va, video branch, 2 GPUs), LoRA SFT (t2va, 8 GPUs) |

## Feature support matrix

* ✅ **Recipe-backed** — exercised by a canonical recipe in `scripts/` or a CI test.
* 🟡 **Implemented** — the code path exists, but no shipped recipe or test covers this combination yet.
* ❌ **Not supported** — no working code path for this combination today.

| | SD3.5 | Qwen-Image | Wan2.2 | LTX-2.3 | Cosmos3 |
| - | - | - | - | - | - |
| Flow-GRPO (`policy_loss`) | ✅ | ✅ | ✅ | ✅ | ✅ |
| DiffusionNFT (`nft`) | ✅ | 🟡 | 🟡 | 🟡 | 🟡 |
| SFT (`sft_loss`, `--train-only`) | ❌ | ❌ | ✅ | ❌ | ❌ |
| LoRA + IPC weight sync | ✅ | ✅ | ✅ | 🟡 | ✅ |
| Single-prompt multi-gen (microgroup > 1) | ✅ | ✅ | ✅ | 🟡 | ❌ |
| USP sequence parallelism | ❌ | ❌ | ✅ | ❌ | ❌ |
| Deterministic mode | ✅ | ✅ | ✅ | ✅ | ✅ |

The built-in SFT rollout encodes only Wan2.2 and MiniMax H3 (`miles/rollout/encoder_hub/`); other families need a custom
encoder or rollout.

## Start here

1. **[Installation](/docs/diffusion/getting-started/installation)** — Docker image, pinned dependency versions, bare-metal setup.
2. **[Quick Start](/docs/diffusion/getting-started/quick-start)** — a working Flow-GRPO run on SD3.5 with 2 GPUs.
3. **[Core Concepts](/docs/diffusion/user-guide/concepts)** — the four objects in every miles-diffusion job and the loop that connects
   them.
4. **[Launch Scripts](/docs/diffusion/user-guide/launch-script)** — every argument group in a launch script, annotated.
5. **[Rewards](/docs/diffusion/user-guide/rewards)** — built-in reward models and custom reward hooks.
6. **Model guides** — per-model config and recipes, starting from the [supported models](#supported-models) table above.

## Contribute

* GitHub: [github.com/radixark/miles\_diffusion](https://github.com/radixark/miles_diffusion)
* Miles (LLM RL): [github.com/radixark/miles](https://github.com/radixark/miles)


# Cosmos3
Source: https://miles.radixark.com/docs/diffusion/models/cosmos/cosmos3

Cosmos3-Nano Flow-GRPO and PickScore recipe with token-level conditioning and packed single-sample forwards.

## 1. Model introduction

[Cosmos3](https://huggingface.co/collections/nvidia/cosmos3) is NVIDIA's Mixture-of-Transformers (MoT) omni family: an
**UND** (understanding) tower and a **GEN** (generation) tower over a joint text+vision packed sequence, with the Wan2.2
VAE (4× temporal compression). This guide covers **Cosmos3-Nano**, the variant validated by the canonical recipe and
training curve.

**Key highlights for RL training:**

* **No separate text encoder.** Conditioning is token-level: `CondKwargs` carries `text_ids` / `text_mask` / `fps`
  verbatim, which eliminates the text-replay-consistency failure class other families guard against.
* **UND tower frozen inside the training graph.** The UND tower participates in the packed forward, so it is frozen by
  parameter-name fragments rather than dropped; LoRA targets are GEN attention only (`add_q_proj`, `add_k_proj`,
  `add_v_proj`, `to_add_out`).
* **Packed single-sample forward.** The transformer consumes one packed text+vision sequence per forward — one request
  cannot batch multiple outputs, so recipes run `--rollout-microgroup-size 1` and CFG batching is disabled by
  construction.
* **Shifted flow-sigma grid.** sgl-d builds Cosmos3's sigma grid as a linear grid warped by flow shift 3.0 for T2I
  (through its FlowUniPC serving scheduler, which overrides the checkpoint's scheduler class). Rollout requests do not
  use UniPC: they take that grid verbatim on a FlowMatchEuler scheduler and step it with Euler (ODE steps) or the SDE
  kernel (SDE steps). Step sizes are non-uniform, so SDE
  candidate steps must be derived from the grid.

## 2. Validated variant

Only Nano has a validated recipe and training curve.

| Size | Composition | HF checkpoints | Status |
| - | - | - | - |
| Nano | 16 B (8 B UND + 8 B GEN) | [Cosmos3-Nano](https://huggingface.co/nvidia/Cosmos3-Nano) | **Validated** — canonical recipe |

## 3. Family config

From `miles/backends/fsdp_utils/configs/cosmos3.py`:

| Property | Value | Why |
| - | - | - |
| Timestep dtype | fp32, no scaling | Grid timesteps are non-integer and sgl-d conditions on exact fp32 values — bf16 rounds 977.92 → 976 |
| Cond dtype | pass-through | cond carries only token ids, the text mask and fps; the packed forward casts its own inputs |
| CFG batching | Off (asserted) | Packed forward is single-sample |
| LoRA targets | GEN attention (`add_*_proj`, `to_add_out`) | UND tower and unused sound/action heads stay frozen |
| Frozen params | Name-fragment allowlist (`_GEN_PARAM_FRAGMENTS`) | UND sits inside the graph and cannot be detached |

## 4. Launch

Canonical recipe: `scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py` — train, rollout,
and PickScore colocated on 4 GPUs; T2I (832×480, 1 frame).

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py
```

The recipe sets `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` in the job environment: RL scores raw
samples, so the serving-side guardrail models are skipped.

## 5. Recipe notes

`epoch_global_random_choice` draws two steps from
`--diffusion-sde-candidate-steps 8,9,10,11`, shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups).

On the 16-step rollout grid (linear, T2I flow shift 3.0, taken verbatim from the serving scheduler), steps 0–7 run
σ ≈ 1.00 → 0.79 with `|dt|` ≈ 0.02–0.04, and the candidates 8–11 sit at σ ≈ 0.75 → 0.58 with `|dt|` ≈ 0.05–0.08.
Step numbers are **not transferable across grids**: re-derive candidates from σ and `|dt|` when changing the
model, the step count, or the shift.

## 6. Reference results

The 4-GPU colocated Cosmos3-Nano recipe raises PickScore
(`rollout/reward/raw_mean`) from \~0.77 to \~0.85 over 250 rollouts:

![Cosmos3 PickScore reward mean](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/cosmos3/reward_mean.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=34aa6f5e1711dde983fbd38afb26ac49)

## 7. Pairs well with

* [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora) — GEN-tower LoRA sync.
* [Rewards](/docs/diffusion/user-guide/rewards) — PickScore worker pool configuration.


# FLUX 3 Image Klein
Source: https://miles.radixark.com/docs/diffusion/models/flux3/flux3-image

FLUX 3 Image Klein DiffusionNFT with PickScore (LoRA) and full-finetune Flow-GRPO with HPSv2.1, both fully colocated on 8 GPUs.

## 1. Model introduction

FLUX 3 Image Klein is a flow-matching DiT for text-to-image generation and image editing: per-stream mode blocks for the
image and text tokens, then joint single-stream blocks over `[image | text]`, with 4-axis RoPE and per-token AdaLN
modulation. Miles supports RL on the text-to-image task only (`--diffusion-task t2i`); image-conditioned requests are
not supported. Rollouts run on sglang-diffusion's native FLUX 3 Image Klein pipeline (`Flux3ImagePipeline`) at
1024×1024 in bf16; the trainer runs the `MilesModelBackend` package `miles/backends/fsdp_utils/models/flux3_image/`,
which builds only the text-to-image path.

## 2. Supported variants

| Variant | Checkpoint | Status |
| - | - | - |
| Klein Base (CFG, not step-distilled) | [black-forest-labs/flux-3-image-klein-base](https://huggingface.co/black-forest-labs/flux-3-image-klein-base) (gated: needs `HF_TOKEN`) | Recipes below |

The step-distilled Klein variant samples without CFG through a different denoise loop and is not supported.
The launchers' `--model-path` option also takes a local package directory.

## 3. Family config

From `miles/backends/fsdp_utils/configs/flux3_image.py`:

| Property | Value | Why |
| - | - | - |
| Rollout microgroup | 1 (validated) | The engine's rollout path accepts one batch with one output per request |
| Micro-batch | 1 (validated) | Text length varies per prompt and the DiT takes no padding mask |
| Attention | `--fsdp-attention-backend _flash_3` on Hopper (SM9x), `_flash_4` on SM10x Blackwell; the trainer rejects other GPUs | The engine's default flash kernel on each GPU; its backward is deterministic under `--deterministic-mode` (the Flow-GRPO recipe sets it, the DiffusionNFT recipe does not) |
| CFG off | `--diffusion-guidance-scale 1` requires `--rollout-patch-group flux3_image` | Unpatched, the engine runs the whole CFG pair even at a scale of 1, packed into one forward over both branches: not bitwise the conditional velocity |
| CFG combine | `neg + scale · (pos − neg)`, two batch-1 forwards; `true_cfg_scale` wins when set | Same formula as the engine |
| Full-weight sync | Weight bridge renames `*.{query,key}_norm.scale` to `.weight` | The engine's tensor-update path skips names its module lacks |

The `flux3_image` rollout patch group (`miles/backends/sglang_diffusion_utils/monkey_patches/patch_flux3_image_cfg_off.py`)
runs a conditional-only denoise loop when the CFG scale is 1 or less: one DiT forward per step, the rollout transition
hook intact, CFG-on requests unchanged.

## 4. Launch

Both recipes put training, rollout and the reward model on the same 8 GPUs (`--colocate` plus
`--<reward>-reward-colocate`, one reward worker per engine) and decode rollout responses in parallel parser actors
(`--sglang-server-concurrency 4 --rollout-fetch-in-parser --rollout-parser-num-workers 32`). The recipes target
Hopper; on SM10x Blackwell, append `--extra-args "--fsdp-attention-backend _flash_4"`.

### 4.1 DiffusionNFT + PickScore, LoRA, 8 GPUs colocated

`scripts/run_diffusion_nft_flux3_image_pickscore_8gpu.py` — LoRA rank 32 / alpha 64 synced over CUDA IPC in
`dynamic` merge mode (the engine keeps the adapter unmerged, like the trainer), an EMA rollout policy and deterministic
ODE rollouts at 16 steps without CFG: NFT trains the unguided velocity, so the rollout samples it too. 8 prompts × 8
samples per rollout, lr 3e-4, gradient checkpointing.

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v)

```bash theme={null}
python3 scripts/run_diffusion_nft_flux3_image_pickscore_8gpu.py
```

### 4.2 Flow-GRPO + HPSv2.1, full finetune, 8 GPUs colocated

`scripts/run_diffusion_grpo_flux3_image_hps_8gpu.py` — full finetune (no LoRA) with fp32 master and reduce, bf16
forward, lr 1e-5; full weights sync over CUDA IPC. 10 sampling steps with CFG 3.5 on both the rollout and the trainer,
Flow-SDE noise level 0.7 on one SDE step per rollout drawn from steps 1-2, clip 1e-4, no KL. Old log-probs come from
the rollout: with this noise level the train-rollout drift leaves the ratio at 1 (`approx_kl` at most \~1e-12, `clipfrac` 0).
16 prompts × 16 samples per rollout, gradient checkpointing.

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v)

```bash theme={null}
python3 scripts/run_diffusion_grpo_flux3_image_hps_8gpu.py
```

## 5. Reference results

Both runs use the default seed and the recipe defaults on one 8×H200 node, except that the Flow-GRPO run stops at
150 rollouts (`--num-rollout 150`; the recipe default is 200). Held-out numbers use the first 32 prompts
of each reward's test split, one image per prompt at a fixed seed, at the training step count.

### 5.1 DiffusionNFT + PickScore

The PickScore reward (`rollout/reward/raw_mean`, PickScore / 26) rises from 0.769 (mean of the first 10 rollouts) to
0.857 (mean of the last 10) over 100 rollouts:

![FLUX 3 Image Klein DiffusionNFT PickScore reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/nft_pickscore_reward_mean.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=b9b2e9a650231f09a7be8327aa21b47d)

On held-out PickScore prompts at 16 steps, the training setting. Klein Base's own default is 50 steps with CFG 3.5,
so neither base row shows its default quality:

| Model | Sampling | PickScore / 26 |
| - | - | - |
| Base | CFG 3.5 | 0.833 |
| Base | CFG off | 0.775 |
| NFT, rollout 100 | CFG off | **0.878** |

The trained model without CFG beats the base model without CFG on 32/32 prompts and the base model with CFG 3.5 on
27/32. It also drifts toward PickScore's preferences: brighter, more saturated images, and on some prompts it adds or
beautifies people.

Case studies (base with CFG 3.5, base without CFG, NFT without CFG):

*An impressionist portrait of a shar pei*

![shar pei: base CFG 3.5, base CFG off, NFT](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/nft_shar_pei.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=b15463acd6d1a934cb2a6a0d7421351f)

*A digital illustration of a dragon flying over a castle, fantasy style*

![dragon over a castle: base CFG 3.5, base CFG off, NFT](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/nft_dragon_castle.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=3fc94b2f45316a94dcebf17cdfaa2264)

*an epic angel dressed in blue with white wings*

![angel: base CFG 3.5, base CFG off, NFT](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/nft_angel.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=80b57cd4b76893bd587a91e81b87677a)

*enigmatic black square building on top of a purple hill, smoke stacks*

![black building on a purple hill: base CFG 3.5, base CFG off, NFT](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/nft_black_building.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=726642a0b38a8cee7c74c98d7be4a427)

### 5.2 Flow-GRPO + HPSv2.1, full finetune

The HPSv2.1 reward rises from 0.263 (mean of the first 10 rollouts) to 0.278 (mean of the last 10) over 150 rollouts;
it is flat for the first \~20 rollouts, as expected of a full finetune at lr 1e-5 with clip 1e-4:

![FLUX 3 Image Klein Flow-GRPO HPSv2.1 reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/grpo_hps_reward_mean.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=f5f32ec328ca4d515e0327990524d96c)

On held-out HPDv2 prompts at 10 steps with CFG 3.5 (the training settings, not Klein Base's 50-step default),
HPSv2.1 goes from 0.270 (base) to 0.282 (rollout 150), higher on 29/32 prompts. The gain is small but consistent:
images keep their composition and gain detail and color.

Case studies (base and Flow-GRPO, both with CFG 3.5):

*Anthropomorphic alien plant creature with big eyes and leafy limbs depicted in a detailed painting.*

![alien plant creature: base, Flow-GRPO](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/grpo_alien_plant.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=6385eab8dbd9ea53a50f7e784d7deae5)

*A massive and brightly colored spacecraft in a deserted landscape, depicted in retro 1960s sci-fi art.*

![retro spacecraft: base, Flow-GRPO](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/grpo_retro_spacecraft.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=e2d7d07dcb50b27755b082626f1105d3)

*A cyber girl with demon horns holds a black feather in front of a cybercity with a gloomy expression.*

![cyber girl: base, Flow-GRPO](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/flux3-image/grpo_cyber_girl.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=204f3e61c3c6aa52bca8b79276881ff8)

## 6. Pairs well with

* [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora) — IPC adapter sync and merge modes.
* [Rewards](/docs/diffusion/user-guide/rewards) — PickScore and HPSv2.1.


# MiniMax H3
Source: https://miles.radixark.com/docs/diffusion/models/h3/h3

Text-to-audio-video Flow-GRPO + PickScore on MiniMax H3 — the verified 2-GPU recipe.

## 1. Model introduction

[MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) is MiniMax's omni-modal generative system: one transformer
denoises a packed text + video + audio sequence and emits video with native stereo audio. MiniMax distilled CFG
into the checkpoint, so the official setting has **no classifier-free guidance** — the recipe is unguided
(`--diffusion-guidance-scale 1.0`). The Flow-GRPO recipe trains the **t2va** (text-to-audio-video) workflow, video
branch only — the audio stream is rolled out but not trained.

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| MiniMax-H3 | [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) | t2va video-only Flow-GRPO; t2va LoRA SFT (joint audio/video forward, the recipe supervises video only); ref2va joint audio/video LoRA SFT (`scripts/run_diffusion_sft_h3_ref2va.py`). `fl2va` is not trained yet, and ref2va has no RL path. |

## 3. Family config

From `miles/backends/fsdp_utils/configs/h3.py`:

| Property | Value | Why |
| - | - | - |
| CFG training | Off (asserted) | CFG is distilled into the checkpoint; the forward is unguided |
| Weight sync | Any mode | Rollouts load the native **FL2VA** partition; full weights are converted to its layout by the `h3` ParamTransform in `models/diffusers/h3/weight_bridge.py` (native names, fused QKV in per-head rows, `[gate, value]` FFN), LoRA adapters (IPC or NCCL) in sgl-d's `lora_merge`. Full-weight sync relies on the rollout DiT staying resident, which miles guarantees by pinning sgl-d's `performance_mode` to `manual` (sgl-d's layerwise offload does no TP slicing on tensor updates) |
| LoRA targets | `attn.to_{q,k,v}`, `attn.to_out.0`, `ff.net.0.proj`, `ff.net.2` | Trainer pushes PEFT/diffusers names; sglang-d `lora_merge` maps and fuses QKV / swaps FFN |
| Optimizer state | `audio` allowed missing | GRPO, and SFT with `--fsdp-supervised-streams visual`, never train the audio branch |
| SFT streams | `--fsdp-flow-shift visual=...,audio=...` (asserted) | Every H3 SFT pair carries an audio stream noised on its own schedule — the target's soundtrack, or zero latents for a silent clip, which never enter the loss; `--fsdp-supervised-streams visual` keeps it out of the loss |
| Sample micro-batch | 1 (asserted) | One packed sequence per forward |
| Forced sampling params | `task=t2va`, `short_edge=768`, `conditions=[]` | 768 is the only short edge MiniMax verifies (sgl-d warns on any other value), so miles pins it and does not expose it; `task` follows `--diffusion-task`, and H3 rollout serves only `t2va` (asserted), with no conditions |

## 4. Launch

### 4.1 LoRA SFT on (video, prompt) pairs (8 GPUs, no rollout engines)

Recipe: `scripts/run_diffusion_sft_h3_t2va.py`

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v)

```bash theme={null}
python3 scripts/run_diffusion_sft_h3_t2va.py   # downloads the default dataset on first run
# custom data: --extra-args "--prompt-data /abs/train.jsonl"
```

### 4.2 Flow-GRPO + PickScore (2 GPUs)

Canonical recipe: `scripts/run_diffusion_grpo_h3_t2va_2gpu.py` — train, rollout, and PickScore
colocated on 2 GPUs; 16:9 / 4 s (with `short_edge=768` the canvas is 1344×768 / 107 frames).

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
# alignment diagnostic: freeze the weights, so log_prob_mean_abs_diff measures
# train-vs-rollout deviation rather than parameter drift
python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \
  --num-rollout 2 --eval-interval 0 --extra-args "--debug-skip-optimizer-step"

# long run at the resource-limited scale
python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \
  --rollout-batch-size 4 --n-samples-per-prompt 16 --num-rollout 200
```

## 5. Recipe notes

SFT: the encode pool (`miles/rollout/encoder_hub/h3/`) imports the engine's own recipes
(`minimax_h3_packed_sequence`, `minimax_h3_encode_reference_video_rows`,
`minimax_h3_text_only_ids`), so cached train pairs stay bit-compatible with the engine's
rollout conditioning. Dataset clips must already sit on the serving grid (short\_edge=768
canvas, 24 fps, 17n+5 frames); `--sft-offload-encoder` keeps the \~70GB encoders (text encoder + video/audio VAEs) in
host RAM outside encode bursts, and `--train-only` (required by `--loss-type sft_loss`) runs no rollout
engine, so no weights are synced.
`scripts/export_lora.py` converts any saved checkpoint into one safetensors file (+
`adapter_config.json` sidecar) directly loadable by sgl-diffusion's `--lora-path`.

H3 solves `x0 = x + sigma * v`, the opposite sign from the flow-matching convention the shared
train and rollout steps assume, so `compute_noise_pred` returns the negated velocity. Both halves
must agree: a flipped sign still produces plausible latents and trains on garbage.

`--diffusion-num-steps N` runs N DiT forwards on the engine (N + 1 sigma endpoints); `sglang-miles`
before its bump to sglang `main` (sgl-project/sglang#38671) ran N - 1.

The rollout engine needs sglang's `sglang-miles` branch, which
[Installation](/docs/diffusion/getting-started/installation) already builds from
(`SGLANG_DIFFUSION_BRANCH`).

## 6. Reference results

### 6.1 LoRA SFT (8 GPUs)

Recipe defaults plus `--num-epoch 10` (the default is 3) on a 254-window video dataset (70 rollouts, 280 steps): the
aggregate loss is
dominated by which sigmas each 8-sample step happened to draw, so read convergence from the
per-bucket curves (`--log-loss-sigma-bucket`) instead, which trend downward.

![MiniMax H3 SFT loss](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/sft_loss.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=c1d58a1c4e80d88c9f5d700dd83ba9fb)

### 6.2 Flow-GRPO + PickScore (2 GPUs)

Long run (`rollout_batch_size=4`, `n_samples_per_prompt=16`):

* `train/log_prob_mean_abs_diff` stays in **1.4e-5 .. 7.1e-5** (under `1e-4`).
* Eval PickScore (`eval/pickscore_val`, every 10 rollouts): **0.806 → 0.814**, peak **0.816** at rollout 59.
* Train `rollout/reward/raw_mean` sits around **\~0.79** (64 samples / rollout, 32 per optimizer step; noisier than eval).

This is the batch-4 curve; it is not evidence for a larger topology, and no larger canonical H3
topology is committed yet.

![MiniMax H3 log\_prob\_mean\_abs\_diff](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/h3/log_prob_diff.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=bc2e443dea76d41c616d2d57b2cd9670)

![MiniMax H3 eval PickScore](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/h3/reward_eval.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=0e29cca9bbcc70f251a7c1f5d8afbd75)

![MiniMax H3 train rollout PickScore](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/h3/reward_rollout.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=51edbcb711009a19ec45693c09297218)

## 7. Pairs well with

* [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora) — the IPC merge (`--lora-ipc-weight-sync` with `--colocate`) the Flow-GRPO recipe uses.
* [Rewards](/docs/diffusion/user-guide/rewards) — PickScore worker pool configuration.


# Training an H3 LoRA with SFT
Source: https://miles.radixark.com/docs/diffusion/models/h3/lora_sft_guide

Build a MiniMax H3 fire / explosion / smoke LoRA with the 8-GPU t2va SFT recipe — topic selection, data curation, preprocessing, training and results.

Follow this guide to tune **an explosion / fire / smoke reality enhancer** for MiniMax-H3.
The LoRA is always-on: no trigger word, just load it. Alpha is the strength slider
(128 = full, 64 = half).

## 1. Topic selection

The H3 base is already very strong on aesthetics and general quality (high-grade official
data, CFG distilled), so **a purely aesthetic LoRA struggles to add value** — our first
attempt (a lighting LoRA) ended as a negative result. One viable differentiation is physical dynamics, which
the base is *relatively* weak at (though still strong enough compared with other open source models). Topic selection decides the outcome; verify the base model's weakness with a handful of comparison generations before investing in data work.

## 2. Data curation

H3 is strong. Most stock footage looks *worse* than what it generates, and metadata
filters can't tell (under 20% usable yield). So we built a data pipeline: **a VLM watches the
frames for coarse screening, a human GSB blind test makes the fine cut**, and
preprocessing (next section) lands the winners on H3's serving grid.

```
WISA-80K (79,480)
  → keyword recall for splash/smoke/explosion       4,563 clips
  → Qwen2.5-VL-7B judges 6 frames per clip          top 500 (180/160/160)
  → H3 generations from the same prompt + seed      500 counterparts
  → human GSB blind test, 500 pairs                 198 dataset wins
  → visual dedup (frame correlation > 0.9)          193 clips → 254 training windows
```

* **VLM rubric**: five-field JSON — `category_ok / spectacle / quality / watermark /
  mostly_static`; total = \`0.5*spectacle + 0.3*quality + 2*category\_ok - 4*watermark
  * 3\*mostly\_static\`. About 5.4 s/clip on 8 GPUs — the full pool screens overnight.
* **The GSB blind test is the final quality gate**: randomized sides per pair, opaque
  video URLs (untraceable even from devtools), source revealed after each vote. It
  answers one question — "is this clip actually better than what H3 generates itself?" —
  and only winners enter the training set. Of the 500 pairs, 198 went to the dataset, 187 to
  H3, and 115 were draws; a draw does not enter either, so only \~40% of screened candidates
  survive (splash 40% / smoke 41% / explosion 38%). On the pairs that were decided the
  dataset only just edges H3 out, 51%.
* Long clips yield up to two non-overlapping windows (one 10 s source = 2 training
  samples sharing a caption).

## 3. Preprocessing

Data must land exactly on H3's serving grid (the encoder rejects anything off-grid):

| Item | Value |
| - | - |
| Canvas | short\_edge=768, 16:9 → 1344×768 |
| Frame rate | constant 24 fps (strictly validated: every frame within 2 ms of its 24 fps slot) |
| Frame count | 17n+5 lattice; the recipe uses 107 frames ≈ 4.46 s (`--diffusion-output-num-frames 107`) |

One ffmpeg pass: `fps=24, scale=...:flags=lanczos, crop=1344:768, setsar=1`, cut to 107 frames
(the SFT reader rejects clips with non-square pixels). The training manifest is a
jsonl, one line per sample:

```json theme={null}
{"prompt": "...", "target": {"visual": "clips/clip.mp4"}}
```

Relative paths are anchored at the jsonl's own directory — the dataset works wherever it
is downloaded. The published jsonl also repeats each path under `metadata.video`, the field
releases before `target` read.

The output of sections 2–3 is published as
[rockdu/WISA-80K-Practical-Dynamics-254](https://huggingface.co/datasets/rockdu/WISA-80K-Practical-Dynamics-254)
(254 windows + train.jsonl, fully curated **and already on the serving grid**) — to train
on it you skip both sections and go straight to training below.

## 4. Training

```bash theme={null}
export WANDB_API_KEY=...                       # without it the recipe submits no wandb flags at all
python3 scripts/run_diffusion_sft_h3_t2va.py   # zero arguments: downloads the dataset above on first run and trains
# custom data:      --extra-args "--prompt-data /abs/train.jsonl"
# delivery config:  --num-epoch 10   (recipe default is 3)
```

Watch the run's log in wandb (project `miles-diffusion-sft`).

The recipe defaults are exactly the validated optimum: **lr 3e-5, weight decay 0.01,
LoRA rank 64 / alpha 128, rollout batch 32 (1 sample per GPU per optimizer step)**,
`--fsdp-flow-shift visual=12,audio=3`, `--fsdp-supervised-streams visual`,
`--diffusion-guidance-scale 1.0` (H3 has no CFG), `--sft-offload-encoder`.

miles-diffusion started with RL and later grew to cover SFT. Read a rollout here as whatever
produces the next batch of training samples — RL generates them, SFT encodes them. Each batch is
then consumed by `--num-steps-per-rollout` 4 optimizer steps. Both fit the same loop, which is
why `iter_N` counts rollouts rather than steps.

Performance reference (8×H200, 254 windows):

| Stage | Cost |
| - | - |
| Encoding | \~15 s/clip/GPU; the first 8 rollouts (one dataset pass) each stall \~110 s on their own misses, 230 s for rollout 0 which builds the pool; warm from rollout 8 |
| Cache | keyed by the prompt, conditions, target, encoder settings, and each media file's path, size, and mtime (editing a clip in place without changing those reuses the stale entry); ≈2× dataset size, lives in `.sft_cache/` next to the jsonl, reused across runs |
| Rollout overlap | rollout `i+1` is produced while rollout `i` trains. A colocated encode burst still waits for the train step, so a cold epoch stalls on encoding while a warm one hides the cache load entirely; `--rollout-num-gpus` reserves encoder GPUs and overlaps the encoding too |
| Encoder residency | `--sft-offload-encoder`: the \~70 GB encoder sleeps in host RAM and only occupies the GPU during encode bursts |
| Training | \~31 s per optimizer step, `--num-steps-per-rollout` 4 of them per encoded batch |
| Checkpoints | saved every `--save-interval` rollouts (the recipe passes 20), at every epoch boundary, and on the final rollout; iter\_N means N rollouts completed |

End to end, the default 3 epochs (21 rollouts, 84 steps) take 66 min from a cold cache;
10 epochs (70 rollouts, 280 steps) take 2 h 50 min cold, 2 h 31 min on a warm one.

Training leaves DCP checkpoints only. Export the LoRA with one single command:

```bash theme={null}
python3 scripts/export_lora.py --ckpt-dir <run>/ckpt/iter_0000070 \
  --out practical_dynamics_fx.safetensors --lora-rank 64 --lora-alpha 128
```

This writes the safetensors plus an `adapter_config.json` sidecar; sgl-diffusion loads it
directly via `--lora-path`.

Reference loss (8×H200, same figure as the repo docs):

![MiniMax H3 SFT loss](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/sft_loss.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=c1d58a1c4e80d88c9f5d700dd83ba9fb)

## 5. Training tips

**Learning rate is extremely sensitive.** Ablation on identical data and config (rank 64 / alpha 128, 254 windows × 10 epochs):

| lr | Outcome |
| - | - |
| 3e-4 | collapses |
| 1e-4 | visible frame degradation |
| 5e-5 | overfitting signs within the first epoch |
| **3e-5** | **best** ✅ |
| 1e-5 | underfits, barely differs from base |

Other tips for DiT LoRA SFT:

* **Watch bucketed loss, not the aggregate**: DiT loss magnitude varies hugely with the
  noise level, so the aggregate curve mostly reflects which sigmas a step happened to
  draw — not training progress. The per-bucket curves from `--log-loss-sigma-bucket`
  (default 10 buckets) compare like with like.
* **Quality beats quantity**: \~200 clips that survived the GSB gate converge in
  10 epochs and outperform larger, looser pools.
* **Don't train audio on noisy soundtracks**: real-footage audio tracks are noisy, so the t2va
  recipe passes `--fsdp-supervised-streams visual`. The soundtrack is still encoded, noised and
  fed to the DiT as its native joint input, but excluded from the loss; drop the flag to train audio.

## 6. Results

### Showcase

![forest fire: base vs LoRA](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/lora_forest_fire.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=43f18c727906785df0329df4723fb442)

Beyond looking better, the LoRA also fixes physical badcases of the base model — below:
incense burning upside down, and smoke passing straight through glass.

![incense: base vs LoRA](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/lora_incense.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=44443f3aa7e0629735d68d6b523f67c7)

![smoke clipping: base vs LoRA](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/lora_smoke_clip.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=5dc9bc8fd823d27d2cd5d5a57182eb8f)

### Known boundaries

Generalization beyond the dataset's distribution is limited. On ten out-of-distribution
detailed prompts related to fire / explosions / smoke, the LoRA scored 6 wins / 3 losses / 1 tie against base — the reliable
wins stay close to the training distribution.

![powder keg: base vs LoRA](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/h3/lora_powder_keg.webp?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=d21c078cd6cd0918624dd9605480262c)

Lower alpha or unload for:

| Scenario | What goes wrong at full strength |
| - | - |
| Multi-beat scripted events (e.g. fuse burns, then explodes) | later beats get consumed to amplify the current dynamic |
| Restrained macro shots | over-energizes what should stay subtle |

### Artifact

`practical_dynamics_fx.safetensors` (rank 64, alpha 128, \~1.4 GB) + `adapter_config.json`.

```bash theme={null}
sglang serve --model-path MiniMaxAI/MiniMax-H3 \
  --lora-path practical_dynamics_fx.safetensors
```


# Krea-2
Source: https://miles.radixark.com/docs/diffusion/models/krea/krea2

Krea-2-Raw DiffusionNFT with OCR, synchronous on 2 or 8 colocated GPUs and asynchronous on 6 + 2 GPUs.

## 1. Model introduction

[Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) is a text-to-image flow-matching DiT. Its rollout runs on
sglang-diffusion at 1024×1024 in bf16.

## 2. Supported variants

| Variant | HF checkpoint | Status |
| - | - | - |
| Raw | [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (gated: needs `HF_TOKEN`) | Recipes below |

## 3. Family config

From `miles/backends/fsdp_utils/configs/krea2.py`:

| Property | Value | Why |
| - | - | - |
| Rollout microgroup | 1 (validated) | The engine serves one output per request |
| Cond collation | Pad text to the batch max; one latent grid per batch | Position ids of the grid depend only on resolution |
| Text mask | Dropped when every token is valid | Matches the rollout DiT's maskless fast path |
| CFG combine | `neg + scale · (pos − neg)`, `true_cfg_scale` wins when set | Same formula as the engine |

## 4. Launch

All recipes train DiffusionNFT with the OCR reward (no reward GPU), LoRA rank 32, an EMA rollout policy and
deterministic ODE rollouts (10 steps, no CFG).

### 4.1 Synchronous, 2 GPUs colocated

`scripts/run_diffusion_nft_krea2.py` — FSDP DP=2 sharing the GPUs with two sglang engines; LoRA adapters sync over
CUDA IPC. 8 prompts × 8 samples per step, micro-batch 2 with gradient checkpointing.

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v)

```bash theme={null}
python3 scripts/run_diffusion_nft_krea2.py
```

### 4.2 Synchronous, 8 GPUs colocated

`scripts/run_diffusion_nft_krea2_8gpu.py` — runs `train_diffusion.py`: FSDP DP=8 shares the GPUs with eight sglang
engines, and the LoRA adapters sync over CUDA IPC. 6 prompts × 8 samples per step, micro-batch 2 with gradient
checkpointing, since the engines take part of every GPU.

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v)

```bash theme={null}
python3 scripts/run_diffusion_nft_krea2_8gpu.py
```

### 4.3 Asynchronous, 6 training + 2 rollout GPUs

`scripts/run_diffusion_nft_krea2_async.py` — runs `train_diffusion_async.py`: each batch is sampled under the EMA while
the previous batch trains, on separate pools; LoRA adapters sync to the rollout engines over NCCL. Same batch as §4.2
(6 × 8 per step), micro-batch 4, no gradient checkpointing.

**Status:** [🧩 PG — Proxy gated](/docs/diffusion/user-guide/recipe-verification#pg) by
`tests/e2e/short/test_krea2_nft_ocr_async_8xGPU_4xGPU_proxy.py`, which runs `--four-gpu-ci`: half of each pool and of
the batch (3 training + 1 rollout GPUs, 3 × 8 per step), so every training GPU keeps the production micro-batches.

```bash theme={null}
python3 scripts/run_diffusion_nft_krea2_async.py
```

## 5. Reference results

Both runs use the default seed (1234) and 100 rollouts. `eval/ocr_test` is the held-out OCR accuracy over 1018 prompts
at 52 denoising steps, measured after the last rollout.

### 5.1 Synchronous, 8 GPUs colocated

The OCR reward (`rollout/reward/raw_mean`) rises from \~0.28 (mean of the first 10 rollouts) to \~0.94 (mean of the last
10\), at \~94 s per rollout; `eval/ocr_test` ends at 0.923:

![Krea-2 8-GPU sync DiffusionNFT OCR reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/krea2/sync8_ocr_reward_mean.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=edf41d1254f42750f4c03543d8b1d95f)

### 5.2 Asynchronous, 6 training + 2 rollout GPUs

The OCR reward rises from \~0.29 to \~0.88, at \~79 s per rollout; `eval/ocr_test` ends at 0.875:

![Krea-2 async DiffusionNFT OCR reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/krea2/async_ocr_reward_mean.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=2f4c34fb131edde8a2aa74937502a72d)

NFT varies a lot from seed to seed. Over seeds 1234, 4321 and 2026, the mean OCR reward of the last 10 rollouts is
0.942, 0.905 and 0.907 for §5.1 and 0.885, 0.827 and 0.940 for §5.2.

## 6. Pairs well with

* [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora) — IPC and NCCL adapter sync.
* [Rewards](/docs/diffusion/user-guide/rewards) — the OCR reward.


# LTX-2
Source: https://miles.radixark.com/docs/diffusion/models/ltx/ltx2

Video GRPO on LTX-2.3 — native model package, CPS-SDE, unguided forward.

## 1. Model introduction

[LTX-2](https://github.com/Lightricks/LTX-2) is Lightricks' audio-video DiT. LTX-2.3 is the
variant miles-diffusion trains, and the framework's most unusual family: like FLUX 3 Image Klein it
does **not** go through diffusers (it trains a native model package), it is the only one on the **CPS**
SDE kernel, and, like MiniMax H3, it is trained **unguided**.

**Key highlights for RL training:**

* **Native model package, not diffusers.** `miles/backends/fsdp_utils/models/ltx/` supplies
  loading, modeling, attention, and the FSDP plan directly against `ltx_core`; FSDP wraps
  `BasicAVTransformerBlock`. Requires `ltx-core` (pinned in `requirements.txt`, baked into the
  Docker image).
* **Unguided training.** `supports_cfg_training = False`: any `--diffusion-guidance-scale` other
  than `1.0`, or a `--diffusion-negative-prompt`, is a parse-time error. The training forward is
  a single velocity pass — there is no negative branch to combine.
* **Video branch only.** The audio stream is loaded but never trained
  (`optimizer_state_allowed_missing = ["audio"]`), and `--rollout-patch-group ltx` disables
  audio-video cross-attention on the engine to match.
* **No sequence parallelism.** `sequence_parallel_plan` raises `NotImplementedError`; leave
  `--sequence-parallel-size` at `1`.

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| LTX-2.3 | [`Lightricks/LTX-2.3`](https://huggingface.co/Lightricks/LTX-2.3) | Video branch only |

## 3. Family config

Registered in `miles/backends/fsdp_utils/configs/ltx.py`:

| Property | Value | Why |
| - | - | - |
| Velocity → noise | `forward_velocity` reconstructs x0 via `ltx_core.utils.to_denoised` in fp32 and divides back out | Algebraically an identity, but the fp32 rounding path is what the e2e standards were recorded against |
| Boundary dtypes | `latents` / `cond` cast to the forward dtype, `timestep` passthrough | Element-wise math anchors on `latents.dtype` and rollout runs bf16 — see [Dtype Control](/docs/diffusion/advanced/dtype-control) |
| Precision | bf16 end to end (master, reduce, forward, engine); the trainer runs math-SDPA (`--fsdp-attention-backend sdpa_math`), the engine `torch_sdpa`, where LTX lets PyTorch pick its fused kernels (cuDNN first) | Math-SDPA keeps the training forward deterministic; the e2e standards are recorded with this pairing |
| SDE | CPS kernel | Each train pair carries the σ the engine stepped with, so no scheduler lookup; log-prob is `−(x − μ)²`, without the Gaussian constants or the 1/(2σ²) scale |
| LoRA targets | Attention + FFN: `to_q`, `to_k`, `to_v`, `to_out.0`, `net.0.proj`, `net.2` | |
| Rollout patches | `--rollout-patch-group ltx`: `patch_ltx2_rollout_cond_kwargs`, `patch_ltx2_disable_av_cross` | Video-only training forward needs AV cross-attention off engine-side |

## 4. Launch

Canonical recipe: `scripts/run_diffusion_grpo_ltx23_sglang.py` — a 5-GPU Ray node with 4 GPUs shared by FSDP and
rollout plus 1 separately scheduled PickScore GPU, 57 frames @ 24 fps, 512×768 (height × width).

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
python3 scripts/run_diffusion_grpo_ltx23_sglang.py
```

## 5. Recipe notes

* The step strategy is `epoch_global_random_choice`: 3 SDE steps drawn from candidates 0-9 and
  shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups) — with only 64 samples per rollout, per-request
  randomization would add variance the batch is too small to average out.
* The clip range is `1e-5`, an order of magnitude tighter than the SD3.5 and Qwen-Image recipes' `1e-4`.
* `--pickscore-num-frames 3` scores three evenly-spaced frames per video; scoring all 57 would
  dominate rollout time for a reward signal that barely changes.
* Health checks are relaxed (`--rollout-health-check-interval 120`, router failure threshold 30).
  Both probes (`/health_generate` from miles, `/health` from the router) only check readiness and
  run no generation, so the relaxed values only make a dead engine take longer to detect.
* `--rollout-parser-num-workers 16`: 57-frame trajectory tensors are large enough that one
  deserializer actor starves the engines.

## 6. Reference results

The original bf16-master 4-train-GPU + 1-reward-GPU run, used as the first LTX E2E
reference, raises `rollout/reward/raw_mean` from \~0.68 to \~0.77 over 200 rollouts:

![LTX-2.3 PickScore reward mean](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/ltx/reward_mean.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=0b1397781e58acceeb6f6fccd3b52d36)

A future LTX update will make `train/model_output_mean_abs_diff` **0.0** and switch
`--fsdp-master-dtype` to `fp32`.

## 7. Pairs well with

* [Dtype Control](/docs/diffusion/advanced/dtype-control) — the boundary-dtype policy is why LTX has one.
* [Deterministic Training](/docs/diffusion/advanced/deterministic) — `sdpa_math` is the deterministic-safe
  backend this recipe relies on.


# Qwen-Image
Source: https://miles.radixark.com/docs/diffusion/models/qwen-image/qwen-image

Flow-GRPO with PickScore on Qwen-Image — the flow_grpo-aligned 5-GPU recipe.

## 1. Model introduction

[Qwen-Image](https://huggingface.co/Qwen/Qwen-Image) is Alibaba's MMDiT text-to-image model. In
miles-diffusion it is the **reference image recipe**: the launcher is aligned flag-for-flag with
flow\_grpo's `pickscore_qwenimage` config, so a run here is directly comparable to the published
baseline.

**Key highlights for RL training:**

* **flow\_grpo-aligned.** Per-prompt reward mean with a batch-wide std (`--globalize-reward-std`
  alone) is exactly flow\_grpo's `PerPromptStatTracker` with `global_std=True`; `gaussian` LoRA
  init and `adam-beta2 0.999` match likewise.
* **RoPE caches rebuilt on CUDA.** diffusers builds `QwenEmbedRope`'s frequency tables on CPU
  while sglang-d builds them on CUDA; the fp32 ULP difference drifts every block. In one frozen-weight comparison,
  `noise_pred` mean |Δ| was approximately `2e-2`. `postprocess_model_after_materialize` rebuilds the caches on device
  after FSDP wrapping.
* **Variable-length text.** Conditioning is padded per batch with the attention mask derived
  from `txt_seq_lens`; the trainer rebuilds it and does not use a mask the engine may send for batched requests.

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| Qwen-Image | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | Family key `qwen_image` |

A checkpoint name containing `qwen-image` resolves to this config. A longer match
such as `qwen-image-2.1` resolves to [Qwen-Image 2.1](/docs/diffusion/models/qwen-image/qwen-image-21). For a
differently-named local directory, add `--diffusion-model-family qwen_image`.

## 3. Family config

Registered in `miles/backends/fsdp_utils/configs/qwen_image.py`:

| Property | Value | Why |
| - | - | - |
| Timestep scaling | Trajectory timesteps ÷ 1000 before the DiT (`process_timestep_as_input`) | Matches how sglang-d's Qwen-Image pipeline rescales them |
| CFG combine | `uncond + scale·(cond − uncond)`, rescaled back to `‖cond‖` only when `true_cfg_scale > 1.0` | Mirrors sglang-d's `postprocess_cfg_noise`; the recipe's `true_cfg_scale 4.0` takes the rescaling branch |
| Cond collation | Pad `encoder_hidden_states` to batch max, mask from `txt_seq_lens`; honours `pad_to_len` | Legacy tiling path can reproduce whole-window padding bitwise |
| RoPE caches | Rebuilt on the model's CUDA device after materialize | CPU/CUDA `torch.pow` differ by fp32 ULPs |
| LoRA targets | `to_q to_k to_v to_out.0`, `add_q_proj add_k_proj add_v_proj to_add_out`, img/txt MLP proj | |
| Optimizer state | `transformer_blocks.59` text-branch params allowed missing | The last block's text outputs are discarded, so those params never receive gradient |

## 4. Launch

Canonical recipe: `scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py` — 4 colocate
GPUs + 1 PickScore GPU, 512×512, PickScore reward.

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
python3 scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py
```

Overrides go through the dataclass CLI, e.g.
`--num-rollout 50 --extra-args "--diffusion-kl-beta 0.02"`.

## 5. Recipe notes

* Of the 10 denoising steps, only steps 3-4 are trained: `--diffusion-sde-window-range 3,5` with
  `--diffusion-num-sde-steps 2` leaves the window nowhere to move, so every sample trains the
  same two mid-schedule steps as SDE while the rest run ODE.
* Batch arithmetic: 32 prompts × 16 samples = 512 samples per rollout, split into two optimizer
  steps of 256 samples; the 2 trained SDE steps expand the 512 samples to 1024 train pairs per rollout.
* `--rollout-patch-group qwen_image` applies the diffusers-op-parity patches on the engine. It costs a
  little rollout throughput and buys a much smaller `train/log_prob_mean_abs_diff`.

## 6. Reference results

The 5-GPU recipe raises PickScore (`rollout/reward/raw_mean`) from \~0.84 to \~0.90 over 400 rollouts;
held-out `eval/pickscore_test` goes from 0.859 at the first eval (after rollout 30; the recipe skips the
pre-training eval with `--skip-eval-before-train`) to 0.900:

![Qwen-Image PickScore reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/qwen-image/grpo-pickscore-raw-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=4e4190634a5811e224b9476f19a409ce)

## 7. Pairs well with

* [LoRA weight sync](/docs/diffusion/advanced/lora) — `--lora-ipc-weight-sync` is on in this recipe.
* [Dtype Control](/docs/diffusion/advanced/dtype-control) — why fp32 master + bf16 forward.
* [Deterministic Training](/docs/diffusion/advanced/deterministic) — `--deterministic-mode` is on.


# Qwen-Image 2.1
Source: https://miles.radixark.com/docs/diffusion/models/qwen-image/qwen-image-21

Flow-GRPO with PickScore on Qwen-Image 2.1 — 1-GPU recipe.

## 1. Model introduction

[Qwen-Image 2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) is the unified text-to-image
and image-to-image successor of Qwen-Image. Conditioning is Qwen3-VL (pre-final-norm
hidden states); the DiT is a single-stream block-causal transformer with an RGBA VAE.
CFG is distilled into the checkpoint, so the official default is **no guidance**.

miles-diffusion trains the **T2I** path. Image-conditioned editing is served by the
rollout engine but is not a recipe yet.

Day-0 references:

* diffusers [`#14804`](https://github.com/huggingface/diffusers/pull/14804)
* sglang-diffusion [`#39983`](https://github.com/sgl-project/sglang/pull/39983)

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| Qwen-Image 2.1 | [`Qwen/Qwen-Image-2.1`](https://huggingface.co/Qwen/Qwen-Image-2.1) | Family key `qwen_image21` |

A name containing `qwen-image-2.1`, `qwen-image-21`, `qwenimage21`, or
`qwen_image21` resolves to this family. Matching keeps the longest pattern, so
`Qwen/Qwen-Image-2.1` resolves to `qwen_image21`.

## 3. Family config

Registered in `miles/backends/fsdp_utils/configs/qwen_image21.py`:

| Property | Value | Why |
| - | - | - |
| Timestep scaling | Trajectory timesteps ÷ 1000 | Matches diffusers' `t / 1000`; sglang-d divides inside the DiT |
| CFG combine | `uncond + scale·(cond − uncond)`, **no** cond-norm rescale | Mirrors `pipeline_qwenimage21` |
| Cond collation | Pad `encoder_hidden_states`, rebuild `img_mask` = text zeros + target slots, `img_shapes` from `--diffusion-height/width` | Rollout sends Qwen3-VL text; the train DiT wants the joint mask |
| Noise pred | Slice `[:, -target_tokens:]` | Single-stream forward emits the joint sequence |
| RoPE caches | Rebuilt on CUDA after materialize | Same CPU/CUDA `torch.pow` trap as 1.0 |
| LoRA targets | `to_q/k/v`, `to_out.0`, `img_mlp.{proj,out,gate_layer}` | Single-stream block; no text-branch add-proj |

## 4. Launch

Canonical recipe: `scripts/run_diffusion_grpo_qwenimage21_pickscore_1gpu.py` — 1
colocated GPU, train + sglang + PickScore, 512×512, 16×16 group, lr 1e-4.
Real-weight GRPO needs ≥80GB; 32GB only finished generate.

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v) — 200 rollouts on one H200.
Smoothed `rollout/reward/raw_mean` rises from about 0.82 to about 0.87.

```bash theme={null}
# alignment diagnostic: freeze weights so log_prob_mean_abs_diff is train-vs-rollout
python3 scripts/run_diffusion_grpo_qwenimage21_pickscore_1gpu.py \
  --num-rollout 2 --extra-args "--debug-skip-optimizer-step"

# local checkpoint
python3 scripts/run_diffusion_grpo_qwenimage21_pickscore_1gpu.py \
  --hf-checkpoint /path/to/qwen-image-2.1
```

## 5. Recipe notes

* Official sampling is 40 steps / CFG 1. This recipe uses 10 steps and the
  same SDE window as Qwen-Image 1.0 (`--diffusion-sde-window-range 3,5`, 2 SDE
  steps). Default is 16 prompts × 16 samples, lr 1e-4, 200 rollouts.
* `--rollout-patch-group qwen_image21` turns off the engine's prefix KV cache so
  every SDE step recomputes the full joint sequence, matching the trainer, and keeps
  the RoPE layouts out of the rollout response.
* `--sglang-lora-merge-mode dynamic` keeps PEFT unmerged on the engine, same as 1.0.
* A 32 GB card will not hold the full 2.1 DiT + Qwen3-VL + VAE + PickScore at
  1024² without offload. Raise batch/resolution only after adding GPUs or
  component offload.

## 6. Reference results

The canonical 1-GPU recipe over 200 rollouts (512², 16×16, lr 1e-4) on one H200:

![Qwen-Image 2.1 PickScore reward mean](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/qwen-image/reward_mean.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=1f28ad77d160fbeb634faac314af2ecf)

Smoothed `rollout/reward/raw_mean` moves from about 0.82 to about 0.87.
`train/log_prob_mean_abs_diff` stays near 2e-5 on the same run:

![Qwen-Image 2.1 log-prob mean abs diff](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/qwen-image/log_prob_mean_abs_diff.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=5629b05f877a42213aa8be9f2a12e956)

## 7. Pairs well with

* [Qwen-Image 1.0](/docs/diffusion/models/qwen-image/qwen-image) — the flow\_grpo-aligned 5-GPU reference recipe.
* [LoRA weight sync](/docs/diffusion/advanced/lora) — `--lora-ipc-weight-sync` is on.
* [Deterministic Training](/docs/diffusion/advanced/deterministic) — `--deterministic-mode` is on.


# Stable Diffusion 3 / 3.5
Source: https://miles.radixark.com/docs/diffusion/models/sd3/sd3

SD3 family config, LoRA targets, and launch recipes for Flow-GRPO and DiffusionNFT.

## 1. Model introduction

[Stable Diffusion 3.5 Medium](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium)
is the primary SD3 checkpoint used in miles-diffusion recipes. SD3 uses a
DiT transformer with two text-conditioning inputs (`encoder_hidden_states` +
`pooled_projections`).

**Key highlights for RL training:**

* Single DiT component — weight sync targets `--update-weight-target-module transformer` (default).
* LoRA on the joint-attention `attn` projections of both streams (image `to_*`, text `add_*_proj` / `to_add_out`).
* Supports **Flow-GRPO** (OCR, HPS, or 0.8 OCR + 0.2 PickScore) and **DiffusionNFT** (PickScore) objectives.
* Gated Hugging Face model — requires `HF_TOKEN`.

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| SD3.5 Medium | [stabilityai/stable-diffusion-3.5-medium](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium) | Default in all SD3 scripts |
| SD3 (other sizes) | Any checkpoint whose name contains `stable-diffusion-3` or `sd3` | Auto-detected by family resolver |

Override family detection when the checkpoint name does not match a registered
pattern (e.g. a renamed local directory):

```bash theme={null}
--diffusion-model-family sd3
```

## 3. Environment setup

SD3.5 is gated — export your Hugging Face token:

```bash theme={null}
export HF_TOKEN=<your_hf_token>
```

Optional prefetch (the scripts load the checkpoint by its Hub id, so it goes into the Hugging Face cache):

```bash theme={null}
hf download stabilityai/stable-diffusion-3.5-medium
```

Prompt datasets live under
[`rockdu/miles-diffusion-datasets`](https://huggingface.co/datasets/rockdu/miles-diffusion-datasets):

| Recipe | Subset | Train path |
| - | - | - |
| GRPO + OCR | `flowgrpo_ocr` | `.../flowgrpo_ocr/train.jsonl` |
| GRPO + HPS | `hpdv2` | `.../hpdv2/train.jsonl` |
| NFT + PickScore | `flowgrpo_pickscore` | `.../flowgrpo_pickscore/train.jsonl` |

Launch scripts download the matching subset automatically via
`command_utils.hf_download_dataset`. Override the download root with
`--data-dir` (default `/root/datasets` →
`/root/datasets/miles-diffusion-datasets`).

## 4. Family config

Registered in `miles/backends/fsdp_utils/configs/sd3.py`:

```python theme={null}
@register_train_pipeline_config("sd3")
class SD3TrainPipelineConfig(TrainPipelineConfig):
    hf_ckpt_name_patterns = ("stable-diffusion-3", "sd3")
    lora_target_modules = [
        "attn.to_q", "attn.to_k", "attn.to_v", "attn.to_out.0",
        "attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj", "attn.to_add_out",
    ]
```

Family resolution (`resolve_diffusion_model_family`) matches `--hf-checkpoint`
names against `hf_ckpt_name_patterns`. `--hf-checkpoint` is the single source
for the HF repo / local directory used by both FSDP training and
sglang-diffusion rollout; pass `--diffusion-model-family` only to override the
auto-detected family.

### SD3-specific training behavior

| Property | SD3 | Notes |
| - | - | - |
| Condition inputs | `encoder_hidden_states` + `pooled_projections` | Concatenated from rollout trajectory |
| CFG combine | `neg + scale * (pos - neg)` | Standard classifier-free guidance |
| Weight sync target | `transformer` (single component) | Multi-component models use comma-separated names |
| Rollout patch group | None | No extra sglang monkey patches |

Rollout responses serialize SD3 conditions as **single tensors** (not lists),
unlike some multi-encoder families.

## 5. Launch

All recipes are Python modules under `scripts/`. Each exposes a Typer CLI
(`ScriptArgs` dataclass) and submits training through
`command_utils.execute_train`. Common overrides: `--cuda-visible-devices`,
`--num-rollout`, `--data-dir`, `--extra-args`.

### 5.1 Script inventory

| Script | Reward | GPUs | Algorithm |
| - | - | - | - |
| `run_diffusion_grpo_sd3_ocr_sglang.py` | OCR (CPU) | 2 colocate | Flow-GRPO |
| `run_diffusion_grpo_sd3_hps_sglang.py` | HPS | 2 colocate | Flow-GRPO |
| `run_diffusion_grpo_sd3_ocr_pickscore_sglang.py` | 0.8 OCR + 0.2 PickScore | 2 colocate | Flow-GRPO |
| `run_diffusion_nft_sd3_pickscore.py` | PickScore | 3 (2+1) | DiffusionNFT |

### 5.2 Flow-GRPO + OCR (2 GPU colocate)

Canonical script: `scripts/run_diffusion_grpo_sd3_ocr_sglang.py`

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
export HF_TOKEN=...
python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py \
  --cuda-visible-devices 6,7
```

Walkthrough: [Quick Start](/docs/diffusion/getting-started/quick-start).

E2E test: `tests/e2e/short/test_sd3_ocr_grpo_2xGPU.py`.

### 5.3 Flow-GRPO + HPS (2 GPU colocate)

Canonical script: `scripts/run_diffusion_grpo_sd3_hps_sglang.py`

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v) — 600 rollouts
(2 optimizer steps each, 1,200 in total) on 2×H200; `rollout/reward/raw_mean` 0.284 → 0.349
(mean of the last 100 rollouts, peak 10-rollout moving average 0.362).

```bash theme={null}
export HF_TOKEN=...
python3 scripts/run_diffusion_grpo_sd3_hps_sglang.py \
  --cuda-visible-devices 6,7
```

The recipe uses the SD3 Flow-GRPO SDE, LoRA rank/alpha, IPC sync and precision settings of the OCR
recipe (not its `--sglang-lora-merge-mode merge`, `--rollout-fetch-in-parser`,
`--rollout-parser-num-workers 16`, or `--rollout-return-full-trajectory`), swaps in the deduplicated `hpdv2` prompts and `--rm-type hps`, and colocates one
HPS reward actor with the train and rollout workers. It keeps Flow-GRPO's own KL weight
(`--diffusion-kl-beta 0.01`) and group-wise advantage std instead of the OCR recipe's
`--diffusion-kl-beta 0.04 --globalize-reward-std`.

### 5.4 Flow-GRPO + OCR & PickScore (2 GPU colocate)

Script: `scripts/run_diffusion_grpo_sd3_ocr_pickscore_sglang.py`

**Status:** [📈 V — Verified](/docs/diffusion/user-guide/recipe-verification#v) — 600 rollouts
(2 optimizer steps each) on 2×H200: `ocr_mean` 0.34 → 0.80, `pickscore_mean` 0.80 → 0.83,
`weighted_mean` 0.43 → 0.81 (means of the last 100 rollouts; 10-rollout moving-average peaks
0.88 / 0.85 / 0.87).

```bash theme={null}
export HF_TOKEN=...
python3 scripts/run_diffusion_grpo_sd3_ocr_pickscore_sglang.py \
  --cuda-visible-devices 6,7
```

The OCR recipe with PickScore mixed in through `weighted_mixture_rm`
(`--custom-rm-args ocr=0.8,pickscore=0.2 --reward-key weighted`, the weighting Stepwise-Flow-GRPO
uses for its OCR task) and `--rollout-shuffle`; the PickScore actor is colocated next to the CPU
OCR pool. Every other flag matches the OCR recipe, except that it leaves `--sglang-lora-merge-mode` at
sgl-d's default `auto` instead of `merge`. See [Rewards](/docs/diffusion/user-guide/rewards)
§ Combining rewards for how the mixture is wired.

### 5.5 DiffusionNFT + PickScore (3 GPU)

Script: `scripts/run_diffusion_nft_sd3_pickscore.py`

**Status:** [🛡️ FG — Fully gated](/docs/diffusion/user-guide/recipe-verification#fg)

```bash theme={null}
export HF_TOKEN=...
python3 scripts/run_diffusion_nft_sd3_pickscore.py \
  --cuda-visible-devices 4,5,2
```

Smoke test (1 rollout, OCR dataset, 2 GPUs):

```bash theme={null}
MILES_SCRIPT_SMOKE=1 python3 scripts/run_diffusion_nft_sd3_pickscore.py
```

### Recipe comparison

| | GRPO + OCR | GRPO + HPS | GRPO + OCR & PickScore | NFT + PickScore |
| - | - | - | - | - |
| Script | `run_diffusion_grpo_sd3_ocr_sglang.py` | `run_diffusion_grpo_sd3_hps_sglang.py` | `run_diffusion_grpo_sd3_ocr_pickscore_sglang.py` | `run_diffusion_nft_sd3_pickscore.py` |
| `--loss-type` | `policy_loss` (default) | `policy_loss` (default) | `policy_loss` (default) | `nft` |
| SDE | Full window, noise=0.7, CFG=4.5 | Full window, noise=0.7, CFG=4.5 | Full window, noise=0.7, CFG=4.5 | ODE, noise=0 |
| Reference | LoRA base KL (β 0.04) | LoRA base KL (β 0.01) | LoRA base KL (β 0.04) | None; `π_old` is the EMA (`--use-ema`) |
| Reward placement | CPU OCR | Colocated HPS actor | CPU OCR + colocated PickScore actor | Dedicated PickScore GPU |
| Verification | FG | V | V | FG |

## 6. Recipe configuration

### GPU layout

**GRPO + OCR (default script):** 2 GPUs colocated (`--colocate`); OCR on CPU Ray actors.

**GRPO + HPS:** 2 GPUs colocated (`--colocate --hps-reward-colocate`); one HPS
actor shares the rollout GPUs.

**GRPO + OCR & PickScore:** 2 GPUs colocated (`--colocate --pickscore-reward-colocate`);
OCR on CPU Ray actors, one PickScore actor shares the rollout GPUs.

**NFT + PickScore:**

| GPU role | Count | Flags |
| - | - | - |
| FSDP train + sglang rollout (colocate) | 2 | `--actor-num-gpus-per-node 2`, `--rollout-num-gpus 2`, `--colocate` |
| PickScore reward | 1 | `--pickscore-num-workers 1`, `--pickscore-num-gpus-per-worker 1.0` |
| **Total** | **3** | `--num-gpus-per-node 3` |

PickScore runs as a Ray actor pool on a dedicated GPU. With `--pickscore-reward-colocate`
(not used in the default script), reward workers share the rollout GPUs
instead — useful only when GPU count is tight.

### Algorithm flags

**Flow-GRPO + OCR:**

| Setting | Value |
| - | - |
| Algorithm | `--loss-type policy_loss` (default) |
| Reference | `--diffusion-kl-beta 0.04` |
| Reward | `--rm-type ocr` |
| SDE | Full window (`num_sde_steps=10`, range `0,10`), noise 0.7, CFG 4.5 |
| Step strategy | `miles.rollout.step_strategy_hub.sde_window` |
| Weight sync | `--lora-ipc-weight-sync` (colocate IPC merge) |
| Determinism | `--deterministic-mode` (CI / e2e parity) |

**Flow-GRPO + HPS:** uses the same SDE, LoRA and precision flags as the OCR recipe,
with these settings of its own:

| Setting | Value |
| - | - |
| Reference | `--diffusion-kl-beta 0.01` |
| Advantage | Group-wise std (no `--globalize-reward-std`) |
| Reward | `--rm-type hps --hps-version v2.1` |
| Reward worker | One actor, batch size 8 |
| Placement | `--hps-reward-colocate` |
| Prompt subset | `hpdv2` |

**DiffusionNFT + PickScore:**

| Setting | Value |
| - | - |
| Algorithm | `--loss-type nft` |
| π\_old | `--use-ema --rollout-weights ema` |
| Reward | `--rm-type pickscore` |
| SDE | `--diffusion-sde-type ode --diffusion-noise-level 0.0` |
| LoRA | rank 32, alpha 64, IPC sync |
| Precision | `--diffusion-forward-dtype fp16` |

## 7. LoRA and weight sync

All SD3 recipes use LoRA with IPC weight sync:

```bash theme={null}
--use-lora \
--lora-ipc-weight-sync \
--lora-rank 32 \
--lora-alpha 64 \
--lora-init-weights gaussian \
--update-weight-buffer-size 2147483648
```

See [LoRA weight sync](/docs/diffusion/advanced/lora) for how the weight-sync updater is selected
and the IPC merge internals.

## 8. Precision notes

All SD3 launch scripts use fp16 DiT forward:

```bash theme={null}
--diffusion-forward-dtype fp16 \
--sglang-dit-precision fp16
```

SD3.5 fp16 policy gradients are small enough to underflow without scaling. When
`--diffusion-forward-dtype fp16`, the FSDP actor automatically enables
**`ShardedGradScaler`** around backward / optimizer step (no extra flag — see
`miles/backends/fsdp_utils/actor.py`). bf16/fp32 forward disables the scaler.

Flow-GRPO recipes also set **`--diffusion-clip-range`** (e.g. `1e-4` in the OCR
script) to clip importance ratios during the policy update.

Train/rollout dtype alignment for Flow-GRPO is covered in
[SDE step backend](/docs/diffusion/advanced/sde-backend).

## 9. Reference results

### Flow-GRPO + OCR

Observed `rollout/reward/raw_mean` from one
`scripts/run_diffusion_grpo_sd3_ocr_sglang.py` run (default batch, 600 rollouts): the mean of the first 50 rollouts
is 0.462 and of the last 50 is 0.841. The curve is an example, not a CI acceptance range:

![Flow-GRPO OCR raw reward](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/sd3/grpo-ocr-raw-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=a72d828f6d0724bcd63e31851dbc7c82)

### Flow-GRPO + OCR & PickScore

Observed rewards from one `scripts/run_diffusion_grpo_sd3_ocr_pickscore_sglang.py` run (600 rollouts);
the weighted sum GRPO trains on, then its two components:

![Flow-GRPO OCR + PickScore weighted reward](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/sd3/grpo-ocr-pickscore-weighted-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=7bed3dc3d11548c5575f4ce4c9a42ba3)

![Flow-GRPO OCR + PickScore: OCR component](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/sd3/grpo-ocr-pickscore-ocr-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=59b983b4fd00f5bac6cb3f5b339e7920)

![Flow-GRPO OCR + PickScore: PickScore component](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/sd3/grpo-ocr-pickscore-pickscore-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=43027ac426b87a7d5e54f0e9042ee77b)

### DiffusionNFT + PickScore

Observed `rollout/reward/raw_mean` from one `scripts/run_diffusion_nft_sd3_pickscore.py` run (100 rollouts):

![DiffusionNFT PickScore raw reward](https://mintcdn.com/radixark/sL-lqXHDLHHpg3Xm/diffusion/assets/images/sd3/nft-pickscore-raw-reward.png?fit=max\&auto=format\&n=sL-lqXHDLHHpg3Xm\&q=85\&s=58daeaefc94a5a058f40a5a3af621b0e)

That run observed held-out
**`eval/pickscore_test` 0.850 / 0.859 / 0.863** at rollouts 29 / 59 / 89 with `--eval-interval 30` and 50 denoise steps
at eval time; these values are not asserted by the E2E fixture.

## 10. Pairs well with

* [Quick Start](/docs/diffusion/getting-started/quick-start) — SD3.5 Flow-GRPO OCR walkthrough.
* [Rewards](/docs/diffusion/user-guide/rewards) — OCR, HPS and PickScore scoring.
* [Customization](/docs/diffusion/user-guide/customization) — `--*-path` plug-points.
* [SDE step backend](/docs/diffusion/advanced/sde-backend) — SDE window (GRPO) vs ODE (NFT).
* [LoRA weight sync](/docs/diffusion/advanced/lora) — IPC merge used by all SD3 recipes.


# Wan2.2-T2V-A14B
Source: https://miles.radixark.com/docs/diffusion/models/wan/wan2-2

Dual-expert MoE video model — Flow-GRPO + PickScore recipe, LoRA SFT recipe, and the high/low-noise expert boundary.

## 1. Model introduction

[Wan2.2-T2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) is a text-to-video model with a **dual-expert
MoE DiT**: a high-noise expert (`transformer`) denoises timesteps `t ≥ boundary` and a low-noise expert
(`transformer_2`) handles the rest. Conditioning comes from a UMT5 text encoder; latents go through the Wan VAE (4×
temporal compression).

**Key highlights for RL training:**

* **Two experts, one boundary.** `boundary_ratio = 0.875` — which expert a train pair updates depends only on its
  timestep. Single- and dual-expert training are both supported: `--update-weight-target-module` names the experts to
  load, train, and sync.
* **Two guidance scales.** Rollout denoises low-noise steps with `guidance_scale_2`. Left unset, the engine silently
  substitutes its own default (the model's sampling config, else `guidance_scale`), which the trainer cannot see — so
  training asserts `--diffusion-guidance-scale-2` is set explicitly, because a silent mismatch against rollout would
  corrupt the ratio.
* **USP-ready.** Wan was enabled for Ulysses × Ring sequence parallelism.

## 2. Supported variants

| Model | HF ID | Notes |
| - | - | - |
| Wan2.2-T2V-A14B | [Wan-AI/Wan2.2-T2V-A14B-Diffusers](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) | Default in all three Wan recipes |

## 3. Family config

Registered in `miles/backends/fsdp_utils/configs/wan2_2.py`:

| Property | Value | Notes |
| - | - | - |
| Expert routing | `t ≥ 0.875 × num_train_timesteps` → `transformer`, else `transformer_2` | `component_for_timestep` |
| Guidance routing | high-noise → `--diffusion-guidance-scale`, low-noise → `--diffusion-guidance-scale-2` (required) | `select_guidance_scale` |
| Timestep scaling | None — Wan DiT takes raw scheduler timesteps (0..1000) | |
| Condition inputs | `encoder_hidden_states` only (fixed-length UMT5 embeds) | Concat-collate, no padding needed |
| CFG combine | `neg + scale × (pos − neg)` | Standard |
| CFG batching | Off | |
| FSDP fp32 param map | `norm2` parameters only | `models/diffusers/wan2_2/parallel_plan.py`; the separate rollout patch group aligns selected norm and `scale_shift_table` arithmetic engine-side |

## 4. Launch

### 4.1 Flow-GRPO + PickScore (4 train GPUs + 1 reward GPU)

Canonical recipe: `scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py`

**Status:** [○ NV — Not verified](/docs/diffusion/user-guide/recipe-verification#nv)

```bash theme={null}
python3 scripts/run_diffusion_grpo_wan22_pickscore_5gpu.py
```

### 4.2 Full-finetune Flow-GRPO + PickScore, multi-node (2×8 train GPUs + 1 reward GPU)

Recipe: `scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py`
(full finetune, no LoRA, true on-policy). Start the
[multi-node Ray cluster](/docs/diffusion/user-guide/launch-script#multi-node-training), then run on the head node:

**Status:** [🧩 PG — Proxy gated](/docs/diffusion/user-guide/recipe-verification#pg)

```bash theme={null}
MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py
```

`--four-gpu-ci` scales the recipe to one 4-GPU node for e2e CI.

### 4.3 LoRA SFT on (video, prompt) pairs (4 GPUs, no rollout engines)

Recipe: `scripts/run_diffusion_sft_wan22.py`

**Status:** [○ NV — Not verified](/docs/diffusion/user-guide/recipe-verification#nv)

```bash theme={null}
MILES_SCRIPT_DATA_JSONL=/abs/data.jsonl python3 scripts/run_diffusion_sft_wan22.py
```

Each jsonl row is `{"prompt": "...", "target": {"visual": "/abs/path.mp4"}}`. Prepare the videos offline: each must
already be 832×480 with square pixels and exactly 161 frames (81 frames at `--sft-frame-stride 2`). The encoder
rejects any other video rather than resizing or trimming it.

## 5. Reference results

### 5.1 17-GPU multi-node full-finetune GRPO

#### CFG enabled (`4.0/3.0`)

* `train/model_output_mean_abs_diff`: **0.0** every rollout.
* `rollout/reward/raw_mean`: \~0.77 → \~0.82.
* Final `eval/pickscore_test`: **0.8231** (rollout 199, 28 denoising steps—not 40,
  UniPC, 2048 prompts).

![Wan2.2 CFG PickScore reward mean](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/wan/reward_mean.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=633d3257c774ad2dd9a17a5eaf6a31cd)

#### CFG disabled (`1.0/1.0`)

* `rollout/reward/raw_mean`: \~0.72 → \~0.83.
* Final `eval/pickscore_test`: **0.8243** (rollout 199, 28 denoising steps—not 40,
  UniPC, 2048 prompts).

![Wan2.2 no-CFG PickScore reward mean](https://mintcdn.com/radixark/-43MEB1vZnsTAlRx/diffusion/assets/images/wan/reward_mean_nocfg.png?fit=max\&auto=format\&n=-43MEB1vZnsTAlRx\&q=85\&s=2a376a40a9c67639cfc18b49530791fc)

## 6. Pairs well with

* [Single-Prompt Multi-Generation](/docs/diffusion/advanced/single-prompt-multi-gen) — the microgroup mechanics behind
  `--rollout-microgroup-size 8`.
* [LoRA Training and Weight Sync](/docs/diffusion/advanced/lora) — IPC merge used by the GRPO recipe.
* [SDE Step Backend](/docs/diffusion/advanced/sde-backend) — how the trained SDE step is scored train-side.
* [Rewards](/docs/diffusion/user-guide/rewards) — PickScore worker pool configuration.


# CLI Reference
Source: https://miles.radixark.com/docs/diffusion/user-guide/cli-reference

Every flag train_diffusion.py accepts, grouped by subsystem.

miles-diffusion is configured entirely through flags on `train_diffusion.py`, coming from three
places:

| Source | Where | Notes |
| - | - | - |
| Training-backend flags | `FSDPArgs` in `miles/backends/fsdp_utils/arguments.py` | Every dataclass field becomes `--field-name` automatically |
| miles flags | `miles/utils/arguments.py` | Added on top as an `extra_args_provider` |
| sglang-d passthrough | The rollout engine's own CLI | Re-registered with a `--sglang-` prefix; a short skip list covers what miles sets itself (`model_path`, ports, `base_gpu_id`, `random_seed`, …) |

`--custom-config-path <file.yaml>` loads YAML keys into the namespace at the very end of argument
validation. A key that collides with an existing argument **overrides it** (with a logged warning) —
the YAML wins over the command line, but the override skips every check and derived value computed
before it (a YAML `global_batch_size`, for example, does not re-derive `num_steps_per_rollout`).

`--config <file.yaml>` is different: it only adds YAML keys that are not arguments at all, so it
cannot change any flag.

Reading the prefixes:

| Prefix | Marks |
| - | - |
| `diffusion-` | Diffusion-specific settings — denoising, SDE, CFG, latents, frames, and the diffusion variants of loss knobs (`--diffusion-clip-range`, `--diffusion-kl-beta`). Shared knobs (EMA, LoRA, batching) do not take it |
| `fsdp-` | The training side — compare `--fsdp-flow-shift` (SFT training sigma grid) with `--diffusion-flow-shift` (the engine's generation schedule) |
| `rollout-` / `sglang-` | The engine side |

The prefix does not tell you which argument group a flag lives in — groups follow concern, not
name. `python3 train_diffusion.py --help` is always the ground truth.

***

## Essentials

### The required flags

| Flag | What |
| - | - |
| `--hf-checkpoint` | The diffusers pipeline to train, as an HF repo id or a local directory. |
| `--diffusion-task` | The task to train on that family, e.g. `t2i` for SD3, `t2v` for Wan2.2. |

One value serves three readers: training loads components from it, the sglang-d
engine serves it, and the **model family** is matched from its name. Add
`--diffusion-model-family` when the name carries no family hint — which local weights usually
do not.

### Cluster topology

| Flag | Default | What |
| - | - | - |
| `--actor-num-nodes` | `1` | Nodes for the training actor. |
| `--actor-num-gpus-per-node` | `8` | GPUs per actor node. Train world size is the product. |
| `--rollout-num-gpus` | – | GPUs for rollout-side work. Forced to the train world size under `--colocate`. |
| `--rollout-num-gpus-per-engine` | `1` | GPUs per sglang-d engine (its `num_gpus`). By default they all go to TP (`--sglang-tp-size`); SP is opt-in via `--sglang-sp-degree`, with `tp × sp` equal to this value (checked when both are set). CFG-parallel (`--sglang-enable-cfg-parallel`) is not supported for training — the trainer's CFG combine does not match it. |
| `--num-gpus-per-node` | `8` | Total GPUs the job may use per node. **Set this when using fewer than 8.** |
| `--colocate` | off | Train and rollout share GPUs and sync weights over CUDA IPC. Without it, rollout runs on its own `--rollout-num-gpus` and weights sync over NCCL (required by `train_diffusion_async.py`). |

### Batch sizing

| Flag | Default | What |
| - | - | - |
| `--rollout-batch-size` | unset | Prompts per rollout. **Required** — the parser rejects a run without it. |
| `--n-samples-per-prompt` | `1` | Samples per prompt (GRPO group size). |
| `--global-batch-size` | derived | **Samples** per optimizer step. Must divide by `dp_size`. |
| `--num-steps-per-rollout` | derived | Optimizer steps per rollout. |
| `--micro-batch-size` | `1` | **Train pairs** per DiT forward (flat cut). |
| `--micro-batch-size-sample` × `--micro-batch-size-tstep` | – | 2D tile per DiT forward. Set together; overrides `--micro-batch-size` with their product. |
| `--num-rollout` | – | Total rollout iterations. Set it or `--num-epoch`. |
| `--num-epoch` | – | Alternative to `--num-rollout`; ignored if both are set. |

For how these batch arguments relate to each other, see
[the batch-knob invariant](/docs/diffusion/user-guide/concepts#the-batch-knob-invariant).

### Diffusion sampling

| Flag | Default | What |
| - | - | - |
| `--diffusion-num-steps` | `10` | Denoising steps per rollout sample. |
| `--diffusion-num-sde-steps` | `0` | How many steps become train pairs; read by the step strategy. |
| `--diffusion-step-strategy-path` | – | Which steps, e.g. `sde_window`. Flow-GRPO training needs one. |
| `--diffusion-noise-level` | `0.7` | SDE noise injected during rollout. |
| `--diffusion-guidance-scale` | `4.0` | CFG scale. |
| `--diffusion-height` / `--diffusion-width` | `512` | Output resolution. |
| `--diffusion-output-num-frames` | `1` | Frames — `1` for images. |
| `--rollout-microgroup-size` | `1` | Samples per prompt sent in one `POST /rollout/generate`. |

### Precision

| Flag | Default | What |
| - | - | - |
| `--diffusion-forward-dtype` | `bf16` | Training-side DiT compute — FSDP `param_dtype` and the input cast. It reaches the engine only as the LoRA-merge precision (`fp32` merges in fp32, see [LoRA](/docs/diffusion/advanced/lora)); the engine's DiT precision is `--sglang-dit-precision`, which is not derived from it — set both to the same value. |
| `--fsdp-master-dtype` | `fp32` | Sharded master weights and optimizer state. |
| `--fsdp-reduce-dtype` | `fp32` | Gradient reduce-scatter. |

See [Dtype Control](/docs/diffusion/advanced/dtype-control).

### Algorithm

| Flag | Default | What |
| - | - | - |
| `--loss-type` | `policy_loss` | `policy_loss` (Flow-GRPO), `nft`, `sft_loss`, `custom_loss`. |
| `--advantage-estimator` | `grpo` | Only `grpo` today. |
| `--diffusion-clip-range` | `1e-4` | PPO ratio clip. |
| `--diffusion-adv-clip-max` | `5.0` | Advantage clamp. |
| `--diffusion-kl-beta` | `0.0` | Reference-KL coefficient; `> 0` turns on a reference forward. |
| `--globalize-reward-std` | off | Batch-wide std instead of per-group. |
| `--globalize-reward-mean` | off | Batch-wide mean instead of per-prompt. |

***

## Complete reference

### Cluster

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--actor-num-nodes` | int | `1` | |
| `--actor-num-gpus-per-node` | int | `8` | |
| `--rollout-num-gpus` | int | – | For train-only SFT, unset colocates encoders with training; set it to reserve dedicated encoder GPUs, which also lets encoding overlap training. |
| `--rollout-num-gpus-per-engine` | int | `1` | sgl-d `num_gpus` per engine; defaults to all-TP. |
| `--num-gpus-per-node` | int | `8` | |
| `--colocate` | flag | off | Train and rollout share GPUs and sync weights over CUDA IPC; without it, rollout gets its own GPUs and weights sync over NCCL. Async training requires it off. |
| `--offload` | flag | off | `--offload-train` + `--offload-rollout`. |
| `--offload-train` / `--no-offload-train` | tri-state | – | On by default under `--colocate`. |
| `--offload-rollout` / `--no-offload-rollout` | tri-state | – | On by default under `--colocate`. |
| `--skip-train-actor-gc-collect` | flag | off | The train actor's `clear_memory` after each rollout's training skips `gc.collect()` and only releases the CUDA cache; saves 0.2-0.5 s per rollout on a large actor process. |
| `--distributed-backend` | str | `nccl` | |
| `--distributed-timeout-minutes` | int | `10` | |

### Training backend

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--train-backend` | enum | `fsdp` | Only value. |
| `--fsdp-master-dtype` | enum | `fp32` | `fp32` / `bf16` / `fp16`. Load, shard, and optimizer-state precision. |
| `--fsdp-reduce-dtype` | enum | `fp32` | `bf16` matches flow\_grpo's all-bf16 policy but adds cross-rank add-noise. |
| `--diffusion-forward-dtype` | enum | `bf16` | `bf16` / `fp16` / `fp32`. |
| `--fsdp-cpu-offload` | flag | off | Offloads params, grads, optimizer state; the optimizer then runs on CPU. |
| `--fsdp-cpu-backend` | str | `gloo` | CPU process group for the above. |
| `--dp-replicate-size` | int | `1` | Hybrid sharding: replica count. `dp_shard` takes the ranks left over from this and SP. |
| `--sequence-parallel-size` | int | `1` | USP = Ulysses × Ring. |
| `--ulysses-degree` | int | `0` | `0` = auto (Ulysses fills SP). Ring degree > 1 needs torch ≥ 2.11 and a ring-capable attention backend. |
| `--fsdp-attention-backend` | str | – | diffusers `set_attention_backend` value. |
| `--fsdp-flow-shift` | str | – | SFT training sigma grid shift per latent stream, e.g. `visual=12,audio=3`; a bare number is the visual shift. RL pairs carry their rollout sigmas instead. Distinct from `--diffusion-flow-shift`. |
| `--fsdp-supervised-streams` | str | `visual,audio` | Comma-separated latent streams the SFT loss supervises, among those the pair carries; must include `visual`. A stream left out is still noised and fed to the DiT but adds no loss. RL trajectories carry only the visual stream. |
| `--gradient-checkpointing` | flag | off | |
| `--deterministic-mode` | flag | off | See [Deterministic Training](/docs/diffusion/advanced/deterministic). |
| `--train-env-vars` | JSON | `{}` | Extra env for the training processes. |

### Optimizer and schedule

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--optimizer` | str | `adam` | AdamW. The only value the actor accepts. |
| `--lr` | float | `1e-6` | |
| `--adam-beta1` / `--adam-beta2` | float | `0.9` / `0.999` | β₂ matches flow\_grpo, not the LLM-side `0.95`. |
| `--adam-eps` | float | `1e-8` | |
| `--weight-decay` | float | `0.0` | |
| `--clip-grad` | float | `1.0` | |
| `--lr-decay-style` | str | `constant` | `constant` / `linear` / `cosine` / `inverse-square-root` / `WSD`. |
| `--min-lr` / `--lr-warmup-init` | float | `0.0` | |
| `--lr-warmup-iters` | int | `0` | |
| `--lr-warmup-fraction` | float | – | |
| `--lr-decay-iters` | int | – | |
| `--lr-wsd-decay-iters` / `--lr-wsd-decay-style` | int / str | – | For `--lr-decay-style WSD`; style is `linear` / `cosine` / `exponential` / `minus_sqrt`. |
| `--use-checkpoint-lr-scheduler` | flag | off | No effect in the FSDP backend, which always restores a saved LR-scheduler state from `--load`. |
| `--override-lr-scheduler` | flag | off | No effect in the FSDP backend; rejected together with `--use-checkpoint-lr-scheduler`. |
| `--seed` | int | `1234` | |

### Rollout

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--hf-checkpoint` | str | – | **Required.** Pipeline to train and to serve; also the family source. |
| `--diffusion-model-family` | str | – | Registered family key: `sd3`, `wan2_2`, `ltx`, `qwen_image`, `qwen_image21`, `cosmos3`, `krea2`, `h3`. Overrides name matching. The longest pattern wins, so `Qwen-Image-2.1` resolves to `qwen_image21`. |
| `--diffusion-task` | str | – | **Required** unless the config declares no `tasks`. The task to train, one of the family config's `tasks`: `t2i` (sd3, krea2, qwen\_image, qwen\_image21, cosmos3), `t2v` (wan2\_2, ltx), `t2va` (h3, the only H3 rollout task) or `ref2va` (h3 joint video/audio SFT with ordered references on `transformer_ref`). |
| `--rollout-function-path` | str | `miles.rollout.sglang_rollout.generate_rollout` | Inherited Miles default; that module is not shipped in miles-diffusion, so set it — diffusion recipes pass `miles.rollout.sglang_diffusion_rollout.generate_rollout`. |
| `--train-pipeline-config-path` | str | – | Your own `TrainPipelineConfig` for an unregistered family. Mutually exclusive with `--diffusion-model-family`. |
| `--model-backend-path` | str | – | Override the family's model loader. |
| `--diffusion-num-steps` | int | `10` | |
| `--diffusion-flow-shift` | float | – | Generation-schedule shift, sent to the engine. |
| `--rollout-microgroup-size` | int | `1` | |
| `--diffusion-fps` | float | – | Video only. |
| `--diffusion-output-num-frames` | int | `1` | |
| `--rollout-video-dtype` | enum | `keep` | `keep` / `uint8`: dtype of the decoded video in rollout responses. `uint8` quantizes engine-side by truncation (×255, then clamp), cutting the response \~4x. That matches the PickScore reward path; HPS, OCR and the OpenAI-API reward round instead, so keep `keep` for them and for any consumer that needs the float tensor. |
| `--diffusion-guidance-scale` | float | `4.0` | |
| `--diffusion-guidance-scale-2` | float | – | Wan2.2 low-noise expert; **required** when training it. |
| `--diffusion-true-cfg-scale` | float | – | |
| `--diffusion-negative-prompt` | str | – | Unset sends none, so the engine's per-model default applies; flow\_grpo recipes pass `" "`. |
| `--diffusion-noise-level` | float | `0.7` | |
| `--diffusion-quality` | enum | `exact` | `exact` / `lossless` / `high`: sglang-diffusion request quality level. `exact` is the bit-identical reference path; the looser levels mount rounding-changing fusions. |
| `--diffusion-height` / `--diffusion-width` | int | `512` | Rollout output size; SFT target media must already be this size. |
| `--diffusion-h3-aspect-ratio` | str | `16:9` | MiniMax H3 only: `21:9` / `16:9` / `4:3` / `1:1` / `3:4` / `9:16`. |
| `--diffusion-h3-duration-seconds` | float | `5.0` | MiniMax H3 only: rollout duration, 4.0–15.0 s. |
| `--diffusion-audio-flow-shift` | float | `3.0` | MiniMax H3 only: audio flow shift for rollout. |
| `--diffusion-sde-type` | str | `sde` | `sde` / `cps` / `ode`, checked at startup unless `--sde-step-backend-path` is set. Selects the train-side SDE backend too. |
| `--sde-step-backend-path` | str | – | Custom dynamics. See [SDE backends](#sde-step-backends). |
| `--diffusion-num-sde-steps` | int | `0` | Read by the step strategy. |
| `--diffusion-sde-window-range` | `"lo,hi"` | – | For `sde_window`. Defaults to `[0, num_inference_steps)`. |
| `--diffusion-sde-candidate-steps` | `"1,2,3"` | – | Required by `epoch_global_random_choice`. |
| `--diffusion-step-strategy-path` | str | – | `(args, sample, num_steps, seed) -> (sde, ret)` choosing the SDE and returned steps, e.g. `miles.rollout.step_strategy_hub.sde_window`. Flow-GRPO needs one. |
| `--rollout-return-full-trajectory` | flag | off | Ask the engine for the whole denoising trajectory instead of only the steps the step strategy requests. The train tensors are identical; kept for debugging and A/B runs. |
| `--diffusion-log-prob-no-const` | flag | off | Drop log-prob constants on the engine (pairs with the CPS backend). Forced on for `cps` / `ode`. |
| `--diffusion-generator-device` | str | `cuda` | |
| `--rollout-patch-group` | str | – | Comma-separated numeric-parity patch groups: `qwen_image`, `wan`, `ltx`. |
| `--update-weight-target-module` | str | `transformer` | Modules to train and sync. Wan2.2: `transformer,transformer_2`. |
| `--update-weight-buffer-size` | int | 512 MiB | Weight-sync chunk size in bytes. |
| `--rollout-seed` | int | `42` | |
| `--rollout-shuffle` | flag | off | Shuffle the prompt order once per epoch (seeded by `--rollout-seed`). Off reads the jsonl in file order; the SFT rollout (`miles.rollout.sft_rollout`) shuffles every epoch regardless. |
| `--over-sampling-batch-size` | int | – | Must equal `--rollout-batch-size` today. |
| `--sglang-server-concurrency` | int | `512` | In-flight generate requests are capped at this × the number of engines — one global limit, with each request routed to the least-loaded engine. |
| `--use-distributed-post` | flag | off | Rollout HTTP POSTs go through per-node Ray actors instead of the local client (`--num-gpus-per-node` actors per node). |
| `--use-miles-router` | flag | off | **Required** — the SGLang router is not supported here. |
| `--miles-router-timeout` | float | – | |
| `--miles-router-max-connections` | int | – | |
| `--miles-router-health-check-failure-threshold` | int | `3` | |

### Data and batching

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--prompt-data` | str | – | jsonl, one row per prompt. |
| `--input-key` | str | `input` | |
| `--metadata-key` | str | `metadata` | |
| `--data-source-path` | str | – | Defaults to `RolloutDataSourceWithBuffer`. |
| `--disable-rollout-global-dataset` | flag | off | Manage data yourself. The shipped diffusion rollout asserts the global dataset, so this needs a custom `--rollout-function-path`. |
| `--rollout-batch-size` | int | – | **Required.** |
| `--n-samples-per-prompt` | int | `1` | |
| `--global-batch-size` | int | derived | |
| `--num-steps-per-rollout` | int | derived | |
| `--micro-batch-size` | int | `1` | Flat: train-pair dicts per DiT forward, contiguous within an optimizer window. |
| `--micro-batch-size-sample` | int | – | 2D: samples per DiT-forward tile. |
| `--micro-batch-size-tstep` | int | – | 2D: SDE timesteps per tile. Set with the above. |
| `--diffusion-train-iter-order` | enum | `sample_major` | `sample_major` / `timestep_major`. Tile visit order; only meaningful with the 2D pair. |
| `--train-dp-split-mode` | enum | `contiguous` | `contiguous` lets a micro-batch reproduce a rollout microgroup exactly; `stride` deals round-robin. |
| `--num-rollout` | int | – | Required unless `--num-epoch` is set. |
| `--num-epoch` | int | – | |
| `--start-rollout-id` | int | – | Resumed from `--load` when unset. |
| `--sft-encoder-checkpoint` | str | – | SFT only: tokenizer/text\_encoder/vae source. |
| `--sft-frame-stride` | int | `1` | SFT encode temporal stride; target videos hold exactly `(num_frames - 1) * stride + 1` frames. |
| `--sft-offload-encoder` | flag | off | SFT only: keep the frozen encoder in host RAM, on the GPU only during encode bursts. |

### Evaluation

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--eval-interval` | int | – | Requires configured eval datasets. |
| `--eval-prompt-data` | str+ | – | `<name> <path>` pairs, all after one flag (repeating the flag keeps only the last). |
| `--eval-config` | str | – | OmegaConf YAML/JSON; overrides `--eval-prompt-data`. |
| `--eval-function-path` | str | – | Defaults to `--rollout-function-path`. |
| `--eval-input-key` | str | – | |
| `--n-samples-per-eval-prompt` | int | `1` | |
| `--diffusion-eval-num-steps` | int | – | Defaults to `--diffusion-num-steps`. |
| `--skip-eval-before-train` | flag | off | |
| `--eval-reward-key` | str | – | Defaults to `--reward-key`. |

### Algorithm and loss

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--loss-type` | enum | `policy_loss` | `policy_loss` / `nft` / `sft_loss` / `custom_loss`. |
| `--advantage-estimator` | enum | `grpo` | |
| `--disable-grpo-std-normalization` | flag | off | Subtract the group mean only; skip dividing by the group std. Std normalization is also skipped when `n_samples_per_prompt == 1`, since a one-sample group has no std. With `--globalize-reward-mean` the batch mean is subtracted instead, and the batch-wide std from `--globalize-reward-std` is skipped too. |
| `--globalize-reward-mean` | flag | off | Leave **off** for flow\_grpo parity. |
| `--globalize-reward-std` | flag | off | **On** for flow\_grpo's PickScore recipe. |
| `--diffusion-clip-range` | float | `1e-4` | |
| `--diffusion-adv-clip-max` | float | `5.0` | Under `nft` this also sets the advantage-to-`r` slope. |
| `--diffusion-recompute-old-log-prob` | flag | off | Recompute old log-probs with the trainer forward instead of trusting the rollout's. `policy_loss` only. |
| `--diffusion-kl-beta` | float | `0.0` | KL to the `--ref-mode` model; under `nft` adds `kl_beta * mean((v_θ − v_ref)²)` per pair. |
| `--ref-mode` | enum | – | KL reference: `none` / `lora_base` (actor with adapters disabled) / `ref` (the model from `--ref-load`). Auto: `lora_base` when KL > 0. |
| `--ref-load` | str | – | HF pipeline checkpoint of an independent frozen reference model. Requires `--ref-mode ref`. |
| `--ref-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--ref-load`, one per `--update-weight-target-module` entry, in order. |
| `--ref-cpu-offload` | flag | off | Keep the reference shards in pinned CPU memory between forwards (FSDP `CPUOffloadPolicy`). |
| `--teacher-load` | str | – | HF pipeline checkpoint of a frozen teacher model, read only by custom prepare/loss hooks as `ctx.teacher_models`, so it requires `--custom-prepare-train-batch-path` or `--custom-loss-function-path`. |
| `--teacher-lora-adapter-path` | str+ | – | PEFT adapter directories applied to `--teacher-load`, one per `--update-weight-target-module` entry, in order. |
| `--teacher-cpu-offload` | flag | off | Keep the teacher shards in pinned CPU memory between forwards. |
| `--custom-prepare-train-batch-path` | str | – | Builds DiT inputs. |
| `--custom-loss-function-path` | str | – | Loss **formula** only — the DiT forward stays in the actor. |
| `--diffusion-nft-beta` | float | `1.0` | |
| `--diffusion-nft-timestep-strategy-path` | str | `…nft.drop_random_steps` | `(args, num_steps, generator) -> list[int]` returning the denoising steps each sample trains on, in training order; `generator` is seeded per sample. Built in: `drop_random_steps` (random order), `drop_final_steps` (schedule order, noisiest first). The final clean output (σ = 0) is never a training step. |
| `--diffusion-nft-num-dropped-timesteps` | int | `0` | Steps the built-in strategies drop: random ones per sample for `drop_random_steps`, the final (lowest-noise) ones for `drop_final_steps`. |
| `--no-diffusion-nft-adaptive-weight` | flag | off | |

### Reward

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--rm-type` | str | – | `pickscore` / `hps` / `ocr` / `openai_api` / `custom_api`. Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. |
| `--custom-api-rm-config` | str | – | Required for `--rm-type custom_api`: a YAML file or inline `base64:<payload>` with `actor_class`, `actor_kwargs`, `max_concurrency`. A flat mapping configures the default OpenAI-compatible image actor. |
| `--openai-api-rm-config` | str | – | Required for `--rm-type openai_api`; same format as `--custom-api-rm-config`. |
| `--group-rm` | flag | off | Score a whole prompt group at once. |
| `--custom-rm-path` | str | – | `async def rm(args, samples) -> list[float]`. Batched only; replaces the `--rm-type` dispatch entirely. Shipped: `miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm` (weighted sum of built-in rewards). |
| `--custom-rm-args` | str | – | Opaque config string for the custom RM, read as `args.custom_rm_args`; e.g. `"hps=0.7,pickscore=0.3"` for `rm_hub.weighted_mixture_rm`. |
| `--rm-url` | str | – | URL of a remote reward service, e.g. `http://localhost:8000`; mirrors miles core. No built-in `--rm-type` reads it yet, so read it from a `--custom-rm-path` function. |
| `--reward-key` | str | – | For dict-valued rewards: the entry GRPO trains on. Every entry is also logged as `rollout/reward/<key>_mean` and `eval/<dataset>/<key>`. |
| `--custom-reward-post-process-path` | str | – | Replace advantage normalisation. |
| `--pickscore-model-path` | str | – | Required for `--rm-type pickscore`. |
| `--pickscore-processor-path` | str | – | Required for `--rm-type pickscore`. |
| `--pickscore-num-workers` | int | `1` | |
| `--pickscore-num-gpus-per-worker` | float | `1.0` | Standalone workers only; fractional values allowed. |
| `--pickscore-reward-colocate` | flag | off | Seats `--pickscore-num-workers` actors on the rollout GPUs, one per placement-group bundle, sharing them with the train actor and rollout engine. Requires `--colocate`. |
| `--pickscore-batch-size` | int | `8` | |
| `--pickscore-num-frames` | int | – | Frames scored per video; unset = all. |
| `--hps-num-workers` | int | `1` | |
| `--hps-num-gpus-per-worker` | float | `1.0` | Standalone workers only; fractional values allowed. |
| `--hps-reward-colocate` | flag | off | Seats `--hps-num-workers` actors on the rollout GPUs, one per placement-group bundle, sharing them with the train actor and rollout engine. Requires `--colocate`. |
| `--hps-batch-size` | int | `8` | |
| `--hps-version` | enum | `v2.1` | `v2.0` / `v2.1`. |
| `--hps-checkpoint-path` | str | – | Local checkpoint; unset downloads from Hugging Face. |
| `--ocr-num-workers` | int | `4` | |
| `--rollout-parser-num-workers` | int | `1` | Ray actors deserializing rollout responses. Raise when trajectory tensors are large. |
| `--rollout-fetch-in-parser` | flag | off | Parser actors pick the least-loaded engine through the miles router, then fetch and parse its response in place, skipping the router's data plane and the manager's event loop. Requires `--use-miles-router` and at least as many parser workers as concurrency slots. |

### Rollout customization hooks

Every one takes a dotted path.

| Flag | Replaces |
| - | - |
| `--custom-generate-function-path` | The inner `generate` call in the diffusion rollout. |
| `--dynamic-sampling-filter-path` | Per-group keep/drop (DAPO-style). |
| `--buffer-filter-path` | Buffer selection. |
| `--custom-convert-samples-to-train-data-path` | The whole samples → train-data conversion. |
| `--custom-expand-samples-to-train-pairs-path` | Just the sample → train-pair expansion. |
| `--custom-rollout-log-function-path` | Train rollout logging. |
| `--custom-eval-rollout-log-function-path` | Eval rollout logging. |

### LoRA

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-lora` | flag | off | |
| `--lora-rank` / `--lora-alpha` | int | `64` / `64` | |
| `--lora-adapter-path` | str+ | – | PEFT adapter directories the actor continues from, one per `--update-weight-target-module` entry, in order; `--lora-rank` / `--lora-alpha` must match their configs. |
| `--lora-target-modules` | str+ | – | Defaults per model family. |
| `--lora-init-weights` | str | `gaussian` | `kaiming-uniform` maps to PEFT's default; other PEFT schemes pass through. |
| `--lora-ipc-weight-sync` | flag | off | Push only `lora_A`/`lora_B` over CUDA IPC and let the rollout merge them (`weight_update_mode=lora_merge`); without it, the trainer merges LoRA into the base weights and pushes full weights. Requires `--use-lora`, and `--colocate` whenever weights are synced (not under `--train-only` / `--debug-rollout-only`). Without `--colocate`, LoRA always pushes `lora_A`/`lora_B` over NCCL for the rollout to merge. |

### EMA

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-ema` | flag | off | Maintains an EMA copy as `π_old`. Needs a consumer: `--rollout-weights ema`, or `--loss-type nft`, which trains against it and requires both this flag and `--rollout-weights ema`. |
| `--rollout-weights` | enum | `actor` | `actor` / `ema`: which weights get pushed to rollout. `ema` requires `--use-ema`. |
| `--ema-decay-init` | float | `0.001` | Decay during the flat period, counted in actor optimizer steps. |
| `--ema-decay-ramp` | float | `0.001` | Increase per actor optimizer step after the flat period; the ramp restarts from zero. |
| `--ema-decay-max` | float | `0.5` | Ceiling. |
| `--ema-decay-flat-steps` | int | `0` | |

### Checkpointing

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--save` | str | – | |
| `--save-interval` | int | – | Requires `--save`. |
| `--no-save-optim` | flag | off | Skips AdamW and LR-scheduler state; `--load` still resumes the model, EMA, and RNG with a fresh optimizer. |
| `--load` | str | – | |
| `--ckpt-step` | int | – | Defaults to `latest_checkpointed_iteration.txt`. |
| `--no-load-optim` / `--no-load-rng` | flag | off | |
| `--async-exact-resume` | flag | off | Async training only. Also checkpoints the EMA from before the latest step and resamples the first resumed batch with it, so a resume matches an uninterrupted run. Off: that batch samples the loaded EMA. |

### Logging

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-wandb` | flag | off | |
| `--wandb-project` | str | – | |
| `--wandb-group` / `--wandb-run-id` / `--wandb-team` / `--wandb-host` / `--wandb-key` | str | – | |
| `--wandb-mode` | enum | – | `online` / `offline` / `disabled`. Overrides `WANDB_MODE`. |
| `--wandb-dir` | str | – | Defaults to `./wandb`. |
| `--disable-wandb-random-suffix` | flag | off | Run names include a random suffix by default; pass this flag to disable it. |
| `--wandb-log-num-images` | int | `0` | Images/videos per rollout; `0` disables. |
| `--wandb-log-image-interval` | int | `1` | Send media every N rollouts. |
| `--log-loss-sigma-bucket` | int | `10` | Sigma buckets for the per-bucket loss curves (emitted by the SFT loss): `loss_<stream>_sigma_*` for each supervised stream, bucketed by its own sigma; the total `loss` is never bucketed, since each stream reads the shared grid index through its own shift. `0` disables. |
| `--use-miles-dashboard` | flag | off | Async phase/trajectory telemetry. |
| `--miles-dashboard-workspace` | str | `./miles_dashboard` | |

### Fault tolerance

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-fault-tolerance` | flag | off | Health-check rollout engines and kill one that fails a check; restarting it is not wired into the diffusion training loop. |
| `--rollout-health-check-interval` | float | `30.0` | Seconds between engine probes (`GET /health_generate`, a readiness check that runs no generation); the miles router polls `/health` on the same interval. |
| `--rollout-health-check-timeout` | float | `30.0` | |
| `--rollout-health-check-first-wait` | float | `0` | Grace period for compilation/init. |

### Debugging

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--debug-rollout-only` | flag | off | Rollout, no training. |
| `--train-only` | flag | off | No engines, no weight sync, no eval. (`--debug-train-only` is a legacy alias.) |
| `--debug-skip-optimizer-step` | flag | off | No backward/step — weights never drift. Use it to measure pure forward divergence from the engine. |
| `--save-debug-rollout-data` | path template | – | Fills `{rollout_id}`. |
| `--load-debug-rollout-data` | path template | – | Replays a rollout with **no engines started**. |
| `--load-debug-rollout-data-subsample` | float | – | |
| `--save-debug-train-data` | path template | – | Fills `{rollout_id}` and `{rank}`. |
| `--dump-details` | dir | – | Sets the debug dump paths for post-hoc analysis. |
| `--diffusion-debug-mode` | flag | off | Engine returns per-step debug tensors; enables the `model_output_*_abs_diff` metrics. |

`--debug-skip-optimizer-step` + `--diffusion-debug-mode` is the standard train/rollout alignment
probe: frozen weights, and the metrics report exactly how far the two forwards drift.

### CI

`--ci-test`, `--ci-metric-checker-key`, `--ci-metric-checker-threshold`. Used by `tests/e2e`;
not meant for manual runs.

***

## SDE step backends

`--diffusion-sde-type` picks the train-side backend automatically:

| `--diffusion-sde-type` | Backend | Dynamics |
| - | - | - |
| `sde` | `FlowSdeStepBackend` | Flow-matching SDE over the rollout's recorded sigmas. |
| `ode` | `FlowSdeStepBackend` | Same class; deterministic rollout (used by NFT). |
| `cps` | `CpsSdeStepBackend` | CPS kernel over the rollout's recorded sigmas, log-prob without constants. |

`--sde-step-backend-path` overrides the mapping with your own `SdeStepBackend` subclass.


# Core Concepts
Source: https://miles.radixark.com/docs/diffusion/user-guide/concepts

The four objects that make up every miles-diffusion job and the trajectory-level training loop.

A miles-diffusion training job is a loop over four objects. Once you understand what each one *is* and how data flows
between them, the major flag groups are easier to place.

## The four objects

```mermaid theme={null}
flowchart LR
    P[Prompt dataset] --> R["Rollout<br/>(sglang-diffusion engines)"]
    R -- trajectories --> RM["Scoring<br/>(reward workers)"]
    RM -- timestep pairs --> A["Training<br/>(FSDP actor)"]
    A -- weights --> R
```

| Object | Role | Lives in |
| - | - | - |
| **Prompt dataset** | Source of prompts (plus optional metadata) | JSONL on disk (`--prompt-data`, `--input-key`) |
| **Rollout (sglang-diffusion engines)** | Denoises prompts into images/videos and records the trajectory | One engine per `--rollout-num-gpus-per-engine` GPUs; shipped diffusion recipes explicitly enable the miles router with `--use-miles-router` |
| **Reward workers** | Map `(prompt, generated output) → score` | Built-in `rm_hub` (`--rm-type ocr / pickscore / hps / ...`) or custom (`--custom-rm-path`) — Ray actor pools |
| **Actor (FSDP2 + diffusers or a native package)** | The DiT being trained, usually via LoRA | HF checkpoint (`--hf-checkpoint`), family resolved by `TrainPipelineConfig` |

## The training loop

The whole of `train_diffusion.py` is three calls per iteration:

```python theme={null}
for rollout_id in range(start_rollout_id, num_rollout):
    # 1. Sample + 2. Score: prompts -> scored denoising trajectories,
    #    expanded to (x_t -> x_{t+1}) train pairs and sharded per DP rank
    rollout_data = rollout_manager.generate(rollout_id)

    # 3. Optimize: micro-batch the pairs and step the loss plugin (--loss-type)
    actor_model.async_train(rollout_id, rollout_data)

    # 4. Sync: push updated weights (or LoRA pairs) to rollout engines
    actor_model.update_weights()
```

Every flag in miles-diffusion configures one of these four phases.

`train_sft.py` is the train-only variant of this loop, and the SFT recipes launch it. Its rollout
data does not depend on the trained weights, so it generates rollout `i+1` while rollout `i`
trains; only the first rollout is waited for. Encoders colocated with training hold their encode
burst until the in-flight train step finishes, so a cold cache stalls on encoding while a warm one
hides the cache load and dp split entirely. Reserving encoder GPUs with `--rollout-num-gpus` overlaps the encoding too.

## The batch-knob invariant

In miles-diffusion a sample is a whole denoising **trajectory**, and the fan-out continues below it: a step strategy
picks the trained timesteps, and micro-batching counts **(`x_t → x_{t+1}`) pairs**. One equation per level.

**Trajectory level** — the four-knob invariant, resolved in `miles/utils/arguments.py`. `rollout_batch_size` is
required and `n_samples_per_prompt` defaults to 1; of the right-hand pair, set one and the other is derived by
integer division (set neither and `global_batch_size` defaults to `dp_size`; passing both with contradicting values
aborts). When the left side does not divide evenly, each rollout's samples are trimmed to a multiple of
`global_batch_size` before training:

```bash theme={null}
rollout_batch_size × n_samples_per_prompt
    = global_batch_size × num_steps_per_rollout
```

**Pair level** — what each optimizer step actually trains:

```bash theme={null}
pairs_per_optimizer_step
    = global_batch_size × trained_timesteps_per_trajectory
```

The step strategy picks `trained_timesteps_per_trajectory`. Each rank then trains its `1 / dp_size` share,
`micro_batch_size` pairs per forward, accumulating gradients until the optimizer step.

Two independent knobs control physical batching, one per side of the loop:

| Knob | Side | Governs |
| - | - | - |
| `--rollout-microgroup-size` | rollout | Samples per engine **request** |
| `--micro-batch-size` | training | Train pairs per DiT **forward** |

And one knob controls the trajectory→pair fan-out itself: `--diffusion-step-strategy-path` picks the SDE step subset for
each engine request (the built-in `sde_window` or `epoch_global_random_choice` in `miles/rollout/step_strategy_hub.py`,
or your own function). Training the full 10-step schedule would multiply the pair count by 10× in the example below.
NFT (`--loss-type nft`) instead trains `(x0, t)` pairs at the steps `--diffusion-nft-timestep-strategy-path` picks per sample.

Here is a full example with the numbers:

```mermaid theme={null}
%%{init: {"flowchart": {"wrappingWidth": 800}}}%%
flowchart TD
    P["48 prompts&nbsp;&nbsp;(--rollout-batch-size)"]
    O1["× 16 generations per prompt&nbsp;&nbsp;(--n-samples-per-prompt)"]:::op
    T["768 trajectories&nbsp;&nbsp;— sample level:&nbsp;&nbsp;768 = 2 × 384 ✓"]
    O2["÷ 2 optimizer steps&nbsp;&nbsp;(--num-steps-per-rollout)"]:::op
    G["384 trajectories per optimizer step&nbsp;&nbsp;(--global-batch-size)"]
    O3["× 1 of the 10 SDE timesteps&nbsp;&nbsp;(step strategy)"]:::op
    PR["384 (x_t → x_t+1) pairs per optimizer step"]
    O4["÷ 4 training ranks&nbsp;&nbsp;(dp_size)"]:::op
    RK["96 pairs per rank"]
    O5["÷ 2 pairs per DiT forward&nbsp;&nbsp;(--micro-batch-size)"]:::op
    F["48 gradient-accumulation forwards per rank, then one optimizer step"]
    P --> O1 --> T --> O2 --> G --> O3 --> PR --> O4 --> RK --> O5 --> F
    classDef op fill:#f3f3f3,stroke:none
    style T fill:#fdf6d8,stroke:#b8a23e
```

## Next

* [Launch Scripts](/docs/diffusion/user-guide/launch-script) — a canonical launcher, group by group.
* [Rewards](/docs/diffusion/user-guide/rewards) — `rm_hub`, custom reward functions, prompt data format.
* [CLI Reference](/docs/diffusion/user-guide/cli-reference) — every flag, fully cataloged.
* [SDE Step Backend](/docs/diffusion/advanced/sde-backend) — how train-side log-probs mirror rollout stepping.


# Customization
Source: https://miles.radixark.com/docs/diffusion/user-guide/customization

Plug-points where you can drop in your own Python without forking miles-diffusion.

Most of miles-diffusion's behavior can be replaced with user-supplied Python by
passing a `--*-path` flag (loaded via `miles.utils.misc.load_function` as a
dotted import path). This page lists the diffusion-relevant hooks, the
signatures they expect, and the defaults they replace — in the same spirit as
[Miles customization](https://miles.radixark.com/docs/user-guide/customization).

## At a glance

| Stage | Flag | Replaces |
| - | - | - |
| **Rollout** | `--rollout-function-path` | The whole rollout loop |
| | `--eval-function-path` | The eval rollout |
| | `--data-source-path` | How prompts / buffer are loaded |
| | `--custom-generate-function-path` | One microgroup's generation |
| | `--diffusion-step-strategy-path` | Which denoising steps enter training |
| **Reward** | `--custom-rm-path` | Batched reward computation |
| | `--custom-reward-post-process-path` | Advantage normalization |
| **Filtering** | `--dynamic-sampling-filter-path` | Per-group keep / drop |
| | `--buffer-filter-path` | Buffer dequeue selection |
| **Training** | `--custom-expand-samples-to-train-pairs-path` | Trajectory → train pairs |
| | `--custom-convert-samples-to-train-data-path` | Full sample → train\_data |
| | `--custom-prepare-train-batch-path` | DiT batch preparation |
| | `--custom-loss-function-path` | Loss formula only |
| | `--sde-step-backend-path` | Train-side SDE dynamics |
| **Logging** | `--custom-rollout-log-function-path` | Train-rollout logging |
| | `--custom-eval-rollout-log-function-path` | Eval-rollout logging |
| **Model family** | `--train-pipeline-config-path` | Family `TrainPipelineConfig` |
| | `--model-backend-path` | Model load / FSDP backend |

`--loss-type` (`policy_loss` / `nft` / `sft_loss` / `custom_loss`) is the
shortcut that auto-fills several of the train-side paths; see § Training.

***

## Rollout

### `--rollout-function-path`

Replace the entire train rollout. Diffusion recipes pass
`miles.rollout.sglang_diffusion_rollout.generate_rollout`; the flag's default
(`miles.rollout.sglang_rollout.generate_rollout`, the LLM path) is not shipped in miles-diffusion.

```python theme={null}
def generate_rollout(
    args, rollout_id, data_source, evaluation=False
) -> RolloutFnTrainOutput | RolloutFnEvalOutput:
    ...
```

`miles/ray/rollout.py` passes the data source positionally through `call_rollout_fn`.
Legacy list/dict returns are still wrapped into the corresponding output dataclass by
`miles.rollout.base_types.call_rollout_fn`.

### `--eval-function-path`

Same shape as the rollout function. Defaults to `--rollout-function-path` when
unset.

### `--data-source-path`

**Class** (not a function). Default:
`miles.rollout.data_source.RolloutDataSourceWithBuffer`.

```python theme={null}
class CustomDataSource:
    def __init__(self, args): ...
    def get_samples(self, num_samples) -> list[list[Sample]]: ...
    def add_samples(self, samples) -> None: ...
    def save(self, rollout_id) -> None: ...
    def load(self, rollout_id=None) -> None: ...
```

`RolloutManager` also reads `len(data_source.dataset)` when `--num-rollout` is unset.

### `--custom-generate-function-path`

Replace the inner microgroup generator inside
`generate_and_rm_microgroup`. Default: `generate_microgroup`.

```python theme={null}
async def custom_generate(
    args, microgroup: list[Sample], sampling_params: dict, *, evaluation: bool = False
) -> list[Sample]:
    ...
```

`evaluation=` is optional — if your signature omits it, the caller skips that kwarg.

### `--diffusion-step-strategy-path`

Select which denoising steps contribute SDE log-probs / train pairs, and which
latents the engine ships back for the trainer. Stock implementations live in
`miles/rollout/step_strategy_hub.py`.

```python theme={null}
def strategy(args, sample, num_steps, seed) -> tuple[list[int] | None, list[int] | None]:
    # returns (sde_step_indices, return_step_indices)
    # sde None = no per-step filter: every step uses --diffusion-sde-type
    # return None = the engine returns every step's latent
    ...
```

`--rollout-return-full-trajectory` overrides `return_step_indices` back to all
steps for debugging / A-B runs.

| Stock | Behavior |
| - | - |
| `sde_window` | Random contiguous window (`--diffusion-num-sde-steps`, `--diffusion-sde-window-range`); returns all latents its steps touch |
| `epoch_global_random_choice` | Subset of `--diffusion-sde-candidate-steps`, shared by each block of `--rollout-batch-size` consecutive prompt groups (one rollout, unless a dynamic-sampling filter drops groups); returns all latents its steps touch |
| `ode_and_return_last` | No step filter (every step uses `--diffusion-sde-type`, which the DiffusionNFT recipes set to `ode`); return only the final clean `x0` |

DiffusionNFT picks the timesteps it trains on with a separate hook,
`--diffusion-nft-timestep-strategy-path`: `(args, num_steps, generator) -> list[int]`, built in
`drop_random_steps` (default) and `drop_final_steps` in `miles/ray/data_conversion_hub/nft.py`.

***

## Reward

Built-in scorers (`--rm-type pickscore` / `hps` / `ocr` / `custom_api` / `openai_api`) are documented in
[Rewards](/docs/diffusion/user-guide/rewards). `--custom-rm-path` replaces that dispatch entirely; `--custom-reward-post-process-path`
replaces the normalization that runs after scoring.

### `--custom-rm-path`

```python theme={null}
async def custom_rm(args, samples: list[Sample], **kwargs) -> list[float]:
    ...
```

Wired only through `batched_async_rm` — implement per-sample routing inside your
batched function if needed. Once set, `--rm-type` and `metadata.rm_type` are ignored: the
custom function is the whole dispatch.

```bash theme={null}
--custom-rm-path my_project.rewards.aesthetic_rm
```

`--custom-rm-args` is an opaque string passed through as `args.custom_rm_args` for the
custom RM to parse (requires `--custom-rm-path`).

Shipped custom RMs:

| Path | What |
| - | - |
| `miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm` | Weighted sum of built-in rewards (`hps`, `pickscore`, `ocr`, `custom_api`, `openai_api`), weights from `--custom-rm-args "hps=0.7,pickscore=0.3"`; returns a dict per sample, train on it with `--reward-key weighted`. See [Rewards](/docs/diffusion/user-guide/rewards) § Combining rewards. |

HTTP / remote scoring: `--rm-type openai_api` / `custom_api` (configured by
`--openai-api-rm-config` / `--custom-api-rm-config`) call an external service; otherwise
implement a batched custom RM and read `args.rm_url` (or your own flags). Encode images
from `sample.generated_output` (see `generated_output_to_rgb_hwc_uint8_frames` in
`miles/utils/processing_utils.py`):

```python theme={null}
import aiohttp
from miles.utils.types import Sample

async def api_rm(args, samples: list[Sample], **kwargs) -> list[float]:
    async with aiohttp.ClientSession() as session:
        rewards = []
        for sample in samples:
            payload = {"prompt": sample.prompt, "image_b64": "<your encoding>"}
            async with session.post(args.rm_url, json=payload) as resp:
                rewards.append((await resp.json())["score"])
        return rewards
```

```bash theme={null}
--custom-rm-path my_project.rewards.api_rm \
--rm-url http://localhost:8000/score
```

### `--custom-reward-post-process-path`

Replace GRPO advantage normalization in
`RolloutManager._post_process_rewards`.

```python theme={null}
def post_process(args, samples) -> tuple[list[float], list[float]]:
    # Returns (raw_rewards, normalized_rewards)
    ...
```

Default behavior: reshape to `(-1, n_samples_per_prompt)`, subtract mean
(`--globalize-reward-mean` for batch-level), optionally divide by std
(on by default — `--disable-grpo-std-normalization` turns it off; `--globalize-reward-std`
switches per-group std to batch-wide).

***

## Filtering

### `--dynamic-sampling-filter-path`

Per-group filter after scoring (DAPO-style). Stock:
`miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std`.

```python theme={null}
def filter_function(args, samples: list[Sample], **kwargs):
    # return DynamicFilterOutput(keep=..., reason=...) or a bool
    ...
```

### `--buffer-filter-path`

Select samples when dequeuing from the rollout buffer. Default is `pop_first`
in `miles/rollout/data_source.py`. The stock diffusion rollout never calls
`add_samples`, so this only fires when a custom rollout refills the buffer.

```python theme={null}
def buffer_filter(
    args, rollout_id: int | None, buffer: list[list[Sample]], num_samples: int
) -> list[list[Sample]]:
    ...
```

***

## Training

`--loss-type` picks the default prepare / formula / expand paths:

| `--loss-type` | Expand pairs | Prepare | Loss formula |
| - | - | - | - |
| `policy_loss` (default) | Flow-GRPO built-in | `prepare_flow_grpo_batch` | `flow_grpo_loss_formula` |
| `nft` | `data_conversion_hub.nft.expand_samples_to_train_pairs` | `loss_hub.nft.prepare_nft_batch` | `nft_loss_formula` |
| `sft_loss` | none; convert `sft_rollout.convert_samples_to_train_data` | `loss_hub.sft.prepare_sft_batch` | `sft_loss_formula` |
| `custom_loss` | your expand (or convert) | your prepare | your formula |

`sft_loss` auto-fills nothing: it requires the convert, rollout-log, prepare and loss
paths, plus `--train-only` and a rollout function that runs no engines (e.g.
`miles.rollout.sft_rollout.generate_rollout`; see `scripts/run_diffusion_sft_wan22.py`). `custom_loss` behaves like `policy_loss`:
any hook left unset falls back to the Flow-GRPO built-in.

DiT forward always stays in the FSDP actor; the loss hook only computes the
objective.

### `--custom-expand-samples-to-train-pairs-path`

```python theme={null}
def expand_samples_to_train_pairs(args, samples, rewards, raw_rewards) -> dict:
    ...
```

Default for Flow-GRPO lives under `miles/ray/data_conversion_hub/flow_grpo.py`.

### `--custom-convert-samples-to-train-data-path`

Replace the entire `RolloutManager._convert_samples_to_train_data` (including
reward post-process + expand). Prefer the narrower expand hook unless you need
full control.

```python theme={null}
def convert_samples_to_train_data(args, samples) -> dict:
    ...
```

### `--custom-prepare-train-batch-path`

```python theme={null}
def prepare(ctx, batch, *, pad_to_len=None) -> PreparedBatch:
    ...
```

Builds DiT inputs from train pairs. Defaults:
`miles.backends.fsdp_utils.loss_hub.flow_grpo.prepare_flow_grpo_batch` or the NFT
equivalent under `loss_hub.nft`. `PreparedBatch.latents`, `timesteps` and
`timesteps_for_model` are keyed by latent stream name: `"visual"`, plus e.g.
`"audio"` for joint audio-video models.

### `--custom-loss-function-path`

```python theme={null}
def loss_formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, metrics, **kwargs):
    ...
```

`new_pred`, `old_pred` and `ref_pred` are keyed by the same stream names as `prepared.latents`.
`old_pred` is `None` unless `--loss-type nft`; `ref_pred` is `None` with `--ref-mode none`.

### `--sde-step-backend-path`

**Class** implementing `SdeStepBackend`. Auto-selected from
`--diffusion-sde-type` (`sde` / `ode` → `FlowSdeStepBackend`, `cps` →
`CpsSdeStepBackend`) unless overridden. See
[SDE step backend](/docs/diffusion/advanced/sde-backend).

***

## Logging

```python theme={null}
def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool:
    ...

def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool:
    ...
```

Return a truthy value to skip the default logging; falsy layers on top.

***

## Model family

### `--hf-checkpoint` / `--diffusion-model-family` / `--diffusion-task` / `--train-pipeline-config-path`

`--hf-checkpoint` names the diffusers pipeline for train + rollout. Family is
resolved from the checkpoint name unless you pass `--diffusion-model-family`
(e.g. `sd3`). The longest matching name pattern wins, so `Qwen/Qwen-Image-2.1`
resolves to `qwen_image21`. `--diffusion-task` (e.g. `t2i`) must be one of the
`tasks` that family's `TrainPipelineConfig` declares. For an unregistered family, pass
`--train-pipeline-config-path` to a `TrainPipelineConfig` subclass instead.

### `--model-backend-path`

**Class** for loading components / FSDP. Default comes from the
family config (usually a `DiffusersModelBackend`).

***

## Worked example

Custom reward + post-process on top of the stock SD3 Flow-GRPO recipe:

```bash theme={null}
python3 scripts/run_diffusion_grpo_sd3_ocr_sglang.py \
  --cuda-visible-devices 6,7 \
  --extra-args "--custom-rm-path my_project.rewards.api_rm --rm-url http://localhost:8000/score --custom-reward-post-process-path my_project.rewards.normalize"
```

(`--extra-args` is spliced verbatim into a `bash -c` command line, so keep it on one line; you can also splice the
flags into a forked recipe.)

***

## Pairs well with

* [Rewards](/docs/diffusion/user-guide/rewards) — built-in PickScore / HPS / OCR and prompt JSONL.
* [SDE step backend](/docs/diffusion/advanced/sde-backend) — step strategies and SDE kernels.
* [LoRA weight sync](/docs/diffusion/advanced/lora) — IPC merge path used by most recipes.
* [SD3 model guide](/docs/diffusion/models/sd3/sd3) — end-to-end recipe flags.


# Launch Scripts
Source: https://miles.radixark.com/docs/diffusion/user-guide/launch-script

What a miles-diffusion launch script does when you run it, how it is structured, and the ways to override a recipe.

Every recipe ships as a Python launch script under `scripts/`, and starting a run is one command:

```bash theme={null}
python3 scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py
```

This page explains what that command does and how to change what it runs. For the meaning of
individual flags, see the [CLI Reference](/docs/diffusion/user-guide/cli-reference).

## How a launch script starts a training job

A launch script is a recipe, not the training process. It assembles the full
`train_diffusion.py` command line, starts a local Ray cluster, and submits the command as a Ray
job — submitted rather than run directly, so the driver lives inside the cluster and sees every
node's GPUs.

| Layer | Location | Role |
| - | - | - |
| Launch script | `scripts/run_*.py` | Holds the recipe: the flag blocks and tuned values |
| Command utilities | `miles/utils/external_utils/command_utils.py` | Starts Ray and submits the job |
| Training entrypoint | `train_diffusion.py` | The train loop (\~90 lines), run inside the Ray job |
| SFT entrypoint | `train_sft.py` | The train-only loop with rollout prefetch; SFT recipes pass `train_script="train_sft.py"` |

## The structure of a launch script

Every launcher follows the same layout:

```python theme={null}
@dataclass
class ScriptArgs(U.ExecuteTrainConfig):     # inherits output_dir, num_nodes, ...
    num_rollout: int = 400
    extra_args: str = ""                     # appended verbatim to the command line

def prepare(args) -> str:                    # download the dataset, return its local path
    ...

def execute(args, data_dir) -> None:
    ckpt_args      = "..."                   # named groups of train_diffusion.py flags
    rollout_args   = "..."
    ...
    U.execute_train(train_args=" ".join(groups), num_gpus_per_node=5, config=args)

@U.dataclass_cli
def main(args): execute(args, prepare(args))
```

### ScriptArgs — script options as flags and `MILES_SCRIPT_*` env vars

`@U.dataclass_cli` exposes each `ScriptArgs` field twice: as a `--kebab-case` CLI option
(`--num-rollout`) and as an environment variable with the `MILES_SCRIPT_` prefix
(`MILES_SCRIPT_NUM_ROLLOUT`). A command-line value beats the env var, which beats the field
default. The env form is how wrappers and cluster tooling inject machine-specific values
without editing the script.

Options shared by every launcher, from the `ExecuteTrainConfig` base class and repo convention:

| Option | Default | Purpose |
| - | - | - |
| `--num-nodes` | `$SLURM_JOB_NUM_NODES` or `1` | Carried by the base class but read by no shipped launcher; the node count comes from `--actor-num-nodes` in the flag blocks |
| `--cuda-visible-devices` | inherited from the environment | Physical GPUs the job may use; must list exactly as many GPUs as the recipe runs on. Applied by exporting it for `ray start`. |
| `--output-dir` | `<repo>/logs` | Where checkpoints and dumps are written |
| `--extra-env-vars` | empty | Extra env vars added to the Ray runtime env, as JSON or space-separated `KEY=VALUE` |
| `--cuda-core-dump` | off | Adds the CUDA core-dump-on-exception env vars to the Ray runtime env; dumps go to `--output-dir` |
| `--extra-args` | empty | Extra flags appended to the `train_diffusion.py` command line |

wandb flags are emitted only when `WANDB_API_KEY` is set — logging turns on by exporting the
key, with no script change.

### execute() — the flag groups

The command line is assembled as one string block per concern, then concatenated — the parser
sees a flat list, and the groups follow **concern, not flag prefix** (a `--diffusion-*` flag can
live in `rollout_args`, a `--micro-batch-size-*` flag in `perf_args`).

| Block | What it carries |
| - | - |
| `ckpt_args` | `--hf-checkpoint` (required — also selects the model family), `--diffusion-task` (required — one of the family config's tasks), `--save` / `--load` |
| `rollout_args` | Prompt data, batch shape, sampler: steps, guidance, noise level, SDE step strategy |
| `eval_args` | Eval datasets and cadence; eval requests are pure generation, no trajectory |
| `grpo_args` | Advantage estimator, reward normalisation, clip range |
| `optimizer_args` | LR, betas, weight decay |
| `lora_args` | Rank, targets, `--lora-ipc-weight-sync` |
| `reward_args` | `--rm-type` and the reward worker pool |
| `sglang_args` | Engine side: router, concurrency, `--sglang-*` passthrough |
| `train_backend_args` | The three dtype flags, `--update-weight-target-module` |
| `perf_args` | Gradient checkpointing, micro-batch tiling, parser workers |
| `misc_args` | GPU layout, `--colocate`, `--deterministic-mode` |

## Ways to override a recipe

From lightest to heaviest:

1. **Append flags with `--extra-args`.** The value lands at the end of the command line, and for
   an argparse flag given twice the later occurrence wins — so it overrides any value the recipe
   sets. A switch the recipe turns on (`--colocate`, `--use-lora`, …) cannot be turned off this
   way; edit the script for that:

   ```bash theme={null}
   python3 scripts/run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py \
       --num-rollout 50 --extra-args "--diffusion-kl-beta 0.02"
   ```

2. **Set a script option, as a flag or a `MILES_SCRIPT_*` env var.**

3. **Edit the script.** The launcher is the canonical home of a recipe's hyperparameters; change
   the flag blocks directly for anything you want to keep.

The Typer flags are also the CI interface: an e2e test names the recipe and passes its knobs
(`script="scripts/run_diffusion_grpo_sd3_ocr_sglang.py", args=["--num-rollout", "2"]`), so
launch scripts must stay runnable with no arguments and configurable through `ScriptArgs` — CI
drives them the same way you do. The one exception is the multinode Wan2.2 launcher, which needs
`--four-gpu-ci` or an external Ray cluster (`MILES_SCRIPT_EXTERNAL_RAY=1`).

## What execute\_train runs on your machine

1. Kills stale `sglang` / `sgl_diffusion` / `miles` / `redis` processes, and `ray` unless the
   cluster is external.
2. Starts a fresh cluster with `ray start --head`, prefixed with `export CUDA_VISIBLE_DEVICES=... &&`
   when `--cuda-visible-devices` is set — the device list must be in the raylet's own environment;
   set per job or per actor it never reaches the scheduler, which then places work on excluded GPUs.
3. Builds the Ray runtime env: NCCL socket vars, `MASTER_ADDR`, `PYTHONPATH`, plus the recipe's
   own `extra_env_vars` and anything from `--extra-env-vars`.
4. Submits the job: `ray job submit -- python3 train_diffusion.py <flags>`.

| Env var | Effect |
| - | - |
| `MILES_SCRIPT_EXTERNAL_RAY=1` | A scheduler already built the Ray cluster: skip `ray stop` and `ray start`, and submit into the running cluster. `--cuda-visible-devices` must be empty here — export it before each `ray start` instead. Used by the multi-node recipe below. |
| `MILES_SCRIPT_ENABLE_RAY_SUBMIT=0` | Run everything up to the submission (dataset download, process cleanup, `ray start`) and skip `ray job submit`; the training command line is not printed |
| `MASTER_ADDR` | Where the single-node `ray start --head` binds, default `127.0.0.1`. The training process group does NOT use it: rank 0 negotiates its real node IP and a free port at runtime, so multi-node needs no setting here |

## The batch-size arithmetic

Diffusion adds a dimension LLM RL does not have: one sample expands into **several train pairs**,
one per trained denoising step.

```
samples per rollout   = rollout_batch_size × n_samples_per_prompt
train pairs           = samples × (number of SDE step indices)
```

The trajectory-level knobs are locked by
[the batch-knob invariant](/docs/diffusion/user-guide/concepts#the-batch-knob-invariant): passing both
`--global-batch-size` and `--num-steps-per-rollout` with values that disagree aborts at parse time, and a product
that does not divide evenly is not an error — the extra samples of each rollout are trimmed before training. `global_batch_size` counts samples and must divide by `dp_size`
(= train world size ÷ `--sequence-parallel-size`).

| Knob | Unit | Meaning |
| - | - | - |
| `--num-steps-per-rollout` | optimizer steps | Windows the pairs are cut into per rollout |
| `--micro-batch-size` | train pairs | Flat: pairs per DiT forward |
| `--micro-batch-size-sample` × `--micro-batch-size-tstep` | samples × timesteps | 2D tile per DiT forward; overrides the flat knob with `S × T` |
| `--diffusion-train-iter-order` | — | Tile visit order (2D only) |
| `--train-dp-split-mode` | — | `contiguous` or `stride` dealing of pairs to DP ranks |

Most GRPO recipes use the 2D tile — it states "S samples × T timesteps per forward" directly, which
the flat knob cannot express. The NFT and SFT recipes, Cosmos3, and the 5-GPU Wan2.2 recipe use the
flat knob; Wan2.2 trains one shared SDE step per rollout, so its `--micro-batch-size 2` forwards
never span the high/low-noise expert boundary.

Worked example (the Qwen-Image recipe):

```
32 prompts × 16 samples             = 512 samples per rollout
512 samples × 2 SDE steps           = 1024 train pairs
1024 pairs ÷ 4 DP ranks             = 256 pairs per rank
256 ÷ num_steps_per_rollout(2)      = 128 pairs per optimizer step per rank
tile = 8 samples × 1 timestep       = 8 pairs per micro-batch → 16 micro-batches per step
                                      (each runs the CFG pair as two DiT forwards: the recipe uses CFG)
```

Two failure modes to recognize:

| Error | Fix |
| - | - |
| `num_pairs_shard=N not divisible by num_steps_per_rollout=M` | Adjust `rollout_batch_size`, `n_samples_per_prompt`, or `num_steps_per_rollout` so the per-rank pair count splits evenly |
| `Micro-batch mixes denoising phases` (Wan2.2) | A micro-batch spans the expert boundary; set `--micro-batch-size 1` or a phase-pure tiling |

## Adapting a recipe

Start from the closest existing script and change one group at a time:

| Goal | What to change |
| - | - |
| Different reward | `reward_args` — swap `--rm-type`, or point `--custom-rm-path` at your own `async def rm(args, samples) -> list[float]` |
| Different prompts | `--prompt-data`, `--input-key`, and the `prepare()` download |
| Different checkpoint | `--hf-checkpoint`, plus `--diffusion-model-family` if the name carries no family hint |
| Smaller GPU budget | Lower `--rollout-batch-size`, `--n-samples-per-prompt`, `--rollout-microgroup-size`; keep the batch identity satisfied |
| Train more denoising steps | Raise `--diffusion-num-sde-steps` — this multiplies train pairs |
| Try a flag without editing the file | `--extra-args "--flag value"` |

## Multi-node training

The worked example is `scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py`:
wan2.2-A14B full finetune on 2 nodes × 8 GPUs (train + rollout colocated) plus 1 reward GPU on a
separate node. Multi-node runs submit into a cluster you build yourself; the launcher only submits
(`MILES_SCRIPT_EXTERNAL_RAY=1`).

### Bring up the cluster

Each node runs ONE of the blocks below, chosen by its role. The `ulimit` and `export` lines must be
in the environment of the `ray start` daemon itself, which is why every block repeats them.

On the head node:

```bash theme={null}
ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
ray start --head --port=6379 --num-gpus 8 --dashboard-host 0.0.0.0
```

On every other training node:

```bash theme={null}
ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
ray start --address=<head-ip>:6379 --num-gpus 8
```

On the reward node:

```bash theme={null}
ulimit -n 1000000
export CUDA_VISIBLE_DEVICES=0
ray start --address=<head-ip>:6379 --num-gpus 1
```

* **`ulimit -n`**: a non-interactive ssh shell defaults to 1024 open files. The raylet inherits
  it and, once the FSDP actors and engines connect, dies with `epoll: Too many open files` — but
  the driver log shows a misleading `ActorUnavailableError: ... RpcError: Socket closed`. Verify
  with `grep "Max open files" /proc/$(pgrep -f raylet | head -1)/limits` after `ray start`.
* **`CUDA_VISIBLE_DEVICES`** follows the same rule the launcher automates single-node: it must be
  in the raylet's environment, and with an external cluster the launcher refuses a
  `--cuda-visible-devices` of its own.

`ray status` on the head should show every node and the full GPU count.

### Submit

On the head node (the launcher submits to `127.0.0.1:8265` unless `RAY_ADDRESS` is set):

```bash theme={null}
ulimit -n 1000000   # the driver needs it too
MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py
```

Reward workers (`--pickscore-num-workers 4 --pickscore-num-gpus-per-worker 0.25`, no
`--pickscore-reward-colocate`) are default-scheduled and land on the only free GPU: the reward node.

### Verify the run is healthy

* With `--rollout-patch-group wan` and `--sglang-attention-backend torch_sdpa`,
  `train/model_output_mean_abs_diff` is expected to be exactly `0.0` from the first optimizer step. The 4-GPU proxy E2E
  standard records `0.0`, and the documented 17-GPU runs sustained it over 200 rollouts. Any nonzero value indicates a
  train/rollout parity regression; check the patch group, backend, dtype, and versions on every rank.
* Both training nodes near 100% GPU util during rollout; a 100%-vs-idle split between nodes means
  the reward workers were packed onto one node.

### Multi-node pitfalls

| Pitfall | What to know |
| - | - |
| Stale engine hijacks the port | Killing a driver by name leaves its spawned engine alive; the next driver's health check talks to the old, unpatched server while the new one dies on `EADDRINUSE`. Between runs kill by GPU pid and check `nvidia-smi` shows no compute processes |
| Version strings lie under editable installs | `pip show sglang` reports the install-time commit; trust `git rev-parse HEAD` in the checkout (or the `-e git+...@<sha>` line in `pip freeze`) |
| Weight-sync bucket | `--update-weight-buffer-size 4294967296` (4 GiB). In one documented A14B run, 512 MiB took \~92 s and 4 GiB \~15 s; these are run-specific observations, not a benchmark guarantee. |
| Determinism env vars are trainer-side | `--deterministic-mode` sets `NCCL_DETERMINISTIC` and `CUBLAS_WORKSPACE_CONFIG` on the FSDP actors; check rollout parity with `train/model_output_mean_abs_diff` |

## Next

* [CLI Reference](/docs/diffusion/user-guide/cli-reference) — every flag, grouped.
* [Core Concepts](/docs/diffusion/user-guide/concepts) — what one rollout iteration does, object by object.
* [Dtype Control](/docs/diffusion/advanced/dtype-control) — what the three dtype flags actually do.


# Recipe Verification Levels
Source: https://miles.radixark.com/docs/diffusion/user-guide/recipe-verification

Evidence levels for documented training recipes.

Every recipe is assigned one evidence level. The level applies to the exact script and
topology named in the model guide, not to the model family as a whole.

## 🛡️ FG — Fully gated

* A complete training curve has been run.
* A deterministic E2E runs the canonical recipe itself at least nightly.
* Every registered metric matches the committed standard exactly.

## 🧩 PG — Proxy gated

* A complete training curve has been run.
* A deterministic E2E runs a documented, scaled single-node proxy at least nightly.
* Every registered metric matches the committed standard exactly.

## 📈 V — Verified

A complete training curve has been run, but no deterministic E2E satisfies the FG or PG criteria above.

## ○ NV — Not verified

No complete training curve has been run. Smoke tests and short debugging runs do not
count as verification.

## Current levels

* **🛡️ FG**
  * `run_diffusion_grpo_sd3_ocr_sglang.py` — SD3.5 Flow-GRPO + OCR.
  * `run_diffusion_grpo_ltx23_sglang.py` — LTX-2.3 Flow-GRPO + PickScore.
  * `run_diffusion_nft_sd3_pickscore.py` — SD3.5 DiffusionNFT + PickScore.
  * `run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py` — Cosmos3-Nano
    Flow-GRPO + PickScore.
  * `run_diffusion_grpo_h3_t2va_2gpu.py` — MiniMax H3 t2va Flow-GRPO + PickScore.
  * `run_diffusion_grpo_pickscore_5gpu_flowgrpo_aligned.py` — Qwen-Image
    Flow-GRPO + PickScore.
* **🧩 PG**
  * `run_diffusion_grpo_wan22_pickscore_17gpu_multinode.py` — Wan2.2 17-GPU
    full-finetune Flow-GRPO + PickScore.
  * `run_diffusion_nft_krea2_async.py` — Krea-2 8-GPU async DiffusionNFT + OCR.
* **📈 V**
  * `run_diffusion_grpo_qwenimage21_pickscore_1gpu.py` — Qwen-Image 2.1
    Flow-GRPO + PickScore.
  * `run_diffusion_nft_krea2_8gpu.py` — Krea-2 8-GPU colocated DiffusionNFT + OCR.
  * `run_diffusion_nft_krea2.py` — Krea-2 2-GPU colocated DiffusionNFT + OCR.
  * `run_diffusion_nft_flux3_image_pickscore_8gpu.py` — FLUX 3 Image Klein 8-GPU colocated DiffusionNFT + PickScore.
  * `run_diffusion_grpo_flux3_image_hps_8gpu.py` — FLUX 3 Image Klein 8-GPU colocated full-finetune Flow-GRPO + HPSv2.1.
  * `run_diffusion_sft_h3_t2va.py` — MiniMax H3 8-GPU t2va LoRA SFT.
  * `run_diffusion_grpo_sd3_hps_sglang.py` — SD3.5 Flow-GRPO + HPSv2.1.
  * `run_diffusion_grpo_sd3_ocr_pickscore_sglang.py` — SD3.5 Flow-GRPO + 0.8 OCR + 0.2 PickScore.
* **○ NV**
  * `run_diffusion_grpo_wan22_pickscore_5gpu.py` — Wan2.2 5-GPU LoRA
    Flow-GRPO + PickScore.
  * `run_diffusion_sft_wan22.py` — Wan2.2 4-GPU LoRA SFT.
  * `run_diffusion_sft_h3_ref2va.py` — MiniMax H3 8-GPU Ref2VA LoRA SFT.


# Rewards
Source: https://miles.radixark.com/docs/diffusion/user-guide/rewards

Built-in reward models (PickScore, HPS, OCR, API rewards), rm_hub dispatch, and prompt data format.

Miles-diffusion scores generated images (or video frames) after each rollout
microgroup. Reward computation lives in `miles/rollout/rm_hub/` and is invoked
from `sglang_diffusion_rollout.generate_and_rm_microgroup`.

For `--custom-rm-path`, `--custom-reward-post-process-path`, and other
`--*-path` hooks, see [Customization](/docs/diffusion/user-guide/customization).

## 1. At a glance

| Stage | Flag | Role |
| - | - | - |
| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `hps`, `ocr`, `openai_api`, `custom_api`); ignored when `--custom-rm-path` is set |
| API reward config | `--openai-api-rm-config`, `--custom-api-rm-config` | YAML file or inline `base64:<payload>` for the two API rewards |
| Per-sample override | `metadata.rm_type` in JSONL | Overrides global `--rm-type` |
| Custom reward / norm | see [Customization](/docs/diffusion/user-guide/customization) | `--custom-rm-path`, `--custom-reward-post-process-path` |

## 2. Built-in reward models

### PickScore (`--rm-type pickscore`)

Implementation: `miles/rollout/rm_hub/pickscore.py`.

PickScore scores text–image alignment using a CLIP model pair:

* Processor: `--pickscore-processor-path` (e.g.
  `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`)
* Model: `--pickscore-model-path` (e.g. `yuvalkirstain/PickScore_v1`)

Scoring formula:

```
score = exp(logit_scale) * dot(text_emb, image_emb) / 26.0
```

The `/ 26.0` scaling maps raw PickScore logits (\~0–26) to roughly 0–1.

PickScore runs as a **Ray actor pool** (`PickScoreRewardActor`) with round-robin
batching. For video outputs, frames are uniformly sampled
(`--pickscore-num-frames`) and scores are averaged.

| Flag | Default | Description |
| - | - | - |
| `--pickscore-num-workers` | 1 | Ray actor count |
| `--pickscore-num-gpus-per-worker` | 1.0 | GPU per worker (non-colocate) |
| `--pickscore-batch-size` | 8 | Batch size per actor |
| `--pickscore-processor-path` | — | Required for pickscore |
| `--pickscore-model-path` | — | Required for pickscore |
| `--pickscore-num-frames` | None | Evenly spaced video frames scored per sample (None = every frame) |
| `--pickscore-reward-colocate` | False | Seat the `--pickscore-num-workers` actors on rollout GPUs, one per placement-group bundle (requires `--colocate`) |

Example from `scripts/run_diffusion_nft_sd3_pickscore.py`:

```bash theme={null}
--rm-type pickscore \
--pickscore-num-workers 1 \
--pickscore-num-gpus-per-worker 1.0 \
--pickscore-batch-size 8 \
--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \
--pickscore-model-path yuvalkirstain/PickScore_v1
```

### HPS (`--rm-type hps`)

Implementation: `miles/rollout/rm_hub/hps.py`.

HPSv2 scores text–image alignment using a preference-tuned CLIP model:

* Model: `ViT-H-14`
* Version: `--hps-version` (`v2.0` or `v2.1`)
* Checkpoint: `--hps-checkpoint-path` (optional; defaults to `xswu/HPSv2`)

Scoring formula:

```
score = diagonal(image_features @ text_features.T)
```

The diagonal pairs each image with its corresponding prompt in the batch.

HPS runs as a **Ray actor pool** (`HPSRewardActor`) with round-robin batching.
It currently accepts image outputs only (`generated_output` with one frame).
Use a custom RM when defining video frame aggregation semantics.

| Flag | Default | Description |
| - | - | - |
| `--hps-num-workers` | 1 | Ray actor count |
| `--hps-num-gpus-per-worker` | 1.0 | GPU per worker (non-colocate) |
| `--hps-batch-size` | 8 | Batch size per actor |
| `--hps-version` | `v2.1` | `v2.0` or `v2.1` checkpoint |
| `--hps-checkpoint-path` | None | Local checkpoint; unset downloads from `xswu/HPSv2` |
| `--hps-reward-colocate` | False | Seat the `--hps-num-workers` actors on rollout GPUs, one per placement-group bundle (requires `--colocate`) |

Example from `scripts/run_diffusion_grpo_sd3_hps_sglang.py`:

```bash theme={null}
--rm-type hps \
--hps-num-workers 1 \
--hps-batch-size 8 \
--hps-version v2.1 \
--hps-reward-colocate
```

### Reward placement

Every GPU reward pool is placed one of two ways:

* `--<rm>-reward-colocate`: every worker takes one **slot** on a rollout placement-group bundle,
  sharing that GPU with the train actor and the rollout engine. Colocated pools share one slot ledger, so `hps` and `pickscore` can both colocate without
  overlapping; more workers than bundles is rejected at parse time. Requires `--colocate`.
* Otherwise the pool is **standalone**: default-scheduled at `--<rm>-num-gpus-per-worker`,
  which only lands on GPUs outside every placement group (placement groups reserve their
  GPUs, so Ray never packs these onto rollout GPUs).

`RolloutManager` seats the colocated pools before the first rollout; standalone pools are
built on first use.

### Combining rewards

`--custom-rm-path` receives `(args, samples)` and can call the built-in scorers
directly; `--custom-rm-args` is an opaque string the framework hands to that function
through `args`, so the function owns its own config grammar. The shipped example
`miles/rollout/rm_hub/weighted_mixture_rm.py` reads `name=weight,name=weight`, with names
from `hps`, `pickscore`, `ocr`, `custom_api`, `openai_api`:

```bash theme={null}
--custom-rm-path miles.rollout.rm_hub.weighted_mixture_rm.weighted_mixture_rm \
--custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted \
--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \
--pickscore-model-path yuvalkirstain/PickScore_v1 \
--hps-reward-colocate --pickscore-reward-colocate   # each reward keeps its own placement flags
```

The example returns a dict per sample (`{"hps": ..., "pickscore": ..., "weighted": ...}`).
`--reward-key` picks the entry GRPO trains on, and every entry of a dict reward gets its own
`rollout/reward/<name>_mean` panel, so the components stay visible while the sum is optimized.

Weights apply to raw scores (HPSv2.1 ≈ 0.25–0.35, PickScore/26 ≈ 0.8–0.9, OCR ∈ \[0, 1], default
OpenAI API rubric ∈ \[0, 4]), so pick them with the scales in mind. Colocated pools share one slot ledger, so several rewards
can colocate without overlapping. Rewards receive `generated_output` itself, and every reward actor
quantises it to uint8 on its own terms.

A shipped recipe uses it: `scripts/run_diffusion_grpo_sd3_ocr_pickscore_sglang.py` trains SD3.5 on
0.8 OCR + 0.2 PickScore with `--rollout-shuffle`; see [SD3](/docs/diffusion/models/sd3/sd3) § 5.4 for the
recipe and numbers and § 9 for the curves. Shuffling matters more than usual there: with 8 prompts per rollout one hard
batch moves the per-rollout mean visibly.

### OCR (`--rm-type ocr`)

Implementation: `miles/rollout/rm_hub/ocr.py`.

OCR reward compares PaddleOCR output against target text embedded in the
prompt. The target is the string between the first pair of double quotes.
Simplified scoring:

```python theme={null}
target = prompt.split('"')[1]
reward = 1 - levenshtein_distance(recognized, target) / len(target)
```

The implementation also lowercases both sides, strips spaces, treats a
substring hit as perfect (`dist=0`), and caps `dist` at `len(target)` — see
`miles/rollout/rm_hub/ocr.py`.

OCR runs on **CPU** Ray actors (`--ocr-num-workers`, default 4). Used by the
SD3 Flow-GRPO OCR and OCR & PickScore recipes, the three Krea-2 DiffusionNFT recipes, and the
smoke mode of the SD3 DiffusionNFT recipe.

### API rewards (`--rm-type openai_api` / `custom_api`)

Implementation: `miles/rollout/rm_hub/openai_api.py` and `miles/rollout/rm_hub/api.py`.

Each sample is one HTTP request, sent by a single zero-GPU Ray actor that runs up to
`max_concurrency` (default 8) requests at once. The config is a YAML file or inline
`base64:<payload>` with `actor_class`, `actor_kwargs`, and `max_concurrency`; a flat mapping is
taken as the `actor_kwargs` of the default actor.

| `--rm-type` | Config flag | Actor |
| - | - | - |
| `openai_api` | `--openai-api-rm-config` | `OpenAIImageRewardActor` (or a subclass) |
| `custom_api` | `--custom-api-rm-config` | Any `ApiRewardActor` subclass implementing `build_request` / `send_request` / `parse_response`; defaults to `OpenAIImageRewardActor` |

`OpenAIImageRewardActor` sends the generation prompt and the image (PNG data URL) to an
OpenAI-compatible Chat Completions endpoint, with the config `prompt` as the system rubric, and asks
for a JSON `score`. Its kwargs are `model` and `api_key_env` (required), `base_url`, `prompt` or
`prompt_path` (resolved relative to the config file, not allowed in an inline `base64:` config; with neither, a
built-in 0–4 rubric is used), `score_min` / `score_max` (default 0 / 4), and `timeout_s`; scores outside that range
raise. It accepts image outputs only (one frame).

### Remote RM (`--rm-type remote_rm`)

The CLI exposes `--rm-url` for a remote reward service, but **`rm_hub` does not
implement `remote_rm` today** — selecting it raises `NotImplementedError`.
To call an external service, use the API rewards above or a `--custom-rm-path`
function that reads `args.rm_url`.

## 3. Call chain

```
generate_and_rm_microgroup()
  → batched_async_rm(args, microgroup)          # rm_hub/__init__.py
    → custom_rm_path?  user batched function
    → all pickscore?   pickscore_rm (batched)
    → all hps?         hps_rm (batched)
    → all ocr?         ocr_rm (batched, one image per actor call)
    → all custom_api?  custom_api_rm (batched, one request per actor call)
    → all openai_api?  openai_api_rm (batched, one request per actor call)
    → else             per-sample async_rm → ocr / pickscore / hps / custom_api / openai_api / NotImplementedError
  → sample.reward = score

# after the rollout, in the RolloutManager Ray actor:
RolloutManager._convert_samples_to_train_data()
  → _post_process_rewards()                     # GRPO advantage normalization
```

After rollout, `RolloutManager._post_process_rewards` subtracts the mean and
optionally divides by std to produce normalized advantages for training.
Override that path with `--custom-reward-post-process-path` — see
[Customization](/docs/diffusion/user-guide/customization) § Reward.

## 4. Prompt data

### JSONL format

Training prompts are loaded from `.jsonl` files via `miles/utils/diffusion_data.py` and read in
file order; `--rollout-shuffle` permutes them once per epoch (seeded by `--rollout-seed`):

```json theme={null}
{"input": "A photo of a cat wearing sunglasses"}
{"input": "A logo saying \"Miles\"", "metadata": {"rm_type": "ocr"}}
```

| Field | CLI mapping | Notes |
| - | - | - |
| Prompt text | `--input-key input` | Non-empty string; rows without one are skipped |
| Per-sample metadata | `--metadata-key metadata` | Optional dict; supports `rm_type` override |

### Dataset subsets

Dataset repo: [`rockdu/miles-diffusion-datasets`](https://huggingface.co/datasets/rockdu/miles-diffusion-datasets)

| Subset | Used by |
| - | - |
| `hpdv2/` | SD3 HPS Flow-GRPO |
| `flowgrpo_pickscore/` | PickScore recipes (SD3 NFT, Qwen-Image, Wan2.2, LTX, Cosmos3, Krea-2, MiniMax H3) |
| `flowgrpo_ocr/` | SD3 OCR and OCR & PickScore Flow-GRPO, Krea-2 OCR NFT, NFT smoke test |

### Per-sample rm\_type override

JSONL `metadata.rm_type` overrides the global `--rm-type` for that sample:

```python theme={null}
# rm_hub/__init__.py
metadata.get("rm_type") or args.rm_type
```

Mixed rm\_types within one microgroup fall back to per-sample dispatch (no
batched fast path).


# VLM Single-Turn RL (FSDP & Megatron)
Source: https://miles.radixark.com/docs/examples/geo3k-vlm

Training VLMs with FSDP using GRPO on the GEO3K dataset.

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

Use the CUDA 13 Miles image from the [installation guide](/docs/getting-started/installation) and keep its bundled cuDNN version.

![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 theme={null}
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 theme={null}
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/train_data_conversion.py` (`_post_process_rewards`) 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`


# VLM Multi-Turn (geo3k dataset)
Source: https://miles.radixark.com/docs/examples/geo3k-vlm/multi-turn

The same dataset over multiple turns, with the model cropping images through an interactive environment.

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

Use the CUDA 13 Miles image from the [installation guide](/docs/getting-started/installation) and keep its bundled cuDNN version.

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.

> **Environment API (geo3k)**
>
> * `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 theme={null}
# 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.


# Examples
Source: https://miles.radixark.com/docs/examples/index

These examples are runnable starting points for your own RL workflow.

A few are purely demonstrative, but most are verifiable against a concrete performance score.

## Recipes

End-to-end training workflows — the place to start.

* **[geo3k\_vlm](/docs/examples/geo3k-vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset.
  * **[multi\_turn](/docs/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](/docs/examples/multi-lora)**: Serve concurrent LoRA fine-tuning clients on one shared base model through the Tinker protocol.
  * **[harbor\_tinker](/docs/examples/multi-lora/harbor-tinker)**: Terminal-Bench 2 training with Harbor sandboxes and the official Tinker cookbook RL loop.
  * **[tinker\_agent](/docs/examples/multi-lora/tinker-agent)**: Agentic training with your own sandbox or harness through the official Tinker cookbook capture recipe.
* **[multi\_policy](/docs/examples/multi-policy)**: Two policies in one run — a solver answering gsm8k, and a verifier scored on ruling correctly about the solver's answers.
* **[on\_policy\_distillation](/docs/examples/on-policy-distillation)**: Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop.
  * **[qwen3\_5\_35b\_selfdistill](/docs/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](/docs/examples/ppo)**: Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs.
* **[retool\_v2](/docs/examples/retool-v2)**: Tool-enabled language model generation with sandboxed Python code execution interleaved with thinking.
* **[swe-agent-harbor-docker](/docs/examples/swe-agent-harbor-docker)**: Trains coding and terminal agents with Harbor-managed local Docker sandboxes and verifier rewards.

## [Infra Features](/docs/examples/infra-features)

Runtime and infrastructure plumbing rather than training recipes — how miles moves
data and weights around.

* **[fully\_async](/docs/examples/infra-features/fully-async)**: Demonstrates fully asynchronous rollout generation for higher efficiency.
* **[hot\_restart](/docs/examples/infra-features/hot-restart)**: Replaces the orchestration script and rollout executor of a live run, keeping trainers and engines up.
* **[low\_precision](/docs/examples/infra-features/low-precision)**: Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability.
* **[p2p\_weight\_transfer](/docs/examples/infra-features/p2p-weight-transfer)**: Point-to-point weight transfer between training and rollout engines.
* **[random\_async](/docs/examples/infra-features/random-async)**: Dataset-free stress test of the async rollout ↔ trainer loop.
* **[score\_centering](/docs/examples/infra-features/score-centering)**: Center off-policy policy gradients using the sampler's top-k probabilities, with optional TIS or MIS weights.
* **[split\_deployment](/docs/examples/infra-features/split-deployment)**: Installs one run as several helm releases — trainer, engines and orchestration script apart.
* **[train\_infer\_mismatch\_helper](/docs/examples/infra-features/train-infer-mismatch-helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS).
* **[true\_on\_policy](/docs/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.

* **[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.


# Infra Features
Source: https://miles.radixark.com/docs/examples/infra-features

These examples exercise runtime and infrastructure behaviour rather than a training recipe.

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.


# Fully Asynchronous Rollout Example
Source: https://miles.radixark.com/docs/examples/infra-features/fully-async

Demonstrates fully asynchronous rollout generation for higher efficiency.

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.

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

## Files

* `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, `examples.infra_features.fully_async.external_eval_fn.ExternalSglangEvalFn`).
* `run_qwen3_30b_a3b_fully_async.py`: the same pattern on a 30B MoE — `tp=8`, `ep=8`, one 8-GPU rollout engine.
* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it.

## Quick Start

Each launcher downloads its own checkpoint and converts it, then submits the job:

```bash theme={null}
python examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py
```

You should see log lines like:

```
Started fully-async rollout worker
```

## On a single GPU

[`examples/retool_v2`](/docs/examples/retool-v2) runs the same flag colocated with the trainer:
`--fully-async --num-gpus-per-node 1` trains Qwen3-4B on a multi-turn tool-call workload
on one GPU, with the engines and the trainer taking turns on it.

## At a larger scale

[`examples/experimental/openenv/glm52_tbench2`](https://github.com/radixark/miles/tree/main/examples/experimental/openenv/glm52_tbench2) runs
the same flag on a frontier-scale agentic workload: GLM-5.2 744B-A40B on terminal-bench-2,
16 GB300 nodes split 8 training / 8 inference, one Daytona sandbox per episode.


# Hot Restart Example
Source: https://miles.radixark.com/docs/examples/infra-features/hot-restart

Replaces the orchestration script and rollout executor of a live run, keeping trainers and engines up.

Replace the orchestration script and the rollout executor of a running job without taking down
the trainers or the inference engines.

## Quick Start

Start the run under a name the same command can be relaunched with:

```bash theme={null}
python examples/infra_features/hot_restart/run_qwen3_0_6b_hot_restart.py \
    --cluster-backend kubernetes --namespace <your-namespace> --run-id my-run
```

While it trains, hot restart it with the same command plus one flag:

```bash theme={null}
python examples/infra_features/hot_restart/run_qwen3_0_6b_hot_restart.py \
    --cluster-backend kubernetes --namespace <your-namespace> --run-id my-run \
    --hot-restart orchestration,rollout_executor
```

Both invocations name the same W\&B run, so the relaunched orchestration script resumes reporting
into the run the original launch opened.

## TODO

* Show a hot restart that changes the rollout executor's args and resumes under them; today the
  command is repeated verbatim.


# FP8 training examples
Source: https://miles.radixark.com/docs/examples/infra-features/low-precision

Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability.

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 theme={null}
   --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 theme={null}
   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.

### MoE weight updates

With `--megatron-to-hf-mode raw` and `--expert-tensor-parallel-size 1`, EP ranks convert
and quantize their complete local routed experts in parallel. Expert-DP replicas
split that work into contiguous ranges of whole experts, so each expert is
processed once; individual expert tensors are never split across expert DP.

With expert tensor parallelism greater than one, updates use the existing path:
gather unquantized ETP shards and EP experts before conversion and quantization.
Both paths support unquantized updates.

For 256 experts with EP64 and EDP8, each EP rank has four local experts: four
replicas process one expert each, and four process none. All replicas participate
in gathering the converted weights and scales.

Gathering reuses Megatron's existing tensor/data/context-parallel group, which
covers EP and full expert DP at ETP1, including context-parallel replication.
Pipeline stages are gathered first when required by the transfer protocol.

Quantization exclusions still apply. Shared experts and nonexpert layers retain
their existing gathering and conversion behavior. The Bridge exporter is unchanged.

## 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 theme={null}
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 theme={null}
RUNTIME_ENV_JSON="{
  \"env_vars\": {
    ...
    \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\",
    \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\"
  }
}"
```

**Launch Commands:**

```bash theme={null}
# 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.


# P2P Weight Transfer Examples
Source: https://miles.radixark.com/docs/examples/infra-features/p2p-weight-transfer

Point-to-point weight transfer between training and rollout engines.

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](/docs/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 <VARIANT> [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 theme={null}
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 theme={null}
# 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 theme={null}
# 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 theme={null}
# 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.


# Random fully-async example
Source: https://miles.radixark.com/docs/examples/infra-features/random-async

Dataset-free stress test of the async rollout ↔ trainer loop.

Minimal sibling of `examples/infra_features/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 theme={null}
# 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`.


# Score Centering
Source: https://miles.radixark.com/docs/examples/infra-features/score-centering

Center off-policy policy gradients using the sampler's top-k probabilities, with optional TIS or MIS weights.

Correct off-policy score drift using the sampler's top-k probabilities, with optional truncated or masked importance weights.

This implements [Score Centering Stabilizes Off-policy Reinforcement Learning](https://arxiv.org/abs/2609.20807), including the efficient top-k approximation in Appendix A.

## Enable it

Add these arguments to an existing text-only GRPO training recipe:

```bash theme={null}
--loss-type score_centering \
--advantage-estimator grpo \
--rollout-top-logprobs-num 128 \
--score-centering-is none \
--rollout-temperature 1.0 \
--rollout-top-p 1.0 \
--rollout-top-k -1 \
--use-rollout-logprobs \
--disable-grpo-std-normalization \
--calculate-per-token-loss
```

Selecting `--loss-type score_centering` automatically enables `--sglang-enable-fp32-lm-head` for rollouts. This avoids rounding LM-head logits to lower precision before computing sampling probabilities. It can increase serving compute and memory costs, depending on the SGLang version.

Keep reward mean subtraction enabled. Disabling standard-deviation normalization gives the paper's group-centered rewards. This is a separate REINFORCE-style loss: PPO clipping parameters do not apply. Existing batch size and update scheduling still control how many updates consume a rollout batch; choose them explicitly when reproducing an experiment.

Choose `--score-centering-is tis` for weights clipped at `--score-centering-tis-clip` (default 2). Choose `mis` to retain ratios in `[--score-centering-mis-low, --score-centering-mis-high]` (defaults 0.5 and 5), setting other weights to zero. These weights are centered together with the score. Use these options instead of `--use-tis` or a custom TIS function.

The existing entropy and reference-KL loss options remain available. They are separate regularizers; the score-centering identity applies to the policy-gradient term.
With filtered sampling, sampling-support replay rules prohibit reference KL; entropy is computed over the recorded support.
The loss keeps full-vocabulary actor scores for reference KL, matching the reference forward even if called directly with replayed candidates. This does not remove the startup restriction above.
Reference-KL tokens with non-finite probabilities or absolute log-probability ratios above 40 are excluded before exponentiation to keep gradients finite. With unbiased KL, this also applies to the train/rollout ratio. `train/kl_invalid_fraction` reports the excluded fraction.

Training logs include `train/train_rollout_logprob_abs_diff` and `train/train_rollout_kl`. The latter uses the same masked, sampled-token k3 estimator of KL(rollout || train) as the policy loss. It is a detached diagnostic and is emitted even when reference-KL regularization is disabled.

## How it works

When the rollout distribution differs from the current trainer, even a constant reward can produce an unwanted average policy update. Score centering subtracts the expected weighted score under the rollout distribution.

Let `p` be the current trainer distribution, `q` the distribution that actually sampled the token, `H` the stored candidates, `A` the detached advantage, and `f` the selected importance-weight function. The implementation computes:

```text theme={null}
rho   = max(1 - sum(q[H]), 1e-6) / max(1 - sum(p[H]), 1e-6)
alpha = rho * f(1 / rho)
loss  = -A * (stop_gradient(f(p[token] / q[token])) * log(p[token])
             - sum(stop_gradient(q[H] * f(p[H] / q[H]) - alpha * p[H]) * log(p[H])))
```

Outside `H`, the approximation models `q` as `rho * p`. The sampled token always uses its recorded `q[token]`, including when it is outside `H`. All weights and correction coefficients are detached. Full-distribution centering cancels the expected constant-reward gradient; the top-k version approximates the true tail and does not guarantee exact cancellation for an arbitrary tail.

`--rollout-top-logprobs-num K` plays a different role in the two sampling modes:

* **Unfiltered sampling** (`--rollout-top-p 1.0 --rollout-top-k -1`): Miles automatically uses `selected` mode. The sampler draws from the full vocabulary, so Miles records only its top `K` tokens as `H` and the tail model covers the rest. `K` is a real truncation and must be set; a larger `K` leaves less to the tail model.
* **Filtered sampling** (top-p/top-k): Miles automatically uses `support` mode when candidate recording is enabled. The sampler draws only from the realized support, and Miles records that whole support with its post-filter probabilities. `H` is then the full sampling distribution and nothing is approximated. `K` only sizes the arrays: it must be at least `--rollout-top-k` and hold the realized support, which cutoff ties can make larger; a larger support fails validation. The trainer renormalizes its log probabilities over the same support.

SGLang reports three kinds of log probability:

| SGLang field | Distribution | Filtered | Sums to 1 on the support |
| - | - | - | - |
| `output_token_sampling_logprobs` (selected or support mode) | `q'`: renormalized after top-k/top-p filtering | yes | yes |
| `output_token_logprobs`, `output_top_logprobs` | `p_T = softmax(logits / T)` over the full vocabulary | no | no, less than 1 under filtering |
| `output_token_logprobs`, `output_top_logprobs` with `SGLANG_RETURN_ORIGINAL_LOGPROB=1` | `log_softmax(logits)`, without the temperature | no | — |

The first row is the sampler under filtering and the second the sampler without filtering. Miles-managed rollout workers set `SGLANG_RETURN_ORIGINAL_LOGPROB=0`, so the third row does not occur. Each configuration reads:

| Configuration | Sampling-support replay (`append_sampling_metadata`) | Candidates (`append_rollout_topk_logprobs`) | SGLang source |
| - | - | - | - |
| Sampling-support replay only (`--rollout-top-logprobs-num 0`) | stores the support; returns the sampled-token log probability | nothing (`K = 0`) | `output_token_sampling_mask`, `output_token_sampling_logprobs` |
| Score centering, unfiltered | not called (no `return_sampling_mask`) | stores the top `K` candidates | `output_top_logprobs`, or its flat arrays on the session path |
| Score centering, filtered | stores the support; picks the sampled-token log probability from its row | stores `[n, K]`: the same support IDs with their row log probabilities | both read the same `output_token_sampling_mask`, `output_token_sampling_logprobs` |

For support replay without reference KL, the trainer gathers only sampled and support logits, normalizes over that support, and saves activations proportional to the candidate count rather than the vocabulary size. Other paths retain full-vocabulary normalization, reducing normalization scalars and selected logits across tensor-parallel ranks. It excludes padded vocabulary entries and computes probabilities in float32 for BF16/FP16 models. `--log-probs-chunk-size` controls the temporary computation size; `--recompute-loss-function` can trade computation for saved activations. No full vocabulary is gathered across ranks.

## Rollout and data contract

Native SGLang generation, the legacy rollout path, and both session-server versions request candidate probabilities at generation time when `--rollout-top-logprobs-num` is positive. Miles derives the logprob mode at startup from `--rollout-top-p` and `--rollout-top-k`. On these training requests, the configured count and derived mode override any client-supplied `top_logprobs`, `top_logprobs_num`, or `sampling_logprobs_mode`. A count of zero disables candidate recording without requesting support-wide probabilities. Session-server training requests ask for the unfiltered candidates as base64 `[n, K]` arrays (`output_top_logprobs_{val,idx}_flat_b64`) instead of nested `output_top_logprobs` rows; Miles decodes either form. Each `Sample` carries:

* `rollout_topk_token_ids`: int32 array shaped `[response_length, k]`.
* `rollout_topk_log_probs`: float32 array of the same shape.
* `rollout_log_probs`: the actual sampled-token log probabilities.

Evaluation requests skip this collection and may use independent sampling settings, including greedy decoding. The built-in agentic producer marks evaluation sessions when creating them; custom session clients should create them with `POST /sessions` with JSON body `{"evaluation": true}`.

Unused candidate slots and non-trained observation rows contain token ID `-1` and log probability `-inf`. Tool-observation masks, multi-turn merging, retries, trailing-token trimming, and truncation preserve row alignment. Session serialization retains both arrays. For score centering in support mode, training batches share the recorded support IDs when every sample has exactly the same candidate order and prefix padding. A per-row candidate count preserves observation rows and masked generated rows; any mismatch keeps the original arrays for the whole batch. The trainer reconstructs only its context-parallel rows. The source Samples and session payloads still retain both representations.

Custom rollout producers must supply these fields with probabilities from the actual generation call and no repeated non-negative token ID in a row. Conversion always rejects missing candidate IDs, candidate log probabilities, or sampled-token log probabilities, including from custom producers. With `--ci-test`, Miles also validates the complete sample before training, rejecting duplicate candidates, sampling-support mismatches, and disagreeing sampled/candidate probabilities. This full-sample validation is skipped in normal training because sorting and support matching are expensive on long responses. Rescoring old rollouts with newer weights is not a substitute. With the feature disabled, requests and the session wire format remain unchanged.

## Supported configurations and limits

* The shared loss is wired into Megatron and FSDP. Candidate selection supports tensor parallelism, packed (`thd`) and padded (`bshd`) zigzag context parallelism, and packed all-gather context parallelism.
* Sampling requires a fixed positive temperature and `min_p=0` on every call. Filtered sampling, including top-p filtering, automatically uses support mode and requires a positive `top_k`, for example `--rollout-top-p 0.9 --rollout-top-k 64 --rollout-top-logprobs-num 128`. Global filtered rollout settings automatically enable sampling-support replay; per-request overrides are checked when each request is built. See the [sampling-support replay guide](/docs/advanced/sampling-support-replay) for its request and server requirements.
* Filtered sampling requires SGLang with support log probabilities (SGLang PR [#40932](https://github.com/sgl-project/sglang/pull/40932), included in the `sglang-miles` branch by [#41047](https://github.com/sgl-project/sglang/pull/41047), merge commit [`ae04cb14046896b6d453758c5769d639deedd353`](https://github.com/sgl-project/sglang/commit/ae04cb14046896b6d453758c5769d639deedd353) or a descendant containing it). External servers must set `SGLANG_RETURN_ORIGINAL_LOGPROB=0` like Miles-managed workers. OpenAI session responses must expose the SGLang fields above in `choices[0].meta_info`; a generic OpenAI-compatible server without that metadata is insufficient.
* The SGLang router must include [sgl-router-for-miles #21](https://github.com/radixark/sgl-router-for-miles/pull/21) (`df2c790`; see `Git Commit` in `smg --version-verbose`), which raises the chat `top_logprobs` limit from 20 to 128 and forwards `sampling_logprobs_mode`.
* Compact candidate transport on the session path requires SGLang PR [#42890](https://github.com/sgl-project/sglang/pull/42890), included in the `sglang-miles` branch by [#42838](https://github.com/sgl-project/sglang/pull/42838), and the router change [sgl-router-for-miles #24](https://github.com/radixark/sgl-router-for-miles/pull/24). Without either, SGLang returns the larger nested rows, where a `-inf` candidate arrives as `null` and is stored as NaN.
* Constrained/custom sampling, speculative decoding, true-on-policy mode, OPD, multi-LoRA/Tinker losses, sequence masking, custom policy-loss reducers, custom train-data converters, and logprob recomputation via prefill are rejected. Multimodal token expansion is not supported. The initial advantage estimator is GRPO.
* Retaining `k=128` uses about 1 KiB per response position for the two arrays, before transport overhead. Eligible support-mode training batches replace the duplicate ID array with one int32 count per response position, saving `4 * (k - 1)` bytes per position in training transport and storage. Source Samples and session payloads keep the original storage cost. Larger `k` improves the unfiltered tail approximation at additional storage and compute cost.

## Metrics and verification

Training logs include `sc_correction`, `sc_train_head_mass`, `sc_rollout_head_mass`, `sc_tail_ratio`, `sc_importance_weight`, `train_rollout_kl`, and `train_rollout_logprob_abs_diff`, under the usual `train/` namespace. Small head mass means more of the distribution is approximated by the tail model. Large tail ratios indicate a substantial mismatch in remaining mass.
For filtered sampling, the head covers the full support, so both head masses and the tail ratio should be approximately one.

The numerical tests compare gradients against an independent dense-distribution oracle for all three weighting modes, including sampled tokens outside the head, constant rewards, and tiny tails. Pipeline tests exercise real native/session producers, serialization, data-parallel splitting, masks, advantage computation, regularization, and checkpointed loss scaling. Run the focused tests in the repository's test environment:

```bash theme={null}
python -m pytest tests/fast/backends/training_utils/test_score_centering.py \
    tests/fast/backends/training_utils/test_score_centering_pipeline.py \
    tests/fast/backends/training_utils/test_score_centering_filtered.py \
    tests/fast/backends/training_utils/test_score_centering_support.py
python -m pytest \
    tests/fast/backends/training_utils/test_score_centering_distributed.py
MILES_TEST_CUDA_DISTRIBUTED=1 python -m pytest \
    tests/fast/backends/training_utils/test_score_centering_distributed.py
```

The distributed tests use four CPU/Gloo processes or four CUDA/NCCL processes (TP=2, CP=2), all four context layouts and all weighting modes, compact support replay with reference KL on and off, plus BF16 selected-probability gradients. Support-scoring tests check saved tensor shapes and float32/64/BF16 gradients against a dense oracle. The independent dense-gradient oracle also runs on CUDA when available. These are correctness tests, not a reproduction of the paper's GPU training results.

For a real SGLang server, also run the opt-in protocol probe:

```bash theme={null}
MILES_LIVE_SCORE_CENTERING_ENDPOINT=http://127.0.0.1:30000 \
MILES_LIVE_SCORE_CENTERING_MODEL=/path/to/model \
MILES_LIVE_SCORE_CENTERING_SERVED_MODEL=your-served-model \
python -m pytest --confcutdir=tests/manual tests/manual/test_score_centering_live.py
```

Set `SGLANG_RETURN_ORIGINAL_LOGPROB=0` on the server before starting it. The probe checks unfiltered native and OpenAI response metadata using the production candidate collector and validator, and checks temperature scaling at 0.7, 1.0 and 1.3. It warms the shared prompt first so that cached and uncached prefills do not confound the temperature comparison. Set `MILES_LIVE_SCORE_CENTERING_ARTIFACT_DIR` to keep the raw responses.


# Split Deployment Example
Source: https://miles.radixark.com/docs/examples/infra-features/split-deployment

Installs one run as several helm releases — trainer, engines and orchestration script apart.

One run installed by hand as one helm release per component: one per trainer, one per inference
engine, and one carrying the orchestration script.

## Example 1: A Single Policy Run

Inside a miles workbench the cluster backend and the namespace are already set; name the run once:

```bash theme={null}
export MILES_SCRIPT_RUN_ID=split-demo
export MILES_SCRIPT_RUN_UUID=$(python -c 'from miles.utils.run_uuid import generate_run_uuid; print(generate_run_uuid())')
```

Then install one component per command — identical but for `--deploy-component` and
`--deploy-instance-id`, with `primary` last since installing it blocks until the run ends:

These scripts import their shared address book as `examples.…`, so run them as modules from the
repository root rather than by path:

```bash theme={null}
SCRIPT=examples.infra_features.split_deployment.run_qwen3_0_6b_split

python -m $SCRIPT --deploy-component trainer
python -m $SCRIPT --deploy-component inference --deploy-instance-id e0
python -m $SCRIPT --deploy-component inference --deploy-instance-id e1
python -m $SCRIPT --deploy-component primary
```

When the run ends, the releases that carry no orchestration script are still up:

```bash theme={null}
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-trainer -n $MILES_SCRIPT_NAMESPACE
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-inference-e0 -n $MILES_SCRIPT_NAMESPACE
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-inference-e1 -n $MILES_SCRIPT_NAMESPACE
```

## Example 2: A Run That Trains Several Policies

`run_solver_verifier_gsm8k_split.py` installs the [multi\_policy](/docs/examples/multi-policy) solver /
verifier demo the same way: one release per policy trainer, one per policy's engines, and the
orchestration script last — five releases for two policies.

Name this run once as well, so the example stands on its own:

```bash theme={null}
export MILES_SCRIPT_RUN_ID=split-multi-policy-demo
export MILES_SCRIPT_RUN_UUID=$(python -c 'from miles.utils.run_uuid import generate_run_uuid; print(generate_run_uuid())')
```

```bash theme={null}
SCRIPT=examples.infra_features.split_deployment.run_solver_verifier_gsm8k_split

python -m $SCRIPT --deploy-component trainer --deploy-instance-id solver-actor
python -m $SCRIPT --deploy-component trainer --deploy-instance-id verifier-actor
python -m $SCRIPT --deploy-component inference --deploy-instance-id solver
python -m $SCRIPT --deploy-component inference --deploy-instance-id verifier
python -m $SCRIPT --deploy-component primary
```

```bash theme={null}
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-trainer-solver-actor -n $MILES_SCRIPT_NAMESPACE
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-trainer-verifier-actor -n $MILES_SCRIPT_NAMESPACE
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-inference-solver -n $MILES_SCRIPT_NAMESPACE
helm uninstall miles-run-$MILES_SCRIPT_RUN_ID-inference-verifier -n $MILES_SCRIPT_NAMESPACE
```


# Rollout Correction Methods
Source: https://miles.radixark.com/docs/examples/infra-features/train-infer-mismatch-helper

Algorithmic methods for rollout correction (e.g., TIS, MIS).

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**).
Specify the **IS/RS configs** with the `--tis-*` and `--rs-*` arguments.

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

Choose from the algorithms above with the command-line 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-tis`: True if apply importance sampling/rejection sampling to loss.

| `use_rollout_logprobs` | `use_tis` | 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-tis` enabled, Algorithm 2 & 3), you may specify the IS modes and applied levels.

### Config keys

The hook registers these command-line arguments from `MisConfig`:

```bash theme={null}
--use-tis
--use-rs
--tis-level token
--rs-level token
--tis-mode truncate
--tis-lower-bound 0.5
--tis-upper-bound 2.0
--rs-veto-threshold 1.0e-4
--tis-batch-normalize
--custom-tis-function-path examples.infra_features.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp
```

`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**, **mask**.
* `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. Unset falls back to the `tis_` bounds.
* `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

These metrics help diagnose policy divergence and guide hyperparameter tuning. Which
ones you get depends on the correction function in use. The built-in one
(`vanilla_tis_function`, used when `--custom-tis-function-path` is unset) reports only
`tis`, `tis_clipfrac` and `tis_abs`. Everything below comes from `mis.py`, so it needs
the `--custom-tis-function-path` wiring shown above.

All names below are logged under the `train/` namespace, and every key `mis.py` produces
carries a `mis_` prefix that its wrapper adds on the way out — so `training_log_ppl`
reaches wandb as `train/mis_training_log_ppl`. The two exceptions, marked in the tables,
come from miles itself and are not prefixed.

### Mismatch Monitoring Metrics

These metrics quantify the difference between training and rollout policies. `mis.py`
computes them whenever `rollout_log_probs` are available, whether or not IS/RS correction
is actually applied.

| Metric Name | Description |
| - | - |
| `mis_training_log_ppl` | Negative mean log probability under training policy: $-\mathbb{E}[\log \pi_{\text{train}}]$ |
| `mis_training_ppl` | Perplexity of training policy: $\exp(-\mathbb{E}[\log \pi_{\text{train}}])$ |
| `mis_rollout_log_ppl` | Negative mean log probability under rollout policy: $-\mathbb{E}[\log \pi_{\text{rollout}}]$ |
| `mis_rollout_ppl` | Perplexity of rollout policy: $\exp(-\mathbb{E}[\log \pi_{\text{rollout}}])$ |
| `mis_kl` | Forward KL divergence estimator: $\mathbb{E}[\log \pi_{\text{rollout}} - \log \pi_{\text{train}}]$ |
| `mis_k3_kl` | K3 KL estimator: $\mathbb{E}[\exp(r) - r - 1]$ where $r = \log \pi_{\text{train}} - \log \pi_{\text{rollout}}$ |
| `mis_log_ppl_diff` | Log perplexity difference |
| `mis_log_ppl_abs_diff` | Absolute log perplexity difference |
| `mis_ppl_ratio` | Perplexity ratio |
| `mis_chi2_token`, `mis_chi2_seq` | Token- and sequence-level $\chi^2$ divergence between the two policies |
| `train_rollout_logprob_abs_diff` (miles, unprefixed) | 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. Under the shared `mis_` prefix, the `tis_` and `rs_` parts say which stage produced the number.

| Metric Name | Description | Emitted when |
| - | - | - |
| `ois` (miles, unprefixed) | On-policy importance sampling ratio: $\exp(\log \pi_{\text{train}} - \log \pi_{\text{old}})$ | `--use-tis` (Algorithm 2 only) |
| `mis_tis_weight_before_bound` | Raw IS weights before any bounding: $\exp(\text{log-ratio})$ | `use_tis` |
| `mis_tis_weight_after_bound` | IS weights after the `tis_mode` bounding | `use_tis` |
| `mis_tis_truncate_fraction` | Fraction of weights truncated | `tis_mode: truncate` |
| `mis_tis_clip_fraction_low` / `mis_tis_clip_fraction_high` | Fraction of weights clipped below / above the bound | `tis_mode: clip` |
| `mis_tis_mask_fraction_low` / `mis_tis_mask_fraction_high` | Fraction of tokens rejected below / above the bound | `tis_mode: mask` |
| `mis_rs_mask_fraction_low` / `mis_rs_mask_fraction_high` | Fraction of tokens rejected by rejection sampling | `use_rs` |
| `mis_rs_catastrophic_token_fraction` | Fraction of catastrophic tokens (below the veto threshold) | `rs_veto_threshold` set |
| `mis_rs_catastrophic_seq_fraction` | Fraction of sequences holding a catastrophic token | `rs_veto_threshold` set |
| `mis_is_ratio_mean_after_tis_rs` | Mean IS weight after both TIS and RS | `use_tis` |
| `mis_batch_norm_factor` | Batch normalization factor applied to weights (1.0 when off) | `use_tis` |
| `mis_is_ratio_mean_final` / `mis_is_ratio_min_final` / `mis_is_ratio_max_final` | Final IS weight statistics actually multiplied into the loss | `use_tis` |

## 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)


# True On-Policy between Training and Inference
Source: https://miles.radixark.com/docs/examples/infra-features/true-on-policy

Ensures strictly equal log probabilities between inference (SGLang) and training engines.

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 theme={null}
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 theme={null}
python scripts/run_qwen3_4b.py full-train --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.


# Multi-LoRA Tinker Gateway
Source: https://miles.radixark.com/docs/examples/multi-lora

Serve concurrent LoRA fine-tuning clients on one shared base model through the Tinker protocol.

> **Read the docs:** [Multi-LoRA training](https://miles.radixark.com/docs/advanced/lora#multi-lora-training).

* `serve_qwen3_30b_a3b_tinker.py`: prepare Qwen3-30B-A3B and launch the gateway.
* `run_multi_tenant_example.py`: check marker memorization for one client or adapter isolation across concurrent tenants.
* `tinker_agent/`: client package for agentic training; see [Agentic training](#agentic-training).

## Layout

One 8-GPU node, disaggregated (multi-LoRA forbids `--colocate`):

* 4 training GPUs: TP2 for the dense layers, EP4 for the 128 routed experts.
* 4 sampling GPUs: two SGLang engines of 2 GPUs each, serving adapter versions by name.
* 4 adapter slots (`--multi-lora-n-adapters`), rank up to 32, covering attention
  (`linear_qkv`, `linear_proj`), the per-expert MoE projections (`linear_fc1`, `linear_fc2`),
  and the output layer (`output_layer`) so the cookbook's default `train_unembed=True` is servable.

The example enables attention, MLP, and output-head training. Client SDK flags
must match the server's selected groups. See [LoRA target selection](/docs/advanced/lora#hf-target-source-of-truth)
for `--target-modules attn,mlp,unembed`; use `attn,mlp` to disable output-head
training. Tinker accepts only these group names and rejects `--exclude-modules`.

## Run

The gateway implements the `tinker==0.26.2` wire schema (newer SDKs renamed protobuf fields); install that exact version on the serving node and the client:

```bash theme={null}
pip install "tinker==0.26.2"
```

Start the gateway:

```bash theme={null}
python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py prepare   # once per node
python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py serve     # Tinker API on :10613
```

Checkpoints default to `<output_dir>/checkpoints/<run_id>`; use `--save-dir` to choose another root.

The launcher sets `--log-probs-chunk-size 1024`, which computes log-probs in 1024-token chunks
to bound peak GPU memory on long sequences. To override it, add
`--extra-args "--log-probs-chunk-size <N>"` to the `serve` command.

Install `tinker` on the client, then run the marker checks:

```bash theme={null}
# one client: train, save for sampler, sample back the marker
python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode single

# four tenants training concurrently on the same prompt with different markers;
# passing means the adapters stayed isolated end to end
python examples/multi_lora/run_multi_tenant_example.py --base-model /root/models/Qwen3-30B-A3B --mode multi --clients 4
```

## Agentic training

The gateway also supports agentic RL with your own sandbox or harness through the official
Tinker cookbook capture recipe. The [tinker\_agent](/docs/examples/multi-lora/tinker-agent) client package
records each model call's exact tokens and log-probs and converts each episode into training
data.

## Context window

Each sampled sequence, prompt plus completion, must fit a single window, defined as the smaller of:

* the trainer's per-sequence limit: `--max-tokens-per-gpu`, rounded down to a multiple of
  TP × `--data-pad-size-multiplier`, times the context-parallel size;
* the inference limit: `--sglang-context-length` (the model's maximum length if unset) minus 7
  tokens, so that SGLang accepts a full-window sequence as a prompt.

With the launcher defaults, the window is 8192 tokens. For each sampling request, the gateway:

* clamps `max_tokens` to the remaining window, without modifying the client's request;
* rejects a prompt that leaves no generation budget with HTTP 400 `prompt is too long`, before
  dispatch to the engine. Agent harnesses use this error to trigger history compaction;
* accepts a full-window sequence for scoring with `compute_logprobs`.

Training Datums are validated against the trainer limit only; a shorter inference window does
not reduce the trainable sequence length.

## Supported inputs

Training accepts text with 1-D loss inputs. 2-D soft targets, including SDFT,
are not supported. Sampling requires a `/sampler_weights/` path returned by
`save_weights_for_sampler()`; `/weights/` training checkpoints cannot be sampled directly.

## Failure handling

A terminal failure of `forward_backward`, `optim_step`, or `load_state` ends
training for that model, including commands already queued behind it.
This includes content validation failures with a valid model and sequence.
Create a new model and restore a saved checkpoint to continue; completed
futures and published checkpoints keep their results.

Known request-local failures of `forward` or sampling leave model training
available. Checkpoint load/save execution failures, including filesystem errors,
invalidate the shared trainer cell and stop the server.
Saving sampler weights commits an immutable directory;
it does not call the inference engines. Sampling loads that snapshot from disk
on demand, including after cache eviction. An engine load failure fails the
sampling request; it leaves the snapshot and training state intact. Unknown
trainer execution failures invalidate the shared trainer cell and stop the server.

This gateway provides failure isolation, not automatic training recovery.
Checkpoints persist; futures, deduplication, and unsaved accumulation do not
survive a server restart.

## Sampler snapshots

Training and inference must use the same base checkpoint. Tinker engines load
that frozen base at startup and serve without trainer weight updates; dummy
loading and `update_weights: true` are rejected. Ordinary full-model and
single-LoRA training continue to use the existing weight updater.

`--tinker-checkpoint-root` must be on storage shared by the trainers, gateway,
and every inference engine. All trainer ranks participate in adapter gathering;
rank 0 writes the tensors and config, then `META.json` after all ranks finish.
Existing sampler versions cannot be overwritten. Saving between `forward_backward`
and `optim_step` neither applies nor discards pending gradients.

Training checkpoint saves and loads are serialized within the gateway. Overwriting
deletes the previous checkpoint before writing the new one; a failed overwrite does
not preserve the previous checkpoint. Checkpoints are ordinary directories;
overwriting does not retain hidden versions.


# Terminal-Bench 2 training with Harbor and Tinker
Source: https://miles.radixark.com/docs/examples/multi-lora/harbor-tinker

Terminal-Bench 2 training with Harbor sandboxes and the official Tinker cookbook RL loop.

This recipe trains a LoRA policy on Terminal-Bench 2 (TB2) with the official Tinker cookbook RL
loop and the Miles Tinker gateway. [Harbor](https://github.com/harbor-framework/harbor) runs
each task's agent in a sandbox. The agent samples through the Tinker capture proxy described in
[Agentic training with Tinker capture](/docs/examples/multi-lora/tinker-agent), and the task's verifier
scores each trial. The cookbook computes group advantages and applies PPO to the captured
tokens.

The recipe accepts any dataset in Harbor's native task format; TB2 is the reference
configuration.

## Requirements

* **Gateway.** An 8-GPU node running the Miles Tinker gateway: 4 training and 4 sampling
  GPUs. The reference configuration serves Qwen3.6-35B-A3B on H200 with a 65536-token window
  (see [Setup](#setup)).
* **Client.** Python 3.12 or newer on the machine that runs the recipe.
* **Sandboxes.** An E2B account and API key; see
  [Sandbox Providers](/docs/user-guide/sandbox-providers). The sandboxes must be able
  to reach the client machine, which hosts the proxy.

## Setup

### 1. Start the gateway

Download the checkpoint and serve it with the gateway launcher, overriding the model, LoRA
layout and context window:

```bash theme={null}
hf download Qwen/Qwen3.6-35B-A3B --local-dir /root/models/Qwen3.6-35B-A3B

python examples/multi_lora/serve_qwen3_30b_a3b_tinker.py serve \
  --hf-checkpoint /root/models/Qwen3.6-35B-A3B --model-type qwen3.6-35B-A3B_lora \
  --lora-rank 16 --lora-alpha 32 --n-adapters 1 --target-modules attn,mlp,unembed \
  --extra-args "--tinker-base-model Qwen/Qwen3.6-35B-A3B \
    --max-tokens-per-gpu 65536 --max-position-embeddings 65536 --sglang-context-length 65536 \
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 --qkv-format thd"
```

* `--tinker-base-model` is the model name clients pass as `config.model_name`.
* `--target-modules attn,mlp,unembed` matches the cookbook client's default LoRA layout, which
  trains the output head (`train_unembed=True`). The client's LoRA rank (16) must not exceed
  `--lora-rank`.
* `--max-tokens-per-gpu 65536` and `--sglang-context-length 65536` set the trainer and inference
  limits; with context-parallel size 1, the sampling window is 65529 tokens. See
  [Context window](/docs/examples/multi-lora#context-window). Tensor parallelism does not raise the
  per-sequence limit.
* The gateway serves on port `10613`.

### 2. Install the client

From a Miles checkout:

```bash theme={null}
uv pip install './examples/multi_lora/tinker_agent[harbor]' \
  'harbor[e2b] @ git+https://github.com/harbor-framework/harbor@917b3c1be7cf6553a7a9fe25de8adb6f8e639287'
```

The `harbor` extra pins the tested Harbor revision; `harbor[e2b]` adds the E2B provider at the
same revision.

### 3. Download the dataset

```bash theme={null}
harbor datasets download terminal-bench@2.0 --output-dir /data/harbor-tasks
export TB2_TASKS_DIR=/data/harbor-tasks/terminal-bench
```

The recipe loads every directory under `TB2_TASKS_DIR` that contains a `task.toml`, in sorted
order. Each task keeps its instruction, environment and verifier.

### 4. Write the trial configuration

`trial.json` is a Harbor `TrialConfig` without the `task` field; the recipe supplies each task.
The recipe substitutes two placeholders: `{proxy_url}`, the trajectory's proxy URL, and
`{proxy_api_key}`, the proxy's auth token. The reference configuration runs mini-swe-agent on
E2B:

```json theme={null}
{
  "trials_dir": "/tmp/harbor-trials",
  "environment": {"type": "e2b"},
  "agent": {
    "name": "mini-swe-agent",
    "model_name": "openai/policy",
    "override_timeout_sec": 5400,
    "env": {
      "OPENAI_BASE_URL": "{proxy_url}/v1",
      "OPENAI_API_KEY": "{proxy_api_key}",
      "MSWEA_COST_TRACKING": "ignore_errors"
    },
    "kwargs": {
      "max_tokens": 8192,
      "max_seq_len": 65536,
      "config": {"model": {"model_kwargs": {"temperature": 0.8}}}
    }
  }
}
```

For `claude-code`, set `agent.env.ANTHROPIC_BASE_URL` to `{proxy_url}` and
`agent.env.ANTHROPIC_API_KEY` to `{proxy_api_key}`. All other fields, including sandbox
credentials, timeouts and verifier settings, follow Harbor's schema.

## Run

```bash theme={null}
export TINKER_API_KEY=tml-<tenant>
export E2B_API_KEY=<e2b-api-key>
export HARBOR_PROXY_HOST=<client-address-reachable-from-sandboxes>

python -m tinker_agent.train \
  config.base_url=http://<gateway-host>:10613 \
  config.model_name=Qwen/Qwen3.6-35B-A3B config.renderer_name=qwen3_5 \
  config.lora_rank=16 config.learning_rate=1e-6 config.max_tokens=8192 config.loss_fn=ppo \
  config.ttl_seconds=None config.log_path=/tmp/harbor-tinker \
  config.dataset_builder.tasks_dir="$TB2_TASKS_DIR" \
  config.dataset_builder.groups_per_batch=4 config.dataset_builder.group_size=8 \
  harbor.trial_config=trial.json \
  harbor.listen_host=0.0.0.0 harbor.advertised_host="$HARBOR_PROXY_HOST"
```

Each batch runs `groups_per_batch × group_size` trials: 4 tasks with 8 trajectories each.
Add `config.max_steps=1` for a single-batch run. Omit `config.base_url` to train on the hosted
Tinker service instead of the Miles gateway.

## Configuration

`config` is the cookbook's `tinker_cookbook.rl.train.Config`, with its native field names and
defaults. Run `python -m tinker_agent.train --help` for the full list.

| Option | Default | Description |
| - | - | - |
| `config.dataset_builder.tasks_dir` | required | Directory of Harbor task directories. |
| `config.dataset_builder.groups_per_batch` | `4` | Tasks per batch. |
| `config.dataset_builder.group_size` | `4` | Trajectories per task; advantages are computed within a group. |
| `config.dataset_builder.epochs` | `1` | Passes over the dataset. |
| `harbor.trial_config` | required | Path to `trial.json`. |
| `harbor.reward_key` | `reward` | Key in the verifier's rewards used as the trial reward. |
| `harbor.listen_host` | `127.0.0.1` | Address the proxy binds to. Use `0.0.0.0` for remote sandboxes. |
| `harbor.advertised_host` | `127.0.0.1` | Address placed in `{proxy_url}`; it must be reachable from the sandboxes. |
| `harbor.max_parallel_trials_per_group` | `4` | Concurrent trials per group. |

The recipe enforces the following:

* `config.rollout_error_tolerance` is set by the recipe; setting it is rejected.
* `config.temperature` must remain `1.0`. The agent issues its own sampling requests, so the
  sampling temperature is set in `trial.json`.
* `config.max_tokens` is the proxy's default for requests that do not set `max_tokens`.
* `config.renderer_name`, if omitted, resolves to the cookbook's recommended renderer for
  `config.model_name`. Qwen3.6 uses `qwen3_5`.
* `config.ttl_seconds=None` is required with the Miles gateway, which does not implement
  checkpoint expiry.

Each group uses a fresh proxy with a random auth token. Rendering and capture run on the client
machine; captured data is stored under `<log_path>/capture`.

## Limitations

* **Trial outcomes.** A completed trial, an agent timeout, a context-window overflow and an
  output-token limit are scored with the verifier's reward. Any other Harbor failure, a missing
  or non-finite reward, a capture loss or a sample without log-probs raises: it cancels the
  group's remaining trials and stops training. Infrastructure errors are never scored as zero.
* **Failed requests.** A failed model call produces no training sample; the agent and Harbor
  decide whether the trial continues.
* **Algorithm.** Training is Tinker LoRA with the cookbook's PPO and group advantages. The agent,
  sandbox, batch size, sampling temperature, token limits and learning rate match the
  [cloud Harbor example](https://github.com/radixark/miles/blob/main/examples/experimental/harbor/README.md), but its full-parameter GRPO, optimizer defaults and KL behavior differ, so
  results are not directly comparable.
* **Recovery.** The recipe provides no partial-episode recovery.


# Agentic training with Tinker capture
Source: https://miles.radixark.com/docs/examples/multi-lora/tinker-agent

Agentic training with your own sandbox or harness through the official Tinker cookbook capture recipe.

The Miles Tinker gateway supports agentic RL with your own sandbox or harness through the
official Tinker cookbook capture recipe. Agents call an OpenAI- or Anthropic-compatible proxy;
each model call is recorded with its exact prompt tokens, sampled tokens and log-probs, and each
episode is read back as a cookbook `Trajectory` for training. Training uses the sampled tokens
directly, with no re-tokenization of the chat transcript.

`tinker_agent` is a standalone client package with no Miles dependency. It runs against the
Miles gateway or the hosted Tinker service.

## Requirements

* A Tinker endpoint: the Miles gateway from [Multi-LoRA Tinker Gateway](/docs/examples/multi-lora) (port
  `10613`), or the hosted Tinker service.

* Python 3.12 or newer, with the client package installed. The package pins `tinker==0.26.2` and
  `tinker-cookbook`:

  ```bash theme={null}
  uv pip install ./examples/multi_lora/tinker_agent
  ```

* The cookbook renderer for the served model: `qwen3` for Qwen3-30B-A3B, the checkpoint the
  gateway launcher serves; `qwen3_instruct` for Qwen3 Instruct-2507 models. See
  `tinker_cookbook.renderers.get_renderer` for all renderers.

## Example

`agent_rl.py` runs one two-turn episode through the proxy, assigns it a reward and applies one
PPO step to the captured tokens. `run_agent` stands in for the agent.

```python theme={null}
import argparse
import asyncio
import tempfile

import httpx
import tinker
from tinker_agent.proxy import serve_proxy
from tinker_agent.trajectory import flush_capture, read_samples, rows_to_trajectory
from tinker_cookbook import renderers
from tinker_cookbook.capture.store import capture_to_store
from tinker_cookbook.rl.data_processing import trajectory_to_data

RUN_ID = "demo"


async def run_agent(proxy_url):
    # Stand-in for your agent: any OpenAI-compatible client with base_url=proxy_url + "/v1" works.
    messages = [{"role": "user", "content": "Reply with the number 1."}]
    async with httpx.AsyncClient(base_url=proxy_url, timeout=300) as http:
        for _ in range(2):
            response = await http.post("/v1/chat/completions", json={"messages": messages, "max_tokens": 32})
            response.raise_for_status()
            messages += [response.json()["choices"][0]["message"], {"role": "user", "content": "Now reply with 2."}]
    return 1.0  # the episode's reward, from your verifier


async def main(base_url, base_model, renderer_name, store_url):
    service = tinker.ServiceClient(base_url=base_url, api_key="tml-demo")
    training = await service.create_lora_training_client_async(base_model=base_model, rank=8)
    saved = await (await training.save_weights_for_sampler_async(name="step-0"))
    sampler = await service.create_sampling_client_async(model_path=saved.path)
    renderer = renderers.get_renderer(renderer_name, sampler.get_tokenizer())

    async with serve_proxy(renderer, sampler, model_label=base_model, default_max_tokens=32) as port:
        reward = await run_agent(f"http://127.0.0.1:{port}/r/run/{RUN_ID}/group/0/traj/0")
    await flush_capture()

    rows = await read_samples(store_url, RUN_ID, group_idx=0, traj_idx=0)
    trajectory = rows_to_trajectory(rows, "completed")
    # With several trajectories per task, pass reward minus the group mean instead.
    datums = trajectory_to_data(trajectory, traj_advantage=reward)
    # "mask" is a client-side field; the cookbook trainer also drops it before forward_backward.
    datums = [
        tinker.Datum(
            model_input=datum.model_input,
            loss_fn_inputs={key: value for key, value in datum.loss_fn_inputs.items() if key != "mask"},
        )
        for datum in datums
    ]
    await (await training.forward_backward_async(datums, loss_fn="ppo"))
    await (await training.optim_step_async(tinker.types.AdamParams(learning_rate=1e-5)))
    print(f"trained on {len(rows)} captured calls")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", required=True)
    parser.add_argument("--base-model", required=True)
    parser.add_argument("--renderer", default="qwen3")
    args = parser.parse_args()
    with capture_to_store(RUN_ID, data_dir=tempfile.mkdtemp(prefix="tinker-capture-")) as captured:
        asyncio.run(main(args.base_url, args.base_model, args.renderer, captured.base_url))
```

Run it against the gateway:

```bash theme={null}
python agent_rl.py --base-url http://<gateway-host>:10613 --base-model /root/models/Qwen3-30B-A3B
```

The [Harbor recipe](/docs/examples/multi-lora/harbor-tinker) applies this workflow to Harbor tasks in sandboxes,
trained with the cookbook's RL loop.

## Workflow

1. **Capture.** `capture_to_store(run_id, data_dir=...)` starts a local capture store and records
   every Tinker sampling call made within its scope. The whole job runs inside it.
2. **Proxy.** `serve_proxy(renderer, sampling_client, model_label=..., default_max_tokens=...)`
   serves the proxy on a free local port and yields the port. `default_max_tokens` applies to
   requests that do not set `max_tokens`.
3. **Episode URL.** Each episode uses its own base URL,
   `http://<host>:<port>/r/run/<run_id>/group/<group_idx>/traj/<traj_idx>`:
   * `run_id` matches the `capture_to_store` run.
   * `group_idx` identifies a set of episodes that are compared with each other, such as several
     attempts at one task. The `group/<group_idx>` segment is optional.
   * `traj_idx` identifies one episode.
   * OpenAI clients append `/v1` to this URL (`/v1/chat/completions`); Anthropic clients use it
     as is (`/v1/messages`).
4. **Flush.** After the episodes finish, exiting `serve_proxy` waits for in-flight calls.
   `flush_capture()` then flushes the exporter and raises if any call was not recorded, so an
   incomplete episode never reaches training.
5. **Read.** `read_samples(store_url, run_id, group_idx=..., traj_idx=...)` returns an episode's
   calls in call order. `rows_to_trajectory(rows, stop_reason)` converts them to a `Trajectory`
   without re-tokenizing; `stop_reason` records how the episode ended, for example
   `"completed"`.
6. **Train.** `trajectory_to_data(trajectory, traj_advantage)` builds the Datums.
7. **Policy update.** After `optim_step`, save new sampler weights and serve the next episodes
   through a new proxy with the updated sampling client.

## Sandboxed agents

The proxy listens on `127.0.0.1` by default. To accept connections from a sandbox, pass a
reachable `host` such as `0.0.0.0` together with an `auth_token`; the proxy refuses a
non-loopback bind without a token. The agent presents the token as its API key, as
`Authorization: Bearer <token>` (OpenAI) or `x-api-key: <token>` (Anthropic). The URL given to
the agent must use an address the sandbox can reach, not `0.0.0.0`.

## Limitations

* **Rewards.** Captured calls carry no reward. The harness computes the reward and passes it as
  the advantage.
* **Failed calls.** Failed calls are not retried and produce no row, so a trajectory contains
  only the successful calls. Retry policy belongs to the agent.
* **Context window.** A prompt that exceeds the gateway's context window fails with HTTP 400
  `prompt is too long`; the agent is expected to compact its history or end the episode. See
  [Context window](/docs/examples/multi-lora#context-window).
* **Multiple samples.** The proxy rejects `n` other than 1. Run each attempt as a separate
  episode with its own `traj_idx` in the same group.
* **Recovery.** The package provides no retries, refusal handling, request rewriting or
  partial-episode recovery.


# Multi Policy Solver and Verifier
Source: https://miles.radixark.com/docs/examples/multi-policy

Two policies in one run — a solver answering gsm8k, and a verifier scored on ruling correctly about the solver's answers.

Two policies trained against each other in one run: a solver answering a gsm8k question and a
verifier ruling on its work. Each rollout yields one sample per policy; each policy has its own
trainer and its own inference engines inside the same job.

## Quick Start

Eight GPUs: 2 trainer GPUs and 2 single-GPU engines per policy. Needs a cluster backend and a
namespace, from `MILES_SCRIPT_CLUSTER_BACKEND` / `MILES_SCRIPT_NAMESPACE` or `--cluster-backend` /
`--namespace`:

```bash theme={null}
python examples/multi_policy/run_solver_verifier_gsm8k.py
```


# On-Policy Distillation Examples
Source: https://miles.radixark.com/docs/examples/on-policy-distillation

Teacher–student distillation on the student's own rollouts, run inside the on-policy training loop.

The canonical OPD documentation lives in
[`docs/advanced/on-policy-distillation.md`](/docs/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.


# Qwen3.5-35B-A3B Self-Distillation on a Single Node (RLVR teacher → OPD)
Source: https://miles.radixark.com/docs/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.

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 theme={null}
# 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 theme={null}
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 theme={null}
# 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 theme={null}
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 theme={null}
# 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 `</think>` 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](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/](https://thinkingmachines.ai/blog/on-policy-distillation/)


# PPO Example
Source: https://miles.radixark.com/docs/examples/ppo

Actor-critic PPO with GAE advantages, where the critic shares the actor's train GPUs.

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 theme={null}
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 theme={null}
--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. Both `--megatron-to-hf-mode raw`
  and `bridge` are supported; in bridge mode the critic's value head is freshly initialized rather
  than loaded from the HF checkpoint, since it has no HF counterpart.
* **`--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 `--indep-dp` (which train fault tolerance implies).** Shared actor/critic
  PPO hands the critic outputs to a single trainer cell as external data.

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


# Retool v2
Source: https://miles.radixark.com/docs/examples/retool-v2

Tool-enabled language model generation with sandboxed Python code execution interleaved with thinking.

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 theme={null}
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.

### Single GPU, fully async and colocated

```bash theme={null}
python examples/retool_v2/run_retool_multi_turn.py --fully-async --num-gpus-per-node 1
```

`--fully-async` switches the run to `train_async.py`, where the rollout engines keep
generating between training steps and share the GPUs with the trainer; see
[fully async](https://github.com/radixark/miles/blob/main/docs/user-guide/fully-async.md).


# SWE-Agent training with Harbor on Docker sandboxes
Source: https://miles.radixark.com/docs/examples/swe-agent-harbor-docker

Trains coding and terminal agents with Harbor-managed local Docker sandboxes and verifier rewards.

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 theme={null}
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://<agent-server>: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 theme={null}
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 theme={null}
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 theme={null}
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://<agent-server>:30000 \
    --session-server-external-host <trainer-host-reachable-from-agent-server> \
    --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.

The launcher resumes from `--save-dir`: if that directory already holds a
checkpoint, training continues from it instead of starting over from
`--ref-load`. Relaunching the same command after a crash or preemption picks up
where the run stopped. Use a fresh `--save-dir` for a new run, because reusing
an old one silently continues that run. Pass `--load-dir` to resume from a
different directory.

`--session-server-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. The launcher starts 32 session-server workers on ports 30000-30031 and the SGLang router on port 31000, so ensure that range and port are reachable end to end; Tailscale is one option when the machines are on different networks.

`--session-server-external-host` publishes every session server at that one host and keeps them all on the head node, so it must reach the head; that is what a multi-node job needs when only the head is reachable from the sandboxes. When every trainer pod has its own reachable address, leave the flag unset and have the deployment set `MILES_NODE_EXTERNAL_IP` in each pod to that pod's own address -- a tailscale sidecar's `tailscale ip -4`, for instance; a DNS name also works. Each node then reports its own, and Miles hands every sandbox the address of the one instance holding its session. Nothing is needed at all when the pod addresses already route from the sandbox network, which is the case whenever the sandbox platform is allowed the trainer's subnet.

## 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/<step>` and `rollout_data/<step>.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.


# Getting Started
Source: https://miles.radixark.com/docs/getting-started/index

Install Miles and run your first RL training job — the two pages you need to go from zero to a working loop.

Two pages take you from a fresh machine to a running RL training job:

* **[Installation](/docs/getting-started/installation)** — Docker (recommended), pip, or building from source on NVIDIA / AMD.
* **[Quick Start](/docs/getting-started/quick-start)** — `docker pull` to a GRPO run on Qwen3-4B in under an hour, on a single 8-GPU node.

After the loop is running, [Models](/docs/models/index) covers the family-specific recipes and the [User Guide](/docs/user-guide/index) walks through concepts, data, monitoring, and customization.


# Installation
Source: https://miles.radixark.com/docs/getting-started/installation

Install Miles on NVIDIA or AMD GPUs. Docker is the recommended path.

There are three ways to install Miles. Docker is recommended because Miles pins patched
versions of SGLang, Megatron-LM, and a few CUDA kernels.

## Method 1: Docker (recommended)

**NVIDIA**

```bash theme={null}
docker pull radixark/miles:latest

docker run --rm \
  --gpus all --ipc=host --shm-size=32g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --network=host \
  -it radixark/miles:latest /bin/bash
```

**AMD**

```bash theme={null}
docker pull rocm/sgl-dev:miles-rocm10-mi35x    # or miles-rocm724-mi35x; miles-rocm10-mi30x on MI300X / MI325X

docker run --rm \
  --device /dev/kfd --device /dev/dri --group-add video --group-add render \
  --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged \
  --shm-size 128G \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --network=host \
  -it rocm/sgl-dev:miles-rocm10-mi35x /bin/bash
```

See [AMD ROCm](/docs/hardware-platforms/amd-gpus) for more details.

The image ships with:

* PyTorch (matching the container's CUDA / ROCm version)
* Megatron-LM, SGLang, FlashAttention-3, DeepGEMM, Apex
* Ray, uv, and Miles installed editable at `/root/miles`

See [Hardware requirements](#hardware-requirements) for per-GPU status.

## Method 2: From source

Clone and install in an existing environment.

```bash theme={null}
git clone https://github.com/radixark/miles.git
cd miles
pip install -r requirements.txt
pip install -e . --no-deps
```

> [!WARNING]
>
> **Patched dependencies.** Miles pins patched versions of SGLang and Megatron-LM. Installing them yourself at
> the wrong commit is the most common source of bug reports — use Docker if you can.

## Method 3: Update an existing container

If you already run a Miles image and want the latest code:

```bash theme={null}
cd /root/miles
git pull --rebase
pip install -e . --no-deps
ray stop && ray start --head --port=6379
```

## Verify

Confirm Miles imports and the GPUs are visible:

```bash theme={null}
python -c "import miles; print('Miles import OK')"
nvidia-smi
```

If either command fails, see [Debugging](/docs/developer/debug).

## Hardware requirements

| Hardware | Status |
| - | - |
| NVIDIA GB300 / GB200 / B300 / B200 | Production |
| NVIDIA H200 / H100 | Production (CI guarded) |
| NVIDIA A100 | Supported — FP8 features disabled |
| AMD MI300X, MI325, MI350X, MI355X | Supported via ROCm |

For multi-node training you also need a high-bandwidth interconnect — InfiniBand, RoCEv2,
or Slingshot — and 200+ GB/s per node. Single-node jobs run fine over NVLink only.

## Next steps

* [Quick Start](/docs/getting-started/quick-start) — run your first training job.
* [Core concepts](/docs/user-guide/concepts) — the mental model behind Miles.
* [Training backends](/docs/user-guide/training-backend) — Megatron vs FSDP.


# Quick Start
Source: https://miles.radixark.com/docs/getting-started/quick-start

Get an RL training job up and running in under an hour.

**What you need**

* A node with 8 GPUs (H100 / H200 / B-series).
* At least 500 GB of free disk.
* Docker with GPU access.

**Pre-flight checks**

```bash theme={null}
# Driver up, all 8 GPUs listed?
nvidia-smi -L
# Docker can hand GPUs to a container?
docker run --rm --gpus all ubuntu nvidia-smi
# Enough free disk where Docker stores data?
df -h $(docker info -f '{{.DockerRootDir}}')
```

**What you will accomplish**

* Launch a GRPO run on Qwen3-4B and watch the reward climb!

Training a different model? The flow is the same — see [Models](/docs/models/index) for
the per-model recipes.

## Step 1: Start the container

On the **host**:

```bash theme={null}
docker pull radixark/miles:latest
docker run --rm \
  --gpus all \
  --ipc=host \
  --shm-size=32g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --network=host \
  -it radixark/miles:latest /bin/bash
```

That drops you into a shell inside the container, with Miles at `/root/miles` and
Megatron-LM at `/root/Megatron-LM`. Refresh the editable install so you run the
latest main:

```bash theme={null}
cd /root/miles && git pull && pip install -e . --no-deps
```

**Everything from here on runs inside the container.**

## Step 2: Download the model and data

Three downloads:

```bash theme={null}
# The model you will train
hf download Qwen/Qwen3-4B --local-dir /root/models/Qwen3-4B
# Training prompts: 17k math problems with checkable answers
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
# Eval benchmark: harder problems, evaluated on but never trained on
hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/datasets/aime-2024
```

## Step 3: Convert to Megatron format

Megatron reads its own sharded checkpoint format, so convert the HuggingFace
weights once:

```bash theme={null}
cd /root/miles
# Load MODEL_ARGS, the Megatron-side description of the architecture
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3-4B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"

# Map the HuggingFace weights into a sharded torch_dist checkpoint
PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Qwen3-4B \
   --save /root/models/Qwen3-4B_torch_dist
```

Keep the original HuggingFace directory too — the rollout engines still read from
it.

## Step 4: Launch training

```bash theme={null}
python scripts/run_qwen3_dense.py --model-name Qwen3-4B
```

That's it — the launcher starts a local Ray cluster and submits the training job!
A few things it already does for you:

* Checkpoints land in `/root/shared_data/checkpoints` every 20 rollouts; `--output-dir`
  moves them.
* The policy is evaluated on AIME-2024 every 20 rollouts.
* If the run dies, relaunch the same command — training resumes from the last
  checkpoint.

Once the engines warm up and the first rollout completes, the log settles into
per-rollout metric lines (values illustrative, keys abridged):

```text theme={null}
perf 0: {'perf/rollout_time': 98.6, ...}
rollout 0: {'rollout/raw_reward': 0.32, 'rollout/log_probs': -0.27, ...}
step 0: {'train/loss': 0.0021, 'train/pg_loss': 0.0018, 'train/grad_norm': 0.62, ...}
perf 0: {'perf/train_wait_time': 101.2, 'perf/actor_train_time': 41.3, ...}
```

`rollout/raw_reward` is the number to watch: the mean reward of the freshly scored
responses, drifting upward as the policy improves. You have a live RL run.

## What's happening

A Miles job combines two engines: [SGLang](https://github.com/sgl-project/sglang)
generates responses from the current policy (the *rollout*), and
[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) updates the policy from those
responses (the *training*). In this recipe both share the same 8 GPUs, taking
turns.

```mermaid theme={null}
flowchart LR
    P[Prompt dataset] --> R[SGLang rollout]
    R --> RM[Reward fn]
    RM --> A[Megatron actor]
    A == weight sync ==> R
    A -. KL .-> Ref[(Reference model)]
```

Every iteration runs the same loop:

1. Sample a batch of prompts and let the SGLang engines generate several candidate
   responses per prompt.
2. Score every response — the reward function checks each final answer against the
   dataset label.
3. Compute the GRPO objective from the scores and step the optimizer.
4. Sync the updated weights back into the SGLang engines, and go again.

The batch-sizing knobs satisfy one identity, and Miles fills in whichever side you
leave unset:

```
rollout_batch_size × n_samples_per_prompt
  = global_batch_size × num_steps_per_rollout
```

In this recipe, 32 prompts × 8 samples = 256 = one optimizer step at global batch
size 256.

### The fine print 🔍

* **The docker flags (Step 1).** `--gpus all` exposes the GPUs, `--ipc=host` and
  `--shm-size=32g` give NCCL and Ray the shared memory they need, and
  `--network=host` lets you reach the Ray dashboard from the host.
* **`MODEL_ARGS` (Step 3).** The Megatron-side description of the architecture
  (layer count, hidden sizes, attention layout). The converter uses it to map the
  HuggingFace weights into a sharded `torch_dist` checkpoint.
* **The launcher (Step 4).** `scripts/run_qwen3_dense.py` builds the flags it passes
  to `train.py` as one group per concern (checkpoint paths, rollout, GRPO, optimizer,
  performance, eval, SGLang) and is meant to be read and edited — it is the canonical
  place to change hyperparameters. `--extra-args` appends flags without touching the
  file, and the same launcher covers Qwen3-32B and the Qwen3.5 / Qwen3.6 dense sizes
  through `--model-name`.
* **Colocation.** The recipe sets `--colocate`: four SGLang engines (2 GPUs each)
  and the Megatron trainer share the same 8 GPUs, alternating between generation
  and training. Sharing GPUs also makes the weight sync local — each rank gathers
  its shards over NCCL and hands them to its engine through IPC, no network
  involved. Disaggregated runs choose a transport with
  `--update-weight-transfer-mode`: `broadcast` (the default, one NCCL broadcast per tensor),
  `broadcast_packed` (one byte-buffer broadcast per bucket; non-colocated Megatron),
  [`p2p`](/docs/advanced/p2p-weight-transfer) (point-to-point RDMA via Mooncake), or
  [`disk-delta`](/docs/advanced/disaggregated-rollout) (versioned deltas through
  shared storage). `p2p` and `disk-delta` are incompatible with `--colocate`.
* **The reward function.** `--rm-type deepscaler` — a rule-based verifier, no
  learned reward model.
* **KL regularization.** The frozen reference model can add a KL term to the loss;
  this recipe sets `--kl-loss-coef 0.00`, so the term is off.
* **The metric lines (Step 4).** The `step` line reports each optimizer step, and
  the two `perf` lines time the generation side (`perf/rollout_time`) and the
  training side (`perf/actor_train_time`) of each iteration.

### Inspecting a run 📊

| Question | Where to look |
| - | - |
| Is the policy learning? | `rollout/raw_reward` in stdout, or wandb |
| Rollout or train bottleneck? | `perf/rollout_time` vs. `perf/actor_train_time` |
| Are GPUs saturated? | `nvidia-smi dmon -s u` |
| SGLang internals? | Ray worker logs under `~/.ray/session_latest/logs/`; raise verbosity with `--sglang-log-level` |
| Ranks crashing? | `~/.ray/session_latest/logs/worker-*.err` |

## Next steps

* [Core concepts](/docs/user-guide/concepts) — the model behind rollout / actor / reference.
* [Launch script](/docs/user-guide/launch-script) — what a launch script does when you run it,
  how it is structured, and the three ways to override a recipe.
* [Training backends](/docs/user-guide/training-backend) — Megatron vs FSDP.
* [Customization](/docs/user-guide/customization) — plug in custom rollout / reward.
* [Models](/docs/models/index) — recipes for Qwen3.5, GLM5.2, DeepSeek V4, Kimi K2.6, and more.

If you hit issues, feel free to open an issue on [GitHub](https://github.com/radixark/miles/issues).


# AMD ROCm
Source: https://miles.radixark.com/docs/hardware-platforms/amd-gpus

Run Miles on AMD MI350X / MI355X and MI300X / MI325X with the ROCm images. Docker is the recommended path.

Miles runs on AMD GPUs through ROCm. The ROCm images ship SGLang, Megatron-LM, and Miles
preinstalled, with `MILES_HARDWARE_PLATFORM=rocm` already set. The recipes and `train.py`
flags are the same as on NVIDIA; what changes is the image, the `docker run` flags, and the
launcher path.

## Images

The images are built daily from `main` by the sgl-project/sglang nightly
workflows and published to Docker Hub under
[`rocm/sgl-dev`](https://hub.docker.com/r/rocm/sgl-dev/tags?name=miles):

| Image | ROCm | GPUs | Notes |
| - | - | - | - |
| `rocm/sgl-dev:miles-rocm10-mi35x` | 10 | MI350X / MI355X | Python 3.12 — the image the nightly tests run on |
| `rocm/sgl-dev:miles-rocm724-mi35x` | 7.2.4 | MI350X / MI355X | Python 3.12 |
| `rocm/sgl-dev:miles-rocm10-mi30x` | 10 | MI300X / MI325X | Python 3.12 — the image the MI300X nightly tests run on |

Each undated tag moves with every build; append `-YYYYMMDD` (e.g.
`miles-rocm10-mi35x-20260916`) to pin one.

To build an image yourself, `docker/Dockerfile.rocm` holds the recipe:

```bash theme={null}
python docker/build.py --variant rocm10-mi35x --image-tag dev    # or rocm724-mi35x, rocm10-mi30x
```

## Start the container

On the **host**:

```bash theme={null}
docker pull rocm/sgl-dev:miles-rocm10-mi35x

docker run --rm \
  --device /dev/kfd --device /dev/dri --group-add video --group-add render \
  --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged \
  --shm-size 128G \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  --network=host \
  -it rocm/sgl-dev:miles-rocm10-mi35x /bin/bash
```

That drops you into a shell inside the container, with Miles at `/root/miles`, Megatron-LM
at `/root/Megatron-LM`, and SGLang at `/sgl-workspace/sglang`.

**Everything from here on runs inside the container.**

## Verify

Confirm Miles imports and the GPUs are visible:

```bash theme={null}
python -c "import miles; print('Miles import OK')"
rocm-smi --showproductname
```

If either command fails, see [Debugging](/docs/developer/debug).

## Launch training

The AMD launchers live under `scripts/amd/` and mirror the CUDA recipes. Download the model
and data and convert the checkpoint as in the [Quick Start](/docs/getting-started/quick-start)
(Steps 2 and 3), then launch:

```bash theme={null}
cd /root/miles
python scripts/amd/run_qwen3_4b.py --hardware MI355X    # or MI350X
```

The other recipes under `scripts/amd/` launch the same way.

## Next steps

* [Quick Start](/docs/getting-started/quick-start) — the same Qwen3-4B run, step by step.
* [Hardware requirements](/docs/getting-started/installation#hardware-requirements) — per-GPU status.
* [Low Precision RL](/docs/advanced/low-precision) — FP8 block-wise on MI350X / MI355X.
* [Docker build](/docs/developer/ci/02-docker-build) — the ROCm Dockerfile, variants, and tags.


# NVIDIA GPUs
Source: https://miles.radixark.com/docs/hardware-platforms/nvidia-gpus

Get started with Miles on NVIDIA GPUs using the installation guide.

See the [Installation Guide](/docs/getting-started/installation) to get started with Miles on NVIDIA GPUs.


# Hardware Platforms
Source: https://miles.radixark.com/docs/hardware-platforms/overview

Platform-specific guides for running Miles on NVIDIA and AMD GPUs.

* [NVIDIA GPUs](/docs/hardware-platforms/nvidia-gpus)
* [AMD GPUs](/docs/hardware-platforms/amd-gpus)


# Welcome
Source: https://miles.radixark.com/docs/index

Miles is an open-source RL framework for large-scale LLM post-training, pairing SGLang rollout with Megatron-LM training at trillion-parameter scale.

Miles is a high-performance, enterprise-ready reinforcement learning framework for
**large-scale model post-training**. It pairs [SGLang](https://github.com/sgl-project/sglang)
for high-throughput rollout with [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) for
scalable training, and ships the precision, stability and observability features an RL run
needs at trillion-parameter scale. A PyTorch FSDP2 backend is available for runs that would
rather train the HuggingFace implementation as-is, though the recipes, the parallelism and
the largest models all live on Megatron-LM. See
[Training Backends](/docs/user-guide/training-backend).

> *"A journey of a thousand miles begins with a single rollout."*

## Core features

### Performance

* **Fully async RL.** Rollout and training workers are decoupled, with configurable on- and
  off-policy schedules, a pipeline tuned for fewer bubbles, and customizable async rollout
  and eval modes. See [Fully Async RL](/docs/user-guide/fully-async).
* **Fast agentic rollout.** Generation runs on [SGLang](https://github.com/sgl-project/sglang)
  behind a router that spreads requests across engines, preserves per-request metadata and
  health-checks the fleet. Tuned for multi-turn agentic workloads.
* **Fast weight updates.** New weights reach the engines in-loop in seconds, even on a
  trillion-parameter model such as Kimi-K2.6, with
  [P2P RDMA](/docs/advanced/p2p-weight-transfer) as the fast path for disaggregated setups.
* **Low-precision training.** [MXFP8 and NVFP4](/docs/advanced/low-precision) training with a
  numerically stable RL recipe that reduces precision-induced divergence. FP8,
  [INT4 QAT](/docs/advanced/int4-qat), BF16 and FP16 are also supported.
* **LoRA and multi-LoRA.** [Low-rank adapters](/docs/advanced/lora) train frontier-scale models
  on a fraction of the GPUs, and the same adapters load straight into SGLang for rollout.

### Correctness and resilience

* **Token-in-token-out (TITO).** Supported for
  [every model and every black-box harness](/docs/user-guide/agentic-rollout), with no
  detokenize and retokenize round-trip between rollout and training.
* **Rollout Routing Replay (R3).** Expert routing recorded during rollout is
  [replayed in the trainer's forward pass](/docs/advanced/miles-router), removing the MoE routing
  mismatch that destabilizes large runs, with compute and communication overlapped to keep
  the cost down.
* **Fault tolerance.** When an SGLang engine dies, Miles
  [recovers it and resumes the run in place](/docs/advanced/fault-tolerance): no restart, no
  pause.

### What Miles runs

* **Day-0 model support.** DeepSeek-V4, Kimi-K3, GLM-5.2, Inkling and Nemotron landed on
  release day. Beyond day 0, nearly every frontier model runs on Miles, including Kimi-K2.6
  and Qwen3.5. See [Supported models](#supported-models).
* **Extensive hardware support.** NVIDIA from H100 through GB300, and AMD MI300X through
  MI355X via ROCm. See [Supported hardware](#supported-hardware).
* **Wide recipe support.** GRPO, GSPO, PPO and REINFORCE++ for RL, plus SFT and
  [on-policy distillation](/docs/advanced/on-policy-distillation).
* **Agentic environments.** Train coding and computer-use agents through connectors for
  Harbor, HUD, NeMo Gym, OpenEnv, Verifiers and more, each plugging into the rollout
  layer that fits it, with task sandboxes on AgentENV, Daytona, E2B or Modal. See
  [Agentic Environments](/docs/user-guide/environments).
* **Diffusion models.** Flow-GRPO, DiffusionNFT and SFT on an sglang-diffusion rollout
  engine and an FSDP2 trainer, in
  [Miles-diffusion](https://github.com/radixark/miles_diffusion).
* **Comprehensive CI.** Unit suites run on every pull request, and tag-triggered end-to-end
  GPU training tests cover the supported model families on both NVIDIA and AMD runners.

## Supported models

Each model name links to its recipe page or launch script. The table is not
exhaustive — it highlights recent releases; many more models run on Miles out
of the box, including older generations of the families below.

| Family | Models |
| - | - |
| **DeepSeek** | [DeepSeek-V4 Pro](/docs/models/deepseek/deepseek-v4-pro) [DeepSeek-V4 Flash](/docs/models/deepseek/deepseek-v4-flash) [DeepSeek-V3.2](/docs/models/deepseek/deepseek-v3-2) [DeepSeek-V3](/docs/models/deepseek/deepseek) |
| **Thinking Machines** | [Inkling](/docs/models/thinkingmachines/inkling) [Inkling-Small](/docs/models/thinkingmachines/inkling-small) |
| **Qwen** | [Qwen3.8-27B](/docs/models/qwen/qwen3-8) [Qwen3.6 MoE](/docs/models/qwen/qwen3-6-moe) [Qwen3.6](/docs/models/qwen/qwen3-6) [Qwen3.5-35B-A3B](/docs/models/qwen/qwen3-5-moe) [Qwen3.5-4B / 9B / 27B](/docs/models/qwen/qwen3-5) |
| **GLM** | [GLM-5.3](/docs/models/glm/glm5-3) [GLM-5.3-Flash](/docs/models/glm/glm5-3-flash) [GLM-5.2](/docs/models/glm/glm5-2) [GLM-5.1](/docs/models/glm/glm5) [GLM-5](/docs/models/glm/glm5) [GLM-4.7-Flash](/docs/models/glm/glm4-7-flash) |
| **Kimi** | [Kimi-K3](/docs/models/kimi/kimi-k3) [Kimi-K2.6](/docs/models/kimi/kimi-k2.5) [Kimi-K2.5](/docs/models/kimi/kimi-k2.5) |
| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/docs/models/nemotron/nemotron-3-ultra) [Nemotron-3-Super-120B-A12B-FP8](/docs/models/nemotron/nemotron-3-super) [Nemotron-3-Nano MoE](/docs/models/nemotron/nemotron-3-nano-moe) [Nemotron-3-Nano](/docs/models/nemotron/nemotron-3-nano) |
| **Gemma** | [Gemma-4 26B-A4B](/docs/models/gemma/gemma-4) [Gemma-4 31B](/docs/models/gemma/gemma-4) |
| **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) |
| **Stable Diffusion** | [SD3.5](/docs/diffusion/models/sd3/sd3) |
| **Qwen-Image** | [Qwen-Image](/docs/diffusion/models/qwen-image/qwen-image) |
| **Wan** | [Wan2.2-T2V-A14B](/docs/diffusion/models/wan/wan2-2) |
| **LTX** | [LTX-2.3](/docs/diffusion/models/ltx/ltx2) |
| **Cosmos** | [Cosmos3-Nano](/docs/diffusion/models/cosmos/cosmos3) |
| **MiniMax** | [MiniMax H3](/docs/diffusion/models/h3/h3) |

See [Models](/docs/models/index) for LLM family guides and [Diffusion](/docs/diffusion) for
diffusion recipes and validation details.

## Supported hardware

* **NVIDIA**: GB300, GB200, B300, B200, H200, H100, A100.
* **AMD**: MI300X, MI325, MI350, MI355X (via ROCm).

See [Installation](/docs/getting-started/installation#hardware-requirements) for per-GPU status
and the container images for each.

## News

* \[2026/08] 🔥 Miles v0.1 is released! Read the blog post here: [Miles v0.1: Production-level Post-training](https://www.lmsys.org/blog/2026-08-18-miles-v0-1).
* \[2026/07] Towards Blackwell-Native 8-bit and 4-bit RL: End-to-End MXFP8 and NVFP4 RL in Miles ([blog](https://www.lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl)).
* \[2026/07] 🔥 SGLang and Miles add day-0 support for Kimi K3 ([blog](https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support)).
* \[2026/07] On-policy distillation lands in Miles ([blog](https://www.lmsys.org/blog/2026-07-18-opd-support-in-miles)).
* \[2026/07] 🔥 SGLang and Miles add day-0 support for Inkling, a frontier multimodal model ([blog](https://www.lmsys.org/blog/2026-07-15-inkling-day0-support)).
* \[2026/07] DeepSeek-V4 Flash RL training comes to AMD Instinct MI355X with Miles ([blog](https://www.lmsys.org/blog/2026-07-10-rocm-miles-dsv4)).
* \[2026/06] SGLang and Miles add day-0 support for NVIDIA Nemotron 3 Ultra ([blog](https://www.lmsys.org/blog/2026-06-04-nvidia-run-nemotron-3-ultra)).
* \[2026/05] No token left behind: token-in-token-out in Miles ([blog](https://www.lmsys.org/blog/2026-05-13-no-token-left-behind)).
* \[2026/04] Updating 1 T parameters in seconds: P2P weight transfer in large-scale distributed RL ([blog](https://www.lmsys.org/blog/2026-04-29-p2p-update)).
* \[2026/04] 🔥 DeepSeek-V4 on day 0: from fast inference to verified RL with SGLang and Miles ([blog](https://www.lmsys.org/blog/2026-04-25-deepseek-v4)).

## Start here

1. **[Installation](/docs/getting-started/installation)** — Docker, bare metal, AMD.
2. **[Quick Start](/docs/getting-started/quick-start)** — a training job up and running in under an hour.
3. **[Core concepts](/docs/user-guide/concepts)** — the four objects in every Miles job.
4. **[Launch script](/docs/user-guide/launch-script)** — what `python scripts/run_*.py` does
   and how to override a recipe.
5. **[Training backends](/docs/user-guide/training-backend)** — Megatron-LM and FSDP: parallelism,
   checkpoints, and hooks.

## Acknowledgment

Miles was forked from [slime](https://github.com/THUDM/slime), and integrates
[SGLang](https://github.com/sgl-project/sglang),
[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) and
[torch\_memory\_saver](https://github.com/fzyzcjy/torch_memory_saver).

Miles is shaped by the teams that build on it and support its development,
from hardware and cloud to model labs, agent infrastructure, and academia:

![Organizations building on, contributing to, and collaborating with Miles](https://mintcdn.com/radixark/-UwsFWJCgDxbWdMf/assets/images/acknowledgment.png?fit=max\&auto=format\&n=-UwsFWJCgDxbWdMf\&q=85\&s=ef555b678c24773c655944b6695bd253)

## Contribute

* GitHub: [github.com/radixark/miles](https://github.com/radixark/miles)
* Slack: [slack.sglang.ai](https://slack.sglang.ai), channel `#miles-rl`
* Contributing: [developer guide](/docs/developer/contributor-guide)


# DeepSeek V3
Source: https://miles.radixark.com/docs/models/deepseek/deepseek

Launch recipe for DeepSeek-V3 (671 B total / 37 B active) via scripts/run_deepseek.py.

## 1. Model Introduction

[DeepSeek-V3](https://huggingface.co/deepseek-ai/DeepSeek-V3) is a large-scale Mixture-of-Experts language model from DeepSeek.

**Key highlights:**

* **Fine-grained MoE architecture**: 671 B total / 37 B active per token, 256 routed experts with top-8 plus 1 shared expert, sigmoid router with bias.
* **MLA attention**: Multi-head Latent Attention with q-LoRA rank 1536, keeping the KV cache compact under long contexts.
* **Long-context capability**: trained at 32 K response length in this recipe; supports extended reasoning and agent-style workflows.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| DeepSeek-V3 | 37 B / 671 B | [deepseek-ai/DeepSeek-V3](https://huggingface.co/deepseek-ai/DeepSeek-V3) |

Passing a `--model-name` that contains `<N>layer` (for example `DeepSeek-V3-0324-5layer`) selects a
layer-pruned checkpoint from the `fzyzcjy` org and the matching `deepseek-v3-<N>layer` Megatron
definition — the single-node smoke-test path.

## 3. Launch

`scripts/run_deepseek.py` splits the pipeline across three commands:

* `prepare` — HF download, FP8 → BF16 cast, HF → Megatron `torch_dist` conversion.
* `train` — the `train.py` submission; each trainer pod first rsyncs both model directories to
  node-local storage.
* `full-train` — `prepare` followed by `train`, which is the whole pipeline in one invocation.

```bash theme={null}
python scripts/run_deepseek.py full-train --num-nodes 16 --num-gpus-per-node 8
```

Single-node smoke test on a pruned checkpoint:

```bash theme={null}
python scripts/run_deepseek.py full-train --model-name DeepSeek-V3-0324-5layer --num-nodes 1
```

Directories default to `--model-dir /root/models` (shared FS), `--data-dir /root/datasets`, and
`--model-local-dir /root/local_data` (node-local copy the training job reads from). `--task`
selects `dapo_aime` (default) or `gsm8k`; `--mode debug_minimal` shortens responses and drops
dynamic sampling and eval.

### Multi-node fan-out

The `torch_dist` conversion and the node-local rsync fan out across every node of the Ray
cluster, so a multi-node run needs the whole cluster joined **before** the launcher starts, and
the launcher must be told not to replace it. Bring up the head, join the workers, then run with
`MILES_SCRIPT_EXTERNAL_RAY=1`:

```bash theme={null}
# on node 0
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats

# on every other node
ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 \
          --node-ip-address ${WORKER_IP} --disable-usage-stats

# back on node 0
MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_deepseek.py full-train \
   --num-nodes 16 --num-gpus-per-node 8
```

With an MPI-style hostfile (each line `ip slot=8`), fan the workers out from node 0:

```bash theme={null}
for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do
  if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then
    continue
  fi
  ssh root@"${WORKER_IP}" \
    "pkill -9 sglang ; ray stop --force ; pkill -9 miles ; \
     ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 \
               --node-ip-address ${WORKER_IP} --disable-usage-stats" &
done
wait
```

Without `MILES_SCRIPT_EXTERNAL_RAY=1` the launcher runs `ray stop --force` and starts its own
single-node head, which is what the single-node invocation above relies on — the conversion then
runs on that one node.

## 4. Checkpoint conversion

`prepare`, and therefore `full-train`, performs the two conversion steps for you; the equivalent
manual commands are below.

The HF checkpoint ships in block-quantized FP8 — first cast it to BF16:

```bash theme={null}
python tools/fp8_cast_bf16.py \
   --input-fp8-hf-path  /root/models/DeepSeek-V3 \
   --output-bf16-hf-path /root/models/DeepSeek-V3-bf16/
```

Then convert BF16 HF → Megatron `torch_dist`. Run on **4 separate nodes** (`NODE_RANK=0..3`);
`MASTER_ADDR` is the IP of node 0:

```bash theme={null}
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py deepseek-v3)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM/ torchrun \
   --nproc-per-node 8 \
   --master-addr ${MASTER_ADDR} --master-port 12345 \
   --nnodes=4 --node-rank ${NODE_RANK} \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --tensor-model-parallel-size 1 \
   --pipeline-model-parallel-size 8 \
   --expert-tensor-parallel-size 1 \
   --expert-model-parallel-size 4 \
   --decoder-first-pipeline-num-layers 7 \
   --decoder-last-pipeline-num-layers 6 \
   --hf-checkpoint /root/models/DeepSeek-V3-bf16/ \
   --save /root/models/DeepSeek-V3_torch_dist/
```

## 5. Recipe Configuration

All values below come straight from `scripts/run_deepseek.py`.

### 5.1 Parallelism

The layout is chosen from `--num-nodes`:

| Nodes | TP | PP | CP | EP | expert-TP | `decoder-last-pipeline-num-layers` |
| - | - | - | - | - | - | - |
| ≤ 2 | 1 | 1 | 4 | 4 | 1 | — |
| ≤ 4 | 4 | 1 | 4 | 4 | 1 | — |
| > 4 | 4 | 4 | 4 | 16 | 1 | 13 (full model only) |

Recomputation is `full` / `uniform` / 1 layer, with `--use-dynamic-batch-size` and
`--max-tokens-per-gpu 2048`.

> [!WARNING]
>
> `--max-tokens-per-gpu 2048` is a deliberately tiny placeholder in the current script (the
> production value is 16384). Raise it before you read any throughput number off this recipe.

DeepSeek-V3 has 61 layers, which doesn't divide evenly into PP=4 —
`--decoder-last-pipeline-num-layers 13` puts the extra layers on the last stage. With
`--use-dynamic-batch-size`, miles packs samples up to `--max-tokens-per-gpu`; under CP=4, a CP
group shares a `CP × max-tokens-per-gpu` budget. miles always trains with data packing and
per-token loss, so dynamic batch size doesn't change the loss.

### 5.2 Algorithm

GRPO with DAPO-style dynamic sampling:

```
--advantage-estimator grpo
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28

--rm-type math
--num-rollout 3000
--rollout-batch-size 128
--n-samples-per-prompt 8
--rollout-max-response-len 32768
--rollout-temperature 1
--num-steps-per-rollout 4
--balance-data

--over-sampling-batch-size 256
--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
```

`--use-kl-loss` is absent, so the reference model is never loaded and the zero KL coefficient
costs nothing. `--over-sampling-batch-size 256` paired with `check_reward_nonzero_std` is the
DAPO-style setup: oversample, then drop prompts whose reward distribution has zero variance. The
two dynamic-sampling flags and the whole eval block are skipped under `--mode debug_minimal`.

### 5.3 Rollout & SGLang

The engine world size follows the node count: 4 GPUs per engine up to 4 nodes, 64 beyond that,
with attention DP 1 and 8 respectively.

```
--rollout-num-gpus-per-engine 64
--sglang-mem-fraction-static 0.7
--sglang-tp-size 64
--sglang-ep-size 64

# dp attention
--sglang-enable-dp-attention
--sglang-dp-size 8
--sglang-moe-dense-tp-size 1
--sglang-enable-dp-lm-head

# enable deepep for sglang
--sglang-moe-a2a-backend deepep
--sglang-deepep-mode low_latency

# make every dp rank has 128 concurrency
--sglang-server-concurrency 1024
--sglang-max-running-requests 2048
--sglang-chunked-prefill-size 16384
--sglang-cuda-graph-max-bs-decode 256
```

`--rollout-num-gpus-per-engine` corresponds to SGLang's `tp_size`. To exploit large-EP inference,
the recipe sets EP64, DP-attention with DP8, and DeepEP `low_latency`.
`--sglang-server-concurrency` is a miles-specific knob to keep the SGLang HTTP server from being
swamped — default 512, raised to 1024 here so each of the 8 DP ranks gets 128 concurrent requests.
`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is exported to match `--sglang-cuda-graph-max-bs-decode`.

### 5.4 Optimizer

```
--optimizer adam
--lr 1e-6
--lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.98
```

CPU offload (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`) is available but commented out in the script. Enabling it puts
the Adam state on host RAM (\~1.4–1.5 TB / 8-GPU node); if a node runs out of host memory, add
more nodes to widen parallelism rather than swapping.

### 5.5 Notable quirks

* **Online FP8 quantization against the HF config**: `--hf-checkpoint` points at the FP8 HF
  directory (also where the tokenizer is read from). miles applies the quantization config from
  the HF checkpoint, so weights are block-wise quantized before being passed to SGLang. The BF16
  directory produced by the cast is used only as the conversion input.
* **`--decoder-last-pipeline-num-layers 13`** is mandatory under PP=4 (61 layers don't divide
  evenly), and is skipped for layer-pruned variants.
* **Node-local reads**: training reads `--hf-checkpoint` and `--ref-load` out of
  `--model-local-dir`, which `train` populates by rsync from the shared `--model-dir`.
* **`--colocate`** runs actor and rollout on the same GPUs
  (`--actor-num-nodes` × `--actor-num-gpus-per-node`).
* **`--attention-backend flash` is deliberately unset** — MLA models need the default backend.

## 6. Pairs Well With

* [PD Disaggregation](/docs/advanced/pd-disaggregation)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)
* [Fault Tolerance](/docs/advanced/fault-tolerance)


# DeepSeek-V3.2
Source: https://miles.radixark.com/docs/models/deepseek/deepseek-v3-2

Launch recipe for DeepSeek-V3.2 (671 B total / 37 B active) — BF16 training, NSA rollout, 8 training nodes and up.

## 1. Model Introduction

[DeepSeek-V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2) is a 37 B-active / 671 B-total Mixture-of-Experts (MoE) model. It keeps the MoE and Multi-head Latent Attention (MLA) shapes of DeepSeek-V3 and adds DeepSeek Sparse Attention (DSA): a separate indexer scores the preceding tokens, and each query attends only to the highest-scoring ones.

**Key highlights:**

* **Fine-grained MoE**: 61 layers (3 dense, then 58 MoE), 256 routed experts with top-8 plus 1 shared expert, a sigmoid router with expert bias, group-limited routing over 8 groups with top-4 groups, and a routed scaling factor of 2.5.
* **MLA with a sparse indexer**: q-LoRA rank 1536, KV-LoRA rank 512, and 128 attention heads split into a 128-dim NoPE part and a 64-dim RoPE part. The indexer runs 64 heads at head dim 128 and keeps `index_topk=2048` keys per query.
* **YaRN RoPE**: rotary base 10000 with scaling factor 40, giving 163,840 max positions.
* **Block-wise FP8 checkpoint**: the published weights are 128×128 FP8 blocks with `ue8m0` scales.

In miles, V3.2 shares its DSA attention implementation with the GLM-5 family — both select it through `--spec miles_plugins.models.glm5.glm5 get_glm5_spec`. Weight import and export run through `DeepseekV32Bridge` (`miles_plugins/mbridge/deepseek_v32.py`), which adds the indexer tensors on top of the V3 bridge. Training is BF16, so the FP8 checkpoint is cast up before conversion.

The raw path also supports Megatron's native DSA through `--dsa-impl megatron --dsa-kernel-backend cudnn`; `--dsa-impl miles` remains the default. See the shared [DSA training implementation selector](/docs/models/glm/glm5#33-dsa-training-implementation) for conversion flags and checkpoint compatibility.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| DeepSeek-V3.2 | 37 B / 671 B, 61 layers | [deepseek-ai/DeepSeek-V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2) |

A 5-layer pruned Megatron config also ships as `deepseek-v32-5layer`. You can select it with `--megatron-model-type deepseek-v32-5layer` when you have a matching pruned checkpoint and want a single-node smoke test.

## 3. Environment Setup

Run everything inside the `radixark/miles:latest` container at `/root/miles`. The whole recipe is driven by one Typer launcher, `scripts/run_deepseek_v32.py`.

### 3.1 Launcher defaults

| Flag | Default | Use |
| - | - | - |
| `--model-org` / `--model-name` | `deepseek-ai` / `DeepSeek-V3.2` | The HF repository to download, and the stem of every derived directory name. |
| `--model-dir` | `/root/models` | Holds the HF checkpoint, the `-bf16` cast, and the Megatron `_torch_dist` directory as siblings. |
| `--model-local-dir` | `/root/models` | Node-local destination each trainer pod rsyncs into before training; worth changing only when `--model-dir` is on shared storage. |
| `--data-dir` | `/root/datasets` | Where dapo-math-17k and aime-2024 are downloaded. |
| `--megatron-path` | `/root/Megatron-LM` | Added to `PYTHONPATH` for both conversion and training. |
| `--output-dir` | `/root/shared_data` | Training checkpoints land under `{output-dir}/{run-id}/checkpoints`. |
| `--hardware` | `B200` | One of `B200`, `B300`, `GB200`, `GB300`, `H100`, `H200`. |

Every option also binds to an env var named `MILES_SCRIPT_<FIELD_NAME_UPPER>` (for example `MILES_SCRIPT_MODEL_DIR`), with precedence CLI flag > env var > built-in default. Run `python scripts/run_deepseek_v32.py train --help` to see each option's env var name.

The launcher does not pass `--colocate`, so training and rollout occupy disjoint GPUs. The flags `--actor-num-nodes` and `--rollout-num-gpus` have no defaults, and every multi-node invocation has to supply them.

### 3.2 Download and convert

The `prepare` subcommand downloads the model and the datasets, casts the FP8 checkpoint to BF16, and converts it to a Megatron `torch_dist` checkpoint:

```bash theme={null}
python scripts/run_deepseek_v32.py prepare --actor-num-nodes 8
```

The stages it runs are the following:

```bash theme={null}
hf download deepseek-ai/DeepSeek-V3.2 --local-dir /root/models/DeepSeek-V3.2
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/datasets/aime-2024

python tools/fp8_cast_bf16.py \
   --input-fp8-hf-path /root/models/DeepSeek-V3.2 \
   --output-bf16-hf-path /root/models/DeepSeek-V3.2-bf16/
```

The `torch_dist` conversion then runs `tools/convert_hf_to_torch_dist.py` under `torchrun` across the Ray cluster. Both the cast and the conversion detect their own output and skip it, so re-running `prepare` after a failure is cheap. If your checkpoint is already BF16, `--from-bf16-ckpt` downloads it straight into the `-bf16` directory and skips the cast.

Because the conversion fans out over Ray with `ray.init(address="auto")`, the cluster has to be up before you call `prepare` on more than one node.

### 3.3 Multi-node Ray

Start Ray yourself and point the launcher at it:

```bash theme={null}
# on the head node
ray start --head --num-gpus 8 --disable-usage-stats
# on every worker
ray start --address=${HEAD_IP}:6379 --num-gpus 8 --disable-usage-stats

export MILES_SCRIPT_EXTERNAL_RAY=1
export RAY_ADDRESS=http://${HEAD_IP}:8265
```

Without `MILES_SCRIPT_EXTERNAL_RAY=1`, the training stage runs `ray stop --force` and starts a fresh local head, tearing down the cluster that the conversion just used. When `RAY_ADDRESS` is unset the launcher submits to `http://127.0.0.1:8265`.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_deepseek_v32.py full-train \
   --actor-num-nodes 8 --rollout-num-gpus 8
```

The `full-train` subcommand chains download → FP8 → BF16 cast → optional rollout quantization → `torch_dist` conversion → training. Staging onto node-local disk is not a separate stage: `train` rsyncs the HF checkpoint and `torch_dist` into `--model-local-dir` on every trainer pod before the job starts.

### 4.2 Individual stages

```bash theme={null}
# download, cast, and convert only
python scripts/run_deepseek_v32.py prepare --actor-num-nodes 8

# re-run just the Megatron conversion
python scripts/run_deepseek_v32.py prepare-megatron-ckpt --actor-num-nodes 8

# train, assuming the stages above already ran
python scripts/run_deepseek_v32.py train --actor-num-nodes 8 --rollout-num-gpus 8
```

### 4.3 Single-node smoke test

The `--use-single-node` flag pins the run to one node, gives half of its GPUs to training and half to rollout (4 + 4 on an 8-GPU node), sets TP and EP to the training GPU count with PP1, converts the checkpoint on that single node, and runs SGLang with 2-GPU engines. Pair it with a pruned checkpoint — the full 671 B model does not fit on 4 GPUs.

## 5. Recipe Configuration

### 5.1 Parallelism

| Stage | TP | PP | CP | EP | expert-TP | Last PP stage |
| - | - | - | - | - | - | - |
| Training, multi-node | 2 | 4 | 1 | 16 | 1 | 13 layers |
| Training, `--use-single-node` | half the node (4) | 1 | 1 | half the node (4) | 1 | — |
| `torch_dist` conversion, multi-node | 4 | 6 | — | 16 | 1 | 13 layers |

61 layers do not divide evenly into PP=4, so `--decoder-last-pipeline-num-layers 13` splits the training stages 16 / 16 / 16 / 13. Megatron also requires the world size to be divisible by `expert-TP × EP × PP` = 64, so `--actor-num-nodes` has to be a multiple of 8 at 8 GPUs per node.

The rest of the performance arguments are fixed by the launcher:

```bash theme={null}
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1
--use-dynamic-batch-size
--max-tokens-per-gpu 32768
--data-pad-size-multiplier 4096
--log-probs-chunk-size 1024
--allgather-cp  # DSA with context parallel uses the sequential allgather-CP layout
```

### 5.2 Algorithm

Using GRPO as an example, you can configure the algorithm with the following flags:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

Rollout reads dapo-math-17k and scores it with `--rm-type deepscaler`: 32 prompts per rollout step, 8 samples per prompt, global batch size 256, and responses capped at 8192 tokens. Evaluation is off until you pass `--enable-eval`, which adds an aime-2024 pass every 20 rollouts at 16 samples per prompt.

The `--enable-mis` flag turns on truncated importance sampling to correct the train/inference mismatch. It writes a custom config with `use_tis: true`, `tis_mode: truncate`, and bounds `[0.5, 2.0]`, then routes the policy loss through `examples.infra_features.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp`. Rejection sampling rides along by default and is disabled with `--no-tis-use-rs`.

### 5.3 Rollout & SGLang

```bash theme={null}
--sglang-mem-fraction-static 0.8
--sglang-attention-backend nsa
--sglang-dsa-decode-backend flashmla_sparse
--sglang-dsa-prefill-backend flashmla_sparse
--sglang-kv-cache-dtype bf16
--sglang-page-size 64  # the NSA KV cache requires 64 on CUDA
--rollout-num-gpus-per-engine 8  # 2 with --use-single-node
--sglang-tp-size 8
--sglang-dp-size 8
--sglang-enable-dp-attention
--sglang-enable-dp-lm-head
--sglang-cuda-graph-max-bs-decode 256
--sglang-moe-runner-backend flashinfer_trtllm_routed  # triton on H100 / H200
```

The launcher exports three env vars into the Ray runtime environment: `SGLANG_NSA_FORCE_MLA=1`, `SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD=0`, and `NVSHMEM_DISABLE_NCCL=1`.

### 5.4 Optimizer

```bash theme={null}
--optimizer adam
--lr 1e-6
--lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.98
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

The Adam state is offloaded to host memory in every configuration the launcher can produce.

### 5.5 Low-precision options

| Flag | Effect |
| - | - |
| `--rollout-fp8` | Runs `tools/convert_hf_to_fp8.py --strategy block --block-size 128 128` and serves the resulting `-FP8` directory to SGLang, while Megatron stays BF16. |
| `--rollout-mxfp8` | Runs `tools/convert_hf_to_mxfp8.py`, serves the `-MXFP8` directory, and switches SGLang to `--sglang-fp8-gemm-backend flashinfer_trtllm`. |
| `--train-mxfp8` | Trains through Transformer Engine with `--fp8-format e4m3 --fp8-recipe mxfp8`. |

Both MXFP8 options require Blackwell and are rejected on H100 / H200, and the two rollout options are mutually exclusive. The `--rollout-mxfp8` path additionally keeps the MLA up-projections in BF16 — `.kv_b_proj.` on the HF side, `linear_kv_up_proj` / `linear_k_up_proj` / `linear_v_up_proj` on the Megatron side — via `--extra-high-precision-layers-hf`, `--extra-high-precision-layers-megatron`, and a generated Transformer Engine precision config.

### 5.6 Notable quirks

* The launcher always passes `--use-fault-tolerance`, and checkpoints every 20 rollouts unless you pass `--no-save`.
* Combining `--train-mxfp8` with `--fp8-param-gather` raises `NotImplementedError`; MXFP8 parameter all-gather is not wired up yet.
* Passing `--train-mxfp8` without `--rollout-mxfp8` points `--hf-checkpoint` at an `-MXFP8` directory that the prepare stage never builds, because only `--rollout-mxfp8` runs the MXFP8 conversion. Pass both flags together.
* The multi-node conversion branch passes PP=6 with `--decoder-last-pipeline-num-layers 13`, which leaves 48 layers for 5 middle stages and fails Megatron's even-split check. Until that default is fixed, convert by calling `tools/convert_hf_to_torch_dist.py` directly with a layout that divides — PP=4 with the same last-stage size matches the training layout and works.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision) — background for the rollout and training quantization flags above.
* [Fault Tolerance](/docs/advanced/fault-tolerance) — enabled by default in this recipe.
* [GLM-5.2](/docs/models/glm/glm5-2) — the other recipe built on the same DSA attention implementation.
* [Agentic Rollout](/docs/user-guide/agentic-rollout) — V3.2 renders through the `deepseekv32` template.


# DeepSeek-V4.1 Flash
Source: https://miles.radixark.com/docs/models/deepseek/deepseek-v4-1-flash

Launch recipe for DeepSeek-V4.1 with radixark/miles:dev — BF16 train / BF16 rollout, colocated, 4-node GB300 (16 GPUs) with optimizer state streamed to NVMe.

## 1. Model Introduction

[DeepSeek-V4.1 Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) is a sparse-attention Mixture-of-Experts model (`model_type: deepseek_v41`): 40 decoder layers, 384 routed experts at top-6 plus one shared expert, fp8 dense weights at a 32-wide ue8m0 block scale with fp4 routed experts.

**Architecture.** Every layer keeps one 512-wide KV latent per position and uses it as both key and value for all 64 query heads, with no separate value projection. Each layer reads two stores that behave oppositely: compressed latents, produced only at source layers and shared forward, collapsing 2 positions into 1 in layers 2–19 and 1-to-1 from layer 20 on; and a 128-position sliding window, recomputed per layer from that layer's own activations, so it can never be shared but also never grows with context. The residual stream is four parallel copies mixed by a per-token doubly stochastic matrix; each sublayer's mixing coefficients are consumed by the next sublayer.

**Sparse retrieval.** KV source layers and index source layers are different lists — four of the former, eight of the latter. The four extra index layers produce no keys at all; they re-score layer 20's keys with their own query, so retrieval decisions are made twice as often as keys are stored. Each retrieving layer picks a top-512 candidate set, and from layer 20 on a coarser candidate-block pass bounds the positions later layers score.

**Engram.** An additive n-gram hash memory at two layers. Token ids are normalized before hashing, so " The", "the" and "THE" cannot fork into separate rows. Its two fp8 tables are the largest single block of weight in the checkpoint.

On the training side the Megatron plugin lives under `miles_plugins/models/deepseek_v4_1/`, the model definition is `scripts/models/deepseek-v4.1.py`, and the launcher is `scripts/run_deepseek_v4_1.py`. The plugin reproduces the model's fp4/fp8 quantization points in the training forward so the trainer scores what the served model computes, and carries the cross-layer state (mixing coefficients, source latents, top-k, candidate mask) inside the inter-layer hidden tensor so pipeline parallelism and full activation recompute work.

## 2. Supported Variants

| Model | Active / Total | Checkpoint |
| - | - | - |
| DeepSeek-V4.1 Flash | \~5 B dense + 6 of 384 experts / \~750 B (544 B experts + 197 B frozen engram) | [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash), converted to a BF16 HF checkpoint on node-local NVMe (see [§3.2](#32-prepare-the-checkpoint)) |

`model_type` is `deepseek_v41`; the launcher selects the recipe with `--model-name DeepSeek-V4.1`.

## 3. Quick Start

The validated configuration is 4 nodes x 4 GB300 (16 GPUs), colocated rollout engines, BF16 training and BF16 rollout, optimizer state streamed to NVMe. Each step below runs inside the image on every node unless stated otherwise.

### 3.1 Pull the image

```bash theme={null}
docker pull radixark/miles:dev
```

The standard `radixark/miles:dev` image (multi-arch; GB300 pulls arm64) carries everything DeepSeek-V4.1 needs: miles at `/root/miles`, and `sglang-miles` at `/sgl-workspace/sglang/python`, which includes the DeepSeek-V4.1 model and runtime support. No separate image or SGLang fork is needed.

Start one container per node with the GPUs, the host network and a node-local NVMe volume mounted at the same path everywhere (`/scratch` below):

```bash theme={null}
docker run -d --name miles --gpus all --network host --ipc host --shm-size 512g \
   -v /scratch:/scratch radixark/miles:dev sleep infinity
```

### 3.2 Prepare the checkpoint

The trainer loads a **BF16 HF checkpoint** directly (no `torch_dist` conversion). The served checkpoint is fp8 dense (32x32 ue8m0 scales) with fp4 experts, and has to be cast once:

* dense fp8 and expert fp4 weights dequantized to bf16;
* the two engram tables kept as fp8 bits plus e8m0 scales (about 197 GB; they are never upcast);
* MTP, DSpark and vision tensors dropped;
* a flattened `config.json` with `model_type: deepseek_v41`.

`tools/fp8_cast_bf16.py` covers the fp8 dense format only; the cast tool that also handles the fp4 experts and the engram layout (`dsv41_cut_cast.py --src <fp8 ckpt> --dst <bf16 ckpt> --config <flattened config.json> --layers 0,...,39`) is kept with the bring-up scripts. The result is about 1.2 TB and must be present on **every training node's local NVMe** at the same path, for example `/scratch/models/DeepSeek-V4.1-bf16`. Reserve a further \~1 TB per node for the streamed optimizer state.

### 3.3 Prepare the data

Once, on the head node:

```bash theme={null}
cd /root/miles
python scripts/run_deepseek_v4_1.py prepare-data --task dapo_aime --data-dir /scratch/datasets
```

This downloads `dapo-math-17k` and `aime-2024`. Copy or re-run on the other nodes if `/scratch` is not shared.

### 3.4 Bring up the Ray cluster

```bash theme={null}
# on node 0
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats
# on every other node
ray start --address=${MASTER_ADDR}:6379 --node-ip-address ${WORKER_IP} --num-gpus 4 --disable-usage-stats
```

`ray status` on the head must list 16 GPUs before launching.

### 3.5 Launch

On the head node:

```bash theme={null}
cd /root/miles
MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_deepseek_v4_1.py train \
   --model-name DeepSeek-V4.1 \
   --hf-checkpoint /scratch/models/DeepSeek-V4.1-bf16 \
   --model-dir /scratch/models --data-dir /scratch/datasets \
   --num-nodes 4 --num-gpus-per-node 4 --hardware GB300 \
   --task dapo_aime --mode debug_minimal \
   --num-rollout 400 --rollout-batch-size 16 --n-samples-per-prompt 8 \
   --rollout-max-response-len 2048 --max-tokens-per-gpu 2048 \
   --load-from-hf --disk-offload --offload-disk-dir /scratch/offload \
   --colocate-memory-peak-device cpu \
   --rollout-gpus-per-engine 8 --sglang-mem-fraction-static 0.7 \
   --pp-size 2 --recompute full \
   --sglang-engram-host-table --sglang-cuda-graph --sglang-radix-cache \
   --no-check-weight-update \
   --extra-env-vars "NCCL_SOCKET_IFNAME=eth0 GLOO_SOCKET_IFNAME=eth0" \
   --extra-args "--stream-optimizer-state-moment-dtype bf16"
```

Replace `eth0` with the cluster interface. Add `--wandb-team`, `--wandb-project` and `--wandb-group` through `--extra-args` to log to Weights & Biases.

Three flags are not optional in this configuration:

| Flag | Why |
| - | - |
| `--sglang-engram-host-table` | Colocated engines release all GPU state when they sleep; the engram tables are constants that the weight update does not carry. Without the shared host table they come back empty and `train/train_rollout_logprob_abs_diff` jumps from \~0.02 to \~10. |
| `--disk-offload --offload-disk-dir` | Full-parameter Adam state for the full model does not fit in GPU or host memory on 16 GPUs; it is streamed to node-local NVMe between steps. |
| `--load-from-hf` | The BF16 HF checkpoint is mapped through the mbridge at startup, so no multi-node `torch_dist` conversion is needed. |

### 3.6 What to expect

| Signal | Expected |
| - | - |
| Step 0 `train/train_rollout_logprob_abs_diff` | \~0.02 nats (0.0246 / 0.0247 in the two validated runs) |
| `train/train_rollout_kl` | 0.0012-0.0017, flat over training |
| `rollout/raw_reward` | \~0.5 at step 0, rising above 0.75 within 80 steps |
| Step time | \~11 min per step (2 min weight sync + rollout, \~8 min train + log-prob) |
| Trainer peak GPU memory | \~218 GB of 277 GB per GPU |

`rollout/rewards` is the group-normalised advantage and is 0 by construction; watch `rollout/raw_reward`.

### 3.7 Launcher path defaults

| Flag | Default | Use |
| - | - | - |
| `--hf-checkpoint` | required | BF16 HF checkpoint; with `--load-from-hf` it is also the reference model (`--ref-load`) |
| `--data-dir` | `/root/datasets` | HF datasets (dapo-math-17k, aime-2024, gsm8k) |
| `--model-dir` | `/root/models` | parent of `DeepSeek-V4.1_torch_dist` when not loading from HF |
| `--save-dir` | `/root/models` | training checkpoints under `{save-dir}/{run-id}/checkpoints/` (`--no-skip-saving`) |
| `--offload-disk-dir` | unset | node-local NVMe directory for the streamed optimizer state and the trainer's weight backups |

Every option also binds to `MILES_SCRIPT_<FIELD_NAME_UPPER>` (CLI flag > env var > default). At TP 8 the launcher adds `--make-vocab-size-divisible-by 32` automatically. On GB300 it exports `NCCL_CUMEM_ENABLE=1`; with it off the expert-parallel all-to-all fails across nodes.

## 4. Recipe Configuration

### 4.1 Megatron parallelism

Tensor, pipeline, context, expert and sequence parallelism are all supported. Full-model layouts validated on 4 x 4 GB300 (16 GPUs, colocated, 8-GPU engines):

| TP | PP | EP | Recompute | Decode CUDA graphs | Step time | Trainer peak GPU (of 276.6 GB) | Notes |
| - | - | - | - | - | - | - | - |
| 4 | 1 | 16 | selective (moe, mlp) | off | 15.7 min | 249.5 GB | 120+ steps validated |
| 4 | 2 | 8 | full (uniform, 1 layer) | on | 11.2 min | 218.1 GB | recommended; identical step-0 parity (abs\_diff 0.0246 vs 0.0247) |

The second layout is the one in [§3.5](#35-launch). Its speedup comes from the rollout side (decode CUDA graphs cut the rollout wait from 493 s to 200 s per step); train time is unchanged because PP 2 halves the layers per rank while full recompute adds a forward.

Tensor, pipeline, context, expert and sequence parallelism can be combined freely; TP/CP/PP/EP combinations were cross-validated against each other and agree to the bf16 floor.

An 8 x 8 H200 bring-up (TP 8, 16-GPU engines) also trains the full model; on SM90 the DeepGEMM paged indexer kernel is unavailable and the engines fall back to the torch indexer for decode, and the engines need data-parallel attention, which is on the H200 branch of the SGLang fork rather than the mainline fork.

`--ep-size` defaults to `actor_nodes * gpus_per_node / pp_size`; `--cp-size > 1` adds `--allgather-cp`; expert TP is always 1.

### 4.2 Algorithm

```bash theme={null}
--advantage-estimator grpo
--eps-clip 0.2
--eps-clip-high 0.28
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--rollout-temperature 0.8
--use-rollout-routing-replay   # R3, on by default (--no-enable-r3 to drop)
--moe-router-freeze-gate --freeze-e-score-correction-bias   # required
```

`--mode debug_minimal` (the default) drops over-sampling and the dynamic-sampling filter; `--mode normal` adds `--over-sampling-batch-size 512` with `check_reward_nonzero_std`. `--task dapo_aime` uses dapo-math-17k prompts with the thinking chat template and 4096-token responses by default; the validated runs cap responses at 2048.

The indexer top-k is **not** replayed into the trainer by default (`--enable-indexer-replay` exists; it pins \~300 GB of host memory per engine rank for the replay capture and measured no parity gain, so the residual 0.02 nats gap is kernel numerics).

### 4.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 8   # 2-node engines on GB300; = GPUs per node otherwise
--sglang-tp-size 8 --sglang-dp-size 1 --sglang-ep-size 8
--sglang-attention-backend dsv4
--sglang-moe-runner-backend auto
--sglang-cuda-graph-max-bs-decode 128   # with --sglang-cuda-graph; else --sglang-disable-cuda-graph
--sglang-max-running-requests 128
--sglang-mem-fraction-static 0.7
--sglang-disable-radix-cache            # unless --sglang-radix-cache
```

Environment set by the launcher: `SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1` (with `--sglang-engram-host-table`), `SGLANG_SKIP_CHECKPOINT_LOAD_CHECK=1`, `NCCL_CUMEM_ENABLE=1`, `SGLANG_DSV4_FP4_EXPERTS=0`, `SGLANG_HEALTH_CHECK_TIMEOUT=900`, `SGLANG_DG_CACHE_DIR_PER_PROCESS=1`, `SGLANG_OPT_FP8_WO_A_GEMM=0`, `SGLANG_OPT_FUSE_WQA_WKV=0`, `SGLANG_DISABLE_MULTIMEM_AG=1`, `TORCHINDUCTOR_COMPILE_THREADS=1`, `CUDA_DEVICE_MAX_CONNECTIONS=1`. With `--train-deterministic` (default) the trainer also gets `--deterministic-mode`, `NCCL_ALGO=Ring`, `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`, `CUBLAS_WORKSPACE_CONFIG=:4096:8`.

**Engram host table.** In a colocated run the engines release all GPU memory during `sleep`; parameters come back through the weight update, but the engram tables are constants that the update does not carry. `--sglang-engram-host-table` keeps one copy of each table in pinned host memory shared by the tensor-parallel ranks of an engine, so a wake-up restores them. On hosts where the shared layout is not auto-selected (H200), set `SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=shared` as well. The weight-update equality check skips `engram_hasher.` and `engram.embed.` for the same reason (`--check-weight-update` re-enables the check; the validated runs pass `--no-check-weight-update`).

**SGLang version.** The rollout side uses the `sglang-miles` build that `radixark/miles:dev` carries at `/sgl-workspace/sglang/python`; it validates the V4.1 feature set from `server_args` at startup.

### 4.4 Optimizer and memory

```bash theme={null}
--optimizer adam
--lr 1e-6 --lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9 --adam-beta2 0.98
--accumulate-allreduce-grads-in-fp32   # --grad-reduce-bf16 switches to --grad-reduce-in-bf16
--attention-softmax-in-fp32
--micro-batch-size 1 --max-tokens-per-gpu 2048
--train-memory-margin-bytes 3221225472
```

Full-parameter training of the full model on 16 GPUs only fits because the optimizer state leaves the GPU. `--disk-offload` adds `--stream-optimizer-state-to-disk --offload-train-target disk --offload-train-disk-dir <dir>`; `--stream-optimizer-state-moment-dtype bf16` halves the streamed moments. Point `--offload-disk-dir` at node-local NVMe with \~1 TB free per node. `--colocate-memory-peak-device cpu` places the colocation peak on the host; on GB300 devboxes the container memory limit (about 626 GB) is below the GPU memory of the node, so keep host residency under \~560 GB per node. `--optimizer-offload` (CPU Adam) is the alternative when host memory allows.

If the run OOMs: lower `--max-tokens-per-gpu`, then add `--grad-reduce-bf16`; `--recompute full` with `--pp-size 2` is already part of the recommended launch.

## 5. Results

DAPO on dapo-math-17k, 16 GB300 GPUs, 2K response cap, 16 prompts x 8 samples per step, first layout of [§4.1](#41-megatron-parallelism). The reward rises from 0.51 to 0.78 (5-step means) over 80 steps while the trainer-vs-rollout policy gap stays flat: per-token KL 0.0012-0.0017 and mean |delta log p| 0.017-0.025 nats. The run completed 120+ steps without a failure; the recommended layout reproduces steps 0-1 within noise (abs\_diff 0.0246 / 0.0220, reward 0.52 / 0.59) and reached 0.77 by step 47.

![Raw reward and trainer-vs-rollout policy mismatch over 80 DAPO steps](https://mintcdn.com/radixark/FalnLh-BcYsayjs5/assets/images/dsv41-rl-training.png?fit=max\&auto=format\&n=FalnLh-BcYsayjs5\&q=85\&s=94a38765284d3c250df309494979c627)

| Metric | 80-step run |
| - | - |
| `rollout/raw_reward` | 0.51 -> 0.78 (5-step means) |
| `train/train_rollout_logprob_abs_diff` | 0.017-0.025 nats, no drift |
| `train/train_rollout_kl` | 0.0012-0.0017 |
| `train/grad_norm` at step 0 | 0.13-0.18 |

## 6. Pairs Well With

* [DeepSeek-V4 Flash](/docs/models/deepseek/deepseek-v4-flash) - the parent architecture and the `torch_dist` conversion flow.
* [Architecture Support](/docs/advanced/architecture-support) - the plugin lives under `miles_plugins/models/deepseek_v4_1/` (`deepseek_v4_1.py`, `engram.py`, `ops/{compressor,indexer,kvnorm}.py`); its RoPE, fake-quant and hyper-connection kernels are in `miles/kernels/{position,quant,hyper_connection}`.
* [Low Precision RL](/docs/advanced/low-precision) - the fake-quantization points the plugin reproduces are the same ones an fp8/fp4 rollout would expose.


# DeepSeek-V4 Flash
Source: https://miles.radixark.com/docs/models/deepseek/deepseek-v4-flash

Launch recipe for DeepSeek-V4-Flash (284 B) — FP8 rollout / BF16 train, 8-node H200 (64 GPUs).

DeepSeek V4 training tracking issue: [`radixark/miles#1046`](https://github.com/radixark/miles/issues/1046).

## 1. Model Introduction

[DeepSeek-V4-Flash](https://huggingface.co/sgl-project/DeepSeek-V4-Flash-FP8) is a 13 B-active / 284 B-total MoE model with a substantially different attention stack from V3/R1. It ships in the `radixark/miles:latest` image. The larger [DeepSeek-V4-Pro](/docs/models/deepseek/deepseek-v4-pro) shares the same V4 architecture family at Pro scale.

**Key highlights:**

* **Hybrid Attention (CSA + HCA)**: combines **Compressed Sparse Attention** (light compression) and **Heavily Compressed Attention** (heavy compression) layers — DeepSeek's official V4 name (see [HF model card §Introduction](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash#introduction)). Implementation uses low-rank Q (`q_lora_rank=1024`), single-head latent KV (`head_dim=512`), grouped output projection (8 groups, LoRA rank 1024). A learned topk **indexer** (`index_topk=512`, 64 heads × 128 dim) picks 512 KV per query at runtime, inheriting V3.2's DSA-style design.
* **KV compressors**: 44-element compression schedule `compress_ratios = [0, 0, 4, 128, 4, 128, …, 4, 0]` — first / last few layers are uncompressed (ratio 0), middle layers alternate 4× (CSA) and 128× (HCA). The compressor RoPE has its own base (`compress_rope_theta=160000`), separate from the main attention RoPE.
* **Hyper-connection (HC) routing**: each layer expands hidden state into `hc_mult=4` parallel streams and recombines via sinkhorn-normalized mixing. Pipeline-parallel buffers are 4-D `[s, b, hc_mult, d]` instead of 3-D.
* **YaRN RoPE on main attention**: `rope_theta=10000`, YaRN `factor=16`, `original_max_position_embeddings=65536` → effective context length **1,048,576 tokens (1 M)**. Per-head learnable attention sinks (one scalar per head, added to softmax denominator).
* **FP8 weights with simulated FP8 QAT** on indexer and compressor activations.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| DeepSeek-V4-Flash | 13 B / 284 B | [sgl-project/DeepSeek-V4-Flash-FP8](https://huggingface.co/sgl-project/DeepSeek-V4-Flash-FP8) |
| DeepSeek-V4-Flash-0731 | 13 B / 284 B | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) |

`DeepSeek-V4-Flash-0731` is the official release with MXFP4 routed experts. Its architecture is identical
to DeepSeek-V4-Flash; the launcher's `prepare-fp8` stage (`tools/convert_mxfp4_to_fp8.py`) casts the
experts losslessly to blockwise FP8, after which the pipeline is identical to `DeepSeek-V4-Flash-FP8`.

## 3. Quick start

### 3.1 One-line launch

One command runs the full pipeline — dataset download, FP8 → BF16 cast, distributed `torch_dist` conversion, and the training loop:

```bash theme={null}
# Pull the image:
docker pull radixark/miles:latest

# 8-node Flash run (colocated), inside the container
cd /root/miles
python scripts/run_deepseek_v4.py full-train \
   --model-name DeepSeek-V4-Flash-FP8 \
   --num-nodes 8 --num-gpus-per-node 8

# 16-node disaggregated run: 8 actor (training) nodes + 8 dedicated rollout nodes
python scripts/run_deepseek_v4.py full-train \
   --model-name DeepSeek-V4-Flash-FP8 \
   --num-nodes 16 --num-gpus-per-node 8 --rollout-num-nodes 8
```

```bash theme={null}
# Official MXFP4 release (0731) — adds an MXFP4->FP8 cast before the same pipeline.
# 8-node GB300 run (32 GPUs):
python scripts/run_deepseek_v4.py full-train \
   --model-name DeepSeek-V4-Flash-0731 \
   --num-nodes 8 --num-gpus-per-node 4
```

The `full-train` subcommand chains `prepare-download → [prepare-fp8 →] prepare-single → prepare-spmd → train` (`prepare-fp8` only for the MXFP4 `0731` variant). Each stage has a sentinel-based skip so you can re-run safely after the first invocation.

### 3.2 Launcher path defaults

The Python launcher (`scripts/run_deepseek_v4.py`) takes its path arguments from CLI flags. The defaults are:

| Flag | Default | Use |
| - | - | - |
| `--data-dir` | `/root/datasets` | HF datasets (e.g. dapo-math-17k, …) |
| `--model-dir` | `/root/models` | parent directory holding the HF checkpoint and Megatron `_torch_dist` artifacts as separate sibling sub-directories |
| `--model-local-dir` | unset → same as `--model-dir` | local NVMe path on each node; when it differs from `--model-dir`, every trainer rsyncs the HF checkpoint and `_torch_dist` here before training starts, so the trainer reads from local disk instead of shared storage (only worth setting when `--model-dir` is on shared/remote storage) |
| `--save-dir` | `/root/models` | training checkpoints under `{save-dir}/{run-id}/checkpoints/` |

You can override these via the CLI flags above or equivalently via env vars — every launcher option binds to `MILES_SCRIPT_<FIELD_NAME_UPPER>` (e.g. `MILES_SCRIPT_MODEL_DIR`), with precedence CLI flag > env var > built-in default; run `train --help` to see each option's `[env var: …]` name.

### 3.3 Colocated vs. disaggregated rollout

By default the launcher runs **colocated**: training and SGLang rollout share all `--num-nodes × --num-gpus-per-node` GPUs. Pass `--rollout-num-nodes N` (`0 < N < --num-nodes`) to run **disaggregated**: `N` nodes serve rollout, the rest train. The verified parallelism recipes are keyed on the **training** nodes, so the 8-node Flash recipe in disaggregated form is `--num-nodes 16 --rollout-num-nodes 8` (8 train + 8 rollout — the validated layout).

## 4. Script breakdown

In this section, we explain what `full-train` does under the hood, and how to drive each stage manually if you need to debug or run outside the one-line launcher.

### 4.1 Download model + datasets

```bash theme={null}
# inside the radixark/miles:latest container
hf download sgl-project/DeepSeek-V4-Flash-FP8 --local-dir /root/models/DeepSeek-V4-Flash-FP8
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/datasets/aime-2024
```

The Python launcher's `prepare-download` subcommand does the dataset fetch automatically; pass `--hf-checkpoint <path>` to skip the model download when the FP8 weights are already on a shared filesystem.

### 4.2 HF → Megatron `torch_dist` conversion

The conversion happens in two stages — a single-rank FP8 → BF16 cast, followed by a distributed `torch_dist` shard:

```bash theme={null}
cd /root/miles
python tools/fp8_cast_bf16.py \
   --input-fp8-hf-path /root/models/DeepSeek-V4-Flash-FP8 \
   --output-bf16-hf-path /root/models/DeepSeek-V4-Flash-FP8-bf16/

MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py deepseek-v4-flash)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun \
   --nproc-per-node 4 --nnodes 8 \
   --master-addr ${MASTER_ADDR} --master-port 12345 \
   --node-rank ${NODE_RANK} \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --tensor-model-parallel-size 1 \
   --pipeline-model-parallel-size 8 \
   --expert-model-parallel-size 4 \
   --decoder-first-pipeline-num-layers 7 \
   --decoder-last-pipeline-num-layers 6 \
   --hf-checkpoint /root/models/DeepSeek-V4-Flash-FP8-bf16/ \
   --save /root/models/DeepSeek-V4-Flash-FP8_torch_dist/
```

The Python launcher's `prepare-spmd` subcommand drives the same conversion.

### 4.3 Multi-node fan-out

The Python launcher manages Ray internally — start each pod with the `radixark/miles:latest` image and a working shared filesystem mounted at the same path on every node, then on the head node:

```bash theme={null}
ray start --head --num-gpus 8 --disable-usage-stats
# … then on each worker:
ray start --address=${HEAD_IP}:6379 --num-gpus 8 --disable-usage-stats
```

Alternatively, you can set `MILES_SCRIPT_EXTERNAL_RAY=1` and `RAY_ADDRESS=…` to point the launcher at an existing Ray cluster (for example, one that an orchestration layer has already brought up). When `RAY_ADDRESS` is unset, the launcher boots a local Ray head.

### 4.4 Notable quirks

* **`model_type` rewrite for pruned checkpoints.** The older layer-pruning flow wrote `"model_type": "deepseek_ref"` into the local `config.json`. The launcher rewrites it back to `deepseek_v4` in place (`scripts/run_deepseek_v4.py`), so HF's `AutoConfig.from_pretrained` resolves the checkpoint without a patched `transformers`.

## 5. Example Recipe Configuration

### 5.1 Megatron Parallelism

These are the validated layouts shipped with the launcher; All parallelisms are supported, you can supply any other TP / EP / PP / CP combination that fits your compute.

| Hardware | Nodes × GPUs | TP | PP | CP | EP | expert-TP | Pipeline layout |
| - | - | - | - | - | - | - | - |
| H200 | 8 × 8 = 64 | 8 | 8 | 1 | 8 | 1 | first 4 / last 3 layers |
| GB300 | 8 × 4 = 32 | 8 | 4 | 1 | 8 | 1 | first 11 / last 10 layers |
| GB300 | 8 × 4 = 32 | 2 | 8 | 2 | 4 | 1 | first 4 / last 3 layers |

The Nodes × GPUs column counts **actor (training) nodes** — in disaggregated mode (see [§3.3](#33-colocated-vs-disaggregated-rollout)) rollout nodes come on top of these.

### 5.2 Algorithm

Using GRPO as an example, you can configure the algorithm with the following flags:

```bash theme={null}
--advantage-estimator grpo
--eps-clip 0.2
--eps-clip-high 0.28
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
```

The flags `--moe-router-freeze-gate` and `--freeze-e-score-correction-bias` are required and asserted on the `mcore` side — bias-update during RL is forbidden.

### 5.3 Rollout & SGLang

```bash theme={null}
SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 8
   --sglang-tp-size 8
   --sglang-dp-size 8
   --sglang-ep-size 8
   --sglang-enable-dp-attention
   --sglang-attention-backend compressed  # V4 sparse-MLA backend
   --sglang-page-size 256
   --sglang-max-running-requests 64
   --sglang-chunked-prefill-size 8192
   --sglang-mem-fraction-static 0.5  # leave headroom for Megatron during wake_up
   --use-rollout-routing-replay  # MoE routing replay (R3)
)
```

The launcher sets the required env vars for you: `SGLANG_SKIP_CHECKPOINT_LOAD_CHECK=1`, `SGLANG_DSV4_FP4_EXPERTS=0`, `SGLANG_HEALTH_CHECK_TIMEOUT=120`, `SGLANG_DG_CACHE_DIR_PER_PROCESS=1`, and `SGLANG_OPT_FP8_WO_A_GEMM=0`. Because `--train-deterministic` defaults to on, a stock run also gets `--deterministic-mode` plus `NCCL_ALGO=Ring`, `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` and `CUBLAS_WORKSPACE_CONFIG=:4096:8`; pass `--no-train-deterministic` to drop those four.

On the Megatron side, V4 needs `--qkv-format bshd` with CP-aware data slicing. The DSA indexer additionally supports replay via `--use-rollout-indexer-replay` (off by default).

### 5.4 Optimizer

```bash theme={null}
--optimizer adam
--lr 1e-6 --lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9 --adam-beta2 0.98
--accumulate-allreduce-grads-in-fp32
--attention-softmax-in-fp32
--clip-grad 1.0  # Megatron default; not overridden by the launcher
```

The `--low-memory-resume` flag (off by default) puts optimizer states on CPU during ckpt resume to avoid OOM on the very first iteration.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)
* [Architecture Support](/docs/advanced/architecture-support) — the V4 plugin lives under `miles_plugins/models/deepseek_v4/`.


# DeepSeek-V4 Pro
Source: https://miles.radixark.com/docs/models/deepseek/deepseek-v4-pro

Launch recipe for DeepSeek-V4-Pro (1.6 T) — V4-family architecture at Pro scale.

DeepSeek V4 training tracking issue: [`radixark/miles#1046`](https://github.com/radixark/miles/issues/1046).

## 1. Model Introduction

[DeepSeek-V4-Pro](https://huggingface.co/sgl-project/DeepSeek-V4-Pro-FP8) is a 49 B-active / 1.6 T-total MoE that scales up the same sparse-MLA + DSA-indexer + KV-compressor + hyper-connection stack as [V4-Flash](/docs/models/deepseek/deepseek-v4-flash). The architecture family is identical; the deltas are size and a handful of tuned knobs (indexer top-k, output-projection groups, compression schedule). The miles + Megatron-Core integration ships in the same image as Flash and is selected with `--model-name DeepSeek-V4-Pro-FP8`.

**Key highlights** (deltas vs [V4-Flash](/docs/models/deepseek/deepseek-v4-flash#1-model-introduction)):

* **Scaled-up V4 architecture**: 61 layers (vs 43), hidden-size 7168 (vs 4096), 128 attention heads (vs 64), `ffn_hidden_size=3072` and `moe_ffn_hidden_size=3072` (vs 2048). All layers are MoE (same `--moe-layer-freq` pattern). `q_lora_rank=1536` (vs 1024); latent KV (`kv_lora_rank=512`, `qk_head_dim=512`, `v_head_dim=512`) is unchanged across V4.
* **Hybrid Attention with wider indexer and output projection**: `index_topk=1024` (vs Flash's 512) — Pro keeps 64 indexer heads × 128 dim but picks twice as many KV per query. Grouped output projection uses `o_groups=16` (vs 8), keeping `o_lora_rank=1024`.
* **KV compressors start heavily compressed**: 60-element schedule `[128, 128, 4, 128, 4, 128, …, 4, 0]` — Pro skips Flash's two leading uncompressed layers and starts at ratio-128 (HCA) from layer 0. Middle layers still alternate 4× (CSA) and 128× (HCA); only the final layer is uncompressed. Compressor RoPE base (`compress_rope_theta=160000`) is shared with Flash.
* **MoE topology**: 384 routed experts + 1 shared (vs Flash's 256 + 1), top-6. `--moe-router-topk-scaling-factor 2.5` (vs Flash 1.5) compensates for the larger expert pool. The first 3 layers (`num_hash_layers=3`) remain dense-routed via hash buckets.
* **Identical YaRN RoPE and context**: `rope_theta=10000`, YaRN `factor=16`, `original_max_position_embeddings=65536` → effective context length **1,048,576 tokens (1 M)**, same as Flash.
* **Hyper-connection (HC) routing**: `hc_mult=4` parallel streams with sinkhorn-normalized mixing, same as Flash (PP buffers stay 4-D).
* **FP8 weights with simulated FP8 QAT** on indexer and compressor activations; default training is BF16 on the cast checkpoint and default rollout is FP8 in SGLang with `--sglang-attention-backend compressed`.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| DeepSeek-V4-Pro-FP8 | 49 B / 1.6 T | [sgl-project/DeepSeek-V4-Pro-FP8](https://huggingface.co/sgl-project/DeepSeek-V4-Pro-FP8) |

## 3. Quick start

### 3.1 One-line launch

```bash theme={null}
# Pull the image:
docker pull radixark/miles:latest

# Production Pro run, inside the container
cd /root/miles
python scripts/run_deepseek_v4.py full-train \
   --model-name DeepSeek-V4-Pro-FP8 \
   --num-nodes 32 --num-gpus-per-node 8
```

The `full-train` subcommand chains `prepare-download → prepare-single → prepare-spmd → train`. Each stage has a sentinel-based skip so you can re-run safely after the first invocation.

### 3.2 Launcher path defaults

| Flag | Default | Use |
| - | - | - |
| `--data-dir` | `/root/datasets` | HF datasets (e.g. dapo-math-17k, …) |
| `--model-dir` | `/root/models` | parent directory holding the HF checkpoint and Megatron `_torch_dist` artifacts |
| `--model-local-dir` | unset → same as `--model-dir` | local NVMe path on each node; when it differs from `--model-dir`, every trainer rsyncs the HF checkpoint and `_torch_dist` here before training starts (set it when `--model-dir` is on shared/remote storage) |
| `--save-dir` | `/root/models` | training checkpoints under `{save-dir}/{run-id}/checkpoints/` |

Pro uses the same launcher as V4-Flash, so every option above can also be preconfigured via `MILES_SCRIPT_<FIELD_NAME_UPPER>` env vars (precedence: CLI flag > env var > built-in default) — see [V4-Flash §3.2](/docs/models/deepseek/deepseek-v4-flash#32-launcher-path-defaults) for details.

## 4. Script breakdown

The under-the-hood stages are essentially identical to V4-Flash — see the [V4-Flash Script breakdown](/docs/models/deepseek/deepseek-v4-flash#4-script-breakdown) and substitute the Pro model name and path defaults shown above.

## 5. Example Recipe Configuration

### 5.1 Megatron Parallelism

These are the validated layouts shipped with the launcher; All parallelisms are supported, you can supply any other TP / EP / PP / CP combination that fits your compute.

| Hardware | Nodes × GPUs | TP | PP | CP | EP | expert-TP | Pipeline layout |
| - | - | - | - | - | - | - | - |
| H200 | 32 × 8 = 256 | 8 | 8 | 1 | 32 | 1 | first 7 / last 6 layers |

### 5.2 Algorithm

Same as Flash — see [V4-Flash §5.2 Algorithm](/docs/models/deepseek/deepseek-v4-flash#52-algorithm).

### 5.3 Rollout & SGLang

```bash theme={null}
SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 32
   --sglang-tp-size 32
   --sglang-dp-size 32
   --sglang-ep-size 32
   --sglang-enable-dp-attention
   --sglang-attention-backend compressed       # V4 sparse-MLA backend
   --sglang-page-size 256
   --sglang-max-running-requests 64
   --sglang-chunked-prefill-size 8192
   --sglang-server-concurrency 1024
   --sglang-moe-a2a-backend deepep             # DeepEP normal-mode dispatch
   --sglang-cuda-graph-max-bs-decode 8                # see hang caveat below
   --sglang-mem-fraction-static 0.7            # Pro needs a larger dynamic buffer
)
```

Required env vars (the launcher sets these for you): `SGLANG_SKIP_CHECKPOINT_LOAD_CHECK=1`, `SGLANG_DSV4_FP4_EXPERTS=0`, `SGLANG_HEALTH_CHECK_TIMEOUT=120`, `SGLANG_DG_CACHE_DIR_PER_PROCESS=1`, `SGLANG_OPT_FP8_WO_A_GEMM=0`, and the Pro-only pair `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256`, `SGLANG_JIT_DEEPGEMM_PRECOMPILE=0`. Pro shares the V4 launcher, so `--train-deterministic` is on by default here too and brings the same `NCCL_ALGO` / `NVTE_ALLOW_NONDETERMINISTIC_ALGO` / `CUBLAS_WORKSPACE_CONFIG` trio.

Megatron side: `--qkv-format bshd` (V4 needs `bshd` with CP-aware data slicing). The DSA indexer additionally supports replay via `--use-rollout-indexer-replay` (off by default).

> [!WARNING]
>
> 1. **Engine size ≥ 32 GPUs.** Pro needs a single SGLang engine spanning at least 32 GPUs — the launcher hard-codes `--rollout-num-gpus-per-engine 32`. Smaller engines do not leave enough memory after weights, KV cache, indexer state, and DeepEP buffers, and rollout will OOM under load.
> 2. **EP is mandatory; pure TP will not shard the model.** 384 routed experts × `moe_ffn_hidden_size=3072` cannot be partitioned by tensor parallelism alone — the model must use expert parallelism (`--sglang-ep-size 32`) to spread the expert MLPs across ranks. `--sglang-tp-size 32` only covers the attention / embedding paths.
> 3. **DeepEP normal-mode + CUDA graphs can hang at large batch sizes.** When `--sglang-moe-a2a-backend deepep` is on, an overly large `--sglang-cuda-graph-max-bs-decode` makes SGLang hang during graph capture or replay. The launcher pins it to `8` for Pro — raise it only after verifying the engine doesn't deadlock at your target batch.

### 5.4 Optimizer

```bash theme={null}
--optimizer adam
--lr 1e-6 --lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9 --adam-beta2 0.98
--accumulate-allreduce-grads-in-fp32
--attention-softmax-in-fp32
--clip-grad 1.0                       # Megatron default; not overridden by the launcher

# Pro-only — forced on by the launcher (optimizer_offload=True)
--optimizer-cpu-offload
--use-precision-aware-optimizer
--overlap-cpu-optimizer-d2h-h2d
```

Pro selects `--model-name DeepSeek-V4-Pro-FP8`, which flips `optimizer_offload=True` in the launcher (`scripts/run_deepseek_v4.py`) and appends the three CPU-offload flags above. Adam states live on host RAM and are D2H/H2D-overlapped with the backward pass, freeing GPU memory for the 1.6 T weight + KV footprint. The `--low-memory-resume` flag (off by default) additionally puts optimizer states on CPU during ckpt resume to avoid OOM on the very first iteration.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)
* [Architecture Support](/docs/advanced/architecture-support)
* [DeepSeek V4 Flash](/docs/models/deepseek/deepseek-v4-flash) — sibling recipe; shares the V4-family architecture.


# DeepSeek
Source: https://miles.radixark.com/docs/models/deepseek/index

Miles recipes for the DeepSeek family — V4.1, V4 Flash, V4 Pro, and V3.2.

Miles ships recipes for the DeepSeek family across three generations. **DeepSeek-V4.1** extends the V4 stack with cross-layer shared attention sources, two-level candidate selection and an engram n-gram memory, and trains from a BF16 checkpoint with the optimizer state streamed to NVMe. **DeepSeek-V4** pairs sparse multi-head latent attention with a learned indexer, KV compressors, and hyper-connection routing. **DeepSeek-V3.2** keeps the V3 MoE and MLA shapes and adds DeepSeek Sparse Attention (DSA), the same attention implementation the GLM-5 recipes use. **DeepSeek-V3** itself remains available through `scripts/run_deepseek.py`.

## Variants

| Model | Active / Total | HF ID | Recipe |
| - | - | - | - |
| DeepSeek-V4.1 Flash | 6 of 384 experts / \~750 B | `deepseek-ai/DeepSeek-V4.1-Flash`; image `radixark/miles:dev` | [deepseek-v4-1-flash](/docs/models/deepseek/deepseek-v4-1-flash) |
| DeepSeek-V4-Pro | 49 B / 1.6 T | TBA | [deepseek-v4-pro](/docs/models/deepseek/deepseek-v4-pro) |
| DeepSeek-V4-Flash | 13 B / 284 B | `sgl-project/DeepSeek-V4-Flash-FP8` | [deepseek-v4-flash](/docs/models/deepseek/deepseek-v4-flash) |
| DeepSeek-V3.2 | 37 B / 671 B | `deepseek-ai/DeepSeek-V3.2` | [deepseek-v3-2](/docs/models/deepseek/deepseek-v3-2) |
| DeepSeek-V3 | 37 B / 671 B | `deepseek-ai/DeepSeek-V3` | [deepseek](/docs/models/deepseek/deepseek) |

A validated DeepSeek-V4-Pro recipe is not yet available — see [`radixark/miles#1046`](https://github.com/radixark/miles/issues/1046) for tracking.

## Fastest path to train

DeepSeek-V4-Flash needs 8 nodes of 8× H200 and the `radixark/miles:latest` image:

```bash theme={null}
cd /root/miles
python scripts/run_deepseek_v4.py full-train \
   --model-name DeepSeek-V4-Flash-FP8 \
   --num-nodes 8 --num-gpus-per-node 8
```

DeepSeek-V3.2 needs 8 training nodes of 8 GPUs plus separate rollout GPUs:

```bash theme={null}
cd /root/miles
python scripts/run_deepseek_v32.py full-train \
   --actor-num-nodes 8 --rollout-num-gpus 8
```

See the [DeepSeek-V4 Flash](/docs/models/deepseek/deepseek-v4-flash) page for the V4 architecture summary, parallelism layouts, and known workarounds. See the [DeepSeek-V3.2](/docs/models/deepseek/deepseek-v3-2) page for the V3.2 flow — FP8 → BF16 conversion, the TP2 / PP4 / EP16 training layout, the NSA rollout settings, and the FP8 / MXFP8 options; see the [DeepSeek V3](/docs/models/deepseek/deepseek) page for the V3 recipe.

## Pairs well with

* [Low Precision RL](/docs/advanced/low-precision) — both generations ship FP8 checkpoints and optional low-precision rollout.
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer) — amortize weight sync across ranks.
* [Fault Tolerance](/docs/advanced/fault-tolerance) — node failures are inevitable at 8-node scale and above.


# Gemma-4
Source: https://miles.radixark.com/docs/models/gemma/gemma-4

Launch recipes for Gemma-4 26B-A4B-it (MoE) and 31B-it (dense) via the HF to Megatron bridge.

## 1. Model Introduction

[Gemma-4](https://huggingface.co/google) is Google's multimodal model line.
miles trains both released instruction-tuned sizes as language models, on the
base VLM checkpoint directly.

Both go through the HF to Megatron bridge (`--megatron-to-hf-mode bridge`), and
on the rollout side sglang runs `Gemma4ForConditionalGeneration`, which loads
Gemma-4's hybrid `head_dim` weights correctly. There is no offline `torch_dist`
conversion and no LLM-view rewrite of the checkpoint.

**Key highlights:**

* **Two shapes, one recipe family**: 26B-A4B is MoE (128 experts, top-8), 31B is
  dense. They differ mainly in expert parallelism and the token budget.
* **Bridge-mode load** straight from the VLM checkpoint.
* **Tied embeddings**: neither config passes
  `--untie-embeddings-and-output-weights`.
* **Single node**: both recipes target 8 × H200.

## 2. Supported Variants

| Model | Class | Active / Total | Layers | Hidden | HF ID |
| - | - | - | - | - | - |
| Gemma-4 26B-A4B-it | MoE, 128 experts top-8 | 4 B / 26 B | 30 | 2816 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it) |
| Gemma-4 31B-it | Dense | 31 B | 60 | 5376 | [google/gemma-4-31B-it](https://huggingface.co/google/gemma-4-31B-it) |

Both use GQA with `kv_channels=256`, RoPE base 1e6, and a 262144-token vocab.

The 31B recipe requires the `zhichen/gemma4-dense` branch of `radixark/Megatron-Bridge`.

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download google/gemma-4-26B-A4B-it --local-dir /root/models/gemma-4-26B-A4B-it
```

`--model-dir` and `--data-dir` default to `/root/models` and `/root/datasets`.
`prepare` performs these downloads for you.

### 3.2 No `torch_dist` conversion

The bridge reads the HF checkpoint directly, so `--hf-checkpoint` and
`--ref-load` both point at the download:

```bash theme={null}
--hf-checkpoint <model-dir>/<model-name>
--ref-load      <model-dir>/<model-name>
--megatron-to-hf-mode bridge
```

## 4. Launch

```bash theme={null}
cd /root/miles

# MoE, single node
python scripts/run_gemma_4_26b_a4b.py full-train --num-nodes 1

# dense, single node
python scripts/run_gemma_4_31b.py full-train --num-nodes 1
```

Passing `--num-nodes 1` puts the recipe into `debug_minimal` mode, which shortens
`--rollout-max-response-len` to 256 for a quick smoke test. Multi-node runs use
the full 8192.

## 5. Recipe Configuration

### 5.1 Parallelism

| Model | TP | PP | CP | EP | ETP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - | - |
| 26B-A4B (MoE) | 4 | 1 | 1 | 8 | 1 | 1024 | 8 (1 × 8) |
| 31B (dense) | 4 | 1 | 1 | — | — | 512 | 8 (1 × 8) |

Activation checkpointing is on for the MoE recipe
(`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`). The dense 31B runs a smaller token budget because
its 60 dense layers at hidden 5376 cost more activation memory per token than
the MoE's 30 layers at 2816.

### 5.2 Algorithm

GRPO. The MoE recipe adds low-variance KL; the dense one runs without it:

```bash theme={null}
--advantage-estimator grpo
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
--rm-type gemma_math
--balance-data

# 26B-A4B only
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
```

Rollout batch 32 at 8 samples per prompt, global batch 256, `--lr 1e-6`. AIME
evaluation every 20 steps is available behind `--enable-eval` and is off by
default.

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 4
--sglang-mem-fraction-static 0.55   # 0.5 for the dense 31B
```

The MoE recipe pins sglang to conservative kernels:

```bash theme={null}
--sglang-attention-backend triton
--sglang-moe-runner-backend triton
--sglang-disable-custom-all-reduce
--sglang-disable-cuda-graph
--sglang-disable-overlap-schedule
--sglang-disable-radix-cache
--use-rollout-routing-replay
```

`--use-rollout-routing-replay` replays the rollout's expert routing during the
training forward pass, so train log-probs match rollout log-probs. Every
sigmoid- or softmax-routed MoE recipe in miles needs this; the dense 31B does
not.

### 5.4 Notable quirks

* **Trained on the VLM checkpoint.** miles does not strip the vision tower; the
  bridge and sglang both handle the multimodal config, and the RL recipe simply
  trains the language stack.
* `--attention-backend unfused` on the training side for the MoE recipe.
* Routing is softmax with `seq_aux_loss` balancing and the bias update rate at 0
  (`--moe-router-bias-update-rate 0 --moe-aux-loss-coeff 0`), plus
  `--moe-grouped-gemm` and `--moe-router-dtype fp32`.
* The 31B recipe needs the `gemma4-dense` branch of `radixark/Megatron-Bridge`,
  because the dense config is driven straight through `Gemma4VLBridge`.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# Gemma
Source: https://miles.radixark.com/docs/models/gemma/index

Miles recipes for Google's Gemma-4 line, trained on the base VLM checkpoint through the HF to Megatron bridge.

Miles supports Google's Gemma-4 in both released instruction-tuned sizes. Both train as
language models on the base VLM checkpoint, through the HF to Megatron bridge
(`--megatron-to-hf-mode bridge`), so there is no offline `torch_dist` conversion.

## Variants

| Model | Class | Active / Total | HF ID | Recipe |
| - | - | - | - | - |
| Gemma-4 26B-A4B-it | MoE, 128 experts top-8 | 4 B / 26 B | `google/gemma-4-26B-A4B-it` | [gemma-4](/docs/models/gemma/gemma-4) |
| Gemma-4 31B-it | Dense | 31 B | `google/gemma-4-31B-it` | [gemma-4](/docs/models/gemma/gemma-4) |

## Fastest path to train

Both recipes run on a single 8-GPU node:

```bash theme={null}
cd /root/miles
python scripts/run_gemma_4_26b_a4b.py full-train --num-nodes 1
```

`--num-nodes 1` shortens the response length for a smoke test. See
[Gemma-4](/docs/models/gemma/gemma-4) for the full walkthrough.

## Which variant do I pick?

* **Cheaper to train, sparse** → 26B-A4B. Four billion active parameters, and expert
  parallelism carries the width.
* **Dense, no routing to reason about** → 31B. It needs the `gemma4-dense` branch of
  `radixark/Megatron-Bridge`, and runs a smaller per-GPU token budget because its 60 dense
  layers cost more activation memory per token.

## Pairs well with

* [Backends Beyond Megatron](/docs/advanced/architecture-support), the bridge path Gemma rides on.
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# GLM4.5
Source: https://miles.radixark.com/docs/models/glm/glm4-5

Launch recipes for GLM-4.5 (355B-A32B) — the 8-node launcher and the Blackwell launcher.

## 1. Model Introduction

[GLM-4.5](https://huggingface.co/zai-org/GLM-4.5) is Zhipu AI's flagship MoE language model with advanced capabilities in reasoning, function calling, and multi-modal understanding.

**Key highlights:**

* **Sparse MoE architecture**: 355 B / 32 B-active for frontier runs and 106 B / 12 B-active for two-node experimentation.
* **Strong reasoning**: built-in step-by-step reasoning, with FP8 rollout supported on Blackwell hardware.
* **Speculative decoding**: EAGLE/MTP rollout is always on in the 8-node launcher; `run_glm45_355b_a32b.py` exposes `--enable-mtp`.
* **R3 / MIS opt-in**: routing-stability extensions available behind a flag (`--enable-mis`) on `run_glm45_355b_a32b.py`.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| GLM-4.5-355B-A32B | 32 B / 355 B | [zai-org/GLM-4.5](https://huggingface.co/zai-org/GLM-4.5) |
| GLM-4.5-Air (106B-A12B) | 12 B / 106 B | [zai-org/GLM-4.5-Air](https://huggingface.co/zai-org/GLM-4.5-Air) |

The 106B-A12B variant has no launcher under `scripts/`; the canonical recipe is [`examples/infra_features/p2p_weight_transfer/GLM-4.5-Air.sh`](https://github.com/radixark/miles/blob/main/examples/infra_features/p2p_weight_transfer/GLM-4.5-Air.sh) (8-node, P2P weight transfer).

## 3. Environment Setup

### 3.1 Required env vars

The 8-node launcher (`run_glm45_355b_a32b_8node.py`) requires:

```bash theme={null}
export MASTER_ADDR=<ray head IP, reachable from every node>
```

Paths are Typer flags: `--model-dir` (default `/root/models`), `--data-dir` (default `/root/datasets`), `--output-dir` (default `/root/shared_data`) — all three must live on a shared FS reachable from every node. `run_glm45_355b_a32b.py` reads no env vars either.

### 3.2 Download model + datasets

```bash theme={null}
hf download zai-org/GLM-4.5 --local-dir /root/models/GLM-4.5-355B-A32B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.3 HF → Megatron `torch_dist` conversion

The 8-node launcher does **not** convert for you — produce `/root/models/GLM-4.5-355B-A32B_torch_dist/` ahead of time:

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py glm4.5-355B-A32B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/GLM-4.5-355B-A32B \
   --save          /root/models/GLM-4.5-355B-A32B_torch_dist
```

`run_glm45_355b_a32b.py` automates the full flow (download → optional `tools/convert_hf_to_fp8.py` → `convert_checkpoint` → `rsync` to `model_local_dir` → submit).

## 4. Launch

### 4.1 Quick start

```bash theme={null}
# 8-node launcher (8 nodes × 8 GPU)
cd /root/miles
export MASTER_ADDR=...
python scripts/run_glm45_355b_a32b_8node.py

# Blackwell-only variant (_execute_train asserts hardware != "H100")
python scripts/run_glm45_355b_a32b.py train --hardware GB300
```

### 4.2 Multi-node fan-out

`run_glm45_355b_a32b_8node.py` performs Ray fan-out internally, ssh-ing every host of `--ray-hostfile` (default `/root/mpi_rack_hostfile`) into the cluster. Pass `--no-join-ray-workers` when the cluster is already joined.

## 5. Recipe Configuration

### 5.1 Parallelism

| Source | TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - | - |
| `run_glm45_355b_a32b_8node.py` | 8 | 4 | 2 | 16 | 1 | 16384 | 64 (8 × 8) |
| `run_glm45_355b_a32b.py` (`num_nodes ≤ 4`, debug) | 4 | 1 | 1 | 4 | 1 | 16384 | ≤ 32 (≤ 4 × 8) |
| `run_glm45_355b_a32b.py` (`num_nodes == 8`) | 4 | 8 | 2 | 8 | 1 | 16384 | 64 (8 × 8) |

### 5.2 Algorithm

| Source | Advantage | Notable flags |
| - | - | - |
| `run_glm45_355b_a32b_8node.py` | GSPO | `--eps-clip 1e-4 --eps-clip-high 2e-4 --use-tis` |
| `run_glm45_355b_a32b.py` | GRPO | `--eps-clip 1e-4 --eps-clip-high 2e-4 --use-tis` |

Neither launcher enables `--use-rollout-routing-replay` by default. `run_glm45_355b_a32b.py` exposes `--enable-mis` (TIS/RS config) as an opt-in.

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 32
--sglang-mem-fraction-static 0.7
--sglang-enable-dp-attention
--sglang-dp-size 4
--sglang-ep-size 32
--sglang-enable-dp-lm-head
--sglang-moe-dense-tp-size 1

# mtp / EAGLE
--sglang-speculative-algorithm EAGLE
--sglang-speculative-num-steps 1
--sglang-speculative-eagle-topk 1
--sglang-speculative-num-draft-tokens 2
--sglang-enable-draft-weights-cpu-backup
```

Megatron side: `--moe-token-dispatcher-type flex`, `--moe-enable-deepep`.

### 5.4 Optimizer

CPU Adam on:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

### 5.5 Notable quirks

* The 8-node launcher passes neither `--load` nor `--save`, so `--load` defaults to the value of `--ref-load`. `--save` turns checkpointing to `<output-dir>/checkpoints` back on, every 20 steps.
* `run_glm45_355b_a32b.py` is Blackwell-only: `_execute_train` asserts `args.hardware != "H100"`.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)
* [INT4 QAT](/docs/advanced/int4-qat)
* [Rollout Routing Replay (R3)](/docs/advanced/miles-router) — opt-in via `--enable-mis` on `run_glm45_355b_a32b.py`.


# GLM4.7 Flash
Source: https://miles.radixark.com/docs/models/glm/glm4-7-flash

Launch recipes for GLM-4.7-Flash — compact MLA + MoE with R3 enabled by default.

## 1. Model Introduction

[GLM-4.7-Flash](https://huggingface.co/zai-org/GLM-4.7-Flash) is a lightweight, high-speed MoE model in the GLM-4.7 series from Zhipu AI, designed for single-GPU-node deployment.

**Key highlights:**

* **Compact MoE architecture**: 30 B total / 3 B active, sparse activation for efficient inference.
* **MLA attention**: Multi-head Latent Attention with q-LoRA rank 768 and kv-LoRA rank 512.
* **MTP head + EAGLE speculative**: built-in `--mtp-num-layers 1` and EAGLE rollout enabled by default.
* **R3 on by default**: the miles launcher enables `--use-rollout-routing-replay` out of the box.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| GLM-4.7-Flash | 3 B / 30 B | [zai-org/GLM-4.7-Flash](https://huggingface.co/zai-org/GLM-4.7-Flash) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download zai-org/GLM-4.7-Flash --local-dir /root/models/GLM-4.7-Flash
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

The launcher does all three downloads itself into `--model-dir` (default `/root/models`) and `--data-dir` (default `/root/datasets`), so this step is optional.

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py glm4.7-flash)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/GLM-4.7-Flash \
   --save          /root/models/GLM-4.7-Flash_torch_dist
```

The launcher does the conversion automatically.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_glm47_flash.py --rollout-num-gpus-per-engine 4

# pass --hardware B200 on B200
```

Defaults (see `ScriptArgs`): `model_org=zai-org`, `model_name=GLM-4.7-Flash`, `num_gpus_per_node=8`, `hardware=H200`, `sglang_attention_backend=None`, `data_dir=/root/datasets`, `model_dir=/root/models`, `output_dir=/root/shared_data`. The `hardware` CLI also accepts `B200`.

## 5. Recipe Configuration

### 5.1 Parallelism

| TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| 4 | 1 | 1 | 8 | 1 | 32768 | 8 (1 × 8) |

Left to itself the launcher picks `--rollout-num-gpus-per-engine 2` on B200 and 1 on H200; those and the 4 above all divide the model's 20 attention heads. The recipe passes no `--sglang-enable-dp-attention` / `--sglang-dp-size` — the in-source comment notes that DP-attention requires `tp_size % dp_size == 0`.

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --use-kl-loss --kl-loss-coef 0.00`.

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 4
--sglang-mem-fraction-static 0.7

# EAGLE speculative decoding (MTP)
--sglang-speculative-algorithm EAGLE
--sglang-speculative-num-steps 2
--sglang-speculative-eagle-topk 1
--sglang-speculative-num-draft-tokens 3

# R3 — on by default in this recipe
--use-rollout-routing-replay
```

### 5.4 Optimizer

CPU Adam on:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

### 5.5 Notable quirks

* Megatron-side DeepEP / `flex` dispatcher are not enabled by this recipe.
* R3 (`--use-rollout-routing-replay`) is enabled by default — atypical for the rest of the model lineup.

## 6. Pairs Well With

* [Rollout Routing Replay (R3)](/docs/advanced/miles-router) — already on by default.
* [Low Precision RL](/docs/advanced/low-precision)


# GLM-5 / GLM-5.1
Source: https://miles.radixark.com/docs/models/glm/glm5

Launch recipe for GLM-5 and GLM-5.1 (744 B / 40 B active) — Python launcher, 16+ node config.

## 1. Model Introduction

[GLM-5](https://huggingface.co/zai-org/GLM-5) is the most powerful language model in Zhipu AI's GLM series, scaling to 744 B parameters (40 B active) and integrating DeepSeek Sparse Attention (DSA) for long-context efficiency. [GLM-5.1](https://huggingface.co/zai-org/GLM-5.1) is the next-generation model for agentic engineering on top of GLM-5, sharing the same model architecture. For the newer GLM-5.2 recipe, see [GLM-5.2](/docs/models/glm/glm5-2).

**Key highlights:**

* **Sparse MoE at frontier scale**: 744 B total / 40 B active per token, 256 routed experts top-8 + 1 shared.
* **MLA + DSA attention**: Multi-head Latent Attention (q-LoRA 2048 / kv-LoRA 512) combined with DeepSeek Sparse Attention to keep KV-cache cost low at long context.
* **Speculative decoding**: EAGLE/MTP rollout supported via `--enable-mtp`.
* **PD disaggregation**: prefill/decode disaggregation enabled by default for ≥1 node.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| GLM-5.1 | 40 B / 744 B | [zai-org/GLM-5.1](https://huggingface.co/zai-org/GLM-5.1) |
| GLM-5 | 40 B / 744 B | [zai-org/GLM-5](https://huggingface.co/zai-org/GLM-5) |

## 3. Environment Setup

### 3.1 Download model + datasets

The Python launcher's `prepare` subcommand handles download + dataset staging:

```bash theme={null}
python scripts/run_glm5_744b_a40b.py prepare --model-name GLM-5 --num-nodes 16
```

### 3.2 HF → Megatron `torch_dist` conversion

Also handled by `prepare`. Before conversion the launcher validates, via `_validate_glm_checkpoint`, that the checkpoint uses the native GLM-5 config (`model_type=glm_moe_dsa`, `architectures=[GlmMoeDsaForCausalLM]`) and fails fast if it does not, then converts it to the `glm5-744B-A40B` Megatron model type. Training itself copies the converted checkpoint from shared NFS to each node's local disk before it starts.

### 3.3 DSA training implementation

The raw Megatron path supports `--dsa-impl miles|megatron` for DSA models using the shared GLM-5 / DeepSeek-V3.2 spec. `miles` remains the default. To select Megatron's native DSA without Megatron-Bridge, pass the following to both `tools/convert_hf_to_torch_dist.py` and training:

```bash theme={null}
--megatron-to-hf-mode raw --dsa-impl megatron --dsa-kernel-backend cudnn
```

Training uses packed `--qkv-format thd` and supports sequence parallelism. Native context parallelism uses zigzag token partitioning within each sequence and `--cp-comm-type allgather` for attention communication. Omitting `--cp-comm-type` selects `allgather` for native DSA; explicitly selecting another communication type is rejected. Miles' `--allgather-cp` instead selects contiguous token partitioning and is rejected for native DSA when CP > 1. GLM-5.2's cross-layer index sharing schedule is preserved. Indexer replay is outside the scope of this backend integration.

The fused cuDNN indexer supports CP1, single-sequence CP batches, and complete multi-sequence CP query partitions. TP-local query slices of multi-sequence CP batches are rejected; select the reference implementation explicitly for that layout. Fused execution uses the configured cuDNN frontend directly and does not silently switch to reference scoring.

`--miles-dsa-topk-backend` selects top-k for both implementations and is accepted by conversion, training, and the standalone `run_megatron` debug worker. The default `canonical` returns the exact top-k with ties at the k-th score kept by the smaller key (the rule of SGLang's DSA top-k) in ascending key order; native Megatron DSA runs it as its exact `torch` top-k, and `torch` keeps the previous Miles selector. This runtime selection does not change the checkpoint layout. The W4A16 test keeps `flashinfer` and `SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK=large`, preserving its top-k backend and tie policy when selecting native Megatron DSA.

Native DSA defaults to `--dsa-indexer-loss-coeff 0` and freezes its indexer parameters, because native top-k selection runs without gradients when the auxiliary objective is disabled. This excludes unused parameters from DDP and optimizer weight decay. A positive native loss coefficient keeps the indexer trainable and requires explicitly selecting the reference implementation with `--dsa-kernel-backend none` or `--attention-backend unfused` when using an external top-k backend.

Convert into a separate `torch_dist` directory when changing implementations: the native attention and indexer parameter names differ from the Miles checkpoint layout. Use the same implementation for conversion, checkpoint loading, and training. This selector applies to DSA training; SGLang rollout backend flags remain independent.

With BF16 training, native cuDNN uses BF16 head weights; Miles produces FP32 head weights from BF16 projection operands. [cuDNN Frontend PR #1311](https://github.com/NVIDIA/cudnn-frontend/pull/1311) adds FP32 input support to forward scoring only. This wiring remains BF16; FP32 projection and auxiliary-loss recompute/backward support are separate work. Head-weight precision affects scores independently of the top-k backend and tie policy.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
python scripts/run_glm5_744b_a40b.py full-train --model-name GLM-5 --num-nodes 16
```

The Typer app exposes three subcommands:

```bash theme={null}
python scripts/run_glm5_744b_a40b.py full-train --model-name GLM-5 --num-nodes <N>

# Just download model + datasets and convert to Megatron
python scripts/run_glm5_744b_a40b.py prepare    --model-name GLM-5 --num-nodes <N>

# Train only (assumes prepare done); copies the checkpoint to each node's local disk first
python scripts/run_glm5_744b_a40b.py train      --model-name GLM-5 --num-nodes <N>
```

The recipe is tested on **H200 / B200 / GB300**; the `--hardware` flag accepts exactly these three values.

## 5. Recipe Configuration

### 5.1 Parallelism

Verbatim from `_execute_train`, `--num-nodes ≥ 16` branch:

| TP | PP | CP | EP | expert-TP | `decoder-last-pipeline-num-layers` | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - | - |
| 4 | 4 | 2 | 32 | 1 | 18 | 16384 | ≥ 128 (≥ 16 × 8) |

Plus `--use-dynamic-batch-size`, `--data-pad-size-multiplier 4096`, `--log-probs-chunk-size 1024`, `--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`.

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28`. R3 (`--use-rollout-routing-replay`) is **not** enabled by default.

### 5.3 Rollout & SGLang

Always-on flags:

```bash theme={null}
--sglang-mem-fraction-static 0.70
--sglang-enable-dp-attention
--sglang-ep-size <world_size>
--sglang-dp-size <world_size>
--sglang-moe-dense-tp-size 1
--sglang-enable-dp-lm-head

# DSA / NSA attention
--sglang-page-size 64
--sglang-dsa-decode-backend flashmla_sparse
--sglang-dsa-prefill-backend flashmla_sparse
--sglang-attention-backend nsa

--sglang-max-running-requests 512
--sglang-watchdog-timeout 3600
```

### 5.4 Optimizer

`--enable-optimizer-offload` adds `--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer` (opt-in).

### 5.5 Notable quirks

The launcher exposes these as flags:

* `--fp8-rollout` — runs `tools/convert_hf_to_fp8.py --strategy block --block-size 128 128` and feeds the FP8 directory to SGLang (Megatron stays BF16).
* `--enable-mtp` — adds SGLang EAGLE speculative decoding (`--sglang-speculative-{algorithm,num-steps,eagle-topk,num-draft-tokens}`).
* `--enable-pd` (default `True` for ≥1 node) — enables prefill/decode disaggregation; with PD the launcher uses larger SGLang world sizes (16 for `<16` nodes, 64 for `≥16` nodes).
* `--use-deepep` (default `True`) — enables Megatron-side DeepEP (`--moe-enable-deepep --moe-token-dispatcher-type flex`); falls back to `alltoall`. Forced off on GB300.

## 6. Pairs Well With

* [PD Disaggregation](/docs/advanced/pd-disaggregation) — on by default for `num_nodes ≥ 1`.
* [Low Precision RL](/docs/advanced/low-precision) — opt-in via `--fp8-rollout`.
* [Speculative Decoding](/docs/advanced/speculative-decoding) — opt-in via `--enable-mtp`.


# GLM-5.2
Source: https://miles.radixark.com/docs/models/glm/glm5-2

Launch recipe for GLM-5.2 (744 B / 40 B active) — FP8 KV cache, TIS, 16+ node config.

## 1. Model Introduction

[GLM-5.2](https://huggingface.co/zai-org/GLM-5.2) is the successor to GLM-5 / GLM-5.1 in Zhipu AI's GLM series. It keeps the same 744 B-parameter (40 B active) `glm_moe_dsa` architecture — MoE plus DeepSeek Sparse Attention (DSA) with cross-layer index sharing — and differs from GLM-5 in the checkpoint it loads, the Megatron model args (`--rotary-base 8000000`), and the rollout recipe (FP8 KV cache, `flashmla_kv` decode, optional EAGLE speculative decoding).

**Key highlights:**

* **Sparse MoE at frontier scale**: 744 B total / 40 B active per token, 256 routed experts top-8 + 1 shared, 3 dense + 75 MoE layers.
* **MLA + DSA with cross-layer index sharing**: only the *computing* layers (1, 2, 3, 7, 11, …, 75 in Megatron 1-indexing; `index_topk_freq=4`) carry indexer weights and compute the sparse top-k — the remaining layers reuse the most recent computing layer's indices. This constrains the pipeline split: every PP stage must start on a computing layer.
* **FP8 KV cache rollout**: `fp8_e4m3` KV cache with `flashmla_kv` decode and `flashmla_sparse` prefill.
* **Truncated importance sampling**: the GLM-5.2 recipe enables TIS by default (`--use-tis`).

## 2. Supported Variants

| Model | Active / Total | HF ID | Scale |
| - | - | - | - |
| GLM-5.2 | 40 B / 744 B (78 layers) | [zai-org/GLM-5.2](https://huggingface.co/zai-org/GLM-5.2) | ≥ 16 nodes |
| GLM-5.2\_5layer | 5-layer pruned | Pinaster/GLM-5.2\_5layer | 1 node (smoke test) |

A LoRA variant of the recipe ships as `scripts/run_glm5_2_744b_a40b_lora.py` (see [LoRA](/docs/advanced/lora)).

## 3. Environment Setup

Use the `radixark/miles:dev` docker image.

### 3.1 Download model + datasets

The Python launcher's `prepare` subcommand handles download + dataset staging (dapo-math-17k):

```bash theme={null}
python scripts/run_glm5_2_744b_a40b.py prepare --model-name GLM-5.2 --num-nodes 32
```

### 3.2 HF → Megatron `torch_dist` conversion

Also handled by `prepare`. Before conversion the launcher validates, via `_validate_glm_checkpoint`, that the checkpoint uses the native GLM-5.2 config (`model_type=glm_moe_dsa`, `architectures=[GlmMoeDsaForCausalLM]`, `num_hidden_layers=78`, no `auto_map`) and fails fast if it does not, then converts it to the `glm5.2-744B-A40B` Megatron model type. The full model converts with PP = 4 (18/20 first/last layer split); the pruned model converts on a single GPU, because DSA's cross-layer index sharing forbids a pipeline stage that starts on a skip layer. `train` copies the converted checkpoint from shared NFS to node-local disk on every trainer pod before the job starts.

The shared [DSA training implementation selector](/docs/models/glm/glm5#33-dsa-training-implementation) also applies to GLM-5.2. The 8-GPU W4A16 test, `tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4_w4a16.py`, uses raw Megatron DSA with cuDNN for conversion and training, with 4 actor GPUs and 4 rollout GPUs. Its native checkpoint is stored separately as `GLM-5.2_5layer-megatron-dsa_torch_dist`.

## 4. Launch

### 4.1 Quick start

Single-node smoke test with the 5-layer pruned model:

```bash theme={null}
python scripts/run_glm5_2_744b_a40b.py full-train --model-name GLM-5.2_5layer --num-nodes 1
```

Full model (≥ 16 nodes):

```bash theme={null}
python scripts/run_glm5_2_744b_a40b.py full-train --model-name GLM-5.2 --num-nodes 32
```

The Typer app exposes three subcommands:

```bash theme={null}
python scripts/run_glm5_2_744b_a40b.py full-train --model-name GLM-5.2 --num-nodes <N>

# Just download model + datasets and convert to Megatron
python scripts/run_glm5_2_744b_a40b.py prepare    --model-name GLM-5.2 --num-nodes <N>

# Train only (assumes prepare done)
python scripts/run_glm5_2_744b_a40b.py train      --model-name GLM-5.2 --num-nodes <N>
```

The recipe is tested on **H200 / B200 / GB300**; the `--hardware` flag accepts exactly these three values.

### 4.2 Agentic RL: terminal-bench-2 in Daytona sandboxes (experimental)

Beyond the math recipe above, [`examples/experimental/openenv/glm52_tbench2/`](https://github.com/radixark/miles/tree/main/examples/experimental/openenv/glm52_tbench2) trains GLM-5.2 with fully-async agentic RL on terminal-bench-2: 16 GB300 nodes (4 GPUs each) split into 8 training nodes (TP2 / CP4 / PP4 / EP8, optimizer state streamed to node-local disk) and 8 inference nodes (one 4-GPU dp-attention FP8 SGLang engine per node). Every episode is a multi-turn terminal agent solving one terminal-bench-2 task inside its own Daytona cloud sandbox built from that task's official image; scoring is the task's canonical `tests/test.sh`.

```bash theme={null}
python3 examples/experimental/openenv/glm52_tbench2/run_glm5_2_744b_a40b_daytona.py train --num-nodes 16
```

The recipe's defaults are the reference configuration (100 rollout steps in \~21 h, \~6.5 min/step including evals). See the example's README for the container, Megatron, OpenEnv, and Daytona prerequisites.

## 5. Recipe Configuration

### 5.1 Parallelism

`_execute_train` picks one of three branches:

| Branch | TP | PP | CP | EP | expert-TP | first/last PP layers | `max_tokens_per_gpu` |
| - | - | - | - | - | - | - | - |
| 1 node (5-layer) | 4 | 1 | 1 | = GPUs per node | 1 | — | 2048 |
| ≥ 16 nodes, 4 GPUs/node (GB300) | 8 | 4 | 1 | 16 | 1 | 18 / 20 | 8192 |
| ≥ 16 nodes, 8 GPUs/node | 4 | 8 | 8 | 32 | 1 | 14 / 16 | 8192 |

The uneven first/last pipeline splits are dictated by DSA: every stage must start on a computing layer (e.g. the 14/16 split lands stage starts on layers 1, 15, 23, 31, 39, 47, 55, 63 — all computing).

Plus `--use-dynamic-batch-size`, `--data-pad-size-multiplier 1024`, `--log-probs-chunk-size 16384`, `--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`, and `--allgather-cp` (DSA + context parallel uses the sequential allgather-CP layout).

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28`, plus truncated importance sampling — `--use-tis --tis-clip-low 0.5 --tis-clip 2.0` — which the GLM-5 recipe does not enable. R3 (`--use-rollout-routing-replay`) is **not** enabled by default.

### 5.3 Rollout & SGLang

Always-on flags:

```bash theme={null}
--sglang-mem-fraction-static 0.85   # 0.70 on the 1-node smoke test
--sglang-ep-size <world_size>
--sglang-router-policy consistent_hashing

# DSA / NSA attention with FP8 KV cache
--sglang-kv-cache-dtype fp8_e4m3
--sglang-dsa-decode-backend flashmla_kv
--sglang-dsa-prefill-backend flashmla_sparse
--sglang-attention-backend nsa
--sglang-page-size 64

--sglang-max-running-requests 512   # 256 with --sglang-config balanced
--sglang-watchdog-timeout 3600
```

With PD disaggregation (the multi-node default) the launcher also adds `--sglang-enable-dp-attention --sglang-dp-size <world_size> --sglang-moe-dense-tp-size 1 --sglang-enable-dp-lm-head` and uses SGLang world size 16 (`< 16` nodes) or 64 (`≥ 16` nodes).

Training runs with these env vars: `SGLANG_NSA_FORCE_MLA=1`, `INDEXER_ROPE_NEOX_STYLE=0`, `NVSHMEM_DISABLE_NCCL=1`.

### 5.4 Optimizer

Adam, `--lr 1e-6 --lr-decay-style constant --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98`. `--enable-optimizer-offload` adds `--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer` (opt-in).

### 5.5 Notable quirks

The launcher exposes these as flags:

* `--fp8-rollout` — runs `tools/convert_hf_to_fp8.py --strategy block --block-size 128 128` and feeds the FP8 directory to SGLang (Megatron stays BF16). Combined with `--use-deepep` it also switches SGLang's MoE all-to-all to DeepEP (`--sglang-moe-a2a-backend deepep --sglang-deepep-mode auto`).
* `--enable-mtp` — adds SGLang EAGLE speculative decoding with `--sglang-speculative-draft-attention-backend nsa`; `low-latency` drafts deeper (num-steps 5, draft-tokens 6) than `balanced` (1, 2). Full model only — the MTP layer is pruned away in the 5-layer variant.
* `--enable-pd` (default `True`, forced off on 1 node) — enables prefill/decode disaggregation.
* `--sglang-config {low-latency, balanced}` — `low-latency` (default) runs TP-8 engines with PD; `balanced` is the GLM-5.2 cookbook serving shape: one 4-GPU engine per node with dp-attention + DeepEP. `balanced` is incompatible with PD and requires GPUs per node divisible by 4.
* `--use-deepep` (default `True`) — enables Megatron-side DeepEP (`--moe-enable-deepep --moe-token-dispatcher-type flex`); falls back to `alltoall`. On GB300 you must pass `--no-megatron-use-deepep` (known Megatron DeepEP failure; the launcher asserts).
* On B200/GB300 (without `balanced` or FP8 + DeepEP) the launcher pins `--sglang-moe-runner-backend`: `flashinfer_trtllm_routed` for FP8 rollout, `triton` for BF16.

## 6. Pairs Well With

* [GLM-5.3](/docs/models/glm/glm5-3) — same architecture; runs on the GLM-5.2 LoRA launcher.
* [PD Disaggregation](/docs/advanced/pd-disaggregation) — on by default for multi-node runs.
* [Low Precision RL](/docs/advanced/low-precision) — opt-in via `--fp8-rollout`.
* [Speculative Decoding](/docs/advanced/speculative-decoding) — opt-in via `--enable-mtp`.
* [LoRA](/docs/advanced/lora) — via `scripts/run_glm5_2_744b_a40b_lora.py`.
* [Fully Async Rollout](/docs/examples/infra-features/fully-async) — the terminal-bench-2 agentic example (§4.2) runs fully async.


# GLM-5.3
Source: https://miles.radixark.com/docs/models/glm/glm5-3

GLM-5.3 (744 B / 40 B active) on the GLM-5.2 recipe — same architecture, Megatron config, and launchers.

## 1. Model Introduction

[GLM-5.3](https://huggingface.co/zai-org/GLM-5.3) is the successor to GLM-5.2 in Zhipu AI's GLM series. It keeps the 744 B-parameter (40 B active) `glm_moe_dsa` architecture of GLM-5.2 unchanged — MoE plus DeepSeek Sparse Attention (DSA) with cross-layer index sharing — so Miles trains it with the GLM-5.2 Megatron model args, bridge, and recipe. Only the checkpoint differs.

GLM-5.3 is not GLM-5.3-Flash. The Flash model is a separate KDA + DSA hybrid architecture with its own recipe; see [GLM-5.3-Flash](/docs/models/glm/glm5-3-flash).

## 2. Supported Variants

| Model | Active / Total | HF ID | Notes |
| - | - | - | - |
| GLM-5.3 (BF16) | 40 B / 744 B (78 layers) | [zai-org/GLM-5.3-BF16](https://huggingface.co/zai-org/GLM-5.3-BF16) | Training checkpoint |
| GLM-5.3 (FP8) | 40 B / 744 B (78 layers) | [zai-org/GLM-5.3](https://huggingface.co/zai-org/GLM-5.3) | Official FP8 release, usable as the rollout checkpoint |

Training is BF16, so start from `GLM-5.3-BF16`. For smoke tests, use the 5-layer GLM-5.2 prune (`GLM-5.2_5layer`): the shapes are identical.

## 3. Launch

### 3.1 LoRA

`scripts/run_glm5_2_744b_a40b_lora.py` takes the checkpoint path through `--hf-checkpoint`, so GLM-5.3 runs on the `GLM-5.2` model entry (Megatron model type `glm5.2-744B-A40B_lora`). Stage the checkpoint and dataset yourself and call `train`: `prepare` and `full-train` download the repo named by `--model-name`, which would fetch GLM-5.2.

```bash theme={null}
hf download zai-org/GLM-5.3-BF16 --local-dir /root/models/GLM-5.3-BF16
hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k

python scripts/run_glm5_2_744b_a40b_lora.py train \
  --model-name GLM-5.2 \
  --hf-checkpoint /root/models/GLM-5.3-BF16
```

For `--task dapo-math`, download `zhuzilin/dapo-math-17k` to `/root/datasets/dapo-math-17k` instead.

FP8 rollout works as it does for GLM-5.2: training stays BF16 and SGLang serves `<hf_checkpoint>_fp8`. Point that at the official FP8 release (for example, download `zai-org/GLM-5.3` to `/root/models/GLM-5.3-BF16_fp8`) and add `--fp8-rollout`.

### 3.2 Full-parameter

`scripts/run_glm5_2_744b_a40b.py` does not take a checkpoint override yet: `prepare` downloads `<model-org>/<model-name>` and accepts only the `GLM-5.2` and `GLM-5.2_5layer` names. Use the LoRA launcher above for GLM-5.3 until the full-parameter launcher gains a GLM-5.3 entry.

## 4. Recipe Configuration

Every setting comes from the GLM-5.2 LoRA launcher unchanged: parallelism, DSA kernel backend (`--dsa-attention-backend`), LoRA targets, algorithm, and rollout flags. See [LoRA](/docs/advanced/lora) for the GLM-5.2 Bridge LoRA details and validation results, and [GLM-5.2](/docs/models/glm/glm5-2) for the DSA constraint that every pipeline stage must start on an indexer-computing layer.

## 5. Pairs Well With

* [GLM-5.2](/docs/models/glm/glm5-2) — the recipe GLM-5.3 runs on.
* [LoRA](/docs/advanced/lora) — the GLM-5 / 5.1 / 5.2 Bridge LoRA path also covers GLM-5.3.
* [Low Precision RL](/docs/advanced/low-precision) — FP8 rollout from the official FP8 checkpoint.


# GLM-5.3-Flash
Source: https://miles.radixark.com/docs/models/glm/glm5-3-flash

RL recipe for GLM-5.3-Flash, a KDA + DSA hybrid MoE with mHC hyper-connections and NoPE MLA.

Implementation: [`radixark/miles#2786`](https://github.com/radixark/miles/pull/2786). It goes
with the SGLang
[`sglang-miles-glm53next`](https://github.com/sgl-project/sglang/tree/sglang-miles-glm53next)
branch and [`radixark/Megatron-LM#89`](https://github.com/radixark/Megatron-LM/pull/89); the
image in section 3 pins all three.

## 1. Model Introduction

[GLM-5.3-Flash](https://huggingface.co/zai-org/GLM-5.3-Flash) (`model_type: glm5_next`) is a
**45-layer KDA + DSA hybrid MoE**. It is a different architecture from the 744 B GLM5 and
GLM5.2 flagships, not a smaller cut of them.

* **45 layers, hybrid**: 34 KDA linear-attention layers + 11 DSA sparse-attention layers.
* **288-expert MoE**, sigmoid routing at top-8; the first 3 layers are dense.
* **mHC hyper-connections** at every block.
* **NoPE MLA** — multi-latent attention with the positional half of the QK head empty.
* **kpool-compressed lightning indexer** picks which keys the DSA layers attend.
* Hidden 4096, FFN 12288, 64 attention heads, vocab 154880, rotary base 800000.
* MTP is dropped for training.

## 2. Supported Variants

| Variant | `--model-name` | Layers |
| - | - | - |
| Full | `GLM-5.3-Flash` | 45 |
| 4-layer slice | `GLM-5.3-Flash-4layer` | 4 (launcher default) |

## 3. Environment Setup

Use `docker.io/radixark/miles:glm53next` — the rolling `radixark/miles:dev` image with the
three moving parts checked out at the versions this recipe was built against, multi-arch so
the same tag serves GB300 and x86 nodes.

| Component | Pinned at |
| - | - |
| miles | [`#2786`](https://github.com/radixark/miles/pull/2786) `1cd14c00` |
| SGLang | `sglang-miles-glm53next` `9a26e749` |
| Megatron-LM | [`#89`](https://github.com/radixark/Megatron-LM/pull/89) `e8f57451` |

```bash theme={null}
hf download zai-org/GLM-5.3-Flash --local-dir /root/models/GLM-5.3-Flash
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
```

The reference checkpoint has to be converted first — `--ref-load` resolves to
`<--ckpt-dir>/<megatron_model_type>_torch_dist`:

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py glm5.3-flash)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
CONVERT_KEEP_PP1=1 CUDA_DEVICE_MAX_CONNECTIONS=1 PYTHONPATH=/root/Megatron-LM \
  torchrun --nproc-per-node 8 \
    tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \
    --hf-checkpoint /root/models/GLM-5.3-Flash \
    --save          /root/ckpt/glm5.3-flash_torch_dist
```

## 4. Launch

Bring up a ray cluster across the nodes, `export MILES_SCRIPT_EXTERNAL_RAY=1`, then on the
head node:

```bash theme={null}
cd /root/miles
python scripts/run_glm5_3_flash.py train \
  --model-name GLM-5.3-Flash \
  --num-nodes 16 --num-gpus-per-node 4 \
  --num-rollout 20 --rollout-max-response-len 4096
```

Smoke slice on one node:

```bash theme={null}
python scripts/run_glm5_3_flash.py train --num-nodes 1 --num-gpus-per-node 8
```

| Shape | TP | PP | EP | Rollout engine |
| - | - | - | - | - |
| 16 × 4 (full) | 8 | 4 | 16 | 8 GPUs, SGLang TP 8 / EP 8 |
| 8 × 4 (full) | 8 | 4 | 8 | 8 GPUs, SGLang TP 8 / EP 8 |
| 2 × 4 or 1 × 8 (4-layer slice) | 2 | 2 | 2 | 4 GPUs, SGLang TP 4 / EP 4 |

The PP-4 shapes run 11 / 11 / 11 / 12 layers per stage, since 45 does not divide by 4.
GRPO on DAPO-Math-17k, Adam at `lr 1e-6`, `max_tokens_per_gpu 8192`, full uniform recompute.
Rollout is colocated, with the trainer offloaded to disk; both DSA paths run on tilelang and
the KV cache is BF16. Routing replay turns on with `--enable-r3`.

## 5. What a Healthy Run Looks Like

From the validation run in [#2786](https://github.com/radixark/miles/pull/2786) — 16 nodes ×
4 GB300, DAPO on DAPO-Math-17k:

| Metric | Observed |
| - | - |
| `train/train_rollout_logprob_abs_diff` | 0.0068 – 0.0106 across the first 11 rollouts |
| `rollout/raw_reward` | 0.5 → 0.94 within 10 rollouts |
| `train/ppo_kl` | \~2.6e-4 |
| `train/grad_norm` | 0.31 – 0.49 |

`train/train_rollout_logprob_abs_diff` is the one to read first on a fresh bring-up: it
covers the KDA, DSA and hyper-connection paths at once.

## 6. LoRA RL

[`scripts/run_glm5_3_flash_lora.py`](https://github.com/radixark/miles/blob/main/scripts/run_glm5_3_flash_lora.py)
trains a LoRA adapter through native LoRA (`--megatron-to-hf-mode raw`): the trainer is the
glm5\_next model of section 4, loaded from the `torch_dist` base of section 3, and
`miles.backends.megatron_utils.lora.native` attaches the adapters. With `--lora-base-cpu-backup`
SGLang keeps the frozen base across colocated offload and onload, so each step syncs only the
adapter.

| Layer type | Adapted HF projections |
| - | - |
| KDA (34 layers) | `q/k/v/o_proj`, `b_proj`, `forget_gate.f_a/f_b_proj`, `g_a/g_b_proj` |
| DSA (11 layers) | `q_a/q_b_proj`, `kv_a_proj_with_mqa`, `o_proj` |
| MLP | dense and shared-expert `gate/up/down_proj` |

The 288 routed experts stay frozen, since native LoRA has no per-expert adapter. `kv_b_proj` stays
frozen too: the DSA attention absorbs its weight into the query and output without calling the
module, so an adapter on it would never run. The launcher's
`--target-modules attn,mlp --exclude-modules "*.mlp.experts.*,*.self_attn.kv_b_proj"` selects
exactly the table, and selecting a frozen projection fails at argument validation. The KDA
`conv1d`, `A_log`, `dt_bias` and `o_norm`, the DSA indexer and the mHC parameters are not LoRA
targets.

Serving the adapter needs SGLang LoRA on the KDA projections, which `sglang-miles` does not have
yet; it is in [sgl-project/sglang#44862](https://github.com/sgl-project/sglang/pull/44862).

```bash theme={null}
python scripts/run_glm5_3_flash_lora.py prepare --model-name GLM-5.3-Flash-4layer
python scripts/run_glm5_3_flash_lora.py train --model-name GLM-5.3-Flash-4layer \
  --num-nodes 1 --num-gpus-per-node 8

# full model, 3 x 8 GPUs: TP 8, EP 24, PP 1, one colocated 8-GPU engine per node
export MILES_SCRIPT_EXTERNAL_RAY=1   # after `ray start` on every node
python scripts/run_glm5_3_flash_lora.py train --model-name GLM-5.3-Flash \
  --num-nodes 3 --num-gpus-per-node 8
```


# GLM
Source: https://miles.radixark.com/docs/models/glm/index

Miles recipes for the GLM4.5, GLM4.7 Flash, GLM5, GLM5.2, and GLM-5.3 families, plus the KDA + DSA hybrid GLM-5.3-Flash.

Miles ships RL recipes for every GLM generation currently in production: the GLM4.5 MoE at 106 B-A12B and 355 B-A32B, the compact GLM4.7 Flash with 64 routed experts, the 744 B-A40B GLM5, GLM5.2, and GLM-5.3 flagships, and GLM-5.3-Flash, a KDA + DSA hybrid on a different architecture entirely.

## Variants

| Family | Class | Sizes | Recipe |
| - | - | - | - |
| GLM4.5 | MoE | 12 B / 106 B · 32 B / 355 B | [glm4-5](/docs/models/glm/glm4-5) |
| GLM4.7 Flash | MoE (64 experts, top-4) | Compact | [glm4-7-flash](/docs/models/glm/glm4-7-flash) |
| GLM5 | MoE | 40 B / 744 B | [glm5](/docs/models/glm/glm5) |
| GLM5.2 | MoE | 40 B / 744 B | [glm5-2](/docs/models/glm/glm5-2) |
| GLM-5.3 | MoE (same architecture as GLM5.2) | 40 B / 744 B | [glm5-3](/docs/models/glm/glm5-3) |
| GLM-5.3-Flash | MoE (KDA + DSA hybrid) | 45 layers · 288 experts, top-8 | [glm5-3-flash](/docs/models/glm/glm5-3-flash) |

## Fastest path to train

GLM4.7 Flash on a single 8× H100 node — the smallest GLM recipe:

```bash theme={null}
python scripts/run_glm47_flash.py
```

See the [GLM4.7 Flash](/docs/models/glm/glm4-7-flash) page for weight conversion and the full walkthrough.

## Which variant do I pick?

* **Single-node GLM first try** → GLM4.7 Flash ([glm4-7-flash](/docs/models/glm/glm4-7-flash)).
* **MoE on a budget** → GLM4.5-106B-A12B ([glm4-5](/docs/models/glm/glm4-5)).
* **Full MoE scale (multi-node)** → GLM4.5-355B-A32B ([glm4-5](/docs/models/glm/glm4-5)).
* **Compact MoE for routing experiments (R3)** → GLM4.7 Flash ([glm4-7-flash](/docs/models/glm/glm4-7-flash)).
* **Frontier scale (744 B)** → GLM-5.3 ([glm5-3](/docs/models/glm/glm5-3)) or GLM5.2 ([glm5-2](/docs/models/glm/glm5-2)); GLM5/GLM5.1 ([glm5](/docs/models/glm/glm5)) for the previous generation.
* **Hybrid linear + sparse attention with mHC** → GLM-5.3-Flash ([glm5-3-flash](/docs/models/glm/glm5-3-flash)).


# GPT-OSS 20B
Source: https://miles.radixark.com/docs/models/gpt-oss/gpt-oss

Launch recipe for OpenAI's GPT-OSS 20B — Megatron BF16 on a single 8-GPU node, loading the MXFP4 HF checkpoint through mbridge.

## 1. Model Introduction

[GPT-OSS](https://huggingface.co/openai/gpt-oss-20b) is OpenAI's open-weight language model, designed for reasoning, agentic tasks, and developer use cases. miles supports the 20 B variant.

**Key highlights:**

* **Configurable reasoning effort**: low / medium / high reasoning effort selectable per request.
* **Full chain-of-thought**: the reasoning trace is exposed and trainable.
* **MXFP4 native weights**: the HF checkpoint ships in MXFP4 (post-trained) — the BF16 launcher uses mbridge to load HF directly.
* **Sink attention**: requires `--qkv-format bshd` on the Megatron path, which precludes dynamic batch sizing.

## 2. Supported Variants

| Model | HF ID |
| - | - |
| gpt-oss-20b | [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download openai/gpt-oss-20b --local-dir /root/models/gpt-oss-20b
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
```

### 3.2 HF → Megatron `torch_dist` conversion

The launcher needs no `convert_hf_to_torch_dist.py` step: it loads the HF checkpoint directly via `--megatron-to-hf-mode bridge` (mbridge).

## 4. Launch

### 4.1 Quick start

```bash theme={null}
# Megatron BF16 (1 node × 8 GPU)
cd /root/miles
python scripts/run_gpt_oss_20b.py
```

The checkpoint is read from `--model-dir` (default `/root/models`) and the dataset from
`--data-dir` (default `/root/datasets`).

## 5. Recipe Configuration

### 5.1 Parallelism

| Launcher | TP | PP | CP | EP | expert-TP | `micro-batch-size` | GPUs |
| - | - | - | - | - | - | - | - |
| `run_gpt_oss_20b.py` (Megatron) | 8 | 1 | 1 | 8 | 1 | 1 | 8 (1 × 8) |

`--use-dynamic-batch-size` is **not** used on the Megatron BF16 path — the in-source comment explains: `--qkv-format bshd` (required for sink attention with TE) is incompatible with dynamic batch size. Only `--micro-batch-size 1` is set. `--sequence-parallel` is on (required for TP + EP).

### 5.2 Algorithm

The recipe uses GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --entropy-coef 0.00` and `--rm-type math`. **`--use-kl-loss` is not passed** — bridge mode has no Megatron-format reference checkpoint to supply `--ref-load`, which KL loss would need.

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 4
--sglang-dtype bfloat16
--sglang-decode-log-interval 1000
--sglang-mem-fraction-static 0.70
```

### 5.4 Optimizer

The launcher enables CPU Adam (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`).

### 5.5 Notable quirks

Attention setup:

```bash theme={null}
--attention-dropout 0.0
--hidden-dropout 0.0
--qkv-format bshd        # required for TE sink attention (SWA + learnable softmax offset)
--attention-backend fused
```

`--qkv-format bshd` is mandated by the sink-attention pattern; in turn it precludes `--use-dynamic-batch-size`. Don't toggle either flag without the other.

The launcher passes no `--load`, and no `--save` unless you ask for it: `--save` exports the
bf16 HF checkpoint to `<model-dir>/gpt-oss-20b-BF16` every 50 steps, which can then be fed
back in as the `--hf-checkpoint`.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)


# Supported Models
Source: https://miles.radixark.com/docs/models/index

Per-family recipes covering weight conversion, launch flags, and parallelism choices.

Miles ships ready-to-run recipes for every model family listed below. Each page covers
weight conversion, parallelism, and the launch script in the order you'd actually run
them.

## Language Models

Each model name links to its recipe page.

| Family | Models |
| - | - |
| **DeepSeek** | [DeepSeek-V4.1 Flash](/docs/models/deepseek/deepseek-v4-1-flash) [DeepSeek-V4 Pro](/docs/models/deepseek/deepseek-v4-pro) [DeepSeek-V4 Flash](/docs/models/deepseek/deepseek-v4-flash) [DeepSeek-V3.2](/docs/models/deepseek/deepseek-v3-2) [DeepSeek-V3](/docs/models/deepseek/deepseek) |
| **Thinking Machines** | [Inkling](/docs/models/thinkingmachines/inkling) [Inkling-Small](/docs/models/thinkingmachines/inkling-small) |
| **Qwen** | [Qwen3.8-Flash-Next](/docs/models/qwen/qwen3-8-flash-next) [Qwen3.8-27B](/docs/models/qwen/qwen3-8) [Qwen3.6 MoE](/docs/models/qwen/qwen3-6-moe) [Qwen3.6](/docs/models/qwen/qwen3-6) [Qwen3.5-35B-A3B](/docs/models/qwen/qwen3-5-moe) [Qwen3.5-4B / 9B / 27B](/docs/models/qwen/qwen3-5) [Qwen3-Next-80B-A3B-Thinking](/docs/models/qwen/qwen3-next) [Qwen3-30B-A3B / 235B-A22B](/docs/models/qwen/qwen3-moe) [Qwen3-0.6B / 1.7B / 4B / 8B / 14B / 32B](/docs/models/qwen/qwen3) |
| **GLM** | [GLM-5.3](/docs/models/glm/glm5-3) [GLM-5.3-Flash](/docs/models/glm/glm5-3-flash) [GLM-5.2](/docs/models/glm/glm5-2) [GLM-5.1](/docs/models/glm/glm5) [GLM-5](/docs/models/glm/glm5) [GLM-4.7-Flash](/docs/models/glm/glm4-7-flash) [GLM-4.5](/docs/models/glm/glm4-5) |
| **Kimi** | [Kimi-K3](/docs/models/kimi/kimi-k3) [Kimi-K2.6](/docs/models/kimi/kimi-k2.5) [Kimi-K2.5](/docs/models/kimi/kimi-k2.5) [Kimi-K2-Instruct / Thinking](/docs/models/kimi/kimi-k2) |
| **Nemotron** | [Nemotron-3-Ultra-550B-A55B](/docs/models/nemotron/nemotron-3-ultra) [Nemotron-3-Super-120B-A12B-FP8](/docs/models/nemotron/nemotron-3-super) [Nemotron-3-Nano MoE](/docs/models/nemotron/nemotron-3-nano-moe) [Nemotron-3-Nano](/docs/models/nemotron/nemotron-3-nano) |
| **Gemma** | [Gemma-4 26B-A4B](/docs/models/gemma/gemma-4) [Gemma-4 31B](/docs/models/gemma/gemma-4) |
| **JoyAI** | [JoyAI-LLM-Flash](https://github.com/radixark/miles/blob/main/scripts/run_joy_ai_llm_flash.py) |
| **GPT-OSS** | [gpt-oss-20b](/docs/models/gpt-oss/gpt-oss) |
| **MiMo** | [MiMo-V2.6-Flash](/docs/models/mimo/mimo-v2-6-flash) |
| **StepFun** | [Step-5-Preview](/docs/models/stepfun/step5) |
| **Muse Spark** | [Muse Spark Open](/docs/models/muse-spark/muse-spark-open) |
| **Reflection** | [Beam-1.0-501B-A23B](/docs/models/reflection/beam) |

## Diffusion

| Family | Models |
| - | - |
| **Stable Diffusion** | [SD3.5](/docs/diffusion/models/sd3/sd3) |
| **Qwen-Image** | [Qwen-Image](/docs/diffusion/models/qwen-image/qwen-image) |
| **Wan** | [Wan2.2-T2V-A14B](/docs/diffusion/models/wan/wan2-2) |
| **LTX** | [LTX-2.3](/docs/diffusion/models/ltx/ltx2) |
| **Cosmos** | [Cosmos3-Nano](/docs/diffusion/models/cosmos/cosmos3) |
| **MiniMax** | [MiniMax H3](/docs/diffusion/models/h3/h3) |

## How a recipe is structured

Every recipe page follows the same six sections:

1. **Model Introduction** — what the model is and why miles supports it.
2. **Supported Variants** — model sizes + HF links.
3. **Environment Setup** — env vars, downloads, and HF → Megatron conversion.
4. **Launch** — the `scripts/run_<family>.py` invocation.
5. **Recipe Configuration** — parallelism, algorithm, rollout/SGLang, optimizer.
6. **Pairs Well With** — links to the advanced features that complement this recipe.

## Adding a new model

Miles's plugin architecture lets you wrap a HuggingFace implementation as a Megatron
module without patching Megatron core. See
[Backends Beyond Megatron](/docs/advanced/architecture-support) for the workflow.


# Kimi
Source: https://miles.radixark.com/docs/models/kimi/index

Miles recipes for the Moonshot family — Kimi K2.6 / K2.5 (multimodal, 1 T / 32 B-A) and Kimi K2 / K2-Thinking.

Miles supports Moonshot's MoE line from top to bottom. The latest Kimi K2.6 and K2.5 are natively multimodal agentic models at 1 T total / 32 B active per token, and the text-only Kimi K2 (Instruct and Thinking variants) runs at the same 1 T / 32 B scale. K2-Thinking is the canonical INT4 QAT target, and the K2.5 / K2.6 recipe trains an INT4 actor under the same QAT path.

## Variants

| Model | Active / Total | HF ID | Recipe |
| - | - | - | - |
| Kimi-K2.6 | 32 B / 1 T | `moonshotai/Kimi-K2.6` | [kimi-k2.5](/docs/models/kimi/kimi-k2.5) |
| Kimi-K2.5 | 32 B / 1 T | `moonshotai/Kimi-K2.5` | [kimi-k2.5](/docs/models/kimi/kimi-k2.5) |
| Kimi-K2-Instruct | 32 B / 1 T | `moonshotai/Kimi-K2-Instruct` | [kimi-k2](/docs/models/kimi/kimi-k2) |
| Kimi-K2-Thinking | 32 B / 1 T | `moonshotai/Kimi-K2-Thinking` | [kimi-k2](/docs/models/kimi/kimi-k2) |

## Fastest path to train

The single-node Kimi-K2.5 2-layer smoke test, before scaling the full model across many nodes:

```bash theme={null}
python scripts/run_kimi_k25.py full-train --model-name Kimi-K2.5-2layer --num-nodes 1
```

See the [Kimi K2.5](/docs/models/kimi/kimi-k2.5) page for the full multi-node recipe, or [Kimi K2](/docs/models/kimi/kimi-k2) for the 16-node K2-Thinking recipe (including the one-line `model_type` patch that lets Miles treat K2 as a DeepSeek-V3-shaped architecture).

## Which variant do I pick?

* **Latest multimodal agentic model** → Kimi-K2.6 or Kimi-K2.5 ([kimi-k2.5](/docs/models/kimi/kimi-k2.5)).
* **Frontier-scale instruction-tuned MoE** → Kimi-K2-Instruct ([kimi-k2](/docs/models/kimi/kimi-k2)).
* **Reasoning-style training, INT4 QAT target** → Kimi-K2-Thinking ([kimi-k2](/docs/models/kimi/kimi-k2)).


# Kimi K2
Source: https://miles.radixark.com/docs/models/kimi/kimi-k2

Launch recipes for Kimi-K2-Instruct and Kimi-K2-Thinking — 32 nodes × 8 GPU.

## 1. Model Introduction

[Kimi-K2](https://moonshotai.github.io/Kimi-K2/) is a state-of-the-art MoE language model from Moonshot AI with 32 B activated parameters and 1 T total parameters.

**Key highlights:**

* **Trillion-parameter MoE**: 1 T total / 32 B active per token, 61 layers (1 dense + the rest MoE), MLA attention shaped like DeepSeek-V3.
* **Instruct and Thinking variants**: Instruct is the general-purpose chat / agentic post-train; Thinking adds step-by-step reasoning with a 256 K context and ships in native INT4.
* **DeepSeek-V3-shaped architecture**: miles loads it through the DeepSeek-V3 path (one `sed` away), reusing the conversion + parallelism plumbing.
* **INT4 QAT target**: Kimi-K2-Thinking is the canonical reference recipe for INT4 QAT in miles.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Kimi-K2-Instruct | 32 B / 1 T | [moonshotai/Kimi-K2-Instruct](https://huggingface.co/moonshotai/Kimi-K2-Instruct) |
| Kimi-K2-Thinking | 32 B / 1 T | [moonshotai/Kimi-K2-Thinking](https://huggingface.co/moonshotai/Kimi-K2-Thinking) |

## 3. Environment Setup

### 3.1 Required env vars

```bash theme={null}
export MASTER_ADDR=<head node IP>
export MILES_SCRIPT_EXTERNAL_RAY=1
```

`MASTER_ADDR` reaches the ray workers as the torch-distributed rendezvous address; `MILES_SCRIPT_EXTERNAL_RAY=1` tells the launcher the cluster is already up (see §4.2). Paths are Typer flags: `--model-dir` (default `/root/models`), `--data-dir` (default `/root/datasets`), `--output-dir` (default `/root/shared_data`) — all three on a shared FS reachable from every node.

### 3.2 Download model + datasets

```bash theme={null}
hf download moonshotai/Kimi-K2-Instruct --local-dir /root/models/Kimi-K2-Instruct
hf download moonshotai/Kimi-K2-Thinking --local-dir /root/models/Kimi-K2-Thinking-fp8

hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.3 HF → Megatron `torch_dist` conversion

Convert across 4 nodes (mirror the DeepSeek-V3 procedure):

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py kimi-k2)" || exit 1   # or kimi-k2-thinking
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM/ torchrun \
   --nproc-per-node 8 \
   --master-addr ${MASTER_ADDR} --master-port 12345 \
   --nnodes=4 --node-rank ${NODE_RANK} \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Kimi-K2-Instruct/ \
   --save          /root/models/Kimi-K2_torch_dist/
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
export MASTER_ADDR=...

MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_kimi_k2.py --model-name Kimi-K2-Thinking   # or Kimi-K2-Instruct
```

Both variants are one launcher selected by `--model-name`. With `MILES_SCRIPT_EXTERNAL_RAY=1` it skips `ray start` and submits to an **already-running Ray cluster** (`ray job submit ...`).

### 4.2 Multi-node fan-out

Bring up Ray on every node before launching:

```bash theme={null}
# head
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
# each worker
ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP}
```

## 5. Recipe Configuration

### 5.1 Parallelism

Identical for both Instruct and Thinking:

| TP | PP | CP | EP | expert-TP | `decoder-last-pipeline-num-layers` | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - | - |
| 8 | 8 | 4 | 32 | 1 | 5 | 16384 | 256 (32 × 8) |

Both variants pass `--actor-num-nodes 32 --actor-num-gpus-per-node 8 --colocate --update-weight-buffer-size 2147483648` to `train.py`; `--num-nodes` overrides the node count.

### 5.2 Algorithm

| `--model-name` | Advantage | TIS |
| - | - | - |
| `Kimi-K2-Instruct` | GRPO (`--eps-clip 0.2 --eps-clip-high 0.28`) | – |
| `Kimi-K2-Thinking` | GRPO (`--eps-clip 0.2 --eps-clip-high 0.28`) | `--use-tis` |

Both use `--use-kl-loss --kl-loss-coef 0.00 --kl-loss-type low_var_kl --entropy-coef 0.00`.

Rollout shape (both):

```bash theme={null}
--rm-type math
--num-rollout 100
--rollout-batch-size 128
--n-samples-per-prompt 8
--rollout-max-response-len 32768   # Instruct
--rollout-max-response-len 16384   # Thinking
--over-sampling-batch-size 256
--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
--num-steps-per-rollout 4
--balance-data
```

DAPO-style dynamic sampling is on by default for both variants.

### 5.3 Rollout & SGLang

Identical for both:

```bash theme={null}
--rollout-num-gpus-per-engine 16
--sglang-mem-fraction-static 0.7

# dp attention
--sglang-enable-dp-attention
--sglang-dp-size 8
--sglang-moe-dense-tp-size 1
--sglang-enable-dp-lm-head

--sglang-ep-size 16

--sglang-server-concurrency 1024
```

SGLang DeepEP (`--sglang-moe-a2a-backend deepep`, `--sglang-deepep-mode`) is not enabled for either variant. Megatron-side `--moe-enable-deepep` and `--moe-token-dispatcher-type flex` are **on** for `Kimi-K2-Instruct` and **off** for `Kimi-K2-Thinking`.

### 5.4 Optimizer

CPU Adam is enabled in both:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

### 5.5 Notable quirks

* Instruct loads the BF16 HF release, Thinking the FP8 one (`Kimi-K2-Thinking-fp8`). Their `torch_dist` reference directories differ too — `Kimi-K2_torch_dist` vs `Kimi-K2-Thinking_torch_dist`.
* Neither variant passes `--global-batch-size`; the batch is driven by `--rollout-batch-size` and `--num-steps-per-rollout`.

## 6. Pairs Well With

* [PD Disaggregation](/docs/advanced/pd-disaggregation)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)
* [Fault Tolerance](/docs/advanced/fault-tolerance)
* [INT4 QAT](/docs/advanced/int4-qat)


# Kimi K2.5 / K2.6
Source: https://miles.radixark.com/docs/models/kimi/kimi-k2.5

Launch recipe for Kimi-K2.5, running full-parameter GRPO on 32 × 8 H200 with an INT4 actor and a BF16 reference.

The reference launcher is [`scripts/run_kimi_k25.py`](https://github.com/radixark/miles/blob/main/scripts/run_kimi_k25.py), which loads the K2.5 model definition from `scripts/models/kimi-k25.py`.

## 1. Model Introduction

[Kimi-K2.5](https://huggingface.co/moonshotai/Kimi-K2.5) is an open-source, natively multimodal agentic model from Moonshot AI. It is built by continual pretraining on roughly 15 T mixed vision and text tokens on top of Kimi-K2-Base, and it pairs a Mixture-of-Experts (MoE) language backbone with a MoonViT vision encoder so a single model handles both image and text inputs. K2.5 keeps the 1 T-total / 32 B-active shape of the K2 family and extends the context window to 256K tokens.

**Architecture at a glance** (from the model card):

| Specification | Value |
| - | - |
| Architecture | Mixture-of-Experts (MoE) |
| Total / activated parameters | 1 T / 32 B |
| Layers (1 dense + 60 MoE) | 61 |
| Attention hidden dimension | 7168 |
| MoE hidden dimension (per expert) | 2048 |
| Attention heads | 64 |
| Routed experts (selected per token) | 384 (top-8) |
| Shared experts | 1 |
| Attention mechanism | Multi-head Latent Attention (MLA) |
| Activation | SwiGLU |
| Vocabulary size | 160K |
| Context length | 256K |
| Vision encoder | MoonViT (400 M) |

**Key features:**

* **Native multimodality.** K2.5 is pretrained on both vision and language tokens, so it covers visual knowledge, cross-modal reasoning, and tool use grounded in images alongside text.
* **Coding with vision.** It generates code from visual specifications such as UI designs and video workflows, and drives tools for visual data processing.
* **Agent swarm.** It decomposes a complex task into parallel sub-tasks run by dynamically instantiated, domain-specific agents, rather than scaling a single agent.

## 2. Supported Variants

The K2.5 launcher expects two checkpoints under `$BASE_DIR`: an INT4 actor checkpoint and a BF16 reference checkpoint.

| Role | Checkpoint | Loaded with |
| - | - | - |
| Actor (trained) | `$BASE_DIR/Kimi-K2.5-int4` | `--hf-checkpoint` |
| Reference | `$BASE_DIR/Kimi-K2.5-bf16` | `--ref-load` |

Both share the K2 family's 1 T-total / 32 B-active MoE + MLA shape inherited from Kimi-K2-Thinking.

## 3. Quick start

### 3.1 Prerequisites

The launcher references two environment variables but never sets them, so you should export them yourself before launch:

```bash theme={null}
export BASE_DIR=<shared FS path, reachable from every node>
export MASTER_ADDR=<head node IP>
```

The `$BASE_DIR` directory must already hold the two K2.5 checkpoints from §2 alongside the DAPO-Math-17k training set (`dapo-math-17k/dapo-math-17k.jsonl`) and the AIME-2024 eval set (`aime-2024.jsonl`).

### 3.2 One-line launch

The launcher runs a `pkill` / `ray stop` cleanup pass first so a failed run can be re-launched cleanly, then starts the ray head and submits with `ray job submit --address http://127.0.0.1:8265`. Export `MILES_SCRIPT_EXTERNAL_RAY=1` to skip the `ray start` and submit to an **already-running Ray cluster** instead.

```bash theme={null}
cd /root/miles
export MASTER_ADDR=...
python scripts/run_kimi_k25.py train
```

`train` submits only. `prepare` does the download and the INT4 → BF16 dequantization; `full-train` runs both plus the training submit in one go, which is how the single-node 2-layer smoke test runs (`--model-name Kimi-K2.5-2layer --num-nodes 1`).

### 3.3 Multi-node fan-out

Bring up Ray on every node before launching, the same way as the other Kimi recipes:

```bash theme={null}
# head
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
# each worker
ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP}
```

## 4. Script breakdown

The launcher builds the flags it passes to `train.py` as one f-string group per concern. The model shape comes from `scripts/models/kimi-k25.py`.

That definition sets the MLA latent ranks (`q_lora_rank=1536`, `kv_lora_rank=512`, `qk_head_dim=128`, `qk_pos_emb_head_dim=64`, `v_head_dim=128`), the MoE routing (384 experts, top-8, sigmoid pre-softmax scoring, FP32 router, `--moe-router-topk-scaling-factor 2.827`), and YaRN RoPE (`--rope-type yarn`, `--rotary-base 50000`, `--rotary-scaling-factor 64.0`, `--original-max-position-embeddings 4096`, `--beta-fast 32`, `--beta-slow 1`). The K2.5 recipe then layers the following on top:

* **`ckpt_args`** wires up the dual checkpoint (INT4 actor via `--hf-checkpoint`, BF16 reference via `--ref-load`) together with `--megatron-to-hf-mode bridge` and `--model-name kimi_k25`.
* **`rollout_args`** and **`eval_args`** configure GRPO sampling and periodic AIME evaluation (covered in §5.2).
* **`perf_args`** sets the parallelism layout and recomputation (§5.1).
* **`grpo_args`** and **`optimizer_args`** set the algorithm and CPU-offloaded Adam (§5.2, §5.4).
* **`sglang_args`** configures the colocated rollout engine (§5.3).

The job runs colocated (`--colocate`) across 32 nodes (`--actor-num-nodes 32 --actor-num-gpus-per-node 8`) with `--update-weight-buffer-size $((4*512*1024*1024))`.

## 5. Example Recipe Configuration

### 5.1 Megatron Parallelism

This is the validated layout shipped with the launcher. All parallelisms are supported, so you can supply any other TP / EP / PP / CP combination that fits your compute.

| Hardware | Nodes × GPUs | TP | PP | CP | EP | expert-TP | `--decoder-last-pipeline-num-layers` | `--max-tokens-per-gpu` |
| - | - | - | - | - | - | - | - | - |
| H200 | 32 × 8 = 256 | 8 | 8 | 4 | 32 | 1 | 5 | 4096 |

Sequence parallelism (`--sequence-parallel`) is on, and the trainer uses dynamic batching (`--use-dynamic-batch-size`) capped at `--max-tokens-per-gpu 4096`. Recomputation is full and uniform over a single layer:

```bash theme={null}
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1
```

### 5.2 Algorithm

The recipe uses GRPO with KL and entropy losses disabled:

```bash theme={null}
--advantage-estimator grpo
--eps-clip 0.2 --eps-clip-high 0.28
--kl-loss-coef 0.00 --kl-loss-type low_var_kl
--entropy-coef 0.00
```

Rollouts draw from DAPO-Math-17k and score with the `deepscaler` reward:

```bash theme={null}
--prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl
--input-key prompt --label-key label
--apply-chat-template
--rollout-shuffle --balance-data
--rm-type deepscaler

--num-rollout 20
--rollout-batch-size 32
--n-samples-per-prompt 8
--rollout-max-response-len 16384
--rollout-temperature 1

--global-batch-size 256
--filter-zero-reward-samples
--use-dynamic-global-batch-size
```

Evaluation runs every 20 steps against AIME-2024, sampling 16 responses per prompt:

```bash theme={null}
--eval-interval 20
--eval-prompt-data aime $BASE_DIR/aime-2024.jsonl
--n-samples-per-eval-prompt 16
--eval-max-response-len 16384
--eval-top-p 1
```

### 5.3 Rollout & SGLang

The rollout engine is colocated with training, spanning 8 GPUs per engine with 8-way expert parallelism:

```bash theme={null}
SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 8
   --sglang-mem-fraction-static 0.7
   --sglang-ep-size 8
   --sglang-server-concurrency 1024
   --sglang-cuda-graph-bs-decode 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128
   --use-rollout-routing-replay
)
```

The `--use-rollout-routing-replay` flag replays the rollout-time MoE routing decisions during training so the two stages stay consistent. On the Megatron side, attention uses the Flash backend (`--attention-backend flash`).

The launcher sets the required env vars for you, including the INT4 QAT pair (`OPEN_TRAINING_INT4_FAKE_QAT_FLAG=1`, `OPEN_TRAINING_INT4_GROUP_SIZE=32`), a long NCCL timeout (`NCCL_TIMEOUT=3600`), `CUDA_DEVICE_MAX_CONNECTIONS=1`, and NVLink-gated NVLS (`NCCL_NVLS_ENABLE` follows the script's NVLink autodetection).

### 5.4 Optimizer

CPU-offloaded Adam is combined with the distributed optimizer:

```bash theme={null}
--optimizer adam
--lr 1e-6 --lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9 --adam-beta2 0.98

--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
--use-distributed-optimizer
```

Adam states live on host RAM and are D2H/H2D-overlapped with the backward pass, freeing GPU memory for the 1 T-parameter weight footprint. Gradients accumulate and all-reduce in FP32 (`--accumulate-allreduce-grads-in-fp32`), and the attention softmax also runs in FP32 (`--attention-softmax-in-fp32`).

## 6. Pairs Well With

* [INT4 QAT](/docs/advanced/int4-qat)
* [PD Disaggregation](/docs/advanced/pd-disaggregation)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)
* [Fault Tolerance](/docs/advanced/fault-tolerance)
* [Kimi K2](/docs/models/kimi/kimi-k2): sibling recipe; K2.5 reuses the K2-Thinking architecture.


# Kimi-K3
Source: https://miles.radixark.com/docs/models/kimi/kimi-k3

LoRA RL recipe for Kimi-K3, a KDA + MLA hybrid with an 896-expert latent MoE, trained colocated with SGLang.

The complete Kimi-K3 LoRA RL implementation is open at the Miles pull request:
[`radixark/miles#1825`](https://github.com/radixark/miles/pull/1825). The scripts and image
below come from that branch. Results and background are in the
[LMSYS day-0 write-up](https://www.lmsys.org/blog/2026-07-27-kimi-k3-day0-support).

## 1. Model Introduction

Kimi-K3 pairs two attention mechanisms in one stack, **KDA and MLA chosen per layer**,
with an **896-expert latent MoE** at top-16. The checkpoint ships in **MXFP4**.

miles trains it with **native LoRA adapters** rather than full fine-tuning, which is what
makes the recipe fit at all: the base weights stay frozen and only the adapters carry
gradients. The adapters are implemented under TP, EP, PP and CP, with **shared-A and
per-expert-B factors** across the 896 experts, and they are exported to the rollout engines
as HF-named chunks over CUDA IPC.

**Key highlights:**

* **Two attention types per layer**: KDA and MLA, with an attention-residual snapshot bank.
* **896-expert latent MoE**, top-16, `moe_latent_size=3584`, plus a shared expert.
* **LoRA RL, not full fine-tuning.** Rank 16 by default, 32 in the validated run.
* **Colocated rollout**: trainer and SGLang share the GPUs, with adapters synced over CUDA IPC.
* **MXFP4 checkpoint** upcast to BF16 once, offline.

## 2. Supported Variants

| `--model-name` | Layers | Purpose | GPUs |
| - | - | - | - |
| `Kimi-K3` | 93 | the release | 64 (16 × 4), validated |
| `Kimi-K3-4layer-64experts` | 4 (1 dense + 3 MoE), 64 routed experts | smoke test and CI, default | one node |
| `Kimi-K3-4layer` | 4 (1 dense + 3 MoE), all 896 experts | rollout TP16 and EP layouts of the release | one node; two for rollout TP16 |

The name sets the checkpoint paths under `--model-dir` and the `megatron_model_type`;
`--train-mode lora|full` picks the recipe.

Architecture, from `scripts/models/kimi-k3.py`: hidden 7168, FFN 33792, 96 attention heads,
`kv_channels=256`, MLA with `q_lora_rank=1536` / `kv_lora_rank=512` /
`qk_head_dim=128` / `qk_pos_emb_head_dim=64` / `v_head_dim=128`, 896 experts at
`moe_ffn_hidden_size=3072`, shared expert 6144, vocab 163840, no position embedding.

## 3. Environment Setup

Use the `radixark/miles:dev` image with the Megatron and SGLang changes from
[radixark/Megatron-LM#94](https://github.com/radixark/Megatron-LM/pull/94) and
[sgl-project/sglang#37704](https://github.com/sgl-project/sglang/pull/37704) until they merge.

The only external asset is the native MXFP4 checkpoint; the BF16 dequantization and the
`torch_dist` conversion derive from it. `scripts/run_kimi_k3.py` names the three by model:
`{model_dir}/{model_name}`, `{model_dir}/{model_name}-bf16` and
`{model_dir}/{model_name}-bf16_torch_dist`, each overridable with `--hf-checkpoint`,
`--bf16-checkpoint` and `--ref-load`.

### 3.1 Four-layer prune (one node)

```bash theme={null}
python scripts/run_kimi_k3.py prepare-download --model-name Kimi-K3-4layer-64experts --task gsm8k
python scripts/run_kimi_k3.py prepare-bf16 --model-name Kimi-K3-4layer-64experts
python scripts/run_kimi_k3.py prepare-torch-dist --model-name Kimi-K3-4layer-64experts
```

`Pinaster/Kimi-K3-4layer` is the first dense layer plus three MoE layers of the release;
`Pinaster/Kimi-K3-4layer-64experts` keeps only the first 64 routed experts of each MoE layer
(the router is sliced to match) so full-parameter training fits one node's host memory. The
`run-ci-model-scripts` recipes train the 64-expert prune; the 896-expert prune is for layouts
that depend on the release's expert count.

### 3.2 Full model

Download the release into `{model_dir}/Kimi-K3` and dequantize it shard by shard across the
nodes (`--shard-rank/--num-shards`, then `--finalize-only` once):

```bash theme={null}
python tools/convert_mxfp4_to_bf16.py --model-dir <native-mxfp4> --save-dir <bf16-hf> --device cuda \
    --shard-rank $RANK --num-shards $NUM_NODES
```

The `torch_dist` conversion runs on 32 ranks; the output re-shards at load, so the conversion
layout does not have to match the training one:

```bash theme={null}
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py kimi-k3)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
torchrun --nnodes=8 --nproc-per-node=4 ... \
    tools/convert_hf_to_torch_dist.py ${MODEL_ARGS[@]} \
    --hf-checkpoint <bf16-hf> --save <torch-dist-dcp> \
    --bf16 --tensor-model-parallel-size 32 --sequence-parallel \
    --pipeline-model-parallel-size 1 --context-parallel-size 1 \
    --expert-model-parallel-size 32 --expert-tensor-parallel-size 1 \
    --megatron-to-hf-mode raw
```

## 4. Launch

The launcher uses `all-linear`, which selects Kimi K3's HF defaults: MLA query/KV down projections, attention output projections, dense/shared MLPs, and routed experts. `attn,mlp` selects the same set. Explicit HF targets can omit `language_model.model.layers.*.block_sparse_moe.experts.*.w2` to leave routed-expert down projections frozen; other layouts remain unsupported by the native backend.

`--train-mode lora` (default) or `full`. Validated on **16 nodes × 4 GPUs**: one container per
node, a ray cluster across them, `export MILES_SCRIPT_EXTERNAL_RAY=1`, then:

```bash theme={null}
python scripts/run_kimi_k3.py train \
  --mode normal --model-name Kimi-K3 --train-mode lora --task dapo-math --reward-model deepscaler \
  --num-nodes 16 --num-gpus-per-node 4 \
  --pipeline-parallel-size 8 --context-parallel-size 2 \
  --rollout-tp-size 16 --rollout-max-concurrency 8 \
  --lora-rank 32 --lora-alpha 64 \
  --num-rollout 1000 --rollout-batch-size 8 --n-samples-per-prompt 8 \
  --rollout-max-response-len 4096 --sglang-max-total-tokens 65536 \
  --global-batch-size 64 --lr 1e-05 --eval-interval 10 \
  --distributed-timeout-minutes 60 \
  --hf-checkpoint <native-mxfp4> --ref-load <torch-dist-dcp> \
  --sglang-path /sgl-workspace/sglang/python \
  --data-dir <datasets> --enable-wandb
```

`--rollout-max-concurrency 8` is passed explicitly: the field default is 64, and the
validated runs pin 8. Off the validated 64 GPUs the full model needs `--tp-size-override`
(and `--ep-size-override`).

For a single-node smoke test use the default `--model-name Kimi-K3-4layer`: the trainer layout
follows the GPU count (TP8/EP8 on 8 GPUs) and the rollout TP/EP default to it; `--rollout-tp-size 16 --rollout-ep-size 1` on two nodes reproduces the TP16 Marlin layout of the full recipe.

## 5. Recipe Configuration

### 5.1 Parallelism

The resolved config at startup should show `expert_model_parallel_size 8`,
`max_tokens_per_gpu 8192`, `colocate_memory_peak_device gpu` and
`lora_base_cpu_backup True`. Checking those four lines is the fastest way to confirm the
run came up in the intended shape.

| Knob | Validated value |
| - | - |
| Pipeline parallel | 8 |
| Context parallel | 2 |
| Expert parallel | 8 (resolved) |
| Rollout TP | 16 |
| `max_tokens_per_gpu` | 8192 |

### 5.2 LoRA

Rank 32 / alpha 64 in the validated run; the script defaults to 16 / 32. Adapters attach to
attention output and both MLA down-projections, the dense MLP, and both expert projections:

```
self_attention.o_proj
self_attention.q_a_proj
self_attention.kv_a_proj_with_mqa
mlp.linear_fc1, mlp.linear_fc2
mlp.experts.linear_fc1, mlp.experts.linear_fc2
```

The 896 experts share one A factor and carry per-expert B factors, which is what keeps the
adapter count tractable at this expert width.

### 5.3 Rollout

Rollout is **colocated**: the trainer and SGLang share GPUs, and adapters sync over CUDA
IPC as HF-named chunks. `lora_base_cpu_backup` keeps a host copy of the frozen base so the
GPU copy can be reclaimed during rollout.

### 5.4 What a healthy run looks like

From the GB300 validation runs:

* 11 to 13 minutes per rollout cycle
* trainer allocated memory returns to about 91 GB after every weight sync
* `rollout/raw_reward` between 0.5 and 0.75 from rollout 0
* `eval/aime` 0.37 to 0.43 at eval\@0 (that spread is temperature-0 nondeterminism), rising
  by at least 0.06 by eval\@9; the measured run went 0.367 to 0.467

The memory figure is the one to watch. If allocated does not return to its baseline after a
weight sync, the adapter export is leaking and the run will die later rather than sooner.

### 5.5 Notable quirks

* The image preloads a small `shm_unlink` shim through `/etc/ld.so.preload`. It tolerates a
  benign PyTorch CUDA-IPC unlink race that otherwise aborts colocated weight sync at scale.
* Weight conversion for K3 lives in
  `miles/backends/megatron_utils/megatron_to_hf/kimi_k3.py`, and the model itself in
  `miles_plugins/models/kimi_k3/`.

## 6. Pairs Well With

* [LoRA](/docs/advanced/lora)
* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Kimi-K2.5](/docs/models/kimi/kimi-k2.5)


# MiMo-V2.6-Flash
Source: https://miles.radixark.com/docs/models/mimo/mimo-v2-6-flash

Launch recipe for Xiaomi's MiMo-V2.6-Flash-RL (309B MoE, hybrid sliding-window attention) — Megatron bridge mode on a BF16 conversion, on two 8-GPU nodes.

## 1. Model Introduction

[MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) is Xiaomi's 309B-parameter MoE model (15B active). miles trains its text decoder.

**Key highlights:**

* **Hybrid attention**: 39 sliding-window layers (window 128, learnable attention sink) and 9 global layers, with different KV head counts (8 vs 4) and RoPE bases (1e4 vs 1e7).
* **Asymmetric heads**: 192-dim Q/K and 128-dim V, V scaled by 0.707, partial RoPE (64 of 192 dims).
* **MoE**: 256 routed experts, top-8 sigmoid routing with a fixed score-correction bias; layer 0 is dense.
* **Compressed checkpoint**: FP8 block-quantized dense weights with a fused, kv-head-interleaved `qkv_proj`, and MXFP4 experts. The Megatron bridge reads a BF16 split-q/k/v conversion instead.

## 2. Supported Variants

| `--model-name` | What it is | HF source |
| - | - | - |
| `MiMo-V2.6-Flash-RL-bf16` | text decoder, 48 layers | [XiaomiMiMo/MiMo-V2.6-Flash-RL](https://huggingface.co/XiaomiMiMo/MiMo-V2.6-Flash-RL) |

The vision, audio and MTP modules stay in the checkpoint but are not trained (text-only RL).

## 3. Environment Setup

### 3.1 Download + BF16 conversion

`prepare` (run by the launcher before training) downloads the official checkpoint to `<model-dir>/MiMo-V2.6-Flash-RL` and converts it unless `<model-dir>/<model-name>` already holds a finished conversion:

```bash theme={null}
hf download XiaomiMiMo/MiMo-V2.6-Flash-RL --local-dir /root/models/MiMo-V2.6-Flash-RL
python tools/convert_mimo_v2_to_bf16.py --model-dir /root/models/MiMo-V2.6-Flash-RL \
    --save-dir /root/models/MiMo-V2.6-Flash-RL-bf16 --device cuda
```

The converter dequantizes each kv-head shard of the fused `qkv_proj` with its own FP8 block scales, splits it into q/k/v in head order, and decodes the MXFP4 experts. The BF16 model is 622 GB.

### 3.2 SGLang

The BF16 engine needs the following MiMo-V2 fixes. The `sglang-miles` branch has them since [`8035002`](https://github.com/sgl-project/sglang/commit/8035002c4ed0360579004791378e5b7d1740503b), and the `radixark/miles:dev` image, built from that branch, includes them:

* upstream [sgl-project/sglang#40448](https://github.com/sgl-project/sglang/pull/40448) (MXFP4 MoE and BF16 router);
* accepting the `split` attention layout of the BF16 conversion;
* allocating the sliding-window KV pools inside the memory-saver region, otherwise a colocated engine cannot release its KV cache;
* passing `layer_id` to the MoE top-k, otherwise `--use-rollout-routing-replay` crashes the CUDA-graph capture.

## 4. Launch

### 4.1 Two nodes

Join the second node to a ray head on the first, then launch from the head:

```bash theme={null}
MILES_SCRIPT_EXTERNAL_RAY=1 MASTER_ADDR=<head ip> python scripts/run_mimo_v2_6_flash.py \
    --mode rl --model-name MiMo-V2.6-Flash-RL-bf16 --num-nodes 2 --train-offload-disk-dir /scratch/offload
```

On two 8×H200 nodes with node-local NVMe (4-drive RAID0), a step with 16 samples of up to 2048 response tokens took about 4 minutes of training, of which about 3 minutes is streaming the optimizer state (72 GB read and 144 GB written per GPU), plus 30 s of weight update. Peak GPU memory was 125 GB and peak host memory 1.05 TB per node.

`--mode sft --prompt-data <chat jsonl>` trains SFT instead. It reads a `messages` column; put a reasoning trace in `reasoning_content`, since the MiMo chat template renders `<think>{reasoning_content}</think>`.

## 5. Recipe Configuration

### 5.1 Parallelism

| `--model-name` | TP | PP | CP | EP | expert-TP | GPUs | rollout engine |
| - | - | - | - | - | - | - | - |
| `MiMo-V2.6-Flash-RL-bf16` | 2 | 2 | 1 | 8 | 1 | 16 (2 × 8) | TP8, `--sglang-mem-fraction-static 0.8` |

TP is capped at 4 by the four global-attention KV heads. Context parallelism is not supported. `--sequence-parallel` is on whenever TP > 1. THD packing with `--use-dynamic-batch-size` is the default; `--qkv-format bshd` runs one sample per micro-batch.

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --entropy-coef 0.00` and `--rm-type math` on dapo-math-17k. No KL loss: bridge mode has no Megatron-format reference checkpoint for `--ref-load`.

### 5.3 Optimizer

Adam at `--lr 1e-6`. The model's Adam state (3.7 TB) fits neither the GPUs nor two hosts' memory, so it streams through node-local NVMe:

```bash theme={null}
--stream-optimizer-state-to-disk
--stream-optimizer-state-moment-dtype bf16
--grad-reduce-in-bf16
--offload-train-target disk
--offload-train-disk-dir <train-offload-disk-dir>
```

### 5.4 Notable quirks

* **Fused attention only**: `--attention-backend fused`. The learnable sink on the sliding-window layers is implemented by TE's FusedAttention (THD needs cuDNN ≥ 9.18) and the unfused path, not by FlashAttention.
* **Sliding window**: a SWA layer attends the query plus `sliding_window_size` (128) previous keys, 129 in total, as SGLang does; the bridge sets `window_size = (sliding_window_size, 0)`. The HF modeling code and Megatron-Bridge's `mimo_v2_flash` attend 128 keys in total (`sliding_window_size - 1` previous), so a log-prob comparison against them differs at positions past the window.
* **Dropout**: the bridge sets `hidden_dropout = 0`. Bridge mode does not copy `--hidden-dropout` onto the provider, so a provider left at Megatron's default (0.1) trains with dropout silently.
* **Frozen towers and `--check-weight-update-equal`**: the trainer never sends the vision/audio weights, so exclude them from the post-update check with `--check-weight-update-skip-list visual. audio_tokenizer. input_local_transformer. speech_embeddings. projection.mlp.`.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)


# Muse Spark Open
Source: https://miles.radixark.com/docs/models/muse-spark/muse-spark-open

Launch recipe for Muse Spark Open (816B omni MoE) — full-parameter GRPO on 64 H200 or 32 B300 GPUs, on math or on image-grounded math.

## 1. Model Introduction

[Muse Spark Open](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B) is an 816 B omni MoE model (42 B active) released in BF16. miles trains its 62-layer language model natively in Megatron.

**Key highlights:**

* **Hybrid attention**: 46 sliding-window layers (window 2048, RoPE) and 16 full-attention layers with no positional encoding; 128 query heads, 16 KV heads, head dim 64.
* **MoE in every layer**: 256 routed experts, top-8, no shared experts.
* **Gated residuals**: a learned per-channel gate mixes the residual stream and each post-normalized sublayer output.
* **Scaled embeddings and logits**: embeddings RMS-normalized and scaled, logits scaled and soft-capped at 20.
* **Frozen towers**: the vision and audio towers stay in the checkpoint but are not trained.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Muse Spark Open | 42 B / 816 B (62 layers) | [meta-models/Muse-Spark-Open-816B-A42B](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B) |

## 3. Environment Setup

### 3.1 Image and storage

```bash theme={null}
docker pull radixark/miles:muse-spark-open
```

Run it on every node — 8 × 8 H200 or 4 × 8 B300, on InfiniBand. `--model-dir`, `--data-dir` and `--output-dir` need shared storage; `--offload-train-disk-dir` needs node-local NVMe with about 1.2 TB free for the streamed Adam state, 2.5 TB on B300.

### 3.2 Download + `torch_dist` conversion

The launcher's `prepare` step downloads the checkpoint and the dataset, then converts the checkpoint on all nodes into `<model-dir>/Muse-Spark-Open-816B-A42B_torch_dist` (about 1.6 TB), skipping a finished conversion. With `--vlm` it prepares `hiyouga/geometry3k` instead of `zhuzilin/dapo-math-17k`.

### 3.3 SGLang

The image carries the [`sglang-miles-muse-spark-open`](https://github.com/sgl-project/sglang/tree/sglang-miles-muse-spark-open) branch, which serves the omni checkpoint natively. Math reasoning is validated on H200 and B300, image-grounded math on H200 only; B200 is untested.

## 4. Launch

### 4.1 Quick start

Start one Ray cluster across the nodes, then launch from node 0:

```bash theme={null}
# node 0
ray start --head --port 6379 --num-gpus 8
# other nodes
ray start --address <node-0-address>:6379 --num-gpus 8

cd /root/miles
# math reasoning on DAPO-Math-17k
MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_muse_spark_64xh200.py

# image-grounded math on Geometry3K
MILES_SCRIPT_EXTERNAL_RAY=1 python scripts/run_muse_spark_64xh200.py \
  --vlm --use-dynamic-sampling-filter --reasoning-effort medium --eval-interval 5
```

`scripts/run_muse_spark_32xb300.py` is the same recipe on 4 × 8 B300 and takes the same options.

| Mode | GPUs | Rollout | Training | Weight sync | `train/train_rollout_logprob_abs_diff` |
| - | - | - | - | - | - |
| Math reasoning | 64 H200 | 2–5.5 min | 3–5.5 min | about 50 s | 0.007 |
| Image-grounded math | 64 H200 | 8–14 min | about 5 min | about 50 s | 0.013 |
| Math reasoning | 32 B300 | 1.7–3.5 min | 8.5–12.5 min (18.5 min at step 0) | about 90 s | 0.004–0.007 |

The log-prob column assumes [R3](/docs/advanced/miles-router) (`--use-rollout-routing-replay`, on in this recipe), which replays the rollout's expert choices so routing noise cannot widen the gap. Reward starts at 0.65–0.75 and reaches 0.8–0.95 within 10 steps on math reasoning; every engine holds bit-identical weights after each sync.

### 4.2 Save and export

`--save-interval N` saves the weights, without optimizer state, to `<output-dir>/<run-id>/checkpoints`, and the launcher converts the last save into a Hugging Face checkpoint at `<output-dir>/<run-id>/hf`. Each save and each conversion takes about 1.6 TB. Save directories count steps from 0, so `iter_0000009` below is the save after step 10:

```bash theme={null}
python tools/convert_muse_spark_torch_dist_to_hf.py \
  --input-dir <output-dir>/<run-id>/checkpoints/iter_0000009 \
  --hf-checkpoint /root/models/Muse-Spark-Open-816B-A42B \
  --output-dir <output-dir>/<run-id>/hf-step-10
```

## 5. Recipe Configuration

### 5.1 Parallelism

| Launcher | TP | PP | CP | EP | expert-TP | first/last stage layers | `max_tokens_per_gpu` |
| - | - | - | - | - | - | - | - |
| `run_muse_spark_64xh200.py` | 8 | 8 | 1 | 8 | 1 | 7 | 4096 |
| `run_muse_spark_32xb300.py` | 8 | 4 | 1 | 8 | 1 | 15 | 8192 |

`--sequence-parallel`, THD packing (`--use-dynamic-batch-size`) and full recomputation are on.

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28`, no KL or entropy term, per-token loss, and R3 (`--use-rollout-routing-replay`). Each step samples 32 prompts × 16 responses of up to 4096 tokens (8192-token context) at temperature 1. `--rm-type math` rewards a matching final `\boxed{}` answer, and `--reasoning-effort` defaults to `xhigh`.

### 5.3 Rollout & SGLang

Rollout runs on 4 colocated engines, up to 64 concurrent requests each on H200 and 128 on B300:

```bash theme={null}
--sglang-ep-size 16                # 16 GPUs per engine on H200; 8 (one node) on B300
--sglang-kv-cache-dtype bfloat16   # FP8 KV widens the rollout/training log-prob gap
```

### 5.4 Optimizer

Adam at `--lr 1e-6`, constant, with weight decay 0.1 and no gradient clipping. The FP32 main weights and moments do not fit the GPUs, so they stream through node-local NVMe:

```bash theme={null}
--stream-optimizer-state-to-disk
--stream-optimizer-state-moment-dtype fp32
--offload-train-target disk
--offload-train-disk-chunk-mb 256
```

`--optimizer muon` switches to Muon (`dist_muon`), which keeps its state and the FP32 main weights in chunked CPU memory (`--chunked-optimizer-state-offload --chunked-muon-master-weights`). Megatron cannot save that state, so a Muon run also needs `--no-load-optim --no-save-optim`.

### 5.5 Image-grounded math (`--vlm`)

`--vlm` trains on Geometry3K. The vision tower stays frozen on the first pipeline stage, outside the optimizer and the checkpoints, and the weight sync re-sends it after every update. `--eval-interval N` scores the Geometry3K test split (601 problems, 4 samples each) and requires `--vlm`.

## 6. Pairs Well With

* [Disk Offload](/docs/advanced/disk-offload)
* [Rollout Routing Replay (R3)](/docs/advanced/miles-router)


# Nemotron
Source: https://miles.radixark.com/docs/models/nemotron/index

Miles recipes for NVIDIA's Nemotron-3 family — Mamba+Attention(+MoE) hybrids loaded via Megatron AutoBridge.

Miles supports NVIDIA's Nemotron-3 line: a Mamba + Attention hybrid that, in the Super tier, adds MoE and ships natively in FP8. All three variants load via the Megatron AutoBridge path, so there is no offline HF → `torch_dist` conversion step.

## Variants

| Model | Active / Total | HF ID | Recipe |
| - | - | - | - |
| Nemotron-3-Nano | 4 B / 4 B (dense) | `nvidia/Nemotron-3-Nano-4B` | [nemotron-3-nano](/docs/models/nemotron/nemotron-3-nano) |
| Nemotron-3-Nano MoE | 3 B / 30 B | `nvidia/Nemotron-3-Nano-30B-A3B` | [nemotron-3-nano-moe](/docs/models/nemotron/nemotron-3-nano-moe) |
| Nemotron-3-Super | 12 B / 120 B (FP8) | `nvidia/Nemotron-3-Super-120B-A12B-FP8` | [nemotron-3-super](/docs/models/nemotron/nemotron-3-super) |
| Nemotron-3-Ultra | 55 B / 550 B | `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16` | [nemotron-3-ultra](/docs/models/nemotron/nemotron-3-ultra) |

## Fastest path to train

Nemotron-3-Nano (dense, 4 B) is the smallest and runs on a single 8-GPU node:

```bash theme={null}
python scripts/run_nemotron_3_nano.py --model-name NVIDIA-Nemotron-3-Nano-4B-BF16
```

See the [Nemotron-3-Nano](/docs/models/nemotron/nemotron-3-nano) page for the dense walkthrough, [Nemotron-3-Nano MoE](/docs/models/nemotron/nemotron-3-nano-moe) for the 30 B MoE variant, and [Nemotron-3-Super](/docs/models/nemotron/nemotron-3-super) for the FP8-native 120 B-A12B recipe.

## Which variant do I pick?

* **Smallest, single-node smoke test** → Nemotron-3-Nano ([nemotron-3-nano](/docs/models/nemotron/nemotron-3-nano)).
* **Mid-scale hybrid MoE** → Nemotron-3-Nano MoE ([nemotron-3-nano-moe](/docs/models/nemotron/nemotron-3-nano-moe)).
* **Frontier-scale FP8-native MoE** → Nemotron-3-Super ([nemotron-3-super](/docs/models/nemotron/nemotron-3-super)).
* **Largest, latent MoE across 16 nodes** → Nemotron-3-Ultra ([nemotron-3-ultra](/docs/models/nemotron/nemotron-3-ultra)).

## Pairs well with

* [Backends Beyond Megatron](/docs/advanced/architecture-support) — the AutoBridge path Nemotron rides on.
* [Low Precision RL](/docs/advanced/low-precision) — Super ships natively in FP8.


# Nemotron-3-Nano
Source: https://miles.radixark.com/docs/models/nemotron/nemotron-3-nano

Launch recipe for the dense NVIDIA Nemotron-3-Nano-4B (Mamba+Attention hybrid) via Megatron AutoBridge.

## 1. Model Introduction

[NVIDIA Nemotron-3-Nano-4B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16)
is a dense `nemotron_h` hybrid model — interleaved Mamba and attention blocks
with squared-relu FFNs, no RoPE, and a 262 144 max position. miles wires it via
the `megatron.bridge` AutoBridge path, so there is **no `torch_dist` conversion
step**: the AutoBridge constructs the full Megatron provider from the HF
`config.json` at load time, including all Mamba-specific fields
(`mamba_num_heads`, `mamba_state_dim`, `hybrid_override_pattern`, etc.).

**Key highlights:**

* **Hybrid architecture**: Mamba + attention layers (`nemotron_h` family).
* **Bridge-mode load**: `--megatron-to-hf-mode bridge` — no separate Megatron checkpoint.
* **No RoPE**: `--position-embedding-type none`.
* **Vocab**: 131 072 tokens, padded to a multiple of 128.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Nemotron-3-Nano-4B | 4 B / 4 B | [nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
hf download nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 --local-dir /root/models/NVIDIA-Nemotron-3-Nano-4B-BF16
```

Those are the `--data-dir` and `--model-dir` defaults; point them elsewhere with the flags.

### 3.2 No `torch_dist` conversion

AutoBridge loads the HF checkpoint directly. Both `--hf-checkpoint` and
`--ref-load` point at the HF directory, and `--megatron-to-hf-mode bridge`
turns on the bridge code path:

```bash theme={null}
--hf-checkpoint <model-dir>/NVIDIA-Nemotron-3-Nano-4B-BF16
--ref-load      <model-dir>/NVIDIA-Nemotron-3-Nano-4B-BF16
--megatron-to-hf-mode bridge
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_nemotron_3_nano.py --model-name NVIDIA-Nemotron-3-Nano-4B-BF16
```

`scripts/run_nemotron_3_nano.py` covers both Nano variants; `--model-name` picks the
dense 4 B here and the MoE 30B-A3B on the [sibling page](/docs/models/nemotron/nemotron-3-nano-moe).
The recipe targets 1 node × 8 GPU (H100/H200), default cell `TP=2 PP=2`, and it is a
10-step smoke test (`--num-rollout`). Checkpoints are written under `--output-dir`
(default `/root/shared_data`).

## 5. Recipe Configuration

### 5.1 Parallelism

The recipe ships a starting cell of `TP=2 PP=2`. Other verified cells (10-step RL
smoke tests, max train/rollout logprob diff): TP=2, TP=4, PP=2, CP=2, TP=2×PP=2.
Change the `tp` / `pp` / `cp` fields of the recipe to switch.

| Cell | TP | PP | CP | EP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| **default (launcher)** | 2 | 2 | 1 | 1 | 9216 | 8 (1 × 8) |
| TP=2 | 2 | 1 | 1 | 1 | 9216 | 8 |
| TP=4 | 4 | 1 | 1 | 1 | 9216 | 8 |
| CP=2 | 1 | 1 | 2 | 1 | 9216 | 8 |

`--sequence-parallel` is not enabled in the dense smoke recipe; activation
checkpointing is also off. Dense Nemotron-3-Nano has no expert parallelism.

### 5.2 Algorithm

GRPO with low-variance KL:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 1
--sglang-mem-fraction-static 0.7
```

### 5.4 Optimizer

GPU Adam in the smoke recipe (no `--optimizer-cpu-offload`). Switch on CPU Adam if
memory pressure rises.

### 5.5 Notable quirks

From `scripts/models/nemotron-3-nano-4b.py` and `scripts/run_nemotron_3_nano.py`:

* **No `--spec`**: the AutoBridge synthesizes the Megatron spec from HF config.
* `--position-embedding-type none` (no RoPE).
* `--vocab-size 131072 --make-vocab-size-divisible-by 128`.
* `--attention-backend auto` (the Mamba layers select their own kernel; flash-only is not safe here).
* Bridge load is required for hybrid `nemotron_h`: the AutoBridge wires `mamba_num_heads`, `mamba_state_dim`, `hybrid_override_pattern`. PP additionally needs miles' PP-unwrap shim (already on the `feat/nemotron-gemma4-rl` branch).

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for the AutoBridge wiring.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# Nemotron-3-Nano MoE
Source: https://miles.radixark.com/docs/models/nemotron/nemotron-3-nano-moe

Launch recipe for NVIDIA Nemotron-3-Nano-30B-A3B (Mamba+Attention+MoE hybrid) via Megatron AutoBridge.

## 1. Model Introduction

[NVIDIA Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16)
is a hybrid Mamba + attention + MoE model. It pairs the `nemotron_h` block
pattern from the dense 4B with a 128-expert sparse layer (top-6 routing,
1 shared expert, DSv3-style sigmoid routing with `routed_scaling_factor=2.5`).

miles loads it through the `megatron.bridge` AutoBridge with a custom
**NemotronH MoE bridge shim** (`miles_plugins/megatron_bridge/nemotron_h.py`) that
wires `routed_scaling_factor`, `n_group`, and `topk_group` onto the Megatron
provider. Without the shim the routed output is silently scaled 1.0× → \~0.28
logprob drift between train and rollout.

**Key highlights:**

* **Hybrid + MoE**: Mamba + attention + sparse MoE in the `nemotron_h` family.
* **128 experts, top-6 routing**, 1 shared expert (3712-dim), aux-free expert-bias load balancing.
* **Sigmoid routing** with `--moe-router-topk-scaling-factor 2.5`.
* **Bridge-mode load**: `--megatron-to-hf-mode bridge` — no `torch_dist` conversion step.
* **No RoPE**: `--position-embedding-type none`.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Nemotron-3-Nano-30B-A3B | 3 B / 30 B | [nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 --local-dir /root/models/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
```

Those are the `--data-dir` and `--model-dir` defaults; point them elsewhere with the flags.

### 3.2 No `torch_dist` conversion

AutoBridge + the NemotronH MoE shim load the HF checkpoint directly. Both
`--hf-checkpoint` and `--ref-load` point at the HF directory:

```bash theme={null}
--hf-checkpoint <model-dir>/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
--ref-load      <model-dir>/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
--save          <output-dir>/checkpoints
--save-interval 20
--megatron-to-hf-mode bridge
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_nemotron_3_nano.py --model-name NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
```

`scripts/run_nemotron_3_nano.py` is shared with the dense
[Nemotron-3-Nano-4B](/docs/models/nemotron/nemotron-3-nano); `--model-name` picks the MoE
recipe here. It targets 1 node × 8 GPU (H200), default cell `TP=2 PP=2 EP=2`, and is a
10-step smoke test (`--num-rollout`).

## 5. Recipe Configuration

### 5.1 Parallelism

Default cell is `TP=2 PP=2 EP=2`. Other verified cells from the upstream PR
(10-step RL smoke, max logprob diff ≈ 0.014):

| Cell | TP | PP | CP | EP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| **default (launcher)** | 2 | 2 | 1 | 2 | 1024 | 8 (1 × 8) |
| EP=4 | 1 | 1 | 1 | 4 | 1024 | 8 |
| TP=2×EP=4+SP | 2 | 1 | 1 | 4 | 1024 | 8 |
| PP=2×EP=4 | 1 | 2 | 1 | 4 | 1024 | 8 |
| CP=2×EP=4 | 1 | 1 | 2 | 4 | 1024 | 8 |
| TP=2×PP=2×EP=2+SP | 2 | 2 | 1 | 2 | 1024 | 8 |

`--sequence-parallel` is enabled by the recipe. Activation checkpointing is on
(`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`).
`--log-probs-chunk-size 128` is required for the smoke memory budget.

### 5.2 Algorithm

GRPO with low-variance KL:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 1
--sglang-mem-fraction-static 0.7
# Replay the exact rollout routing during training forward so
# train logprobs match rollout logprobs (needed for MoE).
--use-rollout-routing-replay
```

The `--use-rollout-routing-replay` flag is what keeps train and rollout
logprobs aligned for the sigmoid-routed MoE — drop it and you'll see the same
\~0.28 drift the bridge shim was added to fix.

### 5.4 Optimizer

GPU Adam in the smoke recipe (no `--optimizer-cpu-offload`). Switch on CPU Adam if
memory pressure rises.

### 5.5 Notable quirks

From `scripts/models/nemotron-3-nano-30b-a3b.py` and `scripts/run_nemotron_3_nano.py`:

* **No `--spec`**: AutoBridge + the NemotronH shim synthesize the Megatron MoE spec from HF config.
* 128 experts, `--moe-router-topk 6`, shared expert (3712-dim).
* Routing: `--moe-router-score-function sigmoid --moe-router-pre-softmax --moe-router-topk-scaling-factor 2.5`.
* Group routing: `--moe-router-num-groups 1 --moe-router-group-topk 1` (no-op for `n_group=1`, kept for parity with HF config).
* Aux-free balancing: `--moe-router-enable-expert-bias --moe-router-load-balancing-type seq_aux_loss --moe-router-bias-update-rate 0 --moe-aux-loss-coeff 0`.
* `--moe-grouped-gemm`, `--moe-router-dtype fp32`.
* `--position-embedding-type none`, `--vocab-size 131072 --make-vocab-size-divisible-by 128`.
* `--attention-backend auto` (Mamba layers select their own kernel).

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for how the bridge
shim layers `routed_scaling_factor` / `n_group` / `topk_group` onto the Megatron provider.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)
* [Low Precision RL](/docs/advanced/low-precision)

## Checkpoint compatibility

Megatron versions exposing `mixer.conv1d_weight` and `mixer.conv1d_bias`
require the corresponding bridge mappings. The Miles Nemotron-H shim
supports these names alongside the older `mixer.conv1d.weight` and
`mixer.conv1d.bias` names. Missing mappings leave these parameters unloaded
and omit them from rollout weight updates.


# Nemotron-3-Super
Source: https://miles.radixark.com/docs/models/nemotron/nemotron-3-super

Launch recipe for NVIDIA Nemotron-3-Super-120B-A12B-FP8 (Mamba+Attention+MoE hybrid, FP8 native) via Megatron AutoBridge.

## 1. Model Introduction

[NVIDIA Nemotron-3-Super-120B-A12B-FP8](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8)
is the Super-tier sibling of Nemotron-3-Nano: the same `nemotron_h` block
pattern (interleaved Mamba and attention blocks, no RoPE, squared-relu FFNs)
scaled to **120 B total / 12 B active** with a sparse MoE FFN, and shipped as
an **FP8-native** checkpoint.

miles loads it through the `megatron.bridge` AutoBridge with the shared
**NemotronH MoE bridge shim** (`miles_plugins/megatron_bridge/nemotron_h.py`)
that wires `routed_scaling_factor`, `n_group`, and `topk_group` onto the
Megatron provider — without the shim the routed output is silently scaled 1.0×,
the same drift class that affects the Nano-MoE recipe.

**Key highlights:**

* **Hybrid + MoE**: Mamba + attention + sparse MoE in the `nemotron_h` family.
* **FP8 native**: weights ship in FP8; load + train without an offline upcast.
* **Sigmoid routing** with per-token group selection, aux-free expert-bias load balancing.
* **Bridge-mode load**: `--megatron-to-hf-mode bridge` — no `torch_dist` conversion step.
* **No RoPE**: `--position-embedding-type none`.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Nemotron-3-Super-120B-A12B-FP8 | 12 B / 120 B | [nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 \
   --local-dir /root/models/NVIDIA-Nemotron-3-Super-120B-A12B-FP8
```

Those are the `--data-dir` and `--model-dir` defaults; on the cluster they are
`/cluster_public/miles_data/datasets` and `/cluster_public/miles_data/models`. `--model-name`
names the checkpoint directory inside `--model-dir` and defaults to the BF16 release, so pass
`--model-name NVIDIA-Nemotron-3-Super-120B-A12B-FP8` for the FP8 one.

### 3.2 No `torch_dist` conversion

AutoBridge + the NemotronH MoE shim load the FP8 HF checkpoint directly. Both
`--hf-checkpoint` and `--ref-load` point at the HF directory:

```bash theme={null}
--hf-checkpoint <model-dir>/<model-name>
--ref-load      <model-dir>/<model-name>
--save          <output-dir>/checkpoints
--save-interval 20
--megatron-to-hf-mode bridge
```

## 4. Launch

### 4.1 Quick start

The two pods take different roles, and each runs exactly one command:

```bash theme={null}
cd /root/miles

# on each worker pod — joins the head's ray cluster and blocks
python scripts/run_nemotron_3_super_120b_a12b.py worker --head-ip <head_ip>

# on the head pod — starts the ray head, waits for both nodes' GPUs, submits
python scripts/run_nemotron_3_super_120b_a12b.py train --head-ip <head_ip>
```

The recipe targets **2 nodes × 8 GPU (H200, FP8)** (`--num-nodes` / `--num-gpus-per-node`).
Default cell is `TP=4 PP=2 EP=8`, and it is a 10-step smoke test (`--num-rollout`).

## 5. Recipe Configuration

### 5.1 Parallelism

Default cell is `TP=4 PP=2 EP=8` on 16 GPUs. The 120B-A12B footprint requires
either a wider EP fan-out or PP=2 to fit the activation memory of the hybrid
Mamba+Attention stack at the FP8 weight resolution.

| Cell | TP | PP | CP | EP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| **default (launcher)** | 4 | 2 | 1 | 8 | 1024 | 16 (2 × 8) |
| TP=4×EP=8 (1 node) | 4 | 1 | 1 | 8 | 1024 | 8 (1 × 8) |
| TP=2×PP=2×EP=8 + SP | 2 | 2 | 1 | 8 | 1024 | 16 (2 × 8) |

`--sequence-parallel` is enabled by the recipe. Activation checkpointing is
on (`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`). `--log-probs-chunk-size 128` keeps the smoke
memory budget intact for FP8.

### 5.2 Algorithm

GRPO with low-variance KL — same defaults as Nano-MoE:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 8     # follows --num-gpus-per-node
--sglang-mem-fraction-static 0.7
# Replay the exact rollout routing during training forward so
# train logprobs match rollout logprobs (needed for sigmoid-routed MoE).
--use-rollout-routing-replay
```

The `--use-rollout-routing-replay` flag keeps train and rollout
logprobs aligned for the sigmoid-routed MoE — the same routing-replay rule that
Nano-MoE uses.

### 5.4 Optimizer

CPU Adam (`--optimizer-cpu-offload`) is the default for the 120B-A12B smoke
recipe; the FP8 weights save GPU memory but the Adam states still dominate at
this scale.

### 5.5 Notable quirks

* **FP8 native load**: the HF checkpoint is FP8; the bridge passes the
  per-block scales through to Megatron — no offline upcast step.
* **No `--spec`**: AutoBridge + the NemotronH shim synthesize the Megatron MoE
  spec from HF config.
* Routing follows the family default: `--moe-router-score-function sigmoid --moe-router-pre-softmax --moe-router-topk-scaling-factor 2.5`.
* Aux-free balancing: `--moe-router-enable-expert-bias --moe-router-load-balancing-type seq_aux_loss --moe-router-bias-update-rate 0 --moe-aux-loss-coeff 0`.
* `--moe-grouped-gemm`, `--moe-router-dtype fp32`.
* `--position-embedding-type none`, `--vocab-size 131072 --make-vocab-size-divisible-by 128`.
* `--attention-backend auto` (Mamba layers select their own kernel).

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for how
the bridge shim layers `routed_scaling_factor` / `n_group` / `topk_group` onto
the Megatron provider, and [Low Precision RL](/docs/advanced/low-precision)
for the FP8 weight format.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Low Precision RL](/docs/advanced/low-precision)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# Nemotron-3-Ultra
Source: https://miles.radixark.com/docs/models/nemotron/nemotron-3-ultra

Launch recipe for NVIDIA Nemotron-3-Ultra-550B-A55B (hybrid Mamba2 + Attention + latent-MoE) via Megatron AutoBridge.

## 1. Model Introduction

[NVIDIA Nemotron-3-Ultra-550B-A55B](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16)
is the Ultra tier of the `nemotron_h` family: the same hybrid block pattern as
Nano and Super, scaled to **550 B total / 55 B active** across 108 layers, with
a **latent MoE** (512 experts, top-22, `moe_latent_size=2048`) and one shared
expert.

miles loads it through the `megatron.bridge` AutoBridge with the shared
NemotronH MoE shim (`miles_plugins/megatron_bridge/nemotron_h.py`), the same
path Super-120B uses. There is no offline `torch_dist` conversion.

**Key highlights:**

* **Hybrid + latent MoE**: Mamba2 and attention blocks with a latent-projection
  MoE FFN, 512 experts at top-22.
* **Bridge-mode load**: `--megatron-to-hf-mode bridge`, straight from the HF
  checkpoint.
* **Sigmoid routing** with aux-free expert-bias load balancing, plus an MTP head
  in the checkpoint.
* **A single-node 4-layer slice** is published alongside the full model, so the
  recipe can be smoke-tested without 16 nodes.

## 2. Supported Variants

| Model | Active / Total | Layers | GPUs | HF ID |
| - | - | - | - | - |
| Nemotron-3-Ultra-550B-A55B-BF16 | 55 B / 550 B | 108 | 128 (16 × 8) | [nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) |
| Nemotron-3-Ultra-550B-A55B-BF16-4layer | slice | 4 | 8 (1 × 8) | pruned slice of the above, for smoke tests |

Tested on H200. Use the `radixark/miles:dev` image.

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 \
   --local-dir /root/models/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16
```

`--model-dir` and `--data-dir` default to `/root/models` and `/root/datasets`.
`--model-name` names the checkpoint directory inside `--model-dir`; passing a
name matching `<N>layer` selects the pruned slice and switches the recipe to its
single-node parallelism automatically.

### 3.2 No `torch_dist` conversion

AutoBridge plus the NemotronH shim read the HF checkpoint directly, so
`--hf-checkpoint` and `--ref-load` both point at the download:

```bash theme={null}
--hf-checkpoint <model-dir>/<model-name>
--ref-load      <model-dir>/<model-name>
--megatron-to-hf-mode bridge
```

## 4. Launch

### 4.1 Single-node smoke test

```bash theme={null}
cd /root/miles
python scripts/run_nemotron_3_ultra_550b_a55b.py full-train \
    --model-name NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-4layer --num-nodes 1
```

### 4.2 Full model

The full 108-layer model needs **16 nodes × 8 GPU**. Bring up the ray cluster
yourself, tell the launcher it is external, and submit from the head — the
launcher has no worker-side subcommand:

```bash theme={null}
# on the head pod
ray start --head --num-gpus 8 --disable-usage-stats
# on every worker pod
ray start --address=${HEAD_IP}:6379 --num-gpus 8 --disable-usage-stats

# then, on the head pod
export MILES_SCRIPT_EXTERNAL_RAY=1
export RAY_ADDRESS=http://${HEAD_IP}:8265
python scripts/run_nemotron_3_ultra_550b_a55b.py train --num-nodes 16
```

Without `MILES_SCRIPT_EXTERNAL_RAY=1` the launcher runs `ray stop --force` and
starts a fresh single-node head, tearing down the cluster the workers joined.

The recipe defaults to a 30-rollout run (`--num-rollout`), rollout batch 32 at 8
samples per prompt, global batch 128.

## 5. Recipe Configuration

### 5.1 Parallelism

| Variant | TP | PP | CP | EP | ETP | GPUs |
| - | - | - | - | - | - | - |
| Full 108-layer | 8 | 4 | 1 | 32 | 1 | 128 (16 × 8) |
| 4-layer slice | 1 | 1 | 1 | 8 | 1 | 8 (1 × 8) |

**Mamba `n_groups=8` caps attention and Mamba tensor parallelism at 8**, because
Megatron requires `n_groups % tp == 0`. That constraint drives the whole layout:
TP cannot grow past 8, so scale comes from PP and EP instead.

The 4-layer slice fits on one node, so it gives every rank to expert parallelism
(512 experts over EP=8 is 64 per rank) and keeps attention and Mamba at TP=1.

`--sequence-parallel` is on when TP > 1. Activation checkpointing is enabled, and
`--log-probs-chunk-size 128` keeps the log-prob pass inside the memory budget.

### 5.2 Algorithm

GRPO with low-variance KL:

```bash theme={null}
--advantage-estimator grpo
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
--rm-type deepscaler
```

### 5.3 Rollout & SGLang

| Variant | GPUs per engine | `--sglang-ep-size` | `--sglang-dp-size` | mem fraction |
| - | - | - | - | - |
| Full 108-layer | 32 | 32 | 4 | 0.7 |
| 4-layer slice | 8 | 8 | 2 | 0.6 |

The 550 B model is roughly 1.1 TB in BF16 and **does not fit one 8-GPU engine**,
so rollout runs 32-GPU engines with EP=32 and DP-attention. The DP size is chosen
so that `attn_tp = gpus_per_engine / dp_size` lands on 8, satisfying the same
Mamba `n_groups` constraint the training side has. The launcher asserts both
divisibility rules rather than letting a bad combination fail deep in SGLang.

### 5.4 Optimizer

CPU Adam, with the host transfer overlapped:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
--lr 1e-6
```

### 5.5 Notable quirks

* **Routing replay is not enabled yet.** `--use-rollout-routing-replay` is off
  for the 108-layer model: the routing capturer needs a fix for per-layer top-22
  under DP-attention. Train and rollout log-probs differ by about 0.01 without
  it. The Super-120B recipe does enable it.
* **No `--spec`**: AutoBridge and the NemotronH shim synthesize the Megatron MoE
  spec from the HF config.
* The shim is what wires `routed_scaling_factor`, `n_group` and `topk_group`
  onto the Megatron provider. Without it the routed output is silently scaled
  1.0×, the same drift class the Nano-MoE and Super recipes call out.
* The checkpoint carries an MTP head (`num_nextn_predict_layers=1`); the RL
  recipe does not train it.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Nemotron-3-Super](/docs/models/nemotron/nemotron-3-super)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# Qwen
Source: https://miles.radixark.com/docs/models/qwen/index

Miles recipes for the full Qwen3, Qwen3.5, and Next line — dense and MoE, through the Qwen4-architecture Qwen3.8-Flash-Next.

Miles ships ready-to-run RL recipes for every generation of the Qwen line: the dense Qwen3 series (0.6 B → 32 B), the Qwen3.5 family with its gated-attention architecture, the Qwen3 and Qwen3.5 MoE variants, the Gated-Delta-Net Qwen3-Next-80B-A3B, and Qwen3.8-Flash-Next — the GDN + QSA hybrid that previews the Qwen4 architecture.

## Variants

| Family | Class | Sizes | Recipe |
| - | - | - | - |
| Qwen3 | Dense | 0.6 B · 1.7 B · 4 B · 8 B · 14 B · 32 B | [qwen3](/docs/models/qwen/qwen3) |
| Qwen3 | MoE | 3 B / 30 B · 22 B / 235 B | [qwen3-moe](/docs/models/qwen/qwen3-moe) |
| Qwen3.5 | Dense | 4 B · 9 B · 27 B | [qwen3-5](/docs/models/qwen/qwen3-5) |
| Qwen3.5 | MoE | 3 B / 35 B | [qwen3-5-moe](/docs/models/qwen/qwen3-5-moe) |
| Qwen3.6 | Dense | 27 B | [qwen3-6](/docs/models/qwen/qwen3-6) |
| Qwen3.6 | MoE | 3 B / 35 B | [qwen3-6-moe](/docs/models/qwen/qwen3-6-moe) |
| Qwen3.8 | Dense | 27 B | [qwen3-8](/docs/models/qwen/qwen3-8) |
| Qwen3-Next | MoE (GDN) | 3 B / 80 B | [qwen3-next](/docs/models/qwen/qwen3-next) |
| Qwen3.8-Flash-Next | MoE (GDN + QSA) | 6 B / 176 B | [qwen3-8-flash-next](/docs/models/qwen/qwen3-8-flash-next) |

## Fastest path to train

Qwen3-4B on a single 8× H100 node — the canonical starter recipe:

```bash theme={null}
cd /root/miles
hf download Qwen/Qwen3-4B --local-dir /root/models/Qwen3-4B
python scripts/run_qwen3_dense.py --model-name Qwen3-4B
```

Dataset is [DAPO-Math-17k](https://huggingface.co/datasets/zhuzilin/dapo-math-17k) at `/root/datasets/dapo-math-17k/dapo-math-17k.jsonl`. See the [Qwen3 Dense](/docs/models/qwen/qwen3) page for the full walkthrough, weight conversion, and variants.

## Which variant do I pick?

* **Learning Miles for the first time** → Qwen3-4B ([qwen3](/docs/models/qwen/qwen3)). Fits one H100 node, fast loop.
* **Need MoE on a single node** → Qwen3-30B-A3B ([qwen3-moe](/docs/models/qwen/qwen3-moe)).
* **Scaling to multi-node** → Qwen3-235B-A22B ([qwen3-moe](/docs/models/qwen/qwen3-moe)).
* **Latest dense architecture (gated attention, A\_log FP32)** → Qwen3.5-4B ([qwen3-5](/docs/models/qwen/qwen3-5)).
* **Hybrid MTP / speculative decoding experiments** → Qwen3.5-35B-A3B ([qwen3-5-moe](/docs/models/qwen/qwen3-5-moe)).
* **Gated-Delta-Net (fla backend, real-CP)** → Qwen3-Next-80B-A3B ([qwen3-next](/docs/models/qwen/qwen3-next)).
* **Qwen4-architecture preview (hyper-connections, QSA sparse attention, PLE)** → Qwen3.8-Flash-Next ([qwen3-8-flash-next](/docs/models/qwen/qwen3-8-flash-next)).


# Qwen3
Source: https://miles.radixark.com/docs/models/qwen/qwen3

Launch recipes for dense Qwen3 models (0.6 B – 32 B).

## 1. Model Introduction

[Qwen3](https://github.com/QwenLM/Qwen3) is the latest generation of Alibaba's Qwen language model series, available in dense and MoE variants with both Instruct and reasoning-enhanced Thinking editions.

**Key highlights:**

* **Stronger general intelligence**: significant improvements in instruction following, logical reasoning, mathematics, science, coding, and tool usage over Qwen2.5.
* **Extended context length**: trained for 256 K-token contexts, useful for long-document reasoning and agentic workflows.
* **Flexible deployment options**: dense sizes from 0.6 B up to 32 B; this page covers the dense recipes (MoE recipes live in [qwen3-moe](/docs/models/qwen/qwen3-moe)).
* **Stronger agent interaction**: improved tool-use and search-based agent performance.

## 2. Supported Variants

| Model | HF ID |
| - | - |
| Qwen3-0.6B | [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) |
| Qwen3-1.7B | [Qwen/Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) |
| Qwen3-4B | [Qwen/Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B) |
| Qwen3-4B-Instruct-2507 | [Qwen/Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507) |
| Qwen3-8B | [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) |
| Qwen3-14B | [Qwen/Qwen3-14B](https://huggingface.co/Qwen/Qwen3-14B) |
| Qwen3-32B | [Qwen/Qwen3-32B](https://huggingface.co/Qwen/Qwen3-32B) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download Qwen/Qwen3-4B --local-dir /root/models/Qwen3-4B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
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 \
   --save          /root/models/Qwen3-4B_torch_dist
```

The converter auto-derives PP from `WORLD_SIZE`; for larger sizes drive it with `torchrun --nproc-per-node 8`. The FSDP launcher (`scripts/run_qwen3_0_6b_fsdp.py`) loads the HF checkpoint directly and skips this step.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_dense.py --model-name Qwen3-4B
```

One launcher covers the whole dense line — pick the variant with `--model-name` (`Qwen3-32B`, and the Qwen3.5 / Qwen3.6 sizes), and it selects the matching `qwen3-XB.py` model config. To run on a slice of a node, add `--num-gpus-per-node 4 --cuda-visible-devices 4,5,6,7`.

Checkpoints are read from `--model-dir` (default `/root/models`), datasets from `--data-dir` (default `/root/datasets`), and `--save`/`--load` point under `--output-dir` (default `/root/shared_data`).

The Qwen3-4B-Instruct-2507 config (`scripts/models/qwen3-4B-Instruct-2507.py`) just calls `qwen3-4B` with `rotary_base=5000000` (`MODEL_ARGS_ROTARY_BASE` still works as an environment override) — load it when converting / launching the Instruct-2507 checkpoint.

## 5. Recipe Configuration

### 5.1 Parallelism

`scripts/run_qwen3_dense.py` holds one recipe per `--model-name`:

| `--model-name` | TP | PP | CP | EP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| `Qwen3-4B` | 2 | 1 | 1 | 1 | 9216 | 8 (1 × 8) |
| `Qwen3-32B` | 8 | 1 | 1 | 1 | 20480 | 8 (1 × 8) |

`Qwen3-4B` also runs on a 4-GPU slice with `--num-gpus-per-node 4`; the parallelism is unchanged.

`--sequence-parallel` is on whenever TP > 1.

### 5.2 Algorithm

GRPO baseline across all dense recipes:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

Rollout uses `--rm-type deepscaler` against `dapo-math-17k`. The SFT recipe (`python scripts/run_qwen3_sft.py --model-name Qwen3-4B-Base`) trains on `/root/datasets/openhermes2_5.parquet`.

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.7
```

`Qwen3-32B` additionally pins `--sglang-cuda-graph-bs-decode 1 2 4 8 $(seq 16 8 256)`. The FSDP variant uses `--attn-implementation flash_attention_3`, SGLang attention backend `fa3`, and adds `--update-weight-buffer-size 536870912 --gradient-checkpointing`.

### 5.4 Optimizer

`Qwen3-32B` enables CPU Adam:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

The 4 B / 8 B / 14 B recipes leave Adam on GPU.

### 5.5 Notable quirks

* **BF16 train + FP8 inference**: download `Qwen/Qwen3-4B-FP8` and pass `--extra-args "--hf-checkpoint /root/models/Qwen3-4B-FP8"` to swap rollout to FP8 while keeping BF16 training. See [Low Precision RL](/docs/advanced/low-precision).
* **FSDP backend**: `python3 scripts/run_qwen3_0_6b_fsdp.py` runs a Qwen3-0.6B recipe with `--train-backend fsdp` (downloads model + datasets itself); no Megatron `torch_dist` conversion needed.
* **AMD ROCm**: `python scripts/amd/run_qwen3_4b.py` mirrors the recipe, with the GPU count per node resolved from `--hardware` (`MI350X` / `MI355X`).

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)
* [Backends Beyond Megatron](/docs/advanced/architecture-support) — for the FSDP variant.


# Qwen3.5
Source: https://miles.radixark.com/docs/models/qwen/qwen3-5

Launch recipes for Qwen3.5-4B / 9B / 27B with attention-output-gate.

## 1. Model Introduction

[Qwen3.5](https://github.com/QwenLM/Qwen3) is the next iteration of the Qwen3 dense series, introducing the gated-attention architecture and an FP32-preserved `A_log` parameter.

**Key highlights:**

* **Attention-output gate**: a learned gate on the attention output, trained alongside attention weights for stronger long-context behavior.
* **Extended rotary base**: `--rotary-base 10000000`, `--rotary-percent 0.25` — wider effective context than the original Qwen3.
* **Larger vocabulary**: 248320 tokens.
* **FP32 `A_log` preservation**: a parameter that must stay in FP32 through Megatron's mixed-precision pipeline; miles handles this via the bridge.

## 2. Supported Variants

| Model | HF ID |
| - | - |
| Qwen3.5-4B | [Qwen/Qwen3.5-4B](https://huggingface.co/Qwen/Qwen3.5-4B) |
| Qwen3.5-9B | [Qwen/Qwen3.5-9B](https://huggingface.co/Qwen/Qwen3.5-9B) |
| Qwen3.5-27B | [Qwen/Qwen3.5-27B](https://huggingface.co/Qwen/Qwen3.5-27B) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.5-4B)" || exit 1   # or qwen3.5-9B / qwen3.5-27B
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.5-4B \
   --save          /root/models/Qwen3.5-4B_torch_dist
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_dense.py --model-name Qwen3.5-4B        # or Qwen3.5-9B / Qwen3.5-27B
```

All three variants are 1 node × 8 GPU. The launcher reads checkpoints from `--model-dir` (default `/root/models`) and datasets from `--data-dir` (default `/root/datasets`), and writes checkpoints under `--output-dir` (default `/root/shared_data`).

## 5. Recipe Configuration

### 5.1 Parallelism

The three sizes are one launcher, `scripts/run_qwen3_dense.py`, selected by `--model-name`:

| `--model-name` | TP | PP | CP | `max_tokens_per_gpu` | SGLang `mem-fraction-static` | CPU Adam | GPUs |
| - | - | - | - | - | - | - | - |
| `Qwen3.5-4B` | 2 | 1 | 1 | 9216 | 0.7 | – | 8 (1 × 8) |
| `Qwen3.5-9B` | 2 | 1 | 1 | 9216 | 0.6 | – | 8 (1 × 8) |
| `Qwen3.5-27B` | 4 | 1 | 1 | 8192 | 0.5 | ✓ | 8 (1 × 8) |

### 5.2 Algorithm

GRPO across all three sizes:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 1
```

All three variants pin `--rollout-num-gpus-per-engine 1` because SGLang TP > 1 produced garbage output for Qwen3.5 on 0.5.9 (in-source comment, [sglang#21039](https://github.com/sgl-project/sglang/issues/21039)). If you bump the SGLang version past the fix, you can raise this back up.

### 5.4 Optimizer

Only the 27 B recipe enables CPU Adam (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`). The 4 B and 9 B recipes leave Adam on GPU.

### 5.5 Notable quirks

From `scripts/models/qwen3.5-4B.py` (and analogous configs for 9 B / 27 B):

* `--spec miles_plugins.models.qwen3_5 get_qwen3_5_spec` — attention-output gate, `A_log` parameter handling.
* `--rotary-base 10000000`, `--rotary-percent 0.25`.
* `--vocab-size 248320`.
* `--apply-layernorm-1p`, `--qk-layernorm`, `--group-query-attention`.
* `--attention-output-gate`.

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for how miles preserves FP32 parameters like `A_log` through Megatron's mixed-precision pipeline.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Low Precision RL](/docs/advanced/low-precision)


# Qwen3.5 MoE
Source: https://miles.radixark.com/docs/models/qwen/qwen3-5-moe

Launch recipe for Qwen3.5-35B-A3B with MTP training and EAGLE speculative rollout.

## 1. Model Introduction

[Qwen3.5-35B-A3B](https://github.com/QwenLM/Qwen3) is the MoE branch of the Qwen3.5 line — 3 B active / 35 B total — combining the gated-attention architecture with a built-in MTP head.

**Key highlights:**

* **Sparse MoE**: 3 B active out of 35 B total parameters.
* **Attention-output gate**: shared with the Qwen3.5 dense series, with FP32-preserved `A_log`.
* **Multi-Token Prediction (MTP)**: `--mtp-num-layers 1` baked into the model config; the recipe trains the MTP head and uses EAGLE speculative decoding at rollout.
* **Single-node footprint**: full recipe fits on 1 × 8 GPU.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Qwen3.5-35B-A3B | 3 B / 35 B | [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.5-35B-A3B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Qwen3.5-35B-A3B \
   --save          /root/models/Qwen3.5-35B-A3B_torch_dist \
   --mtp-num-layers 1
```

`--mtp-num-layers 1` during conversion preserves the MTP layer so it survives into Megatron format. The launcher below also runs the download and conversion itself before submitting, so this step is optional.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_5_35b_a3b_mtp.py \
   --parallelism tp1-ep8 \
   --rollout-max-response-len 8192 \
   --sglang-ep-size 8 \
   --sglang-max-running-requests 512 \
   --enable-eval
```

`--parallelism` also accepts `tp2-cp2-ep8` (the launcher default). Checkpoints are read from `--model-dir` (default `/root/models`) and datasets from `--data-dir` (default `/root/datasets`); checkpoints are written under `--output-dir` (default `/root/shared_data`).

## 5. Recipe Configuration

### 5.1 Parallelism

| TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| 1 | 1 | 1 | 8 | 1 | 8192 | 8 (1 × 8) |

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --use-kl-loss --kl-loss-coef 0.00`. Plus MTP training:

```bash theme={null}
--enable-mtp-training
--mtp-num-layers 1
--mtp-loss-scaling-factor 0.2
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 8
--sglang-mem-fraction-static 0.7
--sglang-ep-size 8
--sglang-cuda-graph-bs-decode 1 2 4 8 16 24 ... 256

# mtp speculative decoding
--sglang-speculative-algorithm EAGLE
--sglang-speculative-num-steps 2
--sglang-speculative-eagle-topk 1
--sglang-speculative-num-draft-tokens 3

--sglang-max-running-requests 512
```

`--sglang-ep-size` and `--sglang-max-running-requests` are launcher options; leave them out and SGLang keeps its own defaults.

### 5.4 Optimizer

CPU Adam is enabled (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`).

### 5.5 Notable quirks

* The Megatron side uses `--moe-token-dispatcher-type flex`; DeepEP isn't enabled here, unlike Qwen3-Next.
* The model config (`scripts/models/qwen3.5-35B-A3B.py`) reuses the Qwen3.5 spec: `--attention-output-gate`, `--rotary-base 10000000`, `--rotary-percent 0.25`, `A_log` kept in FP32 via the bridge. See [Backends Beyond Megatron](/docs/advanced/architecture-support).

## 6. Pairs Well With

* [Speculative Decoding](/docs/advanced/speculative-decoding)
* [Backends Beyond Megatron](/docs/advanced/architecture-support)


# Qwen3.6
Source: https://miles.radixark.com/docs/models/qwen/qwen3-6

Launch recipe for the dense Qwen3.6-27B with attention-output-gate.

## 1. Model Introduction

[Qwen3.6](https://github.com/QwenLM/Qwen3) is the next iteration of Alibaba's
Qwen3 line, focused on agentic-coding workflows and on preserving reasoning
context across long sessions. The family ships two variants — a sparse MoE
(Qwen3.6-35B-A3B) and a dense GDN-backbone model (Qwen3.6-27B) — both with
native hybrid reasoning (thinking by default), built-in tool calling, and
multimodal text / image / video input. Context windows reach 262 K and
extend past 1 M. Weights are Apache 2.0, available in BF16 and FP8.

The dense **Qwen3.6-27B** is the single-GPU-friendly variant. In miles it
reuses the Qwen3.5 Megatron spec
(`miles_plugins.models.qwen3_5.get_qwen3_5_spec`); architecturally it's a
wider, deeper Qwen3.5 with the gated-attention design preserved.

**Key highlights:**

* **Dense GDN backbone**: 27 B parameters, single-GPU friendly footprint.
* **Attention-output gate**: shared with Qwen3.5, trained alongside attention weights.
* **Extended rotary base**: `--rotary-base 10000000`, `--rotary-percent 0.25`.
* **Larger vocabulary**: 248320 tokens.
* **Shape**: `hidden-size 5120`, `ffn-hidden-size 17408`, 64 layers.
* **Long context**: 262 K tokens, extensible past 1 M.

## 2. Supported Variants

| Model | HF ID |
| - | - |
| Qwen3.6-27B | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.6-27B)" || 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.6-27B \
   --save          /root/models/Qwen3.6-27B_torch_dist
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_dense.py --model-name Qwen3.6-27B
```

`scripts/run_qwen3_dense.py` is the shared dense launcher; `--model-name Qwen3.6-27B`
selects this recipe, which targets 1 node × 8 GPU. Checkpoints come from `--model-dir`
(default `/root/models`) and datasets from `--data-dir` (default `/root/datasets`);
checkpoints are written under `--output-dir` (default `/root/shared_data`).

## 5. Recipe Configuration

### 5.1 Parallelism

| TP | PP | CP | EP | `max_tokens_per_gpu` | SGLang `mem-fraction-static` | CPU Adam | GPUs |
| - | - | - | - | - | - | - | - |
| 4 | 1 | 1 | 1 | 8192 | 0.5 | ✓ | 8 (1 × 8) |

`--sequence-parallel` is enabled. Activation checkpointing is on
(`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`).

### 5.2 Algorithm

GRPO with low-variance KL:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 1
--sglang-mem-fraction-static 0.5
```

`--rollout-num-gpus-per-engine 1` follows the Qwen3.5 line; SGLang TP > 1 has been
problematic on this family. If your SGLang version carries the fix for
[sglang#21039](https://github.com/sgl-project/sglang/issues/21039), you can raise it.

### 5.4 Optimizer

CPU Adam is enabled (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`).

### 5.5 Notable quirks

From `scripts/models/qwen3.6-27B.py`:

* `--spec miles_plugins.models.qwen3_5 get_qwen3_5_spec` — Qwen3.6 reuses the Qwen3.5 spec (gated attention, FP32 `A_log`).
* `--rotary-base 10000000`, `--rotary-percent 0.25`.
* `--vocab-size 248320`.
* `--apply-layernorm-1p`, `--qk-layernorm`, `--group-query-attention`.
* `--attention-output-gate`.

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for how miles
preserves FP32 parameters like `A_log` through Megatron's mixed-precision pipeline.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Low Precision RL](/docs/advanced/low-precision)


# Qwen3.6 MoE
Source: https://miles.radixark.com/docs/models/qwen/qwen3-6-moe

Launch recipe for Qwen3.6-35B-A3B with MTP training and EAGLE speculative rollout.

## 1. Model Introduction

[Qwen3.6-35B-A3B](https://github.com/QwenLM/Qwen3) is the sparse MoE branch of
Alibaba's Qwen3.6 line — 35 B total / 3 B active parameters on a Gated Delta
Networks backbone. Like the dense Qwen3.6-27B, it's tuned for agentic-coding
workflows and long-session reasoning, with native hybrid thinking mode,
built-in tool calling, and multimodal text / image / video input. Context
reaches 262 K and extends past 1 M; weights are Apache 2.0 in BF16 and FP8.
Qwen3.6 also ships native Multi-Token Prediction for speculative decoding,
which this recipe trains and serves via EAGLE.

In miles, Qwen3.6-35B-A3B reuses the Qwen3.5 spec
(`miles_plugins.models.qwen3_5.get_qwen3_5_spec`) and bakes in MTP training
plus a shared-expert gate.

**Key highlights:**

* **Sparse MoE on a GDN backbone**: 256 experts, top-8 routing, 3 B active / 35 B total.
* **Attention-output gate**: shared with the Qwen3.5 / 3.6 dense series.
* **Shared expert + gate**: `--moe-shared-expert-intermediate-size 512 --moe-shared-expert-gate`.
* **Multi-Token Prediction (MTP)**: `--mtp-num-layers 1`; trained alongside the policy and served via EAGLE at rollout.
* **Dispatcher**: `--moe-token-dispatcher-type alltoall` for HF→Megatron conversion; runtime uses `flex` (set in the launcher).
* **Long context**: 262 K tokens, extensible past 1 M.
* **Single-node footprint**: full recipe fits on 1 × 8 GPU (H200).

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Qwen3.6-35B-A3B | 3 B / 35 B | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) |

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
hf download Qwen/Qwen3.6-35B-A3B --local-dir /root/models/Qwen3.6-35B-A3B
```

### 3.2 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.6-35B-A3B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Qwen3.6-35B-A3B \
   --save          /root/models/Qwen3.6-35B-A3B_torch_dist \
   --mtp-num-layers 1
```

`--mtp-num-layers 1` during conversion preserves the MTP layer so it survives into Megatron format.

## 4. Launch

### 4.1 Quick start

The launcher is a parametrized Typer script (8 × H200) that exercises arbitrary
(TP, EP, CP, PP, ETP) cells:

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_6_35b_a3b_mtp.py \
   --tp 1 --ep 8 --cp 1 --pp 1 --etp 1 \
   --num-rollout 10
```

Default knobs in the launcher: `--mode debug_minimal`, 8 GPUs, `max_tokens_per_gpu=8192`,
`rollout_batch_size=8`, `n_samples_per_prompt=2`, `global_batch_size=16`,
`rollout_max_response_len=1024`. Override via flags for longer runs.

## 5. Recipe Configuration

### 5.1 Parallelism

The default cell is `TP=1 EP=8 CP=1 PP=1 ETP=1`. Sequence parallelism is on; activation
checkpointing defaults on (`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`)
and can be turned off with `--no-recompute`.

| TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | GPUs |
| - | - | - | - | - | - | - |
| 1 | 1 | 1 | 8 | 1 | 8192 | 8 (1 × 8) |

### 5.2 Algorithm

GRPO with low-variance KL plus MTP training:

```bash theme={null}
GRPO_ARGS=(
   --advantage-estimator grpo
   --use-kl-loss
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --entropy-coef 0.00
   --eps-clip 0.2
   --eps-clip-high 0.28
)

MTP_ARGS=(
   --enable-mtp-training
   --mtp-num-layers 1
   --mtp-loss-scaling-factor 0.2
)
```

### 5.3 Rollout & SGLang

```bash theme={null}
SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 8
   --sglang-mem-fraction-static 0.7
   --sglang-ep-size 8
   --sglang-cuda-graph-bs-decode 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128

   # MTP speculative decoding
   --sglang-speculative-algorithm EAGLE
   --sglang-speculative-num-steps 2
   --sglang-speculative-eagle-topk 1
   --sglang-speculative-num-draft-tokens 3

   --sglang-max-running-requests 256
   --sglang-mamba-radix-cache-strategy extra_buffer
)
```

### 5.4 Optimizer

CPU Adam is enabled (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`).

### 5.5 Notable quirks

From `scripts/models/qwen3.6-35B-A3B.py` and `scripts/run_qwen3_6_35b_a3b_mtp.py`:

* `--spec miles_plugins.models.qwen3_5 get_qwen3_5_spec` — Qwen3.6 reuses the Qwen3.5 spec.
* 256 experts, `--moe-router-topk 8`, `--moe-router-score-function softmax`.
* `--moe-shared-expert-gate` and `--moe-shared-expert-intermediate-size 512`.
* Megatron-side dispatcher overridden to `--moe-token-dispatcher-type flex` at runtime; conversion uses `alltoall`.
* `--moe-grouped-gemm`, `--moe-token-drop-policy probs`, `--moe-router-dtype fp32`, `--moe-permute-fusion`, `--moe-aux-loss-coeff 0`.
* `--attention-output-gate`, `--rotary-base 10000000`, `--rotary-percent 0.25`, `--vocab-size 248320`.

See [Backends Beyond Megatron](/docs/advanced/architecture-support) for FP32 parameter handling and how miles wires the spec.

## 6. Pairs Well With

* [Speculative Decoding](/docs/advanced/speculative-decoding)
* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [P2P Weight Transfer](/docs/advanced/p2p-weight-transfer)


# Qwen3.8
Source: https://miles.radixark.com/docs/models/qwen/qwen3-8

Launch recipe for the dense Qwen3.8-27B, plus where the 2.4T-A95B MoE recipe lives.

## 1. Model Introduction

[Qwen3.8](https://github.com/QwenLM/Qwen3) continues Alibaba's Qwen3 line. The
dense **Qwen3.8-27B** ships the same `config.json` as
[Qwen3.5-27B](/docs/models/qwen/qwen3-5) and [Qwen3.6-27B](/docs/models/qwen/qwen3-6) —
same hybrid GDN backbone, same gated attention, same tokenizer and vocabulary.
It therefore reuses the Qwen3.5 Megatron spec
(`miles_plugins.models.qwen3_5.get_qwen3_5_spec`), and
`scripts/models/qwen3.8-27B.py` is a one-line derivation of the Qwen3.5-27B
model args; the three expand to byte-identical Megatron flags.

The sparse **Qwen3.8-2.4T-A95B** is a different recipe entirely — see
[section 6](#6-qwen38-24t-a95b).

**Key highlights (27 B):**

* **Dense GDN backbone**: 27 B parameters, hybrid linear / full attention (`full_attention_interval 4`).
* **Attention-output gate**: shared with Qwen3.5, trained alongside attention weights.
* **Extended rotary base**: `--rotary-base 10000000`, `--rotary-percent 0.25`.
* **Larger vocabulary**: 248320 tokens.
* **Shape**: `hidden-size 5120`, `ffn-hidden-size 17408`, 64 layers.
* **Multimodal**: `Qwen3_5ForConditionalGeneration`; the RL recipe below trains the text path only.

## 2. Supported Variants

| Model | Class | HF ID | Recipe |
| - | - | - | - |
| Qwen3.8-27B | Dense (GDN) | [Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B) | this page |
| Qwen3.8-2.4T-A95B | MoE | — | [section 6](#6-qwen38-24t-a95b) |

Sections 3–6 cover the dense 27 B.

## 3. Environment Setup

### 3.1 Download model + datasets

```bash theme={null}
hf download Qwen/Qwen3.8-27B --local-dir /root/models/Qwen3.8-27B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.2 HF → Megatron `torch_dist` conversion

Run it on all eight GPUs; the tool shards the 64 layers over the ranks
(`--pipeline-model-parallel-size` is derived from `WORLD_SIZE`) and the output
re-shards at load, so the conversion layout does not have to match the training one:

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.8-27B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   "${MODEL_ARGS[@]}" \
   --hf-checkpoint /root/models/Qwen3.8-27B \
   --save          /root/models/Qwen3.8-27B_torch_dist
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_dense.py --model-name Qwen3.8-27B
```

`scripts/run_qwen3_dense.py` is the shared dense launcher; `--model-name Qwen3.8-27B`
selects this recipe, which targets 1 node × 8 GPU. Checkpoints come from `--model-dir`
(default `/root/models`) and datasets from `--data-dir` (default `/root/datasets`);
checkpoints are written under `--output-dir` (default `/root/shared_data`).

### 4.2 What one step costs

Measured on 1 × 8 H200 with the recipe below, `rollout-batch-size 32`,
`n-samples-per-prompt 8`, `rollout-max-response-len 8192`:

| Phase | Time |
| - | - |
| Rollout (256 sequences) | 307 s |
| Train (ref log-probs 100 s + actor log-probs 33 s + backward/optimizer) | 415 s |
| Weight sync back to SGLang | 2.5 s |
| **Per step** | **≈ 12 min** |

`perf/tokens_per_gpu_per_sec` 478. Mean response length 4585 tokens, 33 % truncated at the
8192 cap. Step 0 reported `rollout/raw_reward` 0.64 and
`train/train_rollout_logprob_abs_diff` 0.010, i.e. the SGLang and Megatron forward passes
agree closely — see [True On-Policy](/docs/examples/infra-features/true-on-policy) for what that
metric does and does not tell you.

## 5. Recipe Configuration

### 5.1 Parallelism

| TP | PP | CP | EP | `max_tokens_per_gpu` | SGLang `mem-fraction-static` | CPU Adam | GPUs |
| - | - | - | - | - | - | - | - |
| 4 | 1 | 1 | 1 | 8192 | 0.8 | ✓ | 8 (1 × 8) |

`--sequence-parallel` is enabled. Activation checkpointing is on
(`--recompute-granularity full --recompute-method uniform --recompute-num-layers 1`).

### 5.2 Algorithm

GRPO with low-variance KL:

```bash theme={null}
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
```

### 5.3 Rollout & SGLang

```bash theme={null}
--rollout-num-gpus-per-engine 1
--sglang-mem-fraction-static 0.8
```

One engine per GPU, inherited from the Qwen3.5 line.

**Why 0.8 and not the 0.5 the rest of the line uses.** The BF16 weights are \~53 GB per
engine. At 0.5 SGLang targets 70 GB of the H200's 140 GB, so after the weights only \~17 GB
is left for the pools — and it reports 69 GB still free. That matters more here than on a
plain transformer: 48 of the 64 layers are linear attention, and their recurrent state is
resident per sequence (48 v-heads × 128 × 128 × FP32 × 48 layers ≈ **151 MB per in-flight
sequence**), an order of magnitude above the KV cost. Measured on one node:

| | `mem-fraction-static 0.5` | `mem-fraction-static 0.8` |
| - | - | - |
| KV pool | 158,018 tok | 517,999 tok |
| GDN state slots (`mamba_cache_size`) | 60 | 198 |
| `#running-req` per engine | 11–12 | 39 |
| decode throughput per engine | 553–690 tok/s | 1855 tok/s |

Decode re-reads all 53 GB of weights every step regardless of batch size, so the small
batch was pure bandwidth waste. Under colocate the training model is offloaded to host RAM
during rollout, and the training-phase peak at 0.8 measured 88–90 GB of 140 GB, so the
higher fraction costs nothing. Two further levers if rollout is still the bottleneck:
`--rollout-num-gpus-per-engine 2` (halves per-GPU weights *and* shards the GDN state), and
`--sglang-mamba-ssm-dtype bfloat16` (halves the state, at a precision cost worth checking
against `train/train_rollout_logprob_abs_diff`).

### 5.4 Optimizer

CPU Adam is enabled (`--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer`).

### 5.5 Notable quirks

From `scripts/models/qwen3.8-27B.py`, which defers to `scripts/models/qwen3.5-27B.py`:

* `--spec miles_plugins.models.qwen3_5 get_qwen3_5_spec` — Qwen3.8 reuses the Qwen3.5 spec (gated attention, FP32 `A_log`).
* `--rotary-base 10000000`, `--rotary-percent 0.25`.
* `--vocab-size 248320`.
* `--apply-layernorm-1p`, `--qk-layernorm`, `--group-query-attention`.
* `--attention-output-gate`.

See [Disk Offload](/docs/advanced/disk-offload) for how miles keeps FP32-marked parameters like
the GDN `A_log` out of the low-precision optimizer path.

## 6. Qwen3.8-2.4T-A95B

The sparse 2.4 T / 95 B-active flagship shares nothing operational with the dense 27 B: the
rollout engines serve the native ModelOpt **NVFP4** experts-only checkpoint
([RadixArk/Qwen3.8-2.4T-A95B-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-2.4T-A95B-NVFP4))
while the Megatron trainer runs **BF16** off a `torch_dist` built from the BF16 parent, and
weight sync ships **LoRA adapters only** — the quantized base never leaves the engines. The
implementation is [radixark/miles#2488](https://github.com/radixark/miles/pull/2488)
("Qwen 3.8 day-0 lora RL support"): `scripts/run_qwen3_8.py` plus the
`scripts/models/qwen3.8-2.4T-A95B{,_4layer,_full}.py` definitions. It is **LoRA-only in
practice** — the native raw-mode `--lora` path is the validated route; treat full-weight
training as unvalidated.

### 6.1 Variants

`--model-variant` selects between them and sets the matching checkpoint paths and
`megatron_model_type` (`scripts/models/qwen3.8-2.4T-A95B{_4layer,_full}.py`):

| Variant | Layers | Purpose | GPUs |
| - | - | - | - |
| `full` | 92 | the real model | 64 (16 × 4 GB300), validated |
| `4layer` | 4 | smoke test, default | single node |

### 6.2 Environment Setup

Use the `docker.io/radixark/miles:dev` image — it is multi-arch, so the same tag serves
GB300 (aarch64) and x86 nodes. Inside the container, multi-node rollout engines need
**NCCL ≥ 2.30.7** (2.28.x deadlocks CUDA-graph-captured cross-node collectives):

```bash theme={null}
pip install nvidia-nccl-cu13==2.30.7
export SGLANG_NCCL_SO_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2
export LD_PRELOAD=$SGLANG_NCCL_SO_PATH
```

The launcher looks checkpoints up under `--model-dir` (default `/root/models`) as
`Qwen3.8-2.4T-A95B-NVFP4_<variant>` (rollout) and
`Qwen3.8-2.4T-A95B-bf16_<variant>` (trainer source):

```bash theme={null}
hf download RadixArk/Qwen3.8-2.4T-A95B-NVFP4 \
    --local-dir /root/models/Qwen3.8-2.4T-A95B-NVFP4_full
# BF16 parent -> /root/models/Qwen3.8-2.4T-A95B-bf16_full
```

Unless `--skip-prepare` is passed, the launcher's prepare step downloads
`zhuzilin/dapo-math-17k` and runs the BF16 → `torch_dist` conversion itself
(`tools/convert_hf_to_torch_dist.py` via the `qwen3.8-2.4T-A95B_<variant>` model args); the
output re-shards at load, so the conversion layout does not have to match the training one.
Training then takes the **NVFP4** directory as `--hf-checkpoint` and the converted
`torch_dist` as `--ref-load` — both wired up by the launcher.

### 6.3 Launch

Single-node smoke test (default `--model-variant 4layer`, 8 GPUs):

```bash theme={null}
python scripts/run_qwen3_8.py
```

Full model, validated on **16 nodes × 4 GB300** (bring up a ray cluster across the nodes,
`export MILES_SCRIPT_EXTERNAL_RAY=1`, then run on the head):

```bash theme={null}
python scripts/run_qwen3_8.py \
  --mode normal --model-variant full --hardware GB300 \
  --num-nodes 16 --num-gpus-per-node 4 \
  --tp 1 --pp 4 --ep 16 --etp 1 --cp 1 \
  --lora --lora-rank 32 --lora-alpha 64 \
  --num-rollout 200 --rollout-batch-size 32 --n-samples-per-prompt 8 \
  --rollout-max-response-len 4096 --global-batch-size 256 \
  --rollout-max-concurrency 8 --recompute --skip-prepare \
  --extra-args "--rollout-num-gpus-per-engine 16 --lr 1e-5"
```

| Side | Parallelism | Notes |
| - | - | - |
| Trainer (BF16) | TP 1 / PP 4 / EP 16 / ETP 1 | `num_query_groups 4` caps native-LoRA TP at 4 |
| Rollout (NVFP4) | 4 engines × TP 16 / EP 1 | `flashinfer_trtllm` MoE runner, triton attention |
| LoRA | rank 32 / alpha 64, `q,k,v,o_proj` | attention projections only; the 69 GDN mixer layers and the routed experts carry no adapter, so the rollout MoE path is untouched |

### 6.4 Memory: the alternating-backup budget

The recipe fits 956 GB hosts only if, at any moment, **exactly one side keeps a host
backup**: during rollout the sleeping trainer's tms backup (\~390 GB/node); during train the
sleeping engines' NVFP4 mirror (\~372 GB/node, `--lora-base-cpu-backup` →
SGLang `enable_weights_cpu_backup`, which is also what lets every weight update skip the
base re-ship). Anything that double-books breaks it. The flags that hold the line:

* `--drop-checkpoint-page-cache-after-load` — a multi-TB DCP read leaves its page cache on
  every node; lazy reclaim loses the race against pinned-allocation spikes at the
  rollout/train handoff.
* `--colocate-memory-peak-device gpu` and no `--use-kl-loss` at `kl-loss-coef 0` — a zero
  coefficient still loads a full ref checkpoint and pins a second per-rank base copy.
* Weight sync must not clone the sleeping side's backup (`get_cpu_backup(zero_copy=True)`
  in #2488); the clone variant retains \~75 GB/rank in malloc arenas per update.

Healthy peaks sit at **600–650 GB per node**; sustained readings above \~750 GB mean a
second backup or the page cache came back.

## 7. Pairs Well With

* [Qwen3.5](/docs/models/qwen/qwen3-5) — same architecture at 4 B / 9 B / 27 B
* [True On-Policy](/docs/examples/infra-features/true-on-policy)
* [Low Precision RL](/docs/advanced/low-precision)


# Qwen3.8-Flash-Next
Source: https://miles.radixark.com/docs/models/qwen/qwen3-8-flash-next

RL recipe for Qwen3.8-Flash-Next, the GDN + QSA hybrid MoE preview of the Qwen4 architecture, with hyper-connections and a host-resident PLE table.

The complete Qwen3.8-Flash-Next RL implementation is open at the Miles pull request:
[`radixark/miles#2777`](https://github.com/radixark/miles/pull/2777). It lands as a set of
three: that branch, the SGLang
[`sglang-miles-qwen38next`](https://github.com/sgl-project/sglang/tree/sglang-miles-qwen38next)
branch, and [`radixark/Megatron-LM#89`](https://github.com/radixark/Megatron-LM/pull/89).
The image in [section 3](#3-environment-setup) pins all three.

## 1. Model Introduction

[Qwen3.8-Flash-Next](https://docs.sglang.io/cookbook/autoregressive/Qwen/Qwen3.8-Flash-Next)
is Qwen's **176 B-parameter (6 B active) GDN + QSA hybrid Mixture-of-Experts preview of the
Qwen4 architecture**. Despite the shared prefix it is not a variant of the dense
[Qwen3.8-27B](/docs/models/qwen/qwen3-8) — it is the next step of the *Next* line that
[Qwen3-Next](/docs/models/qwen/qwen3-next) started, and structurally it has far more in common
with that model than with anything else carrying the Qwen3.8 name.

The block spec is still a uniform GPT decoder, but three components sit outside what a
stock Megatron layer provides, and each one is why a piece of this recipe exists:

**Hyper-connections replace the block layernorms.** The checkpoint ships none: each
hyper-connection's own `hc_norm` *is* the pre-block norm, so the spec drops every block
layernorm and a leftover TE fused norm corrupts the forward pass silently rather than
loudly. Miles fills Megatron's HC `ModuleSpec` slots, but Qwen's output contraction is not
DeepSeek-V4's `learned_output_contract` — it is the same low-rank gated mean as the
per-layer hyper-connections, with the RMS taken per stream rather than over the whole
`n*C` vector. That needed a new Megatron spec slot, `hc_head_contraction`, which is what
Megatron#89 adds; the DeepSeek-V4 default path and its parameter names are untouched.

**QSA (Qwen Sparse Attention) on the full-attention layers.** Twelve of the 48 layers are
full attention. Each projects its own indexer queries and compressed keys, scores them, and
keeps `indexer_budget` key tokens per query; attention then reads only those. Miles
reimplements the indexer and a forward *and backward* Triton sparse-attention kernel rather
than importing SGLang's inference path, with the selection rows built torch-side so the
kernel needs no causal or segment logic of its own.

**A frozen, host-resident PLE table.** The per-layer-embedding n-gram table is \~102 GB. It
lives in host memory, TP-row-sharded, and is deliberately *not* a checkpointed parameter;
token ids reach it over an explicit side channel that raises rather than defaulting when
nothing was published. It is also on the weight-update check's skip list — see
[section 5.5](#55-notable-quirks).

**Key highlights:**

* **48 layers, hybrid**: 36 GDN linear-attention layers + 12 QSA full-attention layers, from
  the released config's `layer_types` (equivalently, every 4th layer is full attention).
* **512-expert MoE at top-10**, `moe-ffn-hidden-size 640`, plus a gated shared expert.
* **Hyper-connections** at every block, with a model-supplied output contraction.
* **PLE** n-gram embeddings, frozen and host-resident.
* **Attention output gate**, `--qk-layernorm`, `--apply-layernorm-1p`.
* **Shape**: hidden 2560, 24 attention heads, 2 query groups, `kv-channels 256`,
  vocab 248320, `--rotary-base 10000000` at `--rotary-percent 0.25`.

Two things the model args deliberately leave out: MTP (`--mtp-num-layers` is omitted, the
MTP tensors are not mapped yet) and the hyper-connection / PLE / QSA fields themselves,
which have no Megatron CLI flags and are derived from the checkpoint by the spec in
`miles_plugins/models/qwen3_8_next/qwen3_8_next.py`.

On the parameter count: Qwen's published headline is 176 B total / 6 B active, while the
docstring in `scripts/models/qwen3.8-flash-next.py` counts 180 B / \~7.4 B from the released
`config.json` shapes. Nothing in the recipe depends on which figure you quote.

## 2. Supported Variants

| Variant | `--model-name` | Layers | Purpose | GPUs |
| - | - | - | - | - |
| Full | `Qwen3.8-Flash-Next` | 48 | the real model | 32 (8 × 4) |
| Smoke slice | `Qwen3.8-Flash-Next-4layer` | 4 | CI and single-node bring-up | 1 × 4 or 1 × 8 |

`--model-name` selects between them and sets the matching `megatron_model_type`
(`qwen3.8-flash-next` / `qwen3.8-flash-next-4layer`), which in turn resolves the model args
and the `torch_dist` path. The launcher asserts the node shape, so a mismatched
`--num-nodes` / `--num-gpus-per-node` fails immediately rather than mid-run.

## 3. Environment Setup

Use the `docker.io/radixark/miles:qwen38next` image. It is the rolling
[`radixark/miles:dev`](/docs/developer/ci/02-docker-build) image with the three moving parts checked out at
the versions this recipe was built against, and nothing else changed — every prebuilt wheel,
TransformerEngine patch and version pin comes from `dev`. It is multi-arch, so the same tag
serves GB300 (aarch64) and x86 nodes.

| Component | Pinned at |
| - | - |
| miles | [`#2777`](https://github.com/radixark/miles/pull/2777) `afd78afd` |
| SGLang | [`sglang-miles-qwen38next`](https://github.com/sgl-project/sglang/tree/sglang-miles-qwen38next) `599d7403` |
| Megatron-LM | [`#89`](https://github.com/radixark/Megatron-LM/pull/89) `e8f57451` |

### 3.1 Download model + dataset

```bash theme={null}
hf download Qwen/Qwen3.8-Flash-Next --local-dir /root/models/Qwen3.8-Flash-Next
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
```

For the smoke slice, download
[`CharyZeng/Qwen3.8-Flash-Next-4layer`](https://huggingface.co/CharyZeng/Qwen3.8-Flash-Next-4layer)
into `/root/models/Qwen3.8-Flash-Next-4layer` instead — the launcher resolves
`--hf-checkpoint` as `<--model-dir>/<--model-name>`, which is the same path CI uses.

### 3.2 HF → Megatron `torch_dist` conversion

Unlike the launcher-prepared recipes, this one takes the converted reference checkpoint as
given: `--ref-load` resolves to `<--ckpt-dir>/<megatron_model_type>_torch_dist`, so the
conversion has to have run first. The output re-shards at load, so the conversion layout
does not have to match the training one:

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3.8-flash-next)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
CONVERT_KEEP_PP1=1 PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   "${MODEL_ARGS[@]}" \
   --hf-checkpoint /root/models/Qwen3.8-Flash-Next \
   --save          /root/ckpt/qwen3.8-flash-next_torch_dist \
   --tensor-model-parallel-size 2 \
   --pipeline-model-parallel-size 1
```

Swap `qwen3.8-flash-next` for `qwen3.8-flash-next-4layer` throughout to convert the smoke
slice.

## 4. Launch

The launcher assumes an **already-running ray cluster**: bring one up across the nodes,
`export MILES_SCRIPT_EXTERNAL_RAY=1`, then run on the head node.

Full model, 8 nodes × 4 GPUs:

```bash theme={null}
cd /root/miles
python scripts/run_qwen3_8_next.py train \
  --model-name Qwen3.8-Flash-Next \
  --num-nodes 8 --num-gpus-per-node 4 \
  --num-rollout 5 --rollout-max-response-len 4096
```

Single-node smoke slice:

```bash theme={null}
python scripts/run_qwen3_8_next.py train \
  --model-name Qwen3.8-Flash-Next-4layer \
  --num-nodes 1 --num-gpus-per-node 8
```

Paths come from `--model-dir` (default `/root/models`), `--data-dir` (default
`/root/datasets`), `--ckpt-dir` (default `/root/ckpt`) and `--megatron-path` (default
`/root/Megatron-LM`). Saving is off by default (`skip_saving`); turning it on writes
checkpoints under `<--save-dir>/<--run-id>/checkpoints` every 10 rollouts, without optimizer
or RNG state.

## 5. Recipe Configuration

### 5.1 Parallelism

| Shape | TP | PP | CP | EP | ETP | Rollout engine |
| - | - | - | - | - | - | - |
| 8 × 4 (full) | 2 | 8 | 1 | 4 | 1 | 8 GPUs, SGLang TP 8 / EP 8 |
| 1 × 4 (4layer) | 2 | 2 | 1 | 2 | 1 | 4 GPUs, SGLang TP 4 / EP 4 |
| 1 × 8 (4layer) | 2 | 2 | 1 | 4 | 1 | 4 GPUs, SGLang TP 4 / EP 4 |

`--sequence-parallel` is on in every shape. Activation checkpointing is full and uniform at
one layer, with `--micro-batch-size 1` and `--max-tokens-per-gpu 8192`.

### 5.2 Algorithm

GRPO, DAPO-Math-17k, thinking mode on:

```bash theme={null}
--advantage-estimator grpo
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
--rollout-batch-size 4
--n-samples-per-prompt 8
--rollout-temperature 0.8
--apply-chat-template-kwargs '{"thinking_mode":"thinking"}'
```

Adam at `--lr 1e-6`, constant schedule, `--weight-decay 0.1`,
`--adam-beta1 0.9 --adam-beta2 0.98`.

### 5.3 Rollout & SGLang

Rollout is **colocated** — trainer and engines share the GPUs — and the trainer offloads to
**disk** rather than host RAM during rollout (`--offload-train-target disk`,
`--offload-train-disk-dir /tmp/train_offload`). Worth keeping in view when sizing hosts: the
PLE table is already holding \~102 GB of host memory for the full model. The 4-layer CI slice
overrides the target back to `--offload-train-target cpu`.

```bash theme={null}
--sglang-linear-attn-prefill-backend flashinfer   # GDN prefill
--sglang-moe-runner-backend triton
--sglang-chunked-prefill-size 8192
--sglang-disable-radix-cache
--sglang-mem-fraction-static 0.7
--linear-attention-backend flashqla               # trainer-side GDN
--qkv-format thd
```

The recipe loosens router
health checking (`--router-health-failure-threshold 40`,
`--router-health-check-interval-secs 15`, `--router-health-success-threshold 1`) and raises
`--rollout-health-check-interval` / `--rollout-health-check-timeout` to 300 s.

### 5.4 What CI watches

`tests/e2e/megatron/model_scripts/test_qwen3_8_next_4layer_ci.py` runs the 4-layer slice on
8 × H200 in `stage-c-8-gpu-h200` with rollout-routing replay on, and gates these metrics:

* `train/grad_norm`
* `train/ppo_kl`
* `train/train_rollout_logprob_abs_diff`
* `train/train_rollout_kl`
* `rollout/raw_reward`

`train/train_rollout_logprob_abs_diff` is the one to read first on a fresh bring-up: it is
the direct measure of whether the SGLang and Megatron forward passes agree, and on this
architecture that covers the GDN, QSA and hyper-connection paths at once. See
[True On-Policy](/docs/examples/infra-features/true-on-policy) for what the metric does and does
not tell you.

### 5.5 Notable quirks

* **Weight-update checking skips two prefixes.** `--check-weight-update-equal` runs with
  `--check-weight-update-skip-list visual. ple_embedding.` — the PLE table is frozen and
  never shipped, so comparing it would fail on a parameter that is working as intended.
* **Every block layernorm is dropped.** The checkpoint has none; each hyper-connection's
  `hc_norm` is the pre-block norm. A leftover TE fused norm corrupts silently.
* **`--moe-aux-loss-coeff 0`.** Routing is not auxiliary-loss balanced here.
* **Triton and Inductor caches are pinned to `/tmp`** (`TRITON_CACHE_DIR`,
  `TORCHINDUCTOR_CACHE_DIR`) with `TORCHINDUCTOR_COMPILE_THREADS=1`.
* **`SGLANG_DISABLE_MULTIMEM_AG=1`** and `SGLANG_SKIP_CHECKPOINT_LOAD_CHECK=1` are set for
  the engines.
* The model lives in `miles_plugins/models/qwen3_8_next/`, and weight conversion in
  `miles/backends/megatron_utils/megatron_to_hf/qwen3_8_next.py`.

## 6. Pairs Well With

* [Qwen3-Next](/docs/models/qwen/qwen3-next) — the previous generation of the same line
* [Qwen3.8](/docs/models/qwen/qwen3-8) — the dense 27 B that shares the name and not the architecture
* [Disk Offload](/docs/advanced/disk-offload)
* [True On-Policy](/docs/examples/infra-features/true-on-policy)


# Qwen3 MoE
Source: https://miles.radixark.com/docs/models/qwen/qwen3-moe

Launch recipes for Qwen3-30B-A3B (single node) and Qwen3-235B-A22B (multi-node).

## 1. Model Introduction

[Qwen3 MoE](https://github.com/QwenLM/Qwen3) is the Mixture-of-Experts branch of the Qwen3 series, available in two sizes: 30 B-A3B (single-node) and 235 B-A22B (multi-node).

**Key highlights:**

* **Sparse MoE architecture**: 30 B / 3 B-active and 235 B / 22 B-active variants, scaling capacity without proportional compute cost.
* **Strong reasoning and coding**: shares the Qwen3 generation's improvements in instruction following, math, and tool usage.
* **Long-context capability**: 256 K-token context inherited from the Qwen3 series.
* **Flexible scaling**: 30 B fits a single 8-GPU node; 235 B is the canonical multi-node target with FP8 rollout.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Qwen3-30B-A3B | 3 B / 30 B | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) |
| Qwen3-235B-A22B | 22 B / 235 B | [Qwen/Qwen3-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) |

## 3. Environment Setup

### 3.1 Required env vars

The 235 B launcher requires:

```bash theme={null}
export MASTER_ADDR=<head node IP>
```

Everything else is a Typer flag: `--model-dir` (default `/root/models`) for the checkpoints, `--data-dir` (default `/root/datasets`) for the datasets, `--output-dir` (default `/root/shared_data`) for what the run writes. Point them at a shared FS path reachable from every node. The 30 B launcher reads no env vars at all.

### 3.2 Download model + datasets

```bash theme={null}
hf download Qwen/Qwen3-30B-A3B --local-dir /root/models/Qwen3-30B-A3B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024

# 235 B (multi-node, FP8 by default)
hf download Qwen/Qwen3-235B-A22B-FP8 --local-dir /root/models/Qwen3-235B-A22B-FP8
```

### 3.3 HF → Megatron `torch_dist` conversion

```bash theme={null}
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3-30B-A3B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Qwen3-30B-A3B \
   --save          /root/models/Qwen3-30B-A3B_torch_dist
```

Drive the conversion across more GPUs / nodes for the 235 B variant; the launcher reads `<model-dir>/Qwen3-235B-A22B_torch_dist` as `--ref-load`.

## 4. Launch

### 4.1 Quick start

```bash theme={null}
# 30 B (1 node × 8 GPU) — launcher handles download + conversion + submit
cd /root/miles
python scripts/run_qwen3_30b_a3b.py

# 235 B (8 actor nodes × 8 GPU + a 64-GPU rollout pool)
export MASTER_ADDR=...
python scripts/run_qwen3_235b_a22b.py
```

### 4.2 Multi-node fan-out

`run_qwen3_235b_a22b.py` ssh-fans-out to workers via `--ray-hostfile` (default `/root/mpi_rack_hostfile`) itself; you only need `MASTER_ADDR` set on the head node. Pass `--no-join-ray-workers` when the cluster is already joined. The 30 B launcher is single-node.

## 5. Recipe Configuration

### 5.1 Parallelism

| Launcher | Backend | TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | Actor GPUs |
| - | - | - | - | - | - | - | - | - |
| `run_qwen3_30b_a3b.py` (H100, 1 node) | Megatron | 4 | 1 | 1 | 8 | 1 | 32768 | 8 (1 × 8) |
| `run_qwen3_235b_a22b.py` | Megatron | 4 | 4 | 2 | 16 | 1 | 16384 | 64 (8 × 8) |
| `run_qwen3_sft.py --model-name Qwen3-235B-A22B` | Megatron | 4 | 1 | 1 | 32 | 1 | 9216 | 32 (4 × 8) |

`run_qwen3_235b_a22b.py` sets `--decoder-last-pipeline-num-layers 22` to balance the layer count across PP=4. Its rollout is disaggregated, so it needs a further 64 GPUs (`--rollout-num-gpus`) on top of the actor pool.

### 5.2 Algorithm

* **30 B launcher**: GRPO with `--eps-clip 0.2 --eps-clip-high 0.28`.
* **235 B launcher**: GSPO (`--advantage-estimator gspo`, `--eps-clip 4e-4`); `--use-kl-loss` is not passed.

### 5.3 Rollout & SGLang

`run_qwen3_30b_a3b.py` (H100, 1 node, BF16 rollout):

```bash theme={null}
--rollout-num-gpus-per-engine 8
--sglang-mem-fraction-static 0.7
--sglang-cuda-graph-max-bs-decode 512
```

`run_qwen3_235b_a22b.py`:

```bash theme={null}
--rollout-num-gpus-per-engine 32
--sglang-mem-fraction-static 0.7
--sglang-enable-dp-attention
--sglang-dp-size 4
--sglang-ep-size 32
--sglang-enable-dp-lm-head
--sglang-cuda-graph-bs-decode 1 2 4 8 16 24 ... 256
--sglang-moe-a2a-backend deepep
--sglang-deepep-mode auto
```

### 5.4 Optimizer

Both `run_qwen3_30b_a3b.py` (H100, 1 node) and `run_qwen3_235b_a22b.py` enable CPU Adam:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

`run_qwen3_30b_a3b.py` removes them when running on Blackwell (`B200/B300/GB200/GB300`) per the hardware match in the launcher.

### 5.5 Notable quirks

* **30 B launcher** supports FP8 / MXFP8 / INT4 rollout, Blackwell hardware, Megatron-bridge mode, and MIS via Typer flags.
* **235 B defaults to the FP8 HF checkpoint** — pass `--no-rollout-fp8` to roll out from the BF16 directory instead.
* **R3 not on by default**; opt-in via `run_qwen3_30b_a3b.py --enable-mis` (TIS / RS) for routing-stability experiments.

## 6. Pairs Well With

* [Low Precision RL](/docs/advanced/low-precision)
* [Rollout Routing Replay (R3)](/docs/advanced/miles-router)


# Qwen3-Next 80B-A3B
Source: https://miles.radixark.com/docs/models/qwen/qwen3-next

Launch recipes for Qwen3-Next-80B-A3B-Thinking on the Megatron backend.

## 1. Model Introduction

[Qwen3-Next](https://huggingface.co/collections/Qwen/qwen3-next) is Alibaba's next-generation Qwen architecture, swapping classical attention for a hybrid Gated DeltaNet + Full Attention design.

**Key highlights:**

* **Hybrid Attention**: combines Gated DeltaNet (linear attention) with Full Attention to handle context lengths up to 262 K tokens efficiently.
* **Highly Sparse MoE**: 80 B total / 3 B active per token — drastically reduces FLOPs per token without sacrificing model capacity.
* **Multi-Token Prediction (MTP)**: built-in MTP layer enables EAGLE-style speculative rollout out of the box.
* **HuggingFace-wrapped Megatron backend**: miles loads the `Qwen/Qwen3-Next-80B-A3B` HF module as a Megatron stage without re-implementing GDN from scratch.

## 2. Supported Variants

| Model | Active / Total | HF ID |
| - | - | - |
| Qwen3-Next-80B-A3B-Thinking | 3 B / 80 B | [Qwen/Qwen3-Next-80B-A3B-Thinking](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Thinking) |

## 3. Environment Setup

### 3.1 Required env vars

```bash theme={null}
export MASTER_ADDR=<head node IP>
```

`--topology 4node` needs `MASTER_ADDR` for the ray fan-out; `--topology single-node` needs nothing. Paths are Typer flags: `--model-dir` (default `/root/models`) for the staged checkpoint, `--data-dir` (default `/root/datasets`) for the datasets, `--output-dir` (default `/root/shared_data`) for what the run writes. On a multi-node run all three must be on a shared FS.

### 3.2 Download model + datasets

```bash theme={null}
hf download Qwen/Qwen3-Next-80B-A3B-Thinking --local-dir /root/models/Qwen3-Next-80B-A3B-Thinking
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024     --local-dir /root/datasets/aime-2024
```

### 3.3 HF → Megatron `torch_dist` conversion

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py qwen3-next-80B-A3B)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
PYTHONPATH=/root/Megatron-LM torchrun --nproc-per-node 8 \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --hf-checkpoint /root/models/Qwen3-Next-80B-A3B-Thinking \
   --save          /root/models/Qwen3-Next-80B-A3B-Thinking_torch_dist
```

## 4. Launch

### 4.1 Quick start

```bash theme={null}
cd /root/miles
export MASTER_ADDR=...
python scripts/run_qwen3_next_80b_a3b.py --topology 4node

# single 8-GPU node (6 GPUs train, 2 serve rollout)
python scripts/run_qwen3_next_80b_a3b.py --topology single-node
```

### 4.2 Multi-node fan-out

`--topology 4node` performs ssh fan-out internally over `/root/mpi_rack_hostfile` — set `MASTER_ADDR` on the head node and the launcher reaches out to the workers. Pass `--no-join-ray-workers` when the ray cluster is already complete. `--topology single-node` never fans out.

## 5. Recipe Configuration

### 5.1 Parallelism

Both layouts are one launcher, `scripts/run_qwen3_next_80b_a3b.py`, selected by `--topology`:

| `--topology` | Backend | TP | PP | CP | EP | expert-TP | `max_tokens_per_gpu` | Actor GPUs |
| - | - | - | - | - | - | - | - | - |
| `4node` | Megatron | 2 | 4 | 2 | 8 | 1 | 8192 | 32 (4 × 8) |
| `single-node` | Megatron | 1 | 6 | 1 | 1 | 1 | 2048 | 6 (1 × 6) |

`4node` colocates the rollout engines on the training GPUs; `single-node` dedicates the remaining 2 GPUs of the node to rollout, which also shrinks every batch dimension.

### 5.2 Algorithm

Both topologies use GSPO (`--advantage-estimator gspo --eps-clip 4e-4`); `--use-kl-loss` is not passed.

### 5.3 Rollout & SGLang

`4node` enables EAGLE speculative rollout:

```bash theme={null}
--rollout-num-gpus-per-engine 8
--sglang-mem-fraction-static 0.8
--sglang-ep-size 8
--sglang-cuda-graph-bs-decode 1 2 4 8 16 24 ... 128

--sglang-speculative-algorithm EAGLE
--sglang-speculative-num-steps 2
--sglang-speculative-eagle-topk 1
--sglang-speculative-num-draft-tokens 3
--sglang-enable-draft-weights-cpu-backup
--sglang-max-running-requests 512
```

`single-node` drops the EAGLE block and uses `--rollout-num-gpus-per-engine 2 --rollout-num-gpus 2 --sglang-mem-fraction-static 0.8 --sglang-ep-size 1`.

### 5.4 Optimizer

Both topologies enable CPU Adam:

```bash theme={null}
--optimizer-cpu-offload
--overlap-cpu-optimizer-d2h-h2d
--use-precision-aware-optimizer
```

### 5.5 Notable quirks

* Gated DeltaNet (GDN) is loaded via the HuggingFace bridge; miles doesn't re-implement GDN in Megatron native code.

## 6. Pairs Well With

* [Backends Beyond Megatron](/docs/advanced/architecture-support)
* [Rollout Routing Replay (R3)](/docs/advanced/miles-router)
* [Speculative Decoding](/docs/advanced/speculative-decoding)


# Beam 1.0
Source: https://miles.radixark.com/docs/models/reflection/beam

Launch recipes for Reflection's Beam 1.0 (501B MoE, 23B active) — full-parameter BF16 GRPO on 96 or 64 H200 GPUs, and LoRA on 32.

The complete Beam RL implementation is open at the Miles pull request:
[`radixark/miles#4149`](https://github.com/radixark/miles/pull/4149), branch
[`beam-support`](https://github.com/radixark/miles/tree/beam-support). The scripts and image
below come from that branch.

## 1. Model Introduction

[Beam-1.0-501B-A23B](https://huggingface.co/reflection/Beam-1.0-501B-A23B) is Reflection's 501B-parameter MoE model (23B active) with a 1M-token context window, released in BF16, FP8 and NVFP4. miles trains the BF16 checkpoint, either the full model or a LoRA adapter on the frozen base, and loads it straight from HF through Megatron-Bridge, with no offline `torch_dist` conversion.

**Key highlights:**

* **Hybrid attention**: 43 sliding-window layers (window 512) with RoPE, and 9 full-attention layers without positional encoding, one every six layers plus the last.
* **SSMax**: the full-attention layers scale each query by a factor that grows with its position.
* **Gated GQA**: 48 query heads over 8 KV heads, QK-norm, and a sigmoid output gate with a per-head bias.
* **Sandwich norms**: every attention and MLP block is normalized on both its input and its output, and its residual add is scaled by depth.
* **MoE**: 256 routed experts with grouped top-8 sigmoid routing and no shared expert; the first 3 layers are dense.
* **FP32 softcap and router**: the logits are soft-capped at 30 in FP32, and the router weights stay in FP32 in the checkpoint, the trainer and SGLang.

## 2. Supported Variants

| Model | Active / Total | Layers | HF ID |
| - | - | - | - |
| Beam-1.0-501B-A23B | 23 B / 501 B | 52 (3 dense + 49 MoE) | [reflection/Beam-1.0-501B-A23B](https://huggingface.co/reflection/Beam-1.0-501B-A23B) |

The Beam bridge loads only this BF16 checkpoint; the FP8 and NVFP4 releases are not supported for training.

## 3. Environment Setup

### 3.1 Checkpoint and data

There is no separate download step: the `full-train` commands in section 4 fetch the checkpoint into `--model-dir` and DAPO-Math-17K and AIME25 into `--data-dir`, then train. Every node must see both at the same path, and the checkpoint needs no Megatron conversion — the trainer reads it through the Beam bridge at startup.

The fetch also writes `aime-2025/aime-2025-answer-format.jsonl`, the AIME25 problems wrapped in the training prompts' answer-format instruction, which the evals read.

To stage the download on its own — to finish it before claiming the cluster, or to share one copy between the recipes — `prepare` does that part alone, and `run_beam_lora.py prepare` is identical.

### 3.2 SGLang

The engine needs the Beam support on SGLang's [`beam-miles`](https://github.com/sgl-project/sglang/tree/beam-miles) branch.

The `radixark/miles:beam` image (multi-arch, built on `radixark/miles:v0.1.2` by `docker/Dockerfile.beam`) carries it and the recipe code:

```bash theme={null}
docker pull radixark/miles:beam
```

## 4. Launch

### 4.1 Bring up the Ray cluster

Run `radixark/miles:beam` on every node, with the checkpoint, the datasets and `--output-dir` on shared storage.

Pass `--join-ray-workers` and the launcher builds the cluster itself, ssh-ing every host in `--ray-hostfile` (default `/root/mpi_rack_hostfile`) into it before submitting the job; skip the rest of this section. Without that flag the launchers only submit, so bring the cluster up yourself:

```bash theme={null}
# on the head node
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
# on every other node
ray start --address=${MASTER_ADDR}:6379 --node-ip-address ${WORKER_IP} --num-gpus 8 --disable-usage-stats
```

`ray status` on the head must list 96, 64 or 32 GPUs before launching, for 4.2, 4.3 and 4.4 respectively.

### 4.2 Full parameter, 12 nodes

`full-train` downloads the checkpoint and datasets if they are not already in place, then runs on twelve 8×H200 nodes: eight train and four host two rollout engines. Each trainer node needs at least 2 TB free on node-local NVMe for the optimizer state, at `/scratch/beam-optimizer` by default (`--optimizer-disk-dir`):

```bash theme={null}
MASTER_ADDR=<head ip> python scripts/run_beam.py full-train --num-rollout 100 \
    --model-dir /root/models/Beam --data-dir /root/datasets \
    --output-dir /root/shared_data/beam-run
```

Training is asynchronous (`train_async.py`): the next rollout generates while the current one trains, so its responses are one update stale; `--no-async-training` switches to synchronous `train.py`. Use `train` instead of `full-train` to skip the download and go straight to training. The run evaluates on AIME25 every five rollouts and saves a checkpoint every twenty.

### 4.3 Full parameter, colocated, 8 nodes

`--colocate` takes the same arguments on eight 8×H200 nodes. Every GPU trains and also hosts one of four rollout engines, which receive the new weights over CUDA IPC, and training is synchronous. Each node's NVMe holds both the optimizer state and the paused trainer:

```bash theme={null}
MASTER_ADDR=<head ip> python scripts/run_beam.py full-train --colocate --num-rollout 100 \
    --model-dir /root/models/Beam --data-dir /root/datasets \
    --output-dir /root/shared_data/beam-colocate-run
```

### 4.4 LoRA, 4 nodes

`run_beam_lora.py` runs on four 8×H200 nodes: two train and two host one rollout engine, asynchronously like `train`. It needs no NVMe, since the adapter's optimizer state stays on the GPUs:

```bash theme={null}
MASTER_ADDR=<head ip> python scripts/run_beam_lora.py full-train --num-rollout 100 \
    --model-dir /root/models/Beam --data-dir /root/datasets \
    --output-dir /root/shared_data/beam-lora-run
```

## 5. Recipe Configuration

### 5.1 Parallelism

| Launcher | TP | PP | CP | EP | expert-TP | GPUs | rollout engines |
| - | - | - | - | - | - | - | - |
| `run_beam.py full-train` | 8 | 1 | 1 | 8 | 8 | 96 (12 × 8): 64 train, 32 rollout | 2 × TP16, `--sglang-mem-fraction-static 0.85` |
| `run_beam.py full-train --colocate` | 8 | 1 | 1 | 8 | 8 | 64 (8 × 8), shared | 4 × TP16, `--sglang-mem-fraction-static 0.7` |
| `run_beam_lora.py full-train` | 8 | 1 | 1 | 16 | 1 | 32 (4 × 8): 16 train, 16 rollout | 1 × TP16, `--sglang-mem-fraction-static 0.85` |

`--sequence-parallel` and full recompute are on in every recipe. Core attention runs unfused on padded `bshd` micro-batches of one sample, up to 5,120 tokens (1,024 prompt + 4,096 response).

### 5.2 Algorithm

GRPO on DAPO-Math-17K with `--eps-clip 0.2 --eps-clip-high 0.28 --eps-clip-c 10`, per-token loss and `--use-tis`; no KL loss and no entropy bonus. Each rollout samples 8 responses for each of 32 prompts at temperature 1.0 with top-p and top-k off, and its 256 responses feed one optimizer step.

### 5.3 Optimizer

Adam at a constant `--lr 1e-6`, no weight decay, gradient clip 1.0, distributed optimizer. The FP32 master weights and BF16 Adam moments live on node-local NVMe and stream through the GPU at each step:

```bash theme={null}
--offload-train-disk-dir /scratch/beam-optimizer
--offload-train-disk-chunk-mb 256
--stream-optimizer-state-to-disk
--stream-optimizer-state-moment-dtype bf16
--grad-reduce-in-bf16
```

`--colocate` adds `--offload-train-target disk`: while the engines generate, the paused trainer is written to the same NVMe.

### 5.4 Rollout

The engines serve BF16 with Triton attention and MoE. Each engine takes its share of the 256 responses, and the trainer sends new weights after every rollout:

```bash theme={null}
--sglang-dtype bfloat16
--sglang-attention-backend triton
--sglang-moe-runner-backend triton
--sglang-chunked-prefill-size 8192
--sglang-max-running-requests 128      # 256 / engines: 128, colocated 64, LoRA 256
--sglang-cuda-graph-max-bs-decode 128  # same as max-running-requests
--update-weight-transfer-mode broadcast
--update-weights-interval 1
```

### 5.5 LoRA

`run_beam_lora.py` keeps the same data, reward, algorithm and rollout, and trains an adapter on every linear layer (attention q/k/v/o and the output gate, the dense MLPs, and each routed expert) while the router, embeddings and LM head stay frozen. What changes:

```bash theme={null}
--lora-rank 16 --lora-alpha 32 --lora-dropout 0.0
--expert-model-parallel-size 16
--expert-tensor-parallel-size 1        # expert LoRA is wrong under expert-TP > 1 in Megatron-Bridge
--lr 1e-4 --min-lr 1e-5 --lr-decay-style cosine --lr-warmup-iters 10 --lr-decay-iters 400
--use-rollout-routing-replay --use-miles-router
```

The adapter's optimizer state stays on the GPUs, so the NVMe flags of 5.3 are dropped, and each weight update sends only the adapter (\~8.7 GiB).

## 6. Pairs Well With

* [LoRA Training and Serving](/docs/advanced/lora)
* [Rollout Routing Replay (R3)](/docs/advanced/miles-router)
* [Disk Offload](/docs/advanced/disk-offload)


# Step-5-Preview
Source: https://miles.radixark.com/docs/models/stepfun/step5

Launch recipe for StepFun's Step-5-Preview (600B MoE, sliding-window + DSA sparse attention) — FP8 rollout, BF16 training, on two 8-GPU B300 nodes.

## 1. Model Introduction

[Step-5-Preview](https://huggingface.co/stepfun-ai/Step-5-Preview) is StepFun's 600B-parameter MoE model (27B active per token). miles trains its text decoder.

**Key highlights:**

* **Hybrid attention**: 69 sliding-window layers (window 512) and 23 full-attention layers with DSA sparse attention: an indexer scores 8-token blocks, and each query attends up to 512 selected blocks plus its current one. The two types use different RoPE (base 1e4 on all dims vs 1e7 on a third).
* **GQA with a gate**: 64 query heads over 4 KV groups, head dim 192, and a head-wise sigmoid gate on the attention output.
* **MoE**: 352 routed experts, top-8 sigmoid routing scaled by 3.0, plus a shared expert; layers 0–2 and 91 are dense.
* **FP8 checkpoint**: routed experts in FP8 E4M3 with 128×128 block scales, except the last three MoE layers (88–90), which ship in BF16. The rollout serves the checkpoint as released; the trainer dequantizes the FP8 experts on load and trains in BF16.

## 2. Supported Variants

| `--model-name` | What it is | HF source |
| - | - | - |
| `Step5` | text decoder, 92 layers | [stepfun-ai/Step-5-Preview](https://huggingface.co/stepfun-ai/Step-5-Preview) |

The vision tower and MTP layers stay in the checkpoint but are not trained (text-only RL).

## 3. Environment Setup

### 3.1 Download

The launcher reads the checkpoint at `<model-dir>/<model-name>`; `prepare` fetches dapo-math-17k:

```bash theme={null}
hf download stepfun-ai/Step-5-Preview --local-dir /root/models/Step5
python scripts/run_step5.py prepare --model-name Step5 --num-nodes 2 --num-gpus-per-node 8
```

The full model needs no Megatron conversion: the trainer reads the FP8 checkpoint through the Step5 bridge at startup and dequantizes the experts on read, since a `torch_dist` copy (1.2 TB) would not fit node disk.

The AIME 2026 evals read `<data-dir>/aime-2026/aime-2026.jsonl` (`prompt` and `label` rows), which `prepare` does not fetch; pass `--no-eval-aime` without it.

### 3.2 SGLang

The engine needs the Step-5-Preview support in [sgl-project/sglang#44356](https://github.com/sgl-project/sglang/pull/44356).

The `radixark/miles:step5` image (multi-arch, built on `lmsysorg/sglang:v0.5.20`) carries it and the recipe code:

```bash theme={null}
docker pull radixark/miles:step5
```

## 4. Launch

### 4.1 Two nodes

Start a ray cluster over both nodes, then launch from the head. `--ray-hostfile` joins the other node over ssh; to start the cluster yourself, set `MILES_SCRIPT_EXTERNAL_RAY=1` instead:

```bash theme={null}
MASTER_ADDR=<head ip> python scripts/run_step5.py train --model-name Step5 \
    --num-nodes 2 --num-gpus-per-node 8 --ray-hostfile /root/hostfile --num-rollout 20
```

On two 8×B300 nodes, with 32 prompts × 8 samples per step, responses capped at 2048 tokens (the recipe default is 32768), DAPO dynamic sampling on and evals off, a step took a median 274 s over 80 steps: 81 s of training and 95 s of weight update, the rest mostly rollout.

## 5. Recipe Configuration

### 5.1 Parallelism

| `--model-name` | TP | PP | CP | EP | expert-TP | GPUs | rollout engine |
| - | - | - | - | - | - | - | - |
| `Step5` | 4 | 1 | 1 | 16 | 1 | 16 (2 × 8) | TP8/EP8 per node, `--sglang-mem-fraction-static 0.75` |

TP is capped at 4 by the four KV groups. Context parallelism is not supported: each layer builds its own RoPE, which would re-slice THD positions per CP rank. `--sequence-parallel` is on whenever TP > 1. A micro-batch holds one prompt beside a full-length response (`--max-tokens-per-gpu 34816` on 2 × 8) with THD packing and `--use-dynamic-batch-size`.

### 5.2 Algorithm

GRPO with `--eps-clip 0.2 --eps-clip-high 0.28 --entropy-coef 0.00` and `--rm-type math` on dapo-math-17k, 8 samples per prompt at temperature 1.0 with top-p and top-k off. No KL loss. The rollout runs FP8 experts against a BF16 trainer, so `--use-rollout-routing-replay` replays the rollout's expert choices and `--use-tis` covers the rest.

### 5.3 Optimizer

Adam at `--lr 1e-6`. With EP 16 on 16 GPUs, expert data parallelism is 1, so the distributed optimizer cannot shard the expert Adam state, and it does not fit HBM beside the BF16 weights and gradients. It stays in pinned host memory (224 GB per GPU) at 6 bytes per parameter, BF16 moments plus the int16 low halves of the FP32 masters, and streams through the GPU in 1 GiB chunks:

```bash theme={null}
--use-precision-aware-optimizer
--main-params-dtype fp32 --exp-avg-dtype bf16 --exp-avg-sq-dtype bf16
--chunked-optimizer-state-offload
--optimizer-state-offload-fraction 1.0
--optimizer-state-offload-chunk-size-mb 1024
--grad-reduce-in-bf16
--train-env-vars '{"PYTORCH_ALLOC_CONF": "pinned_max_cached_size_mb:4"}'
```

`pinned_max_cached_size_mb:4` stops PyTorch from rounding the 12 and 24 MiB expert tensors' pinned blocks up to a power of two, which costs a third more host memory.

### 5.4 Notable quirks

* **`--window-attn-skip-freq` reads backwards**: Megatron treats a 1 as window attention, not a skip. Step5 passes 1 for `layer % 4 != 3`, the 69 sliding-window layers, with `--window-size 511,0`: the query plus 511 previous keys, 512 in total.
* **Frozen DSA indexer**: every parameter but the indexer trains, and the full-attention layers recompute the rollout's region selection with it. The weight sync never sends it; the engine keeps its own copy.
* **`--sglang-max-total-tokens 2097152`** is a correctness bound: an unbounded KV pool overflows the DSA kernels' int32 strides.
* **`--sglang-moe-runner-backend triton`**: routing replay needs the engine's materialized top-k ids.

## 6. Pairs Well With

* [Rollout Routing Replay (R3)](/docs/advanced/miles-router)
* [Low Precision RL](/docs/advanced/low-precision)


# Thinking Machines
Source: https://miles.radixark.com/docs/models/thinkingmachines/index

Miles recipes for Thinking Machines Lab models — Inkling (975 B), a multimodal MoE with short convolution, relative attention, and a shared-expert sink.

Miles ships a native Megatron recipe for **Inkling**, Thinking Machines Lab's 975 B / 41 B-active multimodal mixture-of-experts model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders. The same backend drives both full-parameter and LoRA RL.

## Variants

| Model | Active / Total | Layers | HF ID | Recipe |
| - | - | - | - | - |
| Inkling | 41 B / 975 B | 66 | [thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling) | [Inkling](/docs/models/thinkingmachines/inkling) |
| Inkling-Small | 12 B / 276 B | 42 | [thinkingmachines/Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small) | [Inkling-Small](/docs/models/thinkingmachines/inkling-small) |

## Fastest path to train

Inkling needs 16 nodes of 4× GB300 and the `radixark/miles:inkling` image:

```bash theme={null}
cd /root/miles
python scripts/run_inkling.py train \
   --model-name Inkling --train-mode full --task dapo_math \
   --num-nodes 16 --num-gpus-per-node 4
```

See the [Inkling](/docs/models/thinkingmachines/inkling) page for the architecture summary, HF → Megatron conversion, validated parallelism layouts, training attention backends, LoRA RL, and multimodal RL.


# Inkling
Source: https://miles.radixark.com/docs/models/thinkingmachines/inkling

Launch recipe for Inkling (975 B), Thinking Machines' multimodal MoE with short convolution, relative attention, and a shared-expert sink.

## 1. Model Introduction

[Inkling](https://huggingface.co/thinkingmachines/Inkling) is a mixture-of-experts transformer released by Thinking Machines Lab, with 975 B total parameters and 41 B active, a context window of up to 1 M tokens, and pretraining on 45 trillion tokens of text, images, audio and video. Its architecture introduces short convolution, attention with relative positional embedding, and a novel MoE design with a shared-expert sink. Miles implements Inkling as a native Megatron model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders, and the same backend drives both full-parameter and LoRA RL.

**Key highlights**:

* **ShortConv**: a short causal convolution with a residual connection, applied on the K and V streams and on the attention and MLP/MoE outputs. Miles implements it with fused, packing-aware Triton kernels.
* **Relative attention**: a learned relative-position bias replaces positional embeddings; the stack mixes sliding-window and full-attention layers, with context length up to 1 M tokens.
* **Shared-expert sink MoE**: sigmoid top-k routing where the router also scores the shared experts and renormalizes their weights together with the selected routed experts'.
* **Train–inference consistency by construction**: customized relative-attention score-mod, ShortConv, and FP32 SwiGLU/combine kernels, while Rollout Routing Replay (R3) replays the rollout's routed expert IDs, including over media-expanded multimodal sequences.

## 2. Supported Variants

| Model | Active / Total | Layers | HF ID |
| - | - | - | - |
| Inkling | 41 B / 975 B | 66 | [thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling) |

## 3. Quick start

### 3.1 One-line launch

```bash theme={null}
# Pull the image
docker pull radixark/miles:inkling

# Full-parameter GRPO on 16 nodes x 4 GB300, inside the container
cd /root/miles
python scripts/run_inkling.py train \
   --model-name Inkling --train-mode full --task dapo_math \
   --num-nodes 16 --num-gpus-per-node 4

# LoRA GRPO (rank 32, all-linear), same cluster
python scripts/run_inkling.py train \
   --model-name Inkling --train-mode lora --task dapo_math \
   --num-nodes 16 --num-gpus-per-node 4
```

### 3.2 Launcher path defaults

| Flag | Default | Use |
| - | - | - |
| `--model-dir` | `/root/models` | parent directory holding the HF checkpoint and Megatron `_torch_dist` artifacts |
| `--hf-checkpoint` | `{model-dir}/{model-name}` | released HF weights |
| `--torch-dist` | `{model-dir}/{model-name}_torch_dist` | converted Megatron torch-dist checkpoint |
| `--data-dir` | `/root/datasets` | HF datasets (dapo-math-17k, geo3k, …) |
| `--save-dir` | unset | training checkpoints |

Every option can also be preconfigured via `MILES_SCRIPT_<FIELD_NAME_UPPER>` env vars (precedence: CLI flag > env var > built-in default). Pass any extra Miles / Megatron / SGLang flags through `--extra-args`.

## 4. Script breakdown

This section explains what the launcher does under the hood, and how to drive each stage manually.

### 4.1 Download model + datasets

```bash theme={null}
# inside the container
hf download thinkingmachines/Inkling --local-dir /root/models/Inkling
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/datasets/dapo-math-17k
hf download --repo-type dataset hiyouga/geometry3k --local-dir /root/datasets/geo3k
```

Pass `--hf-checkpoint <path>` to the launcher when the weights are already on a shared filesystem.

### 4.2 HF → Megatron `torch_dist` conversion

Inkling ships in BF16, so conversion is a single distributed `torch_dist` shard (no precision cast). The model definition comes from `scripts/models/inkling.py`:

```bash theme={null}
cd /root/miles
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py inkling)" || exit 1
read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}"
CONVERT_KEEP_PP1=1 PYTHONPATH=/root/Megatron-LM torchrun \
   --nproc-per-node 4 --nnodes 4 \
   --master-addr ${MASTER_ADDR} --master-port 12345 \
   --node-rank ${NODE_RANK} \
   tools/convert_hf_to_torch_dist.py \
   ${MODEL_ARGS[@]} \
   --tensor-model-parallel-size 4 \
   --pipeline-model-parallel-size 1 \
   --expert-model-parallel-size 4 \
   --hf-checkpoint /root/models/Inkling \
   --save /root/models/Inkling_torch_dist/
```

`CONVERT_KEEP_PP1=1` keeps the conversion at PP1: without it the converter auto-bumps PP toward the rank count, which is incompatible with `--tensor-model-parallel-size 4` at 16 ranks.

The saved `torch_dist` checkpoint is parallelism-agnostic: training can load it under any validated TP / PP / EP layout. Point the launcher at the result with `--torch-dist`.

### 4.3 Multi-node fan-out

Start each pod with the image and a shared filesystem mounted at the same path on every node, then bring up Ray:

```bash theme={null}
# head node
ray start --head --num-gpus 4 --disable-usage-stats
# each worker
ray start --address=${HEAD_IP}:6379 --num-gpus 4 --disable-usage-stats
```

Set `MILES_SCRIPT_EXTERNAL_RAY=1` to point the launcher at this existing Ray cluster. When it is unset, the launcher boots a local Ray head itself.

## 5. Example Recipe Configuration

### 5.1 Megatron Parallelism

These are the validated layouts shipped with the launcher. Other TP / EP / PP / CP combinations that fit your compute can be supplied via `--extra-args` (the launcher raises on untested GPU counts to keep you honest).

| Hardware | GPUs | TP | SP | PP | EP | expert-TP | Notes |
| - | - | - | - | - | - | - | - |
| GB300 | 64 | 4 | on | 4 | 16 | 1 | `--decoder-last-pipeline-num-layers 15` |

### 5.2 Algorithm

GRPO with truncated importance sampling. The launcher defaults: global batch size 32, group size 8, maximum response length 4096 with truncation, Adam with `lr 1e-6`, constant decay. Rollout Routing Replay is on by default (`--use-rollout-routing-replay`). It replays only the routed expert IDs; the continuous routed *and* shared weights are recomputed from the current router under one common normalization, so gradients still flow through both.

### 5.3 Training attention backends

The `MILES_INKLING_ATTN_BACKEND` environment variable selects the training-side attention implementation:

| Backend | Role | 8K packed, rel-extent 1024, GB300 |
| - | - | - |
| `flex` (default) | block-sparse FlexAttention, differentiable fwd+bwd | \~34 ms fwd+bwd, 10.9 GB peak |
| `te` | Transformer Engine DPA reference | \~189 ms, 46 GB peak |
| `fa4` | serving-bit-identical FA4 forward + TE-recompute backward | \~195 ms, 46 GB peak |

### 5.4 Rollout & SGLang

```bash theme={null}
SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 16
   --sglang-attention-backend fa4
   --sglang-moe-runner-backend triton
   --sglang-mamba-radix-cache-strategy extra_buffer
   --sglang-enable-multimodal
   --sglang-context-length 8192

   # full-parameter
   --sglang-mem-fraction-static 0.6
   --sglang-max-running-requests 64
   --sglang-max-total-tokens 327680

   # LoRA (replaces the three flags above)
   --sglang-ep-size 16
   --no-offload-rollout --no-offload-train
   --sglang-mem-fraction-static 0.65
   --sglang-max-running-requests 32
   --sglang-lora-backend triton
   --sglang-lora-use-virtual-experts
   --sglang-max-loras-per-batch 1
)
```

Weight updates stream Megatron shards into SGLang's tensor layout in bounded buckets over CUDA IPC (colocated). In LoRA mode only the adapter is synchronized.

### 5.5 Optimizer

```bash theme={null}
--optimizer adam
--lr 1e-6 --lr-decay-style constant
--attention-softmax-in-fp32

# gradient reduction: bf16 under --fully-async, fp32 otherwise
--grad-reduce-in-bf16  # --fully-async
--accumulate-allreduce-grads-in-fp32  # colocated

# full-parameter only, set by the launcher
--micro-batch-size 1
--offload-train-disk-dir /tmp/train_offload
--offload-train-target disk  # colocated: back the paused actor with node-local NVMe
--stream-optimizer-state-to-disk  # --fully-async: stream the optimizer state instead
--offload-train-disk-chunk-mb 256  # --fully-async
```

Within a single GB300 rack the 975 B policy, gradients, and FP32 optimizer state exceed GPU memory, so Miles streams Megatron `DistributedOptimizer` state between a bounded GPU working set and node-local NVMe. The offload changes storage placement, not the update math. The paused training actor's weights are additionally disk-backed through `torch_memory_saver`. LoRA training skips both offloads and uses dynamic batching (`--use-dynamic-batch-size --max-tokens-per-gpu 4096`) instead of the fixed micro-batch.

> [!WARNING]
>
> 1. **Fixed micro-batches for full-parameter runs.** Dynamic token packing exposes a PP-p2p × EP-all-to-all NCCL launch-order race on varlen shapes; the launcher pins `--micro-batch-size 1` for full-parameter training. Keep it pinned until the pad-to-fixed fix lands.
> 2. **One LoRA adapter per engine.** RL serving uses `--sglang-max-loras-per-batch 1` with the triton LoRA backend and virtual experts: the engine holds exactly the current policy's adapter (see [LoRA RL](#6-lora-rl)).

## 6. LoRA RL

`--train-mode lora` switches the same backend and parallel stack to adapter-only training. The base model stays frozen in both runtimes, and GRPO updates only the low-rank factors:

$$
y = Wx + \frac{\alpha}{r}\,B(Ax).
$$

The adapter follows Inkling's released LoRA schema, covering the attention, dense MLP, MoE, and LM-head projections (defaults: rank 32, `alpha = 32`, all-linear). The routed experts use a shared-outer factorization: one factor is shared across experts while the expert-specific factors follow EP sharding.

After each optimizer step, Miles exports a serving-ready adapter directly from the distributed training state and hands it to the colocated SGLang worker over CUDA IPC; the frozen base is never re-transferred. This cuts the weight update from 49.4 s to 2.5 s (20×) and brings training-step time to \~85 % of full-parameter training, with no optimizer offload needed. Released Inkling adapters in safetensors format support warm starts via `--lora-adapter-path`, and per-rank adapter checkpoints allow exact training resume.

## 7. Multimodal RL

`--task geo3k` switches to the vision-language recipe: structured multimodal rollouts, Megatron-side vision and audio towers, and routing replay preserved across the media-expanded sequence. Each image or audio sentinel in the rendered prompt expands into its patch/frame positions before packing, and the R3 trace is indexed over the expanded sequence, so replay stays aligned even though SGLang made its routing decisions after expansion. Production recipes keep the vision and audio towers frozen and train the language model (or its adapter) around their embeddings; training the towers is available as an experimental option. The same path supports audio inputs and both full-parameter and LoRA training.


# Inkling-Small
Source: https://miles.radixark.com/docs/models/thinkingmachines/inkling-small

Launch recipe for Inkling-Small (276 B), the compact sibling of Inkling — same architecture, 4-node H200 footprint.

## 1. Model Introduction

[Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small) is the compact member of Thinking Machines Lab's Inkling family: a 276 B-total / 12 B-active-parameter, 42-layer multimodal MoE (256 routed + 2 shared experts, top-6 sigmoid routing) that matches — and on some benchmarks exceeds — the flagship 975 B model (e.g. SWEBench Verified 80.2 vs 77.6) at a fraction of the deployment footprint. The architecture is the same as [Inkling](/docs/models/thinkingmachines/inkling) — ShortConv, local/global relative attention, and the shared-expert-sink MoE — so everything on the Inkling page (architecture summary, attention backends, R3 routing replay, LoRA schema, multimodal RL) applies unchanged. This page only covers what differs: the model registry entry and the validated small-cluster launch profile.

## 2. Supported Variants

| Model | Active / Total | Layers | HF ID | Recipe |
| - | - | - | - | - |
| Inkling-Small | 12 B / 276 B | 42 | [thinkingmachines/Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small) | this page |

## 3. Quick start

Validated on 4 nodes × 8 H200 (TP4 SP PP8 EP4, DP1):

```bash theme={null}
cd /root/miles

# Full-parameter GRPO. 276 B fits with the CPU-offloaded optimizer -
# no NVMe streaming needed (unlike the 975 B recipe).
python scripts/run_inkling.py train \
   --model-name Inkling-Small --train-mode full --task dapo_math \
   --num-nodes 4 --num-gpus-per-node 8 \
   --lr 5e-5 --rollout-batch-size 64 --global-batch-size 128 \
   --sglang-context-length 4096 --rollout-max-response-len 2048 \
   --extra-args "--offload-train-target cpu --sglang-mem-fraction-static 0.65 \
      --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer"

# LoRA GRPO (rank 32, all-linear), same cluster; adapter-only weight
# sync swaps in ~4 s per rollout.
python scripts/run_inkling.py train \
   --model-name Inkling-Small --train-mode lora --task dapo_math \
   --num-nodes 4 --num-gpus-per-node 8 \
   --lr 2e-4 --rollout-batch-size 64 --global-batch-size 128 \
   --sglang-context-length 4096 --rollout-max-response-len 2048
```

The model definition lives in `scripts/models/inkling-small.py` (`MODEL_ARGS_NUM_LAYERS` overrides the layer count for sliced smoke/parity checkpoints). HF → `torch_dist` conversion uses the same tool as Inkling with this recipe file — a single 8-GPU node (TP8 EP8) converts it in one pass.

## 4. Validated parallelism

| Hardware | GPUs | TP | SP | PP | EP | expert-TP | Notes |
| - | - | - | - | - | - | - | - |
| H200 | 32 | 4 | on | 8 | 4 | 1 | `--decoder-last-pipeline-num-layers 7` (42 = 7×5 + 7) |

Batch shape is configurable from the launcher (`--rollout-batch-size`, `--global-batch-size`; defaults 32/64). The validated Small runs used 64/128: full with `--lr 5e-5`, LoRA with `--lr 2e-4` — both produce a steadily rising dapo-math reward curve. At the launcher's conservative LoRA default (5e-6) the zero-initialised B factors take hundreds of rollouts to accumulate a visible delta-W, which reads as "not learning".


# Subscribe
Source: https://miles.radixark.com/docs/subscribe

Stay up to date with Miles releases and announcements.

* GitHub releases: [radixark/miles/releases](https://github.com/radixark/miles/releases)
* Twitter/X: [@radixark](https://twitter.com/radixark)


# Agentic Rollout (TITO)
Source: https://miles.radixark.com/docs/user-guide/agentic-rollout

Configure an OpenAI-compatible agent loop with Token-In-Token-Out trajectory assembly.

Multi-turn agentic rollout in Miles runs through the Token-In-Token-Out (TITO)
session server. Your agent exchanges OpenAI-compatible chat messages, while Miles
preserves the exact token IDs, logprobs, and routed experts produced during
inference and assembles them into training samples. For the design rationale, see
[No Token Left Behind](https://lmsys.org/blog/2026-05-13-no-token-left-behind/).

This page owns the agentic path: wrapper setup, the custom agent contract,
session behavior, token ownership, model-family selection, and verification.
Use [Generate Endpoint](/docs/user-guide/generate-endpoint) for the lower-level,
stateless `/generate` interface.

> [!WARNING]
>
> **No VLM support yet.** Currently the TITO session path cannot carry image or video inputs. For vision-language models, use the
> [Generate Endpoint](/docs/user-guide/generate-endpoint) path instead.

## Configure the wrapper

Select `agentic_tool_call.generate` as the custom generate function. The wrapper
registers `--custom-agent-function-path` and `--max-seq-len`, creates a TITO
session for each rollout, invokes your agent, and collects the resulting samples.

```bash theme={null}
AGENTIC_ARGS=(
   --custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate
   --custom-agent-function-path    my_agent.run
   --use-session-server
   --hf-checkpoint                 Qwen/Qwen3-4B
   --tito-model                    qwen3
)
```

> [!WARNING]
>
> **Do not apply the chat template to prompt data manually.** Do not pass
> `--apply-chat-template`: `Sample.prompt` must remain a `messages` list. The
> session server renders the first turn and incrementally appends later turns with
> the selected `--tito-model` implementation.

## Write the agent loop

Use `--custom-agent-function-path` to name an async function with this contract:

```python theme={null}
async def run_agent(
    base_url: str,
    prompt,
    request_kwargs: dict,
    metadata: dict,
    **kwargs,
) -> dict | None:
    ...
```

Send OpenAI-compatible chat requests to the session-scoped endpoint:

```python theme={null}
from miles.utils.http_utils import post


async def run_agent(base_url, prompt, request_kwargs, metadata, **kwargs):
    payload = {"model": "default", "messages": prompt, **request_kwargs}
    await post(f"{base_url}/v1/chat/completions", payload)
    return None
```

* `base_url` already includes `/sessions/<id>`; do not append the session path.
* `prompt` is the input sample's OpenAI `messages` list.
* `request_kwargs` contains the rollout sampling settings in
  `ChatCompletionRequest`-compatible form. For example, Miles maps
  `max_new_tokens` to `max_tokens`.
* `metadata` contains the sample metadata, session identifiers, and configured
  `max_seq_len`. Forward only the fields your environment needs.
* Return a dictionary to merge rewards, reports, or metrics into each output
  sample's metadata, or return `None` when there is nothing to add.

For structured parsing, the payload may use SGLang's
`ChatCompletionRequest`-compatible fields, which extend the OpenAI format.

### Where the agent runs

By default (`--custom-agent-function-mode subproc`) every agent call runs in a
child process of its own, started by the rollout process. One episode that blocks
on a slow step, or crashes, then affects only itself; on the rollout's single
event loop it would stall every other episode, long enough under many episodes
for their sandbox requests to time out. The contract changes in three ways:

* The arguments and the return value cross a process boundary, so both must be
  picklable.
* Module-level state lasts one call. A module-level semaphore no longer limits
  anything; use `rollout_semaphore` or `rollout_lock` from
  `miles.rollout.agentic.rollout_limits`, which hold across the rollout's agent
  processes. The number of concurrent calls is already bounded by the rollout's
  batch settings.
* Cleanup the agent leaves in a thread must be a non-daemon thread. The call
  returns without waiting for it, and the process waits up to ten minutes for it
  before exiting.

Cancelling a call cancels the agent inside its process, so cleanup such as closing
a sandbox still runs; an agent still running two minutes later is stopped.

`--custom-agent-function-mode inline` runs the agent on the rollout's event loop,
as before; use it for debugging and for agents that keep state across calls.

### Optional teardown hook

The module named by `--custom-agent-function-path` may expose an `abort` function
alongside the agent entry point:

```python theme={null}
async def abort(args) -> None:
    ...  # cancel this agent's in-flight external work
```

Miles calls this hook in the rollout process during oversampling abort after it
stops in-flight SGLang generation, so under `subproc` it cannot reach the agent
calls' module state; tell the external backend to stop, as below. Use it when the
agent drives an external sandbox or agent server that would otherwise keep issuing
completion requests until its own length limit or timeout. The hook is optional;
modules without it continue to work.

See [`swe_agent_function.abort`](https://github.com/radixark/miles/blob/main/examples/swe-agent-harbor-docker/swe_agent_function.py)
for an implementation that flushes the Harbor agent server.

## TITO

### Leave token ownership to Miles

Send the full `messages` history on every turn. On the first request, the
session server renders the selected template into `input_ids`. After a
successful completion, it checkpoints those prompt IDs together with the output
token IDs and logprobs returned by SGLang.

On later requests, the server reuses the deepest applicable checkpoint,
tokenizes only the appended suffix, and sends the joined `input_ids` to SGLang.
During collection, Miles aligns the turn outputs against the accumulated TITO
sequence, trims model-specific boundary tokens, and builds the training sample.

> [!WARNING]
>
> **Do not set TITO control fields.** Leave `input_ids`, `routed_experts_start_len`, and `logprob_start_len` unset. The server selects `lora_path` and rejects conflicting control values with HTTP 400. It forces `logprobs=True`, `return_meta_info=True`, and the response metadata needed for TITO.

Chat responses from both session-server versions omit `choices[*].meta_info`. Ordinary JSON replies retain the message and `usage`, and carry standard `logprobs` only when the request set `logprobs`. On training sessions that record score-centering candidates, their `top_logprobs` lists are empty: the candidates travel as arrays in `meta_info`. The server keeps the full metadata and logprobs in session records for training sample collection and inspection through `GET /sessions/{id}`. Streaming chunks omit both.

### Choose template options per session

* Ordinary `chat_template_kwargs` use request > continued turn > launch defaults. Fields required to preserve the reused prompt retain their recorded values; fixed model settings override conflicts. Qwen3.8 selects `reasoning_effort` on a new root and preserves its value or absence on continuation. When the field is absent, the template defaults to `xhigh`.
* Omitted `tools` inherit on continuation; incompatible changes return HTTP 400 by default. Sampling parameters remain per-request. New roots resolve without turn history.
* Each successful turn records the full resolved request as `turn_args`. Exported session, tree-node and sample metadata omit `input_ids` and `messages` from this snapshot; the stored history remains complete.

### Choose the session behavior

Use `{"evaluation": true}` in `POST /sessions` for evaluation; omitting it defaults to training. The agentic generator sets this automatically, and the purpose stays fixed across turns, retries, and branches.

`temperature`, `top_p`, and `top_k` in `POST /sessions` provide defaults for omitted or `null` chat fields. Training requests must match the registered temperature or receive HTTP 400; eval temperature and `top_p`/`top_k` remain overridable. The agentic generator registers the sample's resolved values automatically.

Evaluation forces `return_sampling_mask`, `return_routed_experts`, and `return_indexer_topk` off and ignores `routed_experts_start_len`. TITO, logprobs, and sample collection still apply; engine-internal capture may remain enabled.

History handling depends on the selected server version:

* **v1 is linear.** Each request must extend the previous messages at the tail.
  Retrying the latest turn may roll back one assistant checkpoint, including to
  an empty session when retrying the first turn. Earlier divergence or a larger
  rollback is rejected.
* **v2 (Experimental) is an append-only tree.** A request attaches to the deepest checkpoint
  whose complete message path prefixes the request. Any unmatched suffix creates
  a branch, and existing branches are never deleted. A path whose last generation
  ended with `finish_reason=length` cannot be extended. By default, a leaf becomes
  no sample when a later request re-sent its exact prompt tokens; a later request
  that differs is a separate branch and its own sample
  (`--session-sample-picker-path`).

Whether a replayed message counts as "the same" as the stored one is decided by
`--session-message-matcher` (default `strict`); see
[Choose replay matching](#choose-replay-matching).

The v1 wrapper returns one `Sample`. The v2 wrapper returns a `list[Sample]`, one
for each selected tree leaf. Both versions reject `--partial-rollout`. With R3
replay, every pause mode except `retract` requests only additional R3 rows.
`retract` returns full R3 data on every turn and emits a warning because the
payloads can become very large.

Set `--max-seq-len` to cap the context length. Miles also includes this value in the
metadata passed to your agent so an external environment can stop early.

### Pick your `--tito-model`

There is no auto-detection. Pick the family matching your model. Each named
family resolves a maintainer-verified `FIXED_TEMPLATE` registration from
`--tito-model` alone. The registration owns the bundled Jinja or
HuggingFace-native template, fixed template arguments, and the bundled SGLang
reasoning and tool-call parsers.

A named family rejects `--chat-template-path` overrides and conflicting fixed
arguments. Use `--tito-model default` for a custom or checkpoint-native renderer,
but treat it as best-effort until it passes the checks below.

| Your model | `--tito-model` |
| - | - |
| Qwen3 | `qwen3` |
| Qwen3.5 | `qwen35` |
| Qwen3.6 | `qwen36` |
| Qwen3.8-27B | `qwen38small` |
| Qwen3.8-Flash-Next | `qwen4exp` |
| Qwen3-Thinking-2507 / Qwen3-Next | `qwennext` |
| GLM-4.7 / 5 / 5.2 | `glm47` |
| GLM-5.3 / GLM-5.3-Flash (text sessions) | `glm53` |
| NVIDIA Nemotron 3 Nano / Super / Ultra | `nemotron3` |
| Kimi K2.5 / K2.6 | `kimi25` / `kimi26` |
| MiniMax M2.5 / M2.7 | `minimax_m25` / `minimax_m27` |
| DeepSeek-V3.2 / V4 | `deepseekv32` / `deepseekv4` |
| Inkling / Inkling-Small | `inkling` |
| Unregistered model or custom template (best-effort) | `default` |

More model families and verification history live in
[issue #712](https://github.com/radixark/miles/issues/712).

### Verify a new model TITO

To add a named family, register its `TITOTokenizer` and `FIXED_TEMPLATE` in
[`tito_tokenizer.py`](https://github.com/radixark/miles/blob/main/miles/utils/chat_template_utils/tito_tokenizer.py),
then run both checks. Either failure blocks support.

```bash theme={null}
# CPU / fast: rendered token sequences remain append-only
python scripts/tools/verify_chat_template.py \
    --model <hf-id> --tito-model <family>

# GPU / end to end: the invariant holds under real model inference
python scripts/tools/verify_session_tito_tokenizer.py \
    --hf-checkpoint <hf-id> --tito-model <family> \
    --sglang-reasoning-parser <rp> --sglang-tool-call-parser <tcp> \
    --rollout-num-gpus-per-engine 1
```

### Choose replay matching

Some agent harnesses do not replay model messages verbatim: they may reserialize tool-call arguments, replace empty `arguments` with `"{}"`, or omit `reasoning_content` on the next request. Under the default matcher those replays count as divergence — v1 rolls back (or rejects), v2 branches a new lineage.

`--session-message-matcher` is process-wide and defaults to `strict`. It accepts a built-in selector or a trusted dotted import path.

| Selector | Behavior |
| - | - |
| `strict` | Preserves the existing comparison of `role`, `content`, `reasoning_content`, and `tool_calls`, including empty-value and tool-call `index` normalization. |
| `loose_tool_call` | Accepts everything `strict` accepts, plus equivalent JSON-object representations of `tool_calls[].function.arguments`. Call IDs, types, function names, order, unknown fields, and `reasoning_content` still have to match. |
| `role_content_only` | Compares only normalized `role` and `content`. **High risk:** different tool-call or reasoning histories can collapse into one session lineage. |
| dotted import path | Loads a trusted synchronous custom matcher; see [Customization](/docs/user-guide/customization#session-message-matcher). |

The matcher only decides whether the message a client replays and the message stored at the same position in the session count as the same one.

* On a mismatch, the existing paths apply: v1 rolls back (or rejects), v2 branches.
* On a match, the stored messages and token snapshot stay authoritative inside the reusable prefix; only the suffix beyond it is tokenized anew from the client input.

Miles does not reconcile tool-call IDs across that boundary: deployments choosing `role_content_only` must themselves keep a stored call ID `A` followed by a replayed tool result referencing `B` protocol-compatible.

## Example

[`examples/swe-agent-harbor-docker`](https://github.com/radixark/miles/tree/main/examples/swe-agent-harbor-docker)
wires a multi-turn SWE agent, TITO session server, model-family registration,
reward, length limit, and environment teardown into production launchers.


# Argument Groups
Source: https://miles.radixark.com/docs/user-guide/argument-groups

The launch-script argument groups used by Miles recipes, with links to the flags that belong in each group.

Miles launch scripts are Python (`scripts/run_*.py`). The grouping is deliberately
boring: each script builds one string per operational concern, concatenates them into
`train_args`, and hands that to `execute_train`, which submits `train.py` or
`train_async.py` as a Ray job.

Use this page to decide where a flag belongs. Use the [CLI Reference](/docs/user-guide/cli-reference)
when you need the full default and type for an individual flag. For how a script is laid
out and how to override one, see [Launch Script](/docs/user-guide/launch-script).

| Group | Owns | Where it comes from |
| - | - | - |
| [model args](#model-args) | Architecture constants and plugin specs | `scripts/models/<megatron_model_type>.py`, spliced in by `execute_train` |
| [`ckpt_args`](#ckpt-args) | Actor, reference, HF tokenizer/config, save paths | Launch script |
| [`rollout_args`](#rollout-args) | Prompt data, sampling, reward, train/eval batch flow | Launch script |
| [`eval_args`](#eval-args) | Evaluation datasets and eval-only sampling overrides | Launch script |
| [`perf_args`](#perf-args) | Parallelism, recomputation, dynamic batching | Recipe defaults |
| [`grpo_args`](#grpo-args) | RL objective, KL, clipping, entropy, advantage estimator | Recipe defaults |
| [`optimizer_args`](#optimizer-args) | Learning rate, schedule, weight decay, Adam betas | Recipe defaults |
| [`sglang_args`](#sglang-args) | Rollout engine topology and `--sglang-*` passthrough | Deployment shape |
| [`misc_args`](#misc-args) | GPU layout, colocation, dropout, dashboard | Launch script |

The names above are the local variables every launcher uses; a script that needs no eval
simply leaves `eval_args` empty. Model args are the one group a launcher does not build:
it passes `megatron_model_type` to `execute_train`, which resolves the matching file
under `scripts/models/` and prepends those flags to the command line.

## Model args - architecture constants

Model args tell Megatron what model it is instantiating. Megatron cannot infer all
architecture details from a HuggingFace checkpoint, so each recipe loads a matching
file from `scripts/models/`.

Common entries:

| Flag family | Example |
| - | - |
| Transformer shape | `--num-layers`, `--hidden-size`, `--num-attention-heads` |
| Tokenizer/model dimensions | `--seq-length`, `--max-position-embeddings`, `--vocab-size` |
| Rotary and attention variants | `--rotary-base`, `--rotary-percent`, `--kv-channels` |
| MoE architecture | `--num-experts`, `--moe-router-topk`, `--moe-grouped-gemm` |
| Plugin specs | `--spec miles_plugins.models.qwen3_5 get_qwen3_5_spec` |

Keep these values aligned with the checkpoint's `config.json`. If one checkpoint in a
family changes rotary base, vocab padding, or normalization epsilon, override the
sourced defaults in the launch script.

## `ckpt_args` - checkpoint paths

`ckpt_args` wires the three model roles in a run:

| Role | Flag |
| - | - |
| HuggingFace directory for tokenizer, config, and SGLang boot | `--hf-checkpoint` |
| Frozen reference model for KL anchoring | `--ref-load` |
| Actor resume point | `--load` |
| Actor output directory | `--save` |

`--load` and `--save` usually point to the same directory. If `--load` has no
`latest_checkpointed_iteration.txt`, Miles warm-starts the actor from `--ref-load`.

## `rollout_args` - sampling and reward

`rollout_args` controls data entering the loop and how many samples each rollout
produces.

| Concern | Flags |
| - | - |
| Prompt data | `--prompt-data`, `--input-key`, `--label-key`, `--apply-chat-template` |
| Rollout volume | `--rollout-batch-size`, `--n-samples-per-prompt`, `--num-rollout` |
| Training consumption | `--global-batch-size`, `--num-steps-per-rollout` |
| Sampling | `--rollout-temperature`, `--rollout-top-p`, `--rollout-max-response-len` |
| Reward | `--rm-type`, `--custom-rm-path` |
| Filtering | `--over-sampling-batch-size`, `--dynamic-sampling-filter-path` |

The rollout volume and training consumption must satisfy the
[four-knob invariant](/docs/user-guide/concepts#the-four-knob-invariant).

## `eval_args` - evaluation overrides

Evaluation reuses the rollout stack but usually runs with a different dataset and more
deterministic sampling.

Common entries:

| Concern | Flags |
| - | - |
| Cadence | `--eval-interval` |
| Dataset | `--eval-prompt-data` |
| Eval group size | `--n-samples-per-eval-prompt` |
| Eval-only generation | `--eval-max-response-len`, `--eval-top-p`, `--eval-temperature` |

Flags not set in `eval_args` inherit from `rollout_args`.

## `perf_args` - parallelism and memory

`perf_args` controls how training is sharded and how activation memory is managed.

| Concern | Flags |
| - | - |
| Tensor parallelism | `--tensor-model-parallel-size`, `--sequence-parallel` |
| Pipeline parallelism | `--pipeline-model-parallel-size` |
| Context parallelism | `--context-parallel-size` |
| Expert parallelism | `--expert-model-parallel-size`, `--expert-tensor-parallel-size` |
| Recomputation | `--recompute-granularity`, `--recompute-method`, `--recompute-num-layers` |
| Dynamic batching | `--use-dynamic-batch-size`, `--max-tokens-per-gpu` |

Megatron exposes TP, PP, CP, EP, and ETP, but not every product of those dimensions is
valid or worth using for every model. Start from the recipe's tested combination and
see [parallelism compatibility](/docs/user-guide/training-backend#parallelism-compatibility) before changing
more than one dimension.

## `grpo_args` - RL objective

`grpo_args` controls the policy-gradient objective and the stability terms around it.

| Concern | Flags |
| - | - |
| Algorithm | `--advantage-estimator` |
| KL | `--use-kl-loss`, `--kl-loss-coef`, `--kl-loss-type` |
| Clipping | `--eps-clip`, `--eps-clip-high` |
| Entropy | `--entropy-coef`, `--observe-training-entropy` |
| Loss reduction | `--calculate-per-token-loss` |
| Precision/off-policy safety | `--use-tis` |

Zero-weight KL is recipe-specific. `--use-kl-loss --kl-loss-coef 0.00` still loads the
reference and logs KL; it does not remove the reference model.

## `optimizer_args` - optimizer schedule

`optimizer_args` carries the optimizer choice and scalar schedule.

Common entries:

| Concern | Flags |
| - | - |
| Optimizer | `--optimizer` |
| Learning rate | `--lr`, `--min-lr`, `--lr-decay-style` |
| Adam | `--adam-beta1`, `--adam-beta2`, `--adam-eps` |
| Regularization | `--weight-decay`, `--clip-grad` |

Post-training is sensitive to large updates. Most recipes start near `1e-6` and use a
constant schedule unless the model page says otherwise.

## `sglang_args` - rollout engine passthrough

`sglang_args` configures the inference side. Miles owns
`--rollout-num-gpus-per-engine`; everything prefixed with `--sglang-` is forwarded to
`python -m sglang.launch_server` after removing the prefix.

Common entries:

| Concern | Flags |
| - | - |
| Engine tensor parallelism | `--rollout-num-gpus-per-engine` |
| Engine memory | `--sglang-mem-fraction-static` |
| Context length | `--sglang-context-length` |
| MoE serving | `--sglang-ep-size`, `--sglang-moe-a2a-backend`, `--sglang-enable-dp-attention` |
| Debugging | `--sglang-log-level` |

SGLang parallelism is separate from trainer parallelism. For example,
`--rollout-num-gpus-per-engine` maps to the SGLang server's TP size, not Megatron's
`--tensor-model-parallel-size`.

## `misc_args` - GPU layout and everything else

`misc_args` carries what the other groups do not: how many GPUs the actor gets, whether
it shares them with the rollout engines, the Megatron knobs a recipe pins once and never
tunes, and the optional dashboard.

Common entries:

| Concern | Flags |
| - | - |
| GPU layout | `--actor-num-nodes`, `--actor-num-gpus-per-node`, `--num-gpus-per-node` |
| Colocation | `--colocate` |
| Numerics pinned by the recipe | `--attention-dropout 0.0`, `--hidden-dropout 0.0`, `--attention-softmax-in-fp32`, `--accumulate-allreduce-grads-in-fp32` |
| Attention kernel | `--attention-backend` |
| Observability | `--use-miles-dashboard`, `--dump-details` |

Under `--colocate` the actor and the engines share the same GPUs and take turns, so
`--rollout-num-gpus` is ignored; see
[Training Backends](/docs/user-guide/training-backend#3-choosing-the-gpu-layout).


# CLI Reference
Source: https://miles.radixark.com/docs/user-guide/cli-reference

Every command-line flag Miles accepts, grouped by subsystem.

Miles is configured through command-line flags passed to `train.py` or
`train_async.py`. The Megatron flags (such as `--num-layers`, `--rotary-base`,
`--recompute-granularity`) are inherited via Megatron's argument parser; Miles adds
its own flags through an `extra_args_provider`. Run `python3 train.py --help` against
your installed Megatron source for the canonical list.

This page has two passes.

1. **Essentials** lists the flags most runs actually touch.
2. **Complete reference** lists every Miles flag with type and default.

***

## Essentials

### Cluster topology

| Flag | Default | What |
| - | - | - |
| `--actor-num-nodes` | `1` | Total nodes for the actor. |
| `--actor-num-gpus-per-node` | `8` | GPUs per actor node. |
| `--rollout-num-gpus` | derived | GPUs for SGLang rollout (ignored when `--colocate`). |
| `--rollout-num-gpus-per-engine` | `1` | TP size of each SGLang engine. |
| `--colocate` | off | Share GPUs between actor and rollout. |
| `--run-uuid` | generated | Machine-readable id for this launch; 16 lowercase hex characters. |

See [Training Backends](/docs/user-guide/training-backend) for what `--colocate` flips on under the hood.

### Batch sizing

The four-knob invariant:

```
rollout_batch_size × n_samples_per_prompt
  = global_batch_size × num_steps_per_rollout
```

| Flag | Typical | What |
| - | - | - |
| `--rollout-batch-size` | `16 – 256` | Prompts per rollout. |
| `--n-samples-per-prompt` | `4 – 16` | Responses per prompt (GRPO group size). |
| `--global-batch-size` | derived | Samples per optimizer step. |
| `--num-steps-per-rollout` | `1` | Optimizer steps per rollout. |
| `--num-rollout` | `1000 – 10000` | Total rollout iterations. |

### Memory and throughput

| Flag | Default | What |
| - | - | - |
| `--use-dynamic-batch-size` | off | Pack varlen samples into micro-batches. |
| `--max-tokens-per-gpu` | `–` | Token budget per micro-batch per GPU. Required when dynamic batching is on. |
| `--context-parallel-size` | `1` | Spread a single sample across N CP ranks. |
| `--recompute-granularity` | Megatron default | `full` or `selective`. |
| `--recompute-method` | Megatron default | `uniform` or `block`. |
| `--recompute-num-layers` | Megatron default | Layers per recompute chunk. |

Rule of thumb: start with `max_tokens_per_gpu = rollout_max_response_len / cp_size`,
then push up until you OOM.

### RL algorithm

| Flag | Default | What |
| - | - | - |
| `--advantage-estimator` | `grpo` | `grpo`, `gspo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`. On-policy distillation is not an estimator — enable it with `--use-opd` on top of any of these. |
| `--use-kl-loss` | off | Compute KL against the reference model. |
| `--kl-loss-coef` | `0.0` | Weight of KL in the loss (0 means monitor only). |
| `--kl-loss-type` | `k1` | `k1`, `k2`, `k3`, `low_var_kl`. |
| `--entropy-coef` | `0.0` | Entropy bonus weight. |
| `--observe-training-entropy` | off | Log training entropy even when `--entropy-coef` is `0.0`; detached from backward when the coefficient is zero. |
| `--eps-clip` | `0.2` | PPO/GRPO low clip. |
| `--eps-clip-high` | `–` | Asymmetric high clip (DAPO-style). |
| `--use-tis` | off | Truncated Importance Sampling for train/inference precision mismatch. |

### Sampling

| Flag | Default | What |
| - | - | - |
| `--rollout-temperature` | `1.0` | Sampling temperature. |
| `--rollout-top-p` | `1.0` | Top-p truncation. Values below `1` enable [sampling-support replay](/docs/advanced/sampling-support-replay) and require a positive top-k. |
| `--rollout-top-k` | `-1` | Top-k truncation. Positive values enable [sampling-support replay](/docs/advanced/sampling-support-replay). |
| `--rollout-max-response-len` | `–` | Max tokens per response. |
| `--rollout-stop-token-ids` | model default | Stop token IDs. Override when generations don't stop. |
| `--apply-chat-template` | off | Apply the tokenizer's chat template. |
| `--rollout-shuffle` | off | Shuffle prompts each rollout. |

### Optimizer

| Flag | Default | What |
| - | - | - |
| `--optimizer` | `adam` | `adam`, `sgd`. |
| `--lr` | `1e-6` | Learning rate. Post-training is sensitive to large updates; recipes typically stay near `1e-6`. |
| `--lr-decay-style` | `constant` | `constant`, `linear`, `cosine`. |
| `--weight-decay` | `0.1` | L2 weight decay. |
| `--adam-beta1`, `--adam-beta2` | `0.9, 0.98` | Adam moments. |

### Logging

| Flag | Default | What |
| - | - | - |
| `--use-wandb` | off | Log to Weights and Biases. |
| `--wandb-project` | – | wandb project name. |
| `--log-interval` | `1` | Stdout log cadence (rollouts). |
| `--save-interval` | – | Checkpoint cadence (rollouts). Recipes typically set 20 to 100. |

### SGLang passthrough

Any flag accepted by `python -m sglang.launch_server` is accepted by Miles with the
`--sglang-` prefix:

```bash theme={null}
--sglang-log-level INFO
--sglang-mem-fraction-static 0.8
--sglang-ep-size 8
--sglang-moe-a2a-backend deepep
--sglang-enable-dp-attention
```

See [SGLang docs](https://docs.sglang.io) for the full list.

### Environment variables

Set these in Ray's `env_vars` for multi-node runs:

| Variable | Effect |
| - | - |
| `TORCHINDUCTOR_FORCE_DISABLE_CACHES=1` | Workaround for torch-compile JSONDecodeError. |
| `RAY_DEDUP_LOGS=0` | Don't deduplicate worker logs. |
| `NCCL_DEBUG=INFO` | NCCL diagnostics. |
| `PYTHONPATH=/root/Megatron-LM` | Required when using the Megatron backend. |

***

## Complete reference

Sections mirror the launch-script argument groups.

### Cluster

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--cluster-backend` | enum | `ray` | `ray` launches workers from the driver; `kubernetes` expects them to already exist. `kubernetes` is refused during validation until a later milestone provisions those workers. Under `kubernetes`, `--use-prometheus` is ignored. |
| `--actor-num-nodes` | int | `1` | Total nodes for actor training. |
| `--actor-num-gpus-per-node` | int | `8` | GPUs per actor node. |
| `--rollout-num-gpus` | int | derived | Ignored under `--colocate`. |
| `--rollout-num-gpus-per-engine` | int | `1` | TP size of each SGLang engine. |
| `--colocate` | flag | off | Share GPUs between actor and rollout. Implicitly enables `--offload-train`, `--offload-rollout`, and defaults `--sglang-cuda-graph-backend-prefill=disabled`. |

### Model and checkpoints

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--train-backend` | enum | `megatron` | `megatron` or `fsdp`. |
| `--hf-checkpoint` | path | – | HF model dir. Provides tokenizer, config, and the weights FSDP loads. |
| `--ref-load` | path | – | Reference model in `torch_dist` format (Megatron). |
| `--load` | path | – | Actor checkpoint to resume from. |
| `--save` | path | – | Actor checkpoint write directory. |
| `--save-interval` | int | – | Rollouts between saves. |
| `--async-save` | flag | off | Write Megatron checkpoint shards asynchronously. With `--save-hf`, the native write overlaps the HF export. |
| `--save-trigger-sentinel` | path | – | If this file exists at a save point, save a checkpoint now (regardless of `--save-interval`) and remove the file. |
| `--custom-megatron-post-save-hook-path` | `<module>.<fn>` | – | Rank-0 callback after each checkpoint save. |
| `--model-name` | str | – | Set in multi-node to avoid `transformers` file-system race. |
| `--spec` | `<module> <fn>` | – | Plugin spec for custom architectures (e.g. `miles_plugins.models.qwen3_5 get_qwen3_5_spec`). |

### Rollout: data and batching

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--prompt-data` | str | – | Path to a single JSONL file. |
| `--input-key` | str | `input` | JSONL key to `Sample.prompt`. |
| `--label-key` | str | `None` | JSONL key to `Sample.label`. |
| `--metadata-key` | str | `metadata` | JSONL key to `Sample.metadata`. |
| `--apply-chat-template` | flag | off | Apply tokenizer chat template. |
| `--rollout-shuffle` | flag | off | Shuffle prompts each rollout. |
| `--num-rollout` | int | – | Total rollout iterations. If unset, derived from dataset size. |
| `--rollout-batch-size` | int | – | Prompts per rollout. |
| `--n-samples-per-prompt` | int | `1` | Responses per prompt. |
| `--global-batch-size` | int | derived | Samples per optimizer step. |
| `--num-steps-per-rollout` | int | – | Optimizer steps per rollout. Alternative to `--global-batch-size`; setting one derives the other. |
| `--over-sampling-batch-size` | int | – | Oversample size for dynamic sampling (DAPO). |
| `--balance-data` | flag | off | Balance per-rank token count. |

### Rollout: sampling

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--rollout-max-response-len` | int | – | Max tokens per response. |
| `--rollout-temperature` | float | `1.0` | Sampling temperature. |
| `--rollout-top-p` | float | `1.0` | Top-p truncation. Values below `1` require bounded [sampling-support replay](/docs/advanced/sampling-support-replay). |
| `--rollout-top-k` | int | `-1` | Top-k truncation (`-1` disables). Positive values enable [sampling-support replay](/docs/advanced/sampling-support-replay). |
| `--rollout-stop` | str+ | – | Stop strings. |
| `--rollout-stop-token-ids` | int+ | – | Stop token IDs. |

### Eval

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--eval-prompt-data` | str+ | – | One or more `name path` pairs. |
| `--eval-interval` | int | – | Rollouts between eval runs. |
| `--n-samples-per-eval-prompt` | int | `1` | Responses per eval prompt. |
| `--eval-max-response-len` | int | – | Max eval response length. Inherits from rollout if unset. |
| `--eval-temperature` | float | – | Eval temperature. Inherits from rollout if unset. |
| `--eval-top-p` | float | – | Eval top-p. Inherits from rollout if unset. |
| `--eval-top-k` | int | – | Eval top-k. Inherits from rollout if unset. |
| `--eval-num-gpus` | int | `0` | Dedicated eval fleet size. `0` = shared-engine eval. Requires `train_async.py`. |
| `--eval-num-gpus-per-engine` | int | `1` | Eval engine TP, independent of rollout TP. |
| `--eval-hf-dir` | str | – | Staging dir for per-eval HF snapshots (tmpfs recommended). Unset + `--save-hf` = reuse mode. |
| `--eval-max-in-flight` | int | `2` | Snapshots the trainer may export ahead of the eval backend. Evals are serialized, so this buys lead time, not concurrency — and one more staged snapshot. |
| `--eval-overflow-policy` | str | `backpressure` | At the cap: await the oldest eval, or `skip` the new point (logged as `eval/skipped_busy`). |
| `--eval-keep-snapshots` | int | `2` | Retired snapshots kept under `--eval-hf-dir`; with `--eval-max-in-flight` this bounds the staging dir. `--save-hf` output is never deleted. |
| `--eval-sglang-*` | – | – | Per-field override of any `--sglang-*` setting for the eval fleet only. Unset = inherit the rollout engines' value. Booleans take a `--no-` form (`--no-eval-sglang-enable-dp-attention`) so an inherited `True` can be turned off. `tp_size` is not exposed — use `--eval-num-gpus-per-engine`. |

### Performance

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--tensor-model-parallel-size` | int | `1` | TP. |
| `--pipeline-model-parallel-size` | int | `1` | PP. |
| `--context-parallel-size` | int | `1` | CP. |
| `--expert-model-parallel-size` | int | `1` | EP (MoE). |
| `--expert-tensor-parallel-size` | int | `1` | TP within experts. |
| `--sequence-parallel` | flag | off | Enable Megatron sequence parallel. |
| `--use-dynamic-batch-size` | flag | off | Pack varlen samples. Recommended for varlen workloads. |
| `--max-tokens-per-gpu` | int | – | Token budget per micro-batch per GPU. Required when dynamic batching is on. |
| `--micro-batch-size` | int | `1` | Ignored when dynamic batching is on. |
| `--recompute-granularity` | enum | Megatron default | `full` or `selective`. |
| `--recompute-method` | enum | Megatron default | `uniform` or `block`. |
| `--recompute-num-layers` | int | Megatron default | Recompute chunk size. |
| `--gradient-checkpointing` | flag | off | FSDP equivalent of recompute flags. |
| `--fsdp-cpu-offload` | flag | off | FSDP: offload params, grads, optimizer state to CPU. |
| `--fsdp-cpu-backend` | str | `gloo` | FSDP: CPU backend for hybrid offload. |
| `--dp-replicate-size` | int | `1` | FSDP2 hybrid-shard replica count. |
| `--attn-implementation` | str | `flash_attention_2` | FSDP only: passed to `transformers`, e.g. `flash_attention_2`, `flash_attention_3`, `sdpa`, `eager`. |

### RL algorithm

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--advantage-estimator` | enum | `grpo` | `grpo`, `gspo`, `ppo`, `reinforce_plus_plus`, `reinforce_plus_plus_baseline`. |
| `--use-kl-loss` | flag | off | Compute KL vs. reference. |
| `--kl-loss-coef` | float | `0.0` | KL weight in loss (0 means monitor). |
| `--kl-loss-type` | enum | `k1` | `k1`, `k2`, `k3`, `low_var_kl`. |
| `--entropy-coef` | float | `0.0` | Entropy bonus weight. |
| `--observe-training-entropy` | flag | off | Log detached training entropy when entropy bonus weight is zero. |
| `--eps-clip` | float | `0.2` | PPO/GRPO low clip. |
| `--eps-clip-high` | float | – | Asymmetric high clip. |
| `--use-tis` | flag | off | Truncated Importance Sampling. |
| `--use-routing-replay` | flag | off | Forward/backward routing consistency. |
| `--use-rollout-routing-replay` | flag | off | R3 — capture inference-side expert routing and replay it during training. |
| `--calculate-per-token-loss` | flag | off | Per-token loss reduction. |
| `--no-check-for-nan-in-loss-and-grad` | flag | off | Skip NaN/Inf guard (Megatron flag, debug only). |
| `--true-on-policy-mode` | flag | off | Strict on-policy: reject samples from a prior policy. |

### Optimizer

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--optimizer` | enum | `adam` | `adam`, `sgd`. |
| `--lr` | float | `1e-6` | Learning rate. |
| `--lr-decay-style` | enum | `constant` | `constant`, `linear`, `cosine`. |
| `--lr-warmup-iters` | int | `0` | Warmup steps (Megatron flag). |
| `--min-lr` | float | `0` | Lower LR bound for decay schedules (Megatron flag). |
| `--weight-decay` | float | `0.1` | L2 weight decay. |
| `--adam-beta1` | float | `0.9` | |
| `--adam-beta2` | float | `0.98` | |
| `--clip-grad` | float | `1.0` | Grad clipping (Megatron flag). |
| `--optimizer-cpu-offload` | flag | off | Megatron CPU Adam (Megatron flag). |
| `--overlap-cpu-optimizer-d2h-h2d` | flag | off | Overlap D2H/H2D with compute (Megatron flag). |
| `--use-precision-aware-optimizer` | flag | off | Precision-aware optimizer path (Megatron flag). |

### Reward and filters

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--rm-type` | str | – | Built-in reward: `math`, `dapo`, `deepscaler`, `gemma_math`, `f1`, `gpqa`, `ifbench`, `remote_rm`, `random`, `deterministic_random`. A `boxed_` prefix (e.g. `boxed_math`) extracts `\boxed{}` from the response before grading. |
| `--rm-url` | str | – | Endpoint when `--rm-type remote_rm`. |
| `--group-rm` | flag | off | Batched reward computation. |
| `--custom-rm-path` | str | – | Custom reward function (see [Customization](/docs/user-guide/customization)). |
| `--dynamic-sampling-filter-path` | str | – | Group filter (DAPO-style). |
| `--rollout-submission-granularity` | enum | driver | `group` or `sample`: what frees rollout submission capacity. Unset means `sample` under `--fully-async`, `group` otherwise. |
| `--buffer-filter-path` | str | – | Buffer dequeue filter. |
| `--rollout-sample-filter-path` | str | – | Per-sample filter. |

### SGLang and router

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--sglang-router-ip` | str | – | Must stay unset: external router mode was removed. Miles always starts its own router. |
| `--sglang-router-port` | int | – | Pins the port of the router miles starts; model `i` gets `port + i`. Unset lets it move off a busy port. |
| `--sglang-*` | passthrough | | Any flag accepted by `python -m sglang.launch_server` works with this prefix. |
| `--router-*` | passthrough | | Any flag accepted by the active router works with this prefix. |

Common `--sglang-*` flags:

```bash theme={null}
--sglang-mem-fraction-static 0.8
--sglang-context-length 32768
--sglang-log-level INFO
--sglang-ep-size 8
--sglang-enable-dp-attention
--sglang-moe-a2a-backend deepep
--sglang-moe-runner-backend triton
--sglang-deepep-mode auto
--sglang-cuda-graph-backend-prefill  # prefill graphs default to disabled in colocate mode
```

### Agentic sessions

These flags wire an OpenAI-compatible agent loop through Miles' TITO session
server. See [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout) for the request
contract, session behavior, and model-family selection.

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--custom-generate-function-path` | `<module>.<fn>` | – | Set to `miles.rollout.generate_hub.agentic_tool_call.generate` for the built-in agentic wrapper. |
| `--custom-agent-function-path` | `<module>.<fn>` | – | Async agent-environment loop. Registered after selecting the built-in agentic wrapper. |
| `--custom-agent-function-mode` | `subproc` / `inline` | `subproc` | Where each agent call runs: a child process of its own, or the rollout's event loop. See [Where the agent runs](/docs/user-guide/agentic-rollout#where-the-agent-runs). |
| `--use-session-server` | optional `v1` / `v2` | off | Bare flag (or `v1`) selects the linear append-only server; `v2` selects tree serving. Requires `--hf-checkpoint`. |
| `--tito-model` | enum | `default` | TITO model family. Named families load their registered fixed template; `default` is best-effort with a checkpoint-native or custom template. |
| `--max-seq-len` | int | – | Total tokens per session, including prompts, completions, and environment responses. Registered with the agentic wrapper. |
| `--session-server-ip` | str | router IP | Session-server bind address. |
| `--session-server-external-host` | str | – | Host that peers outside the cluster reach every session server on. Keeps the session servers on the head node. Leave unset when each node sets `MILES_NODE_EXTERNAL_IP`. |
| `--session-server-port` | int | auto | First port for standalone session-server instances. When unset, each worker port is auto-allocated. |
| `--session-server-workers` | int | `32` | Number of instances, at least 1; an explicit `--session-server-port` anchors a consecutive range. |
| `--session-sample-picker-path` | `<module>.<fn>` | `drop_same_prompt_retries` | v2 only: selects leaf samples before post-processing. The default trims identical re-sends, including a re-sent first turn; `drop_rolled_back_leaves` also trims a leaf whose later sibling sent a different request. |
| `--session-sample-postprocessor-path` | `<module>.<fn>` | `default_postprocess` | v2 only: finalizes loss masks and rewards. |

`--use-session-server v2` returns `list[Sample]` and rejects `--group-rm`, `--partial-rollout`, and `--recompute-logprobs-via-prefill`.

### MTP / speculative decoding

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--mtp-num-layers` | int | – | MTP layers the model has; built only with `--enable-mtp-training`. |
| `--enable-mtp-training` | flag | off | Build and train the MTP layers alongside the policy (Megatron only). |
| `--mtp-loss-scaling-factor` | float | `0.2` | Weight of MTP loss. |

### Fault tolerance

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-fault-tolerance` | flag | off | Enable rank-level recovery and heartbeats. |
| `--rollout-health-check-first-wait` | int | `0` | Grace period before heartbeats start. |
| `--rollout-health-check-interval` | int | `30` | Seconds between heartbeats. |
| `--rollout-health-check-timeout` | int | `30` | Heartbeat timeout. |

### Async / partial rollout

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--partial-rollout` | flag | off | Resume aborted rollouts in the next iteration. |

### Logging

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--use-wandb` | flag | off | Enable wandb. |
| `--wandb-project` | str | – | Project name. |
| `--wandb-group` | str | – | Group name. |
| `--log-interval` | int | `1` | Stdout log cadence (rollouts). |
| `--custom-rollout-log-function-path` | str | – | Custom train logger. |
| `--custom-eval-rollout-log-function-path` | str | – | Custom eval logger. |

### Profiling

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--profile-target` | enum+ | `[train_overall]` | Which sub-loop to profile: `train_overall`, `train_actor`, `train_log_probs`. |
| `--use-pytorch-profiler` | flag | off | FSDP only: enable PyTorch profiler. |
| `--profile-step-start` | int | `10` | FSDP only: first step to profile. |
| `--profile-step-end` | int | `12` | FSDP only: last step to profile. |
| `--memory-snapshot-path` | str | `snapshot.pickle` | FSDP only: memory snapshot output. |
| `--tensorboard-dir` | str | – | FSDP only: TensorBoard output dir. |

### Debugging

| Flag | Type | Default | Notes |
| - | - | - | - |
| `--debug-rollout-only` | flag | off | Skip Megatron, only spin up SGLang. |
| `--debug-train-only` | flag | off | Skip SGLang, only spin up Megatron. |
| `--save-debug-rollout-data` | path | – | Pickle every rollout to disk. The template must contain `{rollout_id}`. |
| `--load-debug-rollout-data` | path | – | Replay rollouts from disk (implies `--debug-train-only`). |
| `--deterministic-mode` | flag | off | Megatron deterministic mode. |

### Customization

See [Customization](/docs/user-guide/customization) for the full catalog of `--*-path` flags
that replace or extend Miles's behavior.


# Core Concepts
Source: https://miles.radixark.com/docs/user-guide/concepts

The four objects that make up every Miles job, and how data flows between them.

A Miles training job is a loop over four objects. Once you understand what each one
*is* and how data flows between them, every flag in the system has an obvious home.

## The four objects

```mermaid theme={null}
flowchart LR
    subgraph Rollout
      P[Prompt dataset] --> R[SGLang engines]
    end
    subgraph Scoring
      R --> RM[Reward model]
    end
    subgraph Training
      RM --> A[Actor — Megatron]
      A -. KL .-> RF[(Frozen reference)]
    end
    A -- weights --> R
```

| Object | Role | Lives in |
| - | - | - |
| **Prompt dataset** | Source of input examples | JSONL on disk (or `--data-source-path`) |
| **Rollout (SGLang engines)** | Generates responses given prompts | One or more SGLang servers behind a router |
| **Reward model** | Maps `(prompt, response, label) → score` | Built-in (`--rm-type`) or custom (`--custom-rm-path`) |
| **Actor (Megatron / FSDP)** | The model being trained | Megatron `torch_dist` checkpoint, or HF directory under FSDP |
| **Reference** | Frozen copy of the actor for KL anchoring | Loaded from `--ref-load`, never updated |

## The training loop

```python theme={null}
for it in range(num_rollout):
    # 1. Sample
    prompts   = dataset.sample(rollout_batch_size)
    responses = sglang.generate(prompts, n=n_samples_per_prompt)

    # 2. Score
    rewards   = reward_fn(prompts, responses, labels)

    # 3. Optimize
    for step in range(num_steps_per_rollout):
        batch = pack(prompts, responses, rewards, size=global_batch_size)
        loss  = grpo_loss(actor, ref_model, batch)
        loss.backward(); optimizer.step()

    # 4. Sync
    p2p_weight_transfer(actor → sglang_engines)
```

That's the whole thing. Every flag in Miles configures one of these four steps.

## The four-knob invariant

Two knobs govern the sampling half of the loop, two govern the training half, and they
are locked into a single equation:

```
rollout_batch_size × n_samples_per_prompt
  = global_batch_size × num_steps_per_rollout
```

**Every sample produced by rollout is consumed by training, and every sample consumed
by training was produced by rollout.** Set any three sides; Miles fills in the fourth.
Set all four inconsistently and Miles aborts with a validation error.

## Where every flag goes

Use this map when reading any launch script:

| Argument group | Concerns |
| - | - |
| [model args](/docs/user-guide/argument-groups#model-args) | Architecture constants (layers, hidden size, rotary base, ...) |
| [`ckpt_args`](/docs/user-guide/argument-groups#ckpt-args) | Filesystem paths for the actor / reference / save directory |
| [`rollout_args`](/docs/user-guide/argument-groups#rollout-args) | Prompt dataset, batch knobs, sampling parameters, reward type |
| [`eval_args`](/docs/user-guide/argument-groups#eval-args) | Eval dataset, cadence, sampling overrides for evaluation |
| [`perf_args`](/docs/user-guide/argument-groups#perf-args) | Parallelism (TP/PP/CP/EP/ETP), recomputation, dynamic batching |
| [`grpo_args`](/docs/user-guide/argument-groups#grpo-args) | RL algorithm, KL, clipping, entropy bonus, advantage estimator |
| [`optimizer_args`](/docs/user-guide/argument-groups#optimizer-args) | Learning rate, schedule, weight decay, Adam betas |
| [`sglang_args`](/docs/user-guide/argument-groups#sglang-args) | Engine TP, memory fraction, log level, `--sglang-*` passthrough |
| [`misc_args`](/docs/user-guide/argument-groups#misc-args) | GPU layout, colocation, pinned Megatron numerics, dashboard |

## Next

* [Launch Script](/docs/user-guide/launch-script) — what a launch script does when you run it,
  and how to override a recipe.
* [Training Backends](/docs/user-guide/training-backend) — Megatron-LM and FSDP, parallelism, checkpoints, and hooks.
* [Argument Groups](/docs/user-guide/argument-groups) — where each launch-script array belongs.
* [CLI Reference](/docs/user-guide/cli-reference) — every flag, grouped and fully cataloged.


# Customization
Source: https://miles.radixark.com/docs/user-guide/customization

The plug-points where you can drop in your own Python without forking Miles.

Most of Miles's behavior can be replaced with user-supplied Python by passing a
`--*-path` flag. This page lists every such hook, the function signature it expects,
and the default it replaces.

## At a glance

| Stage | Flag | Replaces |
| - | - | - |
| **Rollout** | `--rollout-function-path` | The whole rollout loop |
| | `--custom-agent-function-path` | The agent-environment loop inside a TITO session |
| | `--custom-generate-function-path` | A single sample's generation |
| | `--data-source-path` | How prompts are loaded |
| | `--eval-function-path` | The eval rollout |
| **Session** | `--session-message-matcher` | Prefix-replay equivalence |
| **Reward** | `--custom-rm-path` | Reward computation |
| | `--custom-reward-post-process-path` | Reward normalization |
| **Filtering** | `--dynamic-sampling-filter-path` | Per-group filter (DAPO) |
| | `--buffer-filter-path` | Buffer dequeue filter |
| | `--rollout-sample-filter-path` | Per-sample loss filter |
| | `--rollout-all-samples-process-path` | Inspect all samples post-rollout |
| | `--rollout-data-postprocess-path` | Mutate samples post-logprob |
| **Training** | `--custom-loss-function-path` | The loss formula |
| | `--custom-tis-function-path` | Importance sampling correction |
| | `--custom-pg-loss-reducer-function-path` | Loss reduction (Dr.GRPO) |
| | `--custom-convert-samples-to-train-data-path` | Sample to tensor batch |
| **Megatron hooks** | `--custom-megatron-init-path` | After Megatron init |
| | `--custom-megatron-before-log-prob-hook-path` | Before logprob compute |
| | `--custom-megatron-before-train-step-hook-path` | Before each train step |
| | `--custom-megatron-post-save-hook-path` | After each checkpoint save |
| **Logging** | `--custom-rollout-log-function-path` | Train-rollout logging |
| | `--custom-eval-rollout-log-function-path` | Eval-rollout logging |
| **Model** | `--custom-model-provider-path` | Megatron model factory |

***

## Rollout

### `--rollout-function-path`

Replace the entire rollout function. Use this only for fundamentally different flows
such as multi-agent co-evolution.

```python theme={null}
def generate_rollout(args, rollout_id, data_source, evaluation=False) \
        -> RolloutFnTrainOutput | RolloutFnEvalOutput:
    ...
```

**Default:** `miles.rollout.inference_rollout.inference_rollout_common.InferenceRolloutFn`; use `miles.rollout.sglang_rollout.generate_rollout` under `MILES_USE_LEGACY_ROLLOUT_V1=1`.

Plain functions with the signature above are wrapped in a legacy adapter. A class-based rollout
function needs to subclass `miles.rollout.base_types.BaseRolloutFn`.

**Reference:** [`examples/experimental/multi_agent/rollout_with_multi_agents.py`](https://github.com/radixark/miles/blob/main/examples/experimental/multi_agent/rollout_with_multi_agents.py).

### `--custom-generate-function-path`

Replace just the generation step inside the default rollout. Most tool-use, RAG, and
multi-turn workflows live here.

```python theme={null}
async def custom_generate(args, sample: Sample, sampling_params: dict) -> Sample:
    ...
```

The hook also accepts the `GenerateFnInput -> GenerateFnOutput` form; both
signatures load through the same adapter. See
[Generate Endpoint](/docs/user-guide/generate-endpoint) for the full contract.

**Reference:** [`examples/experimental/search-r1/generate_with_search.py`](https://github.com/radixark/miles/blob/main/examples/experimental/search-r1/generate_with_search.py).

### `--custom-agent-function-path`

Enabled when you set `--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate`.
Use `--custom-agent-function-path` to specify the async agent or environment loop
that sends OpenAI-compatible chat requests through Miles' TITO session server.

```python theme={null}
async def run_agent(
    base_url: str,
    prompt,
    request_kwargs: dict,
    metadata: dict,
    **kwargs,
) -> dict | None:
    ...
```

See [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout) for the full wiring and
message/token ownership contract.

### `--data-source-path`

```python theme={null}
class CustomDataSource(DataSource):
    def get_samples(self, num_samples) -> list[list[Sample]]: ...
    def save(self, directory: Path) -> None: ...
    def load(self, directory: Path) -> None: ...
```

**Default:** `miles.rollout.data_source.RolloutDataSource`.

### `--eval-function-path`

Same signature as `--rollout-function-path`. Defaults to whatever rollout function is
configured.

***

## Session

### `--session-message-matcher`

Some harnesses do not replay history verbatim — they reserialize tool-call arguments or drop `reasoning_content` — and the default `strict` matcher counts that as divergence (v1 rollback, v2 branching). This flag loosens what "the same message" means during replay: choose a looser built-in selector (see [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout#choose-replay-matching)) or supply your own matcher via a trusted dotted import path:

```python theme={null}
def matcher(stored_message: dict[str, Any], replayed_message: dict[str, Any]) -> bool:
    ...
```

`stored_message` is the authoritative stored message; return `True` to accept the replayed message as the same history at that position. Keep the matcher a fast, synchronous, side-effect-free equivalence check — exceptions or non-`bool` results return HTTP 500, and startup fails if the path does not resolve.

***

## Reward

### `--custom-rm-path`

```python theme={null}
async def custom_rm(args, sample: Sample) -> float:
    ...

# Batched mode (set --group-rm)
async def batched_custom_rm(args, samples: list[Sample]) -> list[float]:
    ...
```

**Built-in `--rm-type` options:** `math`, `dapo`, `deepscaler`, `gemma_math`, `f1`,
`gpqa`, `ifbench`, `remote_rm` (with `--rm-url`), `random`, `deterministic_random`.
Prefixing any of them with `boxed_` (for example `boxed_math`) extracts `\boxed{}`
from the response before grading.

### `--custom-reward-post-process-path`

Hook to normalize rewards differently from the default GRPO normalization.

***

## Filtering

### `--dynamic-sampling-filter-path`

Per-group filter; runs after scoring, before queueing for training.

```python theme={null}
def filter_function(args, samples: list[Sample], **kwargs) -> FilterOutput:
    return FilterOutput(keep=True, reason=None)
```

**Stock implementation:** `miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std`.

### `--buffer-filter-path`

Pops samples from the rollout buffer at dequeue time. The default is
`pop_first` in `miles/rollout/data_source.py`.

```python theme={null}
def buffer_filter(
    args,
    rollout_id: int | None,
    buffer: list[list[Sample]],
    num_samples: int,
) -> list[list[Sample]]:
    ...
```

### `--rollout-sample-filter-path`

Per-sample, in-place. Set `s.remove_sample = True` to exclude a sample from the loss
(advantage normalization still uses it).

The framework passes `data: list[list[Sample]]` — a list of
`n_samples_per_prompt`-size groups — so iterate the outer list once to reach `Sample`
objects:

```python theme={null}
def filter_function(args, data: list[list[Sample]]) -> None:
    for group in data:
        for s in group:
            if not_good(s):
                s.remove_sample = True
```

### `--rollout-all-samples-process-path`

Runs after rollout completes and can see all samples, including filtered ones.
Useful for logging or analysis.

### `--rollout-data-postprocess-path`

Runs after log probabilities have been computed but before training. Useful for
updating loss masks based on per-token logprobs.

***

## Training

### `--custom-loss-function-path`

Replace the GRPO/PPO loss. Requires `--loss-type custom_loss`. Useful for novel
objectives or multi-objective work.

### `--custom-tis-function-path`

Importance sampling correction for off-policy training when train and inference
diverge.

**Reference:** [`examples/infra_features/train_infer_mismatch_helper/mis.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/train_infer_mismatch_helper/mis.py).

### `--custom-pg-loss-reducer-function-path`

```python theme={null}
def get_pg_loss_reducer(
    total_lengths: list[int],
    response_lengths: list[int],
    loss_masks: list[torch.Tensor],
    calculate_per_token_loss: bool = False,
) -> Callable[[torch.Tensor], torch.Tensor]:
    ...
```

Use case: Dr.GRPO divides by a constant instead of effective token count.
**Reference:** [`examples/experimental/DrGRPO/custom_reducer.py`](https://github.com/radixark/miles/blob/main/examples/experimental/DrGRPO/custom_reducer.py).

### `--custom-convert-samples-to-train-data-path`

```python theme={null}
def convert_samples_to_train_data(args, samples) -> dict:
    return {
        "tokens":           [...],
        "response_lengths": [...],
        "rewards":          [...],
        "raw_reward":       [...],
        "truncated":        [...],
        "sample_indices":   [...],
        "loss_masks":       [...],
        # optional
        "round_number":            [...],
        "rollout_log_probs":       [...],
        "rollout_routed_experts":  [...],
        "metadata":                [...],
        "multimodal_train_inputs": [...],
        "teacher_log_probs":       [...],
    }
```

***

## Megatron hooks

| Flag | Signature | |
| - | - | - |
| `--custom-megatron-init-path` | `def custom_init(args) -> None` | |
| `--custom-megatron-before-log-prob-hook-path` | `def custom_hook(args, model, store_prefix) -> None` | |
| `--custom-megatron-before-train-step-hook-path` | `def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler) -> None` | |
| `--custom-megatron-post-save-hook-path` | \`def hook(args, rollout\_id: int, checkpoint\_dir: str, hf\_checkpoint\_dir: str | None) -> None\` |

The Megatron init, log-prob, and train-step hooks give access to the live model
and optimizer, useful for custom probes, weight clipping, or surgical interventions.
The post-save hook runs on rank 0 after checkpoint save completion and receives
the saved checkpoint paths instead of live model objects.

***

## Logging

```python theme={null}
# Training rollouts
def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool:
    ...

# Eval rollouts
def log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool:
    ...
```

Return `True` to suppress Miles's default logging, `False` to layer on top.

***

## Model

### `--custom-model-provider-path`

Replace Megatron's default model factory.

```python theme={null}
def custom_model_provider(
    pre_process: bool,
    post_process: bool,
    vp_stage: int | None = None,
) -> GPTModel:
    ...
```

***

## Worked example

A custom rollout plus a custom reward in one launch script:

```bash theme={null}
ROLLOUT_ARGS+=(
   --custom-generate-function-path my_pkg.search_rollout.generate
   --custom-rm-path                my_pkg.rewards.f1_with_grounding
   --metadata-key metadata
   --rollout-max-response-len 4096
)
```

That is the entire delta from the stock GRPO recipe, with no source changes to Miles.

→ Next: [Server arguments reference](/docs/user-guide/cli-reference)


# Agentic Environments
Source: https://miles.radixark.com/docs/user-guide/environments

How Miles trains on RL environments — datasets with rewards, self-wired environments, and optional external ecosystems.

Miles owns the training side of environment RL — batch orchestration, lossless
token-in/token-out recording, reward hooks, filtering — and is agnostic about
where the environment itself comes from:

* **No environment** — single-turn RLVR: a prompt dataset scored by the
  built-in rule-based rewards (math, ifbench, ...) or a custom reward function.
* **Your own environment** — plug your code into one of the three rollout
  layers described in [Integration shapes](#integration-shapes); most
  environments sit in the agent function, with the session server recording
  tokens (see [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout)).
* **An external ecosystem** — adopt a prebuilt connector from the table below,
  spanning coding, computer-use and tool-calling agents; connectors occupy the
  same three layers.

| Integration | Plugs in at | Guide |
| - | - | - |
| [Harbor](https://github.com/harbor-framework/harbor) | agent function | [guide](/docs/user-guide/harbor) |
| [HUD](https://hud.ai) | generate function³ | [example](https://github.com/radixark/miles/tree/main/examples/experimental/hud) |
| [NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) | agent function | [guide](/docs/user-guide/nemo-gym) |
| [OpenEnv](https://github.com/huggingface/openenv) | agent function | [guide](/docs/user-guide/openenv) |
| [Strands Agents](https://strandsagents.com/) | generate function | [example](https://github.com/radixark/miles/tree/main/examples/experimental/strands_sglang) |
| [Verifiers (Prime Intellect)](https://github.com/PrimeIntellect-ai/verifiers) | rollout function | [guide](/docs/user-guide/verifiers) |
| [τ-bench](https://github.com/sierra-research/tau-bench) | generate function | [example](https://github.com/radixark/miles/tree/main/examples/experimental/tau-bench) |

Sandbox providers are a different axis: they provision the task containers
*inside* a connector rather than occupying a rollout layer. Which providers work
with which connector, and how to set one up, is
[Sandbox Providers](/docs/user-guide/sandbox-providers).

Everything above is experimental, and listed alphabetically.

## Integration shapes

The rollout stack is three nested plug-in layers (see
[Customization](/docs/user-guide/customization)): each column in the table below
wraps the one to its left, so replacing an outer layer also takes over
everything an inner one would. A connector replaces exactly one layer.

✓ = the external framework takes it over; ○ = stays in Miles.

| | Agent function (innermost) | Generate function | Rollout function (outermost) |
| - | :-: | :-: | :-: |
| Plug-point flag | `--custom-agent-function-path` | `--custom-generate-function-path` | `--rollout-function-path` |
| Agent–environment loop | ✓ | ✓ | ✓ |
| Trajectory & token recording | ○ | ✓¹ | ✓¹ |
| Reward pathway (RM hooks, group rewards) | ○² | ○² | ✓ |
| Data source (prompts / taskset) | ○ | ○ | ✓ |
| Batch orchestration (grouping, filtering) | ○ | ○ | ✓ |
| Model, engines & weight updates, advantages, optimizer | ○ | ○ | ○ |

¹ Typically by speaking SGLang's native `/generate` (token IDs in and out)
rather than the session-server chat endpoint Miles' own recording uses.

² The environment may grade an episode itself (Harbor and τ-bench do); the
score still enters training through Miles' `Sample.reward` / RM hooks, and
group-level reward handling stays in Miles.

³ HUD's harness records per-turn token ids and sampling logprobs itself when
the inference server returns them, so the connector's job is stitching those
into one training sequence rather than recording. Computer-use observations
are screenshots, which Miles' session-server recording does not carry yet —
once it does, this connector can also sit in the agent function.


# Fully Async RL
Source: https://miles.radixark.com/docs/user-guide/fully-async

How fully async rollout decouples generation from training, which flags control it, and how to evaluate without stalling the trainer.

Fully async rollout decouples generation from training so that **rollout never waits on
a training step**. The engines **keep generating continuously** while the trainer
consumes finished groups, instead of the two taking turns.

### When to use it

Use fully async when **rollout is slow**, or when its wall time is set by a **long tail
of straggler trajectories** rather than by the average one, which is most often the case
in **long-context, tool-use, and agentic workloads**. The trade-off is that training
goes **off-policy**, because the trainer consumes groups generated under older weights.
Staying synchronous is recommended while you are debugging a recipe or validating loss
and reward code, where an exact on-policy cadence is worth more than throughput.

## Usage & Examples

### Basic usage

Switch the entrypoint from `train.py` to `train_async.py`, enable the class-based
rollout API, and pass `--fully-async`:

```diff theme={null}
- python3 train.py ...
+ python3 train_async.py ...
+   --fully-async
```

### Examples

Four launch scripts show the mode end to end, from a single-node smoke test to a
16-node agentic run:

| Script | What it covers |
| - | - |
| [`run_qwen3_30b_a3b_fully_async.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/run_qwen3_30b_a3b_fully_async.py) | The same pattern on a 30B MoE, with `tp=8`, `ep=8`, and one 8-GPU rollout engine |
| [`run_qwen3_5_4b_fully_async_eval.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py) | Both checkpoint eval backends behind one flag, `--eval-backend fleet` or `--eval-backend external` |
| [`run_glm5_2_744b_a40b_daytona.py`](https://github.com/radixark/miles/blob/main/examples/experimental/openenv/glm52_tbench2/run_glm5_2_744b_a40b_daytona.py) | GLM-5.2 744B-A40B on 16 GB300 nodes, split 8 training and 8 inference, with multi-turn terminal-bench-2 episodes in per-task Daytona sandboxes. It runs 128 in-flight trajectories against a 64-sample train batch and evaluates on the shared rollout engines |

### Customizations

Starting from a working run, the rest of this page covers what you can change:

| To change | See |
| - | - |
| How much generation stays in flight | [Arguments: Scheduling options](#arguments-scheduling-options) |
| Whether the engines and the trainer share GPUs | [Colocate](#colocate) |
| How deep the buffer is, how stale a group may be, which groups reach training, or the buffer implementation itself | [Arguments: Buffer options](#arguments-buffer-options) |
| Where eval runs and where it gets its weights | [Evaluation](#evaluation) |
| Which numbers tell you where the bottleneck is, or where the metrics are logged | [Metrics](#metrics) |

## The fully async schedule

### How generation is scheduled

The implementation lives in
[`miles/rollout/fully_async_rollout.py`](https://github.com/radixark/miles/blob/main/miles/rollout/fully_async_rollout.py)
(the worker) and
[`miles/rollout/submission_scheduler.py`](https://github.com/radixark/miles/blob/main/miles/rollout/submission_scheduler.py)
(the submission scheduler).

Fully async rollout splits Miles into two concurrent loops:

1. A background rollout worker keeps SGLang generation in flight and puts completed
   groups into a data buffer.
2. The trainer drains the buffer, runs optimizer steps, and syncs updated weights back
   to the rollout engines.

The worker is a persistent background task, started on the first training step and
generating from then on. It never stops between steps, so a training step no longer
triggers generation; it only takes finished groups out of the buffer.

The worker keeps a bounded amount of generation in flight. It always submits whole
prompt groups, and a **submission scheduler** decides when a finished unit frees the
slot for the next one. Under fully async the default granularity is `sample`: every
finished trajectory frees its own slot, and a new group goes out as soon as
`n_samples_per_prompt` trajectories have completed, whichever groups they came from.
This keeps the engines full when trajectory lengths vary widely, as in agentic
workloads, where holding each slot until the group's slowest trajectory returns would
leave the engines idle.

Generation is the only thing that runs continuously. Everything else still happens on
the driver's step schedule:

1. The trainer drains `--rollout-batch-size` groups from the buffer, waiting if not
   enough have finished yet.
2. It runs the optimizer step while the worker keeps generating.
3. Every `--update-weights-interval` steps it pauses generation, synchronizes the new
   weights through the configured update mode, and resumes. The
   `--pause-generation-mode` flag decides how in-flight requests survive that pause: the
   default `retract` returns them to the waiting queue and recomputes their KV cache,
   while `in_place` freezes them and resumes on the existing cache. Passing `abort`
   would kill them outright, which is why fully async rejects it.

Because generation spans those weight updates, the samples in one group can carry
different weight versions. The gap between a group's oldest weight version and the
engines' current one is its **staleness**, and it is the reason a group that finished
long ago may no longer be worth training on.

### Arguments: Scheduling options

Three flags control how much generation stays in flight:

| Flag | Effect |
| - | - |
| `--rollout-batch-size` | Groups the trainer consumes per step. It is also the default in-flight cap, so raising it widens both the batch and the concurrency unless `--async-max-concurrent-samples` is set |
| `--async-max-concurrent-samples` | In-flight cap in trajectories rather than groups, floored to `value // n_samples_per_prompt` groups. Use it to decouple generation concurrency from batch size |
| `--rollout-submission-granularity` | Sets when a finished unit frees a submission slot. Under `--fully-async` the default is `sample`, which frees each slot as its own sample completes; `group` holds the slot until the whole group returns |

## Colocate

`--fully-async --colocate` puts the engines and the trainer on the same GPUs, down to a
single one. Generation still runs continuously between training steps, but the two now
take turns on the device instead of holding separate GPUs. Generation stops for the whole
train window, so no samples are produced there. That is the same gap as the pause window
in the disaggregated mode, and `--max-weight-staleness` and the weight-version
bookkeeping are unchanged.

### Launch

```diff theme={null}
  python3 train_async.py ...
    --fully-async
+   --colocate
+   --num-gpus-per-node 1
-   --rollout-num-gpus 4
```

[`examples/retool_v2`](https://github.com/radixark/miles/blob/main/examples/retool_v2)
runs a multi-turn tool-call recipe this way on one GPU:
`python examples/retool_v2/run_retool_multi_turn.py --fully-async --num-gpus-per-node 1`.

## Data path

### The data buffer

The implementation lives in
[`miles/rollout/fully_async_data_buffer.py`](https://github.com/radixark/miles/blob/main/miles/rollout/fully_async_data_buffer.py).

The **data buffer** is the store of finished groups between the two loops, and every
group-level decision lives in it. The producer puts each group in as it completes, the
trainer takes groups back out one at a time, and everything in between — what to keep,
what to discard, what to send back for regeneration — is the buffer's call. It is one
replaceable component with the following interface:

| Method | Called by | Purpose |
| - | - | - |
| `put()` | The rollout worker, once per finished group | Store the group, or reject it |
| `get(num_groups=...)` | The trainer, once per training batch | Return that many groups at once, waiting until the buffer holds them |
| `get_metrics(trainer_model_id)` | The trainer, once per step | Report what the buffer did since the previous step. The trainer model id is always passed, and is `None` in a run of one policy |
| `state_dict()` | Checkpoint save | Return every accepted entry needed to resume without loss or duplication |
| `load_state_dict(state)` | Checkpoint load | Replace the buffer state with a previously returned state |

These methods are the whole interface: the worker and the trainer see nothing
else, and everything inside the box below is the built-in `DefaultDataBuffer`.

```mermaid theme={null}
flowchart LR
    DS[Data source] --> W["Rollout worker<br/>keeps N groups in flight"]
    W -->|"put()"| PF
    subgraph DB["DefaultDataBuffer — the built-in DataBuffer implementation"]
        direction LR
        PF{{"put-time filters"}}
        PF -->|aborted group| U["--async-unused-samples-handler"]
        PF -->|dynamic-sampling filter reject| X[Dropped]
        PF -->|kept| B[("Bounded store")]
        B --> GF{{"get-time filter"}}
        GF -->|"staleness > --max-weight-staleness"| U
        U -->|drop| X
    end
    GF -->|"get()"| T["Trainer drains<br/>rollout_batch_size groups"]
    U -->|retry| DS
    T --> S[Optimizer step, weight sync]
```

Groups are filtered at two points, because the two decisions become available at
different times. Whether a group was aborted, and whether
`--dynamic-sampling-filter-path` keeps it, is fixed the moment generation finishes, so
both are decided on `put()`. Staleness depends on how long the group then sits in the
buffer, so it is decided on `get()`. Once the trainer has a full batch it sorts the
groups by index and applies `--rollout-sample-filter-path` to the assembled batch.

If a sample cancels itself (for example, a sandbox client raises `CancelledError`),
its group is marked aborted after the sibling samples settle. The default buffer
discards the entire group from training, counts it in
`rollout/fully_async/aborted_groups_filtered`, and applies
`--async-unused-samples-handler` to the original prompts. Other groups keep running.
Cancelling the producer itself still propagates, and unexpected non-cancellation
exceptions still fail the worker rather than being silently discarded.

The buffer decouples the two loops. As long as it holds finished groups, the trainer
never waits for generation. If it sits empty, rollout is still the bottleneck and async
cannot hide it.

### Arguments: Buffer options

Buffer capacity bounds how far generation can run ahead of training:

| Flag | Effect |
| - | - |
| `--async-data-buffer-capacity-factor` | Buffer holds `floor(factor * rollout_batch_size)` groups, `2.0` by default. When it is full the producer blocks until training consumes |

Staleness control decides which of those groups training is allowed to see:

| Flag | Effect |
| - | - |
| `--max-weight-staleness` | Maximum gap between a group's oldest weight version and the current engine version. Unset by default, which disables the filter |
| `--async-unused-samples-handler` | What happens to a group training does not use, either aborted or too stale. The default `drop` discards it; `retry` recycles its prompts into the data source for regeneration. Dynamic-filter rejects are always dropped |

When those knobs are not enough, `--custom-async-data-buffer-path` replaces the buffer
itself. This is a larger step than setting any flag above: your `DataBuffer` subclass
takes over the full interface and therefore every group-level decision, and the flags in
this section apply only if your class reads them. The one decision that stays outside is
`--rollout-sample-filter-path`, which runs on the assembled batch rather than on
individual groups.

## Evaluation

Fully async rollout changes one thing about eval: generation is always in flight, so an
eval that runs on the rollout engines costs production time. Two independent choices
follow, which backend runs the eval and where that backend gets its weights.

| Backend | Selected by | Cost |
| - | - | - |
| Shared engines | Neither flag below, which is the default | Rollout production pauses for the eval duration |
| Dedicated fleet | `--eval-num-gpus N` | N GPUs carved out of the job; training never pauses |
| External backend | `--eval-function-path` pointing at a `CheckpointEvalFn` | Nothing from the training job beyond the snapshot export |

The fleet and an external backend each select a backend, so passing both is an error.

### Mode 1: Shared engines

The implementation lives in [`miles/rollout/fully_async_rollout.py`](https://github.com/radixark/miles/blob/main/miles/rollout/fully_async_rollout.py).

Eval runs on the rollout engines. The producer stops submitting for the duration of the
blocking eval and resumes after; in-flight requests finish and buffer, and nothing is
aborted. No extra weight movement is needed, because the engines already carry the
weights the step's `update_weights` broadcast pushed. Production stalls for roughly the
eval duration, which is acceptable for a small debug set and for points that must land
strictly on time.

### Mode 2: Dedicated fleet

The implementation lives in [`miles/ray/rollout/eval_fleet.py`](https://github.com/radixark/miles/blob/main/miles/ray/rollout/eval_fleet.py).

The fleet runs on its own GPUs behind its own router, synced by loading HF checkpoint
snapshots rather than by joining training weight updates:

```bash theme={null}
--eval-num-gpus 1  # dedicated eval engines, behind their own router
--eval-interval K
--eval-hf-dir /dev/shm/miles_eval_hf  # snapshot staging; tmpfs avoids the disk dependency
--eval-prompt-data aime /path/to/aime.jsonl
```

On each eval-due step the trainer hands the snapshot to the fleet, fires the eval
without waiting for it, and keeps training. Your `--eval-function-path` function
generates against the fleet exactly as it would against the training engines, so custom
eval functions work unchanged. The point lands at the right step on the x-axis even when
it completes a few steps later, and `eval/lag_steps` reports how late it was.

The fleet's engines inherit every `--sglang-*` setting from the rollout engines, so by
default they are configured exactly like the engines you already tuned. You can override
any single field with the matching `--eval-sglang-*` flag:

```bash theme={null}
--eval-sglang-mem-fraction-static 0.9  # the eval fleet is not sharing with training
--no-eval-sglang-enable-dp-attention  # booleans take a --no- form to turn an inherited True off
```

Tensor parallelism (TP) is the exception, because it comes from
`--eval-num-gpus-per-engine`, which also places the engines. SGLang ties `dp_size`,
`pp_size`, `ep_size`, and `attn_cp_size` to TP, so when the eval and rollout TP sizes
differ those four default to 1 instead of being inherited; inheriting them across a
different TP produces an engine that fails SGLang's own validation at boot. You can set
them explicitly with `--eval-sglang-*` if the fleet is large enough to want them.

### Mode 3: External backend

The contract lives in [`miles/rollout/checkpoint_eval.py`](https://github.com/radixark/miles/blob/main/miles/rollout/checkpoint_eval.py), with a
reference implementation in [`examples/infra_features/fully_async/external_eval_fn.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/external_eval_fn.py).

Subclass `CheckpointEvalFn` and implement `evaluate_checkpoint(checkpoint_dir, input)`.
The trainer hands over a snapshot path per eval point and owns dispatch, logging, and
garbage collection; raise `EvalSkip(reason)` for an attributable skipped point. Because
the function runs in-job it reads the real training args and logs through the trainer,
with no GPU carve-out.

The reference implementation launches its own SGLang server on spare GPUs or attaches to
an existing one, configured through the env vars documented in the script. A non-SGLang
black box implements the same contract by calling out to its API and mapping the
response into `RolloutFnEvalOutput`.

### The weight snapshot pipeline

The implementation lives in [`miles/ray/rollout/eval_dispatch.py`](https://github.com/radixark/miles/blob/main/miles/ray/rollout/eval_dispatch.py).

The fleet and an external backend both evaluate a checkpoint, so both need
`--eval-interval` and one of these two snapshot sources:

| Source | Set | Cost |
| - | - | - |
| A fresh export per eval point | `--eval-hf-dir`, pointed at tmpfs | `eval/export_time_seconds` per point, plus staging space |
| The periodic HF checkpoints | `--eval-hf-dir` unset and `--save-hf` set | None, but `eval_interval` must be a multiple of `save_interval` |

When snapshot eval reuses `--save-hf`, a fixed path such as `/checkpoints/hf`
automatically becomes `/checkpoints/hf/step_{rollout_id}`. Existing rollout ID
format fields, including `{rollout_id:06d}`, keep their original layout. Save paths
are unchanged when eval uses `--eval-hf-dir` or shared engines, or eval is disabled.

On a real run you are persisting checkpoints anyway, so reuse costs nothing. A test run
that saves no checkpoints needs the export, and a tmpfs staging directory keeps it off
disk.

The export is the one part of a point that is not fire-and-forget: it is a collective
across every train actor and the training loop waits for it. Under the default
`--eval-overflow-policy backpressure` a due point can also wait out the oldest pending
eval; `--eval-overflow-policy skip` drops it instead so a slow eval set can never stall
training.

Size the staging directory before pointing it at tmpfs. Snapshots are retired on every
outcome, but the directory holds up to `--eval-keep-snapshots` retired plus
`--eval-max-in-flight` still-evaluating model-sized directories at once, four by
default, so around 32 GB of `/dev/shm` for a 4B model in bf16. Evals are serialized
inside the trainer, so raising `--eval-max-in-flight` does not run more of them at once;
it lets the trainer export further ahead at the cost of one more snapshot on disk.

## Metrics

### Async rollout metrics

The buffer reports these metrics to wandb on every training step, alongside the
standard rollout metrics:

```text theme={null}
rollout/fully_async/queue_size
rollout/fully_async/aborted_groups_filtered
rollout/fully_async/stale_groups_filtered
rollout/fully_async/avg_staleness, rollout/fully_async/max_staleness
rollout/fully_async/avg_post_generation_staleness, rollout/fully_async/max_post_generation_staleness
rollout/fully_async/avg_generation_version_span, rollout/fully_async/max_generation_version_span
rollout/fully_async/token_weighted_staleness
rollout/fully_async/weight_version_sample_coverage
rollout/fully_async/buffer_avg_staleness, rollout/fully_async/buffer_max_staleness
rollout/dynamic_filter/drop_<reason>
```

Staleness is measured in published rollout weight versions: with
`--update-weights-interval 1`, one version normally corresponds to one optimizer step;
otherwise it does not. The metrics separate three parts of policy provenance:

* `avg_staleness` and `max_staleness` are the gap from the current version to the
  oldest version in each selected group. `stale_groups_filtered` counts groups that
  exceeded `--max-weight-staleness`; rejected groups do not enter these averages.
* `post_generation_staleness` is the gap from the current version to the newest version
  in each selected group. `generation_version_span` is the newest minus oldest version
  within the group. Together they distinguish aging after generation from updates that
  landed while a trajectory was being generated.
* `token_weighted_staleness` averages `current_version - span.version`, weighting each
  numeric version span by its number of generated tokens. It therefore describes the
  typical version gap of stamped generation tokens rather than the oldest token only.

`weight_version_sample_coverage` is the fraction of selected samples with at least one
numeric version span. Missing provenance is excluded from lag averages, not treated as
zero. The `buffer_` pair covers groups still sitting in the buffer when the step drained
it.

A `queue_size` pinned at zero means rollout is the bottleneck, so scale rollout capacity
or lower per-sample generation cost. A `queue_size` pinned at capacity means training is
the bottleneck, and the `buffer_` staleness metrics will climb with it. A rising
`stale_groups_filtered` means groups are aging out faster than the trainer consumes
them. In the logs, a `No completed rollout groups for 30.0s` warning means the drain is
starved.

### Async eval metrics

The metrics are logged in [`miles/ray/rollout/metrics.py`](https://github.com/radixark/miles/blob/main/miles/ray/rollout/metrics.py).

Every skipped point is logged at the step it would have landed on, with the reason:

| Metric | Cause |
| - | - |
| `eval/skipped_busy` | At `--eval-max-in-flight` under `--eval-overflow-policy skip` |
| `eval/skipped_export_failed` | The snapshot export raised |
| `eval/skipped_ckpt_missing` | No `.complete` marker in the snapshot directory |
| `eval/skipped_crashed` | Anything else the eval raised |
| `eval/skipped_<reason>` | A `CheckpointEvalFn` raised `EvalSkip(reason)` |

For a point that did run, `eval/{dataset}/weight_version/mean == eval/step` and
`eval/{dataset}/weight_version/mixed_version_ratio == 0` together confirm it measured
exactly the intended weights. A checkpoint backend measures the actor's exact weights at
that step, whatever the broadcast schedule. Shared engines measure the engines'
last-broadcast version, which equals the actor's current weights when
`--update-weights-interval` is 1.

### Performance metrics

For performance work, use the metrics below as a basic reference:

1. **Engine concurrency.** Watch `sglang_num_running_reqs` across engines. If some
   engines sit far below the others, or concurrency collapses without a weight update to
   explain it, check the router configuration — requests are not being spread evenly.
2. **Prefix cache hit rate.** Watch `sglang_cache_hit_rate` per engine, or the per-step
   `prefix_cache_hit_rate` in the rollout metrics. A coding-agent workload should stay
   above 90%, since every turn re-prefills its session prefix. If it is low, suspect the
   KV cache memory (`--sglang-mem-fraction-static`) and the router configuration.
3. **Where the time goes.** Compare rollout time, train time, and the staleness metrics.
   If rollout is slower, consider more GPUs for rollout, a higher
   `--async-max-concurrent-samples`, and throughput-oriented SGLang settings. If
   training is slower, consider more GPUs for training, a lower concurrency, and
   latency-oriented SGLang settings.

### Arguments: Logging options

Two flags replace the default metric logging, both defined in
[`miles/ray/rollout/metrics.py`](https://github.com/radixark/miles/blob/main/miles/ray/rollout/metrics.py):

| Flag | Signature |
| - | - |
| `--custom-rollout-log-function-path` | `log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time) -> bool` |
| `--custom-eval-rollout-log-function-path` | `log_eval_rollout_data(rollout_id, args, data, extra_metrics) -> bool` |

Returning `True` skips the default logging for that call, so a function that only
forwards metrics elsewhere should return `False` and leave the built-in logging in
place. To change which numbers the buffer reports in the first place, override
`get_metrics()` on a custom buffer instead, as described in
[Arguments: Buffer options](#arguments-buffer-options).


# Generate Endpoint
Source: https://miles.radixark.com/docs/user-guide/generate-endpoint

Write a custom generate function against SGLang's raw /generate endpoint, owning prompt construction, tokens, and loss masks.

Write a token-level generate function when your rollout logic must control
prompt construction and token handling directly. Your code builds the request,
calls SGLang's stateless `/generate` endpoint, and writes tokens, logprobs,
loss mask, and status back onto the `Sample`.

This is one of two styles of custom generation, both selected through
`--custom-generate-function-path`. The other style exchanges OpenAI-compatible
chat messages instead of tokens — see
[Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout).

***

## The generate-function hook

`--custom-generate-function-path` accepts two forms. The difference is only
the signature you write — `load_generate_function`
(`miles/rollout/inference_rollout/compatibility.py`) adapts an old-form
function automatically at load time:

| | New form (recommended) | Old form |
| - | - | - |
| Signature | `async def generate(input: GenerateFnInput) -> GenerateFnOutput` | `async def generate(args, sample, sampling_params) -> Sample` |
| `--custom-generate-function-path` | `miles.rollout.generate_hub.single_turn.generate`, `miles.rollout.generate_hub.multi_turn.generate` | `miles.rollout.sglang_rollout.generate` |
| Example file | [miles/rollout/generate\_hub/single\_turn.py](https://github.com/radixark/miles/blob/main/miles/rollout/generate_hub/single_turn.py), [miles/rollout/generate\_hub/multi\_turn.py](https://github.com/radixark/miles/blob/main/miles/rollout/generate_hub/multi_turn.py) | [miles/rollout/sglang\_rollout.py](https://github.com/radixark/miles/blob/main/miles/rollout/sglang_rollout.py) |
| Status | what all `generate_hub` built-ins use | kept for backward compatibility; adapted automatically at load time |

> [!NOTE]
>
> The class-based rollout path is the default; `MILES_USE_LEGACY_ROLLOUT_V1=1`
> selects the deprecated v1 path. Both generate-function forms work on either.

Either way, your function does the same three things:

1. Builds a request from the prompt.
2. Executes it against SGLang.
3. Updates the `Sample` with tokens, logprobs, loss mask, status.

`GenerateFnInput` / `GenerateFnOutput` live in `miles/rollout/base_types.py`.
The input carries:

* `state`: tokenizer, processor, args, sampling defaults.
* `sample`: the prompt, current tokens, response, status.
* `sampling_params`: `max_new_tokens`, `temperature`, `top_p`, etc.
* `evaluation`: whether this call serves an eval rollout.

Minimal skeleton (new form):

```python theme={null}
from miles.rollout.base_types import GenerateFnInput, GenerateFnOutput
from miles.utils.types import Sample


async def generate(input: GenerateFnInput) -> GenerateFnOutput:
    args = input.args
    sample = input.sample
    sampling_params = input.sampling_params

    # 1) build request from prompt and sampling params
    # 2) call backend
    # 3) update sample.tokens, sample.response, sample.rollout_log_probs,
    #    sample.loss_mask, sample.status

    return GenerateFnOutput(samples=sample)


def _add_arguments(parser):
    parser.add_argument("--your-arg", type=str)


generate.add_arguments = _add_arguments
```

> [!TIP]
>
> **Custom CLI flags.** `generate.add_arguments = _add_arguments` registers extra CLI flags. They are
> parsed into `input.args` and available everywhere in your generator.

Helpers:

* `compute_prompt_ids_from_sample` and `compute_request_payload` from
  `miles/rollout/generate_utils/generate_endpoint_utils.py` build `/generate` requests.
* A generate function can set `GenerateFnOutput.samples` to a `Sample` or `list[Sample]`.

## Reference generators

`miles/rollout/generate_hub/` ships reusable token-level generate functions
that compose with tool use and multi-turn logic:

* **`single_turn.py`**: single-turn generation via `/generate`. Text or multimodal prompts.
* **`multi_turn.py`**: multi-turn tool calling via `/generate`. Adds CLI flags
  `--generate-max-turns`, `--generate-tool-specs-path`, `--generate-tool-call-parser`,
  `--generate-execute-tool-function-path`.
* **`benchmarkers.py`**: forces random output sequence length for benchmarking.

***

## Next

* [Customization](/docs/user-guide/customization): browse every Python hook.
* [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout): the message-level
  style of custom generation.


# Harbor
Source: https://miles.radixark.com/docs/user-guide/harbor

Train agents on mixed task suites (SWE-bench, Terminal-Bench, custom) through the Harbor framework.

[Harbor](https://github.com/harbor-framework/harbor) is an agent-environment
framework from the Laude Institute: agent orchestration and grading are unified
in a single `Trial.run()` call, and a task is fully described by four files
(`instruction.md`, `Dockerfile`, `test.sh`, `task.toml`), so mixed task suites —
SWE-bench, Terminal-Bench, custom tasks — train through one endpoint.

Miles integrates Harbor as an
[agent-function integration](/docs/user-guide/environments): the agent function
hands each session's OpenAI-compatible URL to Harbor, which runs the per-task
sandbox, runs the agent against that URL, and grades the result; the grade
becomes the sample's reward through a custom reward hook.

## Try it

Two execution modes; each example README is the complete guide for its mode:

* Tasks on a **local Docker daemon** → the agent-server mode:
  [`examples/swe-agent-harbor-docker`](https://github.com/radixark/miles/tree/main/examples/swe-agent-harbor-docker).
* Tasks on a **cloud [sandbox provider](/docs/user-guide/sandbox-providers)** → the
  in-process mode, no server in between:
  [`examples/experimental/harbor`](https://github.com/radixark/miles/tree/main/examples/experimental/harbor).


# User Guide
Source: https://miles.radixark.com/docs/user-guide/index

Concepts, launch scripts, customization hooks, and a complete CLI reference.

| Page | What it covers |
| - | - |
| [Core Concepts](/docs/user-guide/concepts) | The four objects in the training loop and the four-knob invariant. |
| [Launch Script](/docs/user-guide/launch-script) | What `python scripts/run_*.py` does, how a launch script is structured, and how to override a recipe. |
| [Argument Groups](/docs/user-guide/argument-groups) | Where model args, `perf_args`, `grpo_args`, and the other launch-script flag groups belong. |
| [Fully Async RL](/docs/user-guide/fully-async) | Continuous generation decoupled from training: the schedule, the data buffer, async eval, and the metrics to watch. |
| [Training Backends](/docs/user-guide/training-backend) | Megatron-LM and FSDP: what each one owns, how to choose, parallelism, checkpoints, and hooks. |
| [Monitoring & Logging](/docs/user-guide/monitoring) | wandb, structured logs, per-source breakdowns, profiling, router metrics. |
| [Customization](/docs/user-guide/customization) | The `--*-path` plug-points for custom Python — rollout, reward, filters, loss, hooks. |
| [Generate Endpoint](/docs/user-guide/generate-endpoint) | Custom generate functions that own tokens and loss masks via the raw `/generate` endpoint. |
| [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout) | Configure an OpenAI-compatible agent loop with TITO trajectory assembly. |
| [Agentic Environments](/docs/user-guide/environments) | Supplying an environment: dataset + reward, your own env via the plug points, or an external ecosystem. |
| [CLI Reference](/docs/user-guide/cli-reference) | Every flag Miles accepts, grouped by subsystem. |

## Which pages do I actually need?

* **Training my first job** — read [Core Concepts](/docs/user-guide/concepts), then [Launch Script](/docs/user-guide/launch-script).
* **Tuning a running job** — [Launch Script](/docs/user-guide/launch-script) in depth + [CLI Reference](/docs/user-guide/cli-reference).
* **Plugging in a custom reward / rollout / filter** — skim [Core Concepts](/docs/user-guide/concepts) for vocabulary, then go to [Customization](/docs/user-guide/customization).
* **Contributor onboarding** — read top to bottom.


# Launch Script
Source: https://miles.radixark.com/docs/user-guide/launch-script

What a Miles launch script does when you run it, how it is structured, and the three ways to override a recipe.

Every supported model ships as a python launch script under `scripts/`, and starting a
training run is one command:

```bash theme={null}
python scripts/run_qwen3_dense.py --model-name Qwen3-4B
```

This page explains what that command does and how to change what it runs. For the
meaning of individual training flags, see [Argument Groups](/docs/user-guide/argument-groups)
and the [CLI Reference](/docs/user-guide/cli-reference).

## How a launch script starts a training job

A launch script is a recipe, not the training process. It assembles the full `train.py`
command line for one model family, starts a local Ray cluster, and submits the command
as a Ray job. The pieces involved:

| Layer | Location | Role |
| - | - | - |
| Launch script | `scripts/run_*.py` | Holds the recipe: per-model values and the flag blocks |
| Model definition | `scripts/models/<type>.py` | Provides the Megatron architecture flags |
| Command utilities | `miles/utils/external_utils/command_utils/` | Starts Ray and submits the job |
| Training entrypoint | `train.py` / `train_async.py` | The actual training process, run inside the Ray job |

When the script starts, it prints its resolved options as a table, then every shell
command it issues with an `EXEC:` prefix — the log is a complete record of what ran.

### Running a v1 launch script

Older Ray launch scripts can keep the free-function API by importing its compatibility module:

```python theme={null}
import miles.utils.external_utils.command_utils.legacy as U
```

It retains `execute_train`, the command and checkpoint helpers, dataset download,
Mooncake startup, and the original `ExecuteTrainConfig`. New scripts build a backend
with `config.create_backend()`; the compatibility API does not support Kubernetes.

## The structure of a launch script

The sections below walk `scripts/run_qwen3_dense.py`; every launcher follows the same
layout. The module docstring states the prerequisites (a converted checkpoint, the
datasets) and a runnable example, then three parts follow in file order.

### Selecting a recipe with `--model-name`

For some models one launcher provides multiple recipes, selected with `--model-name`.
The per-variant values live in a `_RECIPES` table at the top of the file:

```python theme={null}
@dataclass(frozen=True)
class _Recipe:
    megatron_model_type: str
    tensor_model_parallel_size: int
    max_tokens_per_gpu: int
    rollout_num_gpus_per_engine: int
    sglang_mem_fraction_static: float
    optimizer_cpu_offload: bool
    num_rollout: int = 3000
    extra_sglang_args: str = ""

_RECIPES: dict[str, _Recipe] = {
    "Qwen3-4B": _Recipe("qwen3-4B", 2, 9216, 2, 0.7, False),
    "Qwen3.5-27B": _Recipe("qwen3.5-27B", 4, 8192, 1, 0.5, True),
    # ... one entry per supported variant
}
```

Single-recipe launchers skip the table and write their values directly in the flag
blocks.

### ScriptArgs — script options as CLI flags and `MILES_SCRIPT_*` env vars

The script's own options are the fields of a `ScriptArgs` dataclass:

```python theme={null}
@dataclass
class ScriptArgs(U.ExecuteTrainConfig):
    run_id: str = U.create_run_id()
    model_name: _MODEL_NAMES = "Qwen3-4B"
    num_gpus_per_node: int = 8
    enable_eval: bool = True
    save: bool = True
    extra_args: str = ""
    data_dir: str = "/root/datasets"
    model_dir: str = "/root/models"
    megatron_path: str = "/root/Megatron-LM"
```

The `@U.dataclass_cli` decorator exposes each field twice: as a `--kebab-case` command
line option (`--model-dir`) and as an environment variable with the `MILES_SCRIPT_`
prefix (`MILES_SCRIPT_MODEL_DIR`). A value given on the command line beats the
environment variable, which beats the field default. The env form is how launch
wrappers and cluster tooling inject machine-specific values without editing the script.

Options shared by every launcher, from the `ExecuteTrainConfig` base class and repo
convention:

| Option | Default | Purpose |
| - | - | - |
| `--model-dir` | `/root/models` | Where checkpoints (HF and `torch_dist`) are read from |
| `--data-dir` | `/root/datasets` | Where prompt and eval datasets are read from |
| `--output-dir` | `/root/shared_data` | Where checkpoints and dumps are written |
| `--num-nodes` | `$SLURM_JOB_NUM_NODES` or 1 | Number of training nodes |
| `--num-gpus-per-node` | 8 | GPUs used on each node |
| `--extra-args` | empty | Extra flags appended to the `train.py` command line |
| `--extra-env-vars` | empty | Extra env vars added to the Ray runtime env |

### execute() — assembling the train.py flags from grouped blocks

The `execute()` function builds the `train.py` command line as one f-string block per
concern, then concatenates them:

```python theme={null}
ckpt_args = (
    f"--hf-checkpoint {args.model_dir}/{args.model_name} "
    f"--ref-load {args.model_dir}/{args.model_name}_torch_dist "
    f"--load {args.output_dir}/checkpoints "
)

grpo_args = (
    "--advantage-estimator grpo "
    "--use-kl-loss "
    "--kl-loss-coef 0.00 "
    "--eps-clip 0.2 "
    "--eps-clip-high 0.28 "
)

train_args = f"{ckpt_args} {rollout_args} {optimizer_args} {grpo_args} ... {args.extra_args}"
```

Each block maps to a section of [Argument Groups](/docs/user-guide/argument-groups), which
documents the flags themselves:

| Block | Flags documented at |
| - | - |
| `ckpt_args` | [Checkpoint paths](/docs/user-guide/argument-groups#ckpt-args) |
| `rollout_args` | [Sampling and reward](/docs/user-guide/argument-groups#rollout-args) |
| `eval_args` | [Evaluation](/docs/user-guide/argument-groups#eval-args) |
| `perf_args` | [Parallelism and memory](/docs/user-guide/argument-groups#perf-args) |
| `grpo_args` | [RL objective](/docs/user-guide/argument-groups#grpo-args) |
| `optimizer_args` | [Optimizer](/docs/user-guide/argument-groups#optimizer-args) |
| `sglang_args` | [Rollout engine](/docs/user-guide/argument-groups#sglang-args) |
| `misc_args` | [GPU layout and everything else](/docs/user-guide/argument-groups#misc-args) |

One block has no Argument Groups section: the wandb flags come from
`U.get_default_wandb_args`, which returns them only when `WANDB_API_KEY` is set — so
wandb logging turns on by exporting the key, with no script change.

## Three ways to override a recipe

From lightest to heaviest:

1. **Append flags with `--extra-args`.** The value is appended to the end of the
   `train.py` command line, and for a flag given twice the later occurrence wins — so
   this overrides any flag the recipe already sets:

   ```bash theme={null}
   python scripts/run_qwen3_dense.py --model-name Qwen3-4B \
       --extra-args "--lr 2e-6 --num-rollout 100"
   ```

2. **Set a script option, as a flag or an env var.** Anything on `ScriptArgs` can come
   from the command line or from `MILES_SCRIPT_*`:

   ```bash theme={null}
   MILES_SCRIPT_MODEL_DIR=/mnt/models python scripts/run_qwen3_dense.py --model-name Qwen3-4B
   ```

3. **Edit the script.** The launcher is the canonical home of a recipe's
   hyperparameters and is meant to be read and edited — change the `_RECIPES` values or
   the flag blocks directly for anything you want to keep.

## What execute\_train runs on your machine

The launcher hands the assembled flags to `U.execute_train`, which issues the `EXEC:`
commands you see in the log, in order:

1. Kills leftover `sglang`, `miles`, and `redis` processes and stops any previous Ray
   cluster.
2. Starts a fresh cluster with `ray start --head`, using `--num-gpus-per-node` GPUs.
3. Runs the launcher's `before_ray_job_submit` hook, if it has one (used for the ssh
   fan-out below).
4. Builds the Ray runtime env for the job: `PYTHONUNBUFFERED`,
   `CUDA_DEVICE_MAX_CONNECTIONS=1`, `NCCL_NVLS_ENABLE` (your exported value if set,
   otherwise probed with `nvidia-smi`), `MASTER_ADDR`, `no_proxy`, and a `PYTHONPATH`
   containing the repo root and `--megatron-path`, plus anything from
   `--extra-env-vars`.
5. Submits the job: `ray job submit -- python3 train.py <architecture flags> <recipe flags>`.

Two environment variables skip parts of this sequence:

| Env var | Effect |
| - | - |
| `MILES_SCRIPT_EXTERNAL_RAY=1` | The Ray cluster is already running: skip the Ray teardown and `ray start`, only submit |
| `MILES_SCRIPT_ENABLE_RAY_SUBMIT=0` | Run everything except the submission — shows what a launcher would do |

The head-node address defaults to `127.0.0.1` and is taken from `MASTER_ADDR`; export
it on multi-node runs so Ray and torch distributed bind to the right interface.

## Multi-step and multi-node launchers

Two launcher shapes go beyond a single `execute()`.

### Subcommand pipelines: prepare-\* and full-train

Large-model launchers split the pipeline — download, precision cast, `torch_dist`
conversion, training — into subcommands of one script (underscores in the function name
become dashes on the CLI):

```bash theme={null}
python scripts/run_deepseek_v4.py prepare-download --model-name DeepSeek-V4-Flash-FP8
python scripts/run_deepseek_v4.py train --model-name DeepSeek-V4-Flash-FP8 \
    --num-nodes 8 --num-gpus-per-node 8
```

The `full-train` subcommand chains all steps and checks a sentinel file before each
one, so a completed step is skipped on re-run — after an interruption, relaunch the
same command and it resumes where it stopped.

### Joining multiple nodes: per-role subcommands and ssh fan-out

A multi-node run needs every node in the Ray cluster before the job is submitted.
Launchers express this in one of two ways:

* **One subcommand per node role.** `scripts/run_nemotron_3_super_120b_a12b.py` has
  `worker` (joins the head's cluster and blocks) and `train` (starts the head, waits
  until the cluster reports every GPU, then submits); you run one command on each node.
* **ssh fan-out from the head.** With `--join-ray-workers`,
  `scripts/run_qwen3_sft.py` sshes every host of an MPI-style hostfile into the
  cluster (`U.ssh_start_ray_workers` as the `before_ray_job_submit` hook), so the whole
  cluster comes up from a single command on the head node.

## Model architecture definitions in scripts/models/

Megatron cannot read the architecture from a HuggingFace checkpoint, so each
`megatron_model_type` has a file in `scripts/models/` named exactly after it, exposing
one function that returns the architecture flags:

```python theme={null}
def model_args(**kwargs) -> str:
    return """
        --num-layers 36
        --hidden-size 2560
        ...
    """
```

`execute_train` resolves the file by name and splices its output into the `train.py`
command line ahead of the recipe flags. A variant (a layer-pruned debug model, a LoRA
target) derives from its base file with `load_sibling_model_args` instead of copying
it.

> [!WARNING]
>
> **Architecture parameters are not self-validating.** Two checkpoints from the same
> family can ship different `--rotary-base`, vocab padding, or normalization epsilon.
> Diff the checkpoint's `config.json` against the file in `scripts/models/` before a
> first run, and override any drifted value by appending it, e.g.
> `--extra-args "--rotary-base 10000"`.

## Next

* [Argument Groups](/docs/user-guide/argument-groups) — which training flags belong to
  which block, and what they mean.
* [CLI Reference](/docs/user-guide/cli-reference) — every flag Miles accepts.
* [Quick Start](/docs/getting-started/quick-start) — downloading and converting a
  checkpoint, the step before any launcher.
* [Customization](/docs/user-guide/customization) — the `--*-path` plug points for custom
  rollout, reward, and filter code.
* [Models](/docs/models/index) — the per-model recipe pages built on these launchers.


# Monitoring & Logging
Source: https://miles.radixark.com/docs/user-guide/monitoring

wandb, structured logs, profiling, and what to look at when something looks off.

Miles emits per-rollout metrics to stdout and (optionally) Weights & Biases. SGLang and
Ray write their own logs to their default directories.

## What gets logged by default

Each rollout iteration emits a structured row to stdout (illustrative shape — exact
fields depend on backend and config):

```text theme={null}
[trainer] iter 12/3000 | loss=0.412 reward=0.61 kl=0.018
                      | rollout=18.4s train=22.1s p2p=2.1s  (total 42.6s)
                      | grad_norm=0.93 lr=1.0e-06
```

When `--use-wandb` is set, metrics also go to wandb under the `train/`, `rollout/`,
and `perf/` namespaces (see `miles/utils/tracking_utils/wandb_utils.py`).

## Enabling wandb

```bash theme={null}
ray job submit --address=auto -- \
  python3 train.py ... \
    --use-wandb \
    --wandb-project miles \
    --wandb-group qwen3-30b-grpo
```

Available flags: `--use-wandb`, `--wandb-project`, `--wandb-group`. `WANDB_API_KEY`
should be supplied via Ray's `env_vars` rather than baked into the launch script.

## What to watch

| Signal | Healthy pattern | Red flag |
| - | - | - |
| `loss` | Slow decay over hundreds of iterations | Spike → crash within an iteration |
| `raw_reward` | Trending up, with healthy variance | Saturates near a single value (collapse) |
| `kl_loss` | Bounded, drifts up over time | Sudden jump (policy diverged from ref) — only logged when `--use-kl-loss` |
| `train_rollout_logprob_abs_diff` | Stable and small (≪ 1.0) | Climbing without bound → train/inference precision drift |
| `entropy_loss` | Slowly decreasing | Falls to \~0 too fast (mode collapse) |
| `grad_norm` | \< `clip_grad` (1.0 by default) | Repeatedly hitting clip threshold |
| `rollout_time` / `train_time` | Roughly balanced | One ≫ other → resource imbalance |
| `train/pg_clipfrac` | \< 0.2 | > 0.5 means policy is moving fast → drop LR |

Panel names follow what `loss.py` and the rollout logger emit; Miles's wandb metrics
live under `train/`, `rollout/`, `perf/`, `multi_turn/`, `passrate/` namespaces.

## Custom loggers

Replace the default rollout logger with your own to push to internal systems:

```python theme={null}
def my_log(rollout_id, args, samples, extra, rollout_time) -> bool:
    statsd.gauge("miles.reward", mean([s.reward for s in samples]))
    return False   # also keep default logging
```

```bash theme={null}
--custom-rollout-log-function-path my_pkg.logging.my_log
```

## Profiling

| Tool | When |
| - | - |
| `nvidia-smi dmon -s u` | Quick sanity check on GPU utilization |
| `nsys profile` | Deep CUDA-level profiling |
| `py-spy dump --pid <ray worker>` | Find Python-side stalls |
| `ray timeline` | Inspect Ray task scheduling |

### Built-in PyTorch profiler

The PyTorch profiler is wired into Miles via `miles/utils/profile_utils.py`. Flags
differ by backend:

**Megatron** — choose which sub-loop to profile:

```bash theme={null}
--profile-target train_overall    # or train_actor, train_log_probs (multi-arg)
```

**FSDP** — additionally exposes the standard FSDPArgs window:

```bash theme={null}
--use-pytorch-profiler
--profile-step-start 10
--profile-step-end 12
--memory-snapshot-path snapshot.pickle
--tensorboard-dir /data/tb-run-42
```

Open the trace in `chrome://tracing` or [Perfetto](https://ui.perfetto.dev/).

## Where the log files live

| Source | Path |
| - | - |
| Trainer stdout | wherever you redirected `ray job submit` (or Ray dashboard) |
| Ray workers | `~/.ray/session_latest/logs/` |
| wandb local cache | `wandb/run-<id>/files/` |
| FSDP profiler / memory snapshot | `--tensorboard-dir`, `--memory-snapshot-path` |

## Router endpoints

The router exposes a small FastAPI surface used internally by Miles:

| Endpoint | Method | What |
| - | - | - |
| `/add_worker` | POST | Register an SGLang engine |
| `/list_workers` | GET | List registered workers |
| any other path | GET / POST / PUT / DELETE | Proxied to a selected SGLang worker — e.g. `/generate`, `/v1/chat/completions`, `/health`. |


# NeMo Gym
Source: https://miles.radixark.com/docs/user-guide/nemo-gym

Train on NVIDIA NeMo Gym environments through the agent-function extension point.

[NeMo Gym](https://github.com/NVIDIA-NeMo/Gym) is NVIDIA's RL environment
ecosystem: environments are HTTP *resources servers* (code execution, search,
SWE tasks, ...) paired with *agents* that drive an episode end-to-end and
grade it. Task containers run through NeMo Gym's own sandbox provider API
(`nemo_gym.sandbox`) — Docker locally, or Daytona / Apptainer / ECS Fargate /
OpenSandbox — selected by config, no agent changes.

Miles integrates NeMo Gym as an
[agent-function integration](/docs/user-guide/environments): per sample, the agent
function POSTs the task to a NeMo Gym agent server's `/run` endpoint with
`policy_base_url` set to the session's OpenAI-compatible URL. NeMo Gym runs
its agent harness (mini-swe-agent v2 in `mini_swe_agent_2`) against that URL,
so Miles' session server records every turn losslessly (token ids, logprobs,
loss masks — see [Agentic Rollout (TITO)](/docs/user-guide/agentic-rollout)); NeMo Gym
grades the episode itself and the grade enters training through a custom
reward hook reading `sample.metadata["reward"]`.

## Try it

The maintained recipe is **SWE-bench GRPO with mini-swe-agent** in
[`examples/experimental/nemo-gym`](https://github.com/radixark/miles/tree/main/examples/experimental/nemo-gym).
In short:

1. **Environment side** — on any docker-capable host, clone NeMo Gym `main`
   (>= `fcca3a8`) and start the `mini_swe_agent_2` responses-API agent server
   with the docker sandbox provider config.
2. **Data** — convert SWE-bench Verified to Miles prompt data with
   `download_and_process_data.py`; the task instance rides in each sample's
   `metadata`.
3. **Training side** — point `NEMO_GYM_URL` at the agent server and launch
   `run.py`, wiring the chain:

```bash theme={null}
--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate
--custom-agent-function-path nemogym_agent_function.run
--custom-rm-path nemogym_generate.reward_func
--use-session-server
--tito-model qwen3
```

The recipe is validated end-to-end: golden and API-policy scans on a real
docker host, plus a 4-GPU GRPO training smoke whose episodes ran in real task
containers with the SWE-bench harness grading them. Follow the
[recipe README](https://github.com/radixark/miles/blob/main/examples/experimental/nemo-gym/README.md)
for the NeMo Gym server setup, no-GPU validation (golden scan / API-policy
scan), the launch walkthrough, and known limitations (SWE-Gym eval specs,
Qwen3 template soft-mismatch diagnostics).


# OpenEnv
Source: https://miles.radixark.com/docs/user-guide/openenv

Train on Hugging Face OpenEnv environments through the agent-function extension point.

[OpenEnv](https://github.com/huggingface/openenv) is Hugging Face's open
protocol for RL environments: an environment is an HTTP service exposing
`reset` / `step` (and optionally `evaluate`), so any environment speaking the
protocol can serve any trainer.

Miles integrates OpenEnv as an
[agent-function integration](/docs/user-guide/environments): a Miles-side agent
function drives the agentic loop — `reset(task_id)`, repeated `step`s, then
scoring the episode with the task's own tests — against an unmodified OpenEnv
server, and the score becomes the sample's reward through a custom reward
hook.

## Try it

The maintained end-to-end recipe is **Terminal-Bench-2 GRPO** in
[`examples/experimental/openenv`](https://github.com/radixark/miles/tree/main/examples/experimental/openenv).
It gives every episode its own cloud sandbox, built from that task's official
image so no resident infrastructure is left behind, on any of the
[sandbox providers](/docs/user-guide/sandbox-providers). One shared Docker env
server is supported as well,
for running without any sandbox platform.
Follow the
[recipe README](https://github.com/radixark/miles/blob/main/examples/experimental/openenv/README.md)
for prompt-data preparation, environment options, launcher flags, and
operational notes.


# Sandbox Providers
Source: https://miles.radixark.com/docs/user-guide/sandbox-providers

Credentials and endpoints for the cloud sandbox providers a recipe can run its tasks on.

Recipes that give each task its own container get it from a sandbox provider.
Rows are providers, columns the connectors that can use them; a filled cell has
had a real rollout run on it. Each column header links to the connector's
recipe, and a check carries its own link where that pair runs on a different
one.

| Sandbox provider | [Harbor](https://github.com/radixark/miles/tree/main/examples/experimental/harbor) | [HUD](https://github.com/radixark/miles/tree/main/examples/experimental/hud) | [NeMo Gym](https://github.com/radixark/miles/tree/main/examples/experimental/nemo-gym) | [OpenEnv](https://github.com/radixark/miles/tree/main/examples/experimental/openenv) |
| - | :-: | :-: | :-: | :-: |
| [AgentENV](https://github.com/kvcache-ai/AgentENV) | ✓ | | | [✓](https://github.com/radixark/miles/tree/main/examples/experimental/agentenv) |
| [Daytona](https://www.daytona.io/) | ✓ | ✓ | ✓ | ✓ |
| [E2B](https://e2b.dev/) | ✓ | | | ✓ |
| [Modal](https://modal.com/) | ✓ | | | ✓ |

The rest of this page is how to set each provider up. Its SDK is installed per
recipe; the recipe README has that line.

Every provider takes its credential the same two ways: exported in the
environment, or in a key file. On a multi-host cluster the file has to be
readable where the rollout workers run — a shared filesystem, or the same path
on every node — because the launcher hands the workers the file's *path*, never
its contents.

## E2B

E2B Cloud, or any server that speaks the E2B API.

```bash theme={null}
mkdir -p ~/.config/e2b && echo e2b_... > ~/.config/e2b/api_key   # or export E2B_API_KEY
```

With nothing else set, the SDK talks to E2B Cloud. For a server of your own,
point both planes at it:

```bash theme={null}
export E2B_API_URL=http://<server>:8000        # control plane
export E2B_SANDBOX_URL=http://<server>:8000    # data plane
```

[AgentENV](https://github.com/kvcache-ai/AgentENV) is one such server —
Firecracker microVMs behind the E2B API — and deploying one is
[its own guide](https://github.com/radixark/miles/tree/main/examples/experimental/agentenv).

E2B Cloud builds templates as a non-root user, so a task Dockerfile whose `RUN`
layers need root fails to build there. Terminal-Bench-2 tasks ship prebuilt
images and are unaffected.

## Daytona

```bash theme={null}
mkdir -p ~/.config/daytona && echo dtn_... > ~/.config/daytona/api_key   # or export DAYTONA_API_KEY
```

Accounts carry a total-disk quota: keep concurrent sandboxes × per-sandbox disk
under it. A Harbor recipe sizes that disk with `HARBOR_OVERRIDE_STORAGE_MB`.

## Modal

The credential is a token pair, kept in the config file Modal's CLI writes:

```bash theme={null}
uv tool install 'modal>=1.5.5' && modal token new     # writes ~/.modal.toml
```

`MODAL_TOKEN_ID` + `MODAL_TOKEN_SECRET` in the environment work as well; one
half without the other is treated as missing. `MODAL_PROFILE` and
`MODAL_ENVIRONMENT` pick the workspace when the profile's default is not the
one you want. Miles defaults `MODAL_SANDBOX_V2=1`; Modal SDK 1.5.5 and newer
honor that switch through the public `Sandbox.create` API. Set it to `0` only
when temporarily diagnosing a V2-specific provider issue.

For OpenEnv TB2, `cpus` and `memory_mb` from each task's `task.toml` become
Modal resource requests, not hard limits. They preserve the task's guaranteed
resource contract while still allowing the sandbox to burst above the request
when its host has spare capacity.


# Training Backends
Source: https://miles.radixark.com/docs/user-guide/training-backend

The contract Megatron-LM, FSDP and torchtitan all implement, which one to pick, and how to configure parallelism, GPU layout, offload, and checkpoints for each.

In Miles a training backend is one class: a `TrainRayActor` subclass that owns the model on
the GPU. `--train-backend` decides which one `miles/ray/train/actor_factory.py` instantiates
on every trainer rank, and there are three choices.

| Value | Class | What it is | Default |
| - | - | - | - |
| [`megatron`](#megatron-lm) | `MegatronTrainRayActor` | Megatron-LM: five parallel dimensions, `torch_dist` checkpoints | ✅ |
| [`fsdp`](#fsdp) | `FSDPTrainRayActor` | The model's own HuggingFace implementation under PyTorch FSDP2 | |
| [`torchtitan`](#torchtitan) | `TorchtitanTrainRayActor` | torchtitan's own `Trainer`, driven as a black box | |

Whichever you pick, the rest of the job talks to it through the same handful of methods, and
that short list is the whole contract between a backend and everything else in Miles:

| Method | What the backend has to do |
| - | - |
| `init` | Build the model, optimizer and parallel layout after the shared base has set up the process group and the device |
| `train` | Consume one rollout's data and take the optimizer steps for it |
| `update_weights` | Push the freshly trained weights into the SGLang engines |
| `save_model` | Write a checkpoint, in whatever format this backend uses |
| `sleep` / `wake_up` | Move the model and optimizer off the GPU and back, so a colocated SGLang engine can use the memory in between |

FSDP and torchtitan implement `train`, `update_weights` and `sleep` / `wake_up` once, in the
shared `TorchNativeTrainRayActor`; each supplies only its model, optimizer and step runner.
Megatron implements the contract on its own, because its microbatch loop belongs to the
pipeline schedule and its checkpoints to `torch_dist`.

That is also why switching backends does not touch the rest of your launch script. Rollout,
reward, eval, the RL algorithm and the SGLang engine all sit above this line, and so does
the GPU layout: **disaggregated** by default, where trainer and engines own separate GPUs,
or **colocated** with `--colocate`, where they share GPUs and `sleep` / `wake_up` hand the
memory back and forth.

What does change is everything below the line, which is what the rest of this page is about.

## Which one do you want?

**Use Megatron-LM for large models and for anything that needs real parallelism.** It is the
recommended backend, the one every recipe in [Models](/docs/models/index) is tuned for, and the
only one that can split a model *inside* itself. If the model is a 100 B+ MoE, if the job
spans racks, or if fitting it at all depends on tensor / pipeline / expert parallelism, this
is the answer.

**Use FSDP when you want the HuggingFace implementation trained verbatim.** It loads a HF
directory as-is, with no conversion step and no architecture flags to write, which makes it
the fast path for bringing up a new architecture, for checking trainer numerics against the
HF reference, and for models that fit under data parallelism alone.

**Use torchtitan when you want its model implementations and its parallelism.** It brings
full parallelism like Megatron-LM and loads an HF directory like FSDP, but the architecture
has to be one torchtitan itself implements — that list, not a spec you write, is the limit.

The rest follows from that split:

| | Megatron-LM | FSDP | torchtitan |
| - | - | - | - |
| Model splitting | TP × PP × CP × EP × ETP, plus DP | `dp_replicate` × `dp_shard` | TP × PP × CP × EP, plus `dp_replicate` × `dp_shard` |
| Model input | `torch_dist` checkpoint (offline conversion step) | HF directory, loaded as-is | HF directory, loaded as-is |
| Architecture definition | `scripts/models/<megatron_model_type>.py` plus a Megatron spec for anything non-standard | HF `config.json`, plus an optional adaptation spec | a torchtitan model name + flavor |
| Checkpoints written | Megatron `torch_dist` | PyTorch Distributed Checkpoint | torchtitan's own DCP checkpointer |
| Activation recompute | `--recompute-granularity / method / num-layers` | `--gradient-checkpointing` | `--gradient-checkpointing` |
| Optimizer on CPU | `--optimizer-cpu-offload` | `--fsdp-cpu-offload` | Not supported |
| Offload beyond host RAM | `--offload-train-target disk`, `--stream-optimizer-state-to-disk` | Not supported | Not supported |
| LoRA | Supported | Not supported | Not supported |

***

## Megatron-LM

Configuring this backend is a handful of decisions, in this order: what the architecture is,
how to split it, where it sits relative to the rollout engines, how to fit it in memory,
where the weights come from, and what you want to hook into.

### 1. Describing the architecture

You do not re-declare Megatron's flags to Miles. Miles imports Megatron's whole argument
surface at launch:

```python theme={null}
from megatron.training.arguments import parse_args
```

so every Megatron flag your checkpoint needs (`--kv-channels`, `--rotary-base`,
`--moe-grouped-gemm`, and the rest) already works. Miles then threads its own flags in
through an `extra_args_provider` (`get_miles_extra_args_provider` in
`miles/utils/arguments.py`), which is why Miles and Megatron flags share one CLI.

That import is also why you export the Megatron source before launching:

```bash theme={null}
export PYTHONPATH=/root/Megatron-LM
```

In a launch script the architecture flags come from `scripts/models/<family>.py`,
selected by the `megatron_model_type` the script hands to `execute_train`. Most models need nothing beyond the stock
`--num-layers / --hidden-size / ...`. For the ones that do, see
[bringing in a new architecture](#going-deeper-bringing-in-a-new-architecture) below.

### 2. Choosing the parallelism

Megatron exposes five useful parallel dimensions, but you can't combine them in arbitrary
ways. Only a subset of TP × PP × CP × EP × ETP combinations is actually supported, and some
legal combinations are slower than the recipe baseline. **Start from the model recipe's
tested combination, then change one dimension at a time.**

| Dimension | Use it for | Compatibility notes |
| - | - | - |
| TP | Shard dense matrix multiplications inside each layer | When `--tensor-model-parallel-size` is set above 1, also pass `--sequence-parallel` unless the recipe says otherwise. |
| PP | Split layers across pipeline stages | Combines with TP and CP, but changes micro-batch scheduling and checkpoint layout. |
| CP | Split long sequences across ranks | Useful for long context; size token budgets as `CP x max_tokens_per_gpu`. |
| EP | Distribute MoE experts across ranks | MoE-only. Keep trainer EP and SGLang EP as separate choices. |
| ETP | Tensor-parallelize expert MLPs | MoE-only. Use it only when the recipe enables it or when EP alone cannot fit the experts. |

Do not assume TP, CP, EP and ETP can all be raised independently for a new model. The exact
set of supported combinations depends on the Megatron Core kernels and model spec in use.
[Argument Groups](/docs/user-guide/argument-groups#perf-args) lists the flags that belong in
`perf_args`.

### 3. Choosing the GPU layout

Parallelism says how the trainer splits the model. This says where the trainer sits relative
to the SGLang engines, and there are two answers.

**Disaggregated is the default.** The trainer takes `--actor-num-nodes` x
`--actor-num-gpus-per-node` GPUs, the engines take `--rollout-num-gpus` more, and the two
sets do not overlap. Nobody has to move: both halves stay resident on their own GPUs for the
whole run, so `--offload-train` / `--offload-rollout` default off and no phase pays an
offload cost. It is also the layout that lets the two halves actually run at the same time,
which is what [Fully Async Rollout](/docs/user-guide/fully-async) and `train_async.py` are for.
Under the synchronous loop in `train.py` the phases still alternate, so each set of GPUs is
idle while the other works.

```bash theme={null}
# 8 GPUs training, 8 more generating
--actor-num-nodes 1 --actor-num-gpus-per-node 8 \
--rollout-num-gpus 8 --rollout-num-gpus-per-engine 2
```

**Colocated shares one set of GPUs.** `--colocate` puts the engines on the training GPUs and
the two take turns: generate, offload the engine, train, offload the trainer, repeat. It is
the right default when GPUs are the scarce resource, since the same 8 GPUs do both jobs
instead of standing idle during the other phase.

```bash theme={null}
--colocate \
--actor-num-nodes 1 --actor-num-gpus-per-node 8 \
--rollout-num-gpus-per-engine 2 \
--sglang-mem-fraction-static 0.8
```

Three things follow from `--colocate` that are worth knowing before you use it:

* `--rollout-num-gpus` is ignored and reconciled to `actor_num_gpus_per_node x
  actor_num_nodes`, since the engines are on the training GPUs by definition.
* `--offload-train` and `--offload-rollout` both turn on, which is what makes the taking of
  turns possible. That is the memory story in the next section.
* The trainer reserves HBM at init before SGLang starts, so `--sglang-mem-fraction-static`
  has to come down, typically to 0.8 or lower. Miles also defaults
  `--sglang-cuda-graph-backend-prefill=disabled` here to avoid an NVLS OOM.

The layout also decides how `update_weights` gets the weights across. Colocated, the engine
is on the same device, so the actor hands over CUDA IPC handles and nothing crosses the
network. Disaggregated, the weights have to travel, and `--update-weight-transfer-mode`
picks how: `broadcast` (the default) sends each tensor over the training-to-engine process group,
`broadcast_packed` sends one packed byte buffer per bucket for non-colocated Megatron,
`p2p` uses [RDMA point-to-point](/docs/advanced/p2p-weight-transfer), and
[`disk-delta`](/docs/advanced/disaggregated-rollout) publishes only the bytes that changed since
the last sync for each engine to pull.

On a node with fewer than 8 usable GPUs, set `--num-gpus-per-node` too, otherwise the
rollout side still assumes 8. And `--fully-async` cannot be colocated: its whole point is
that rollout keeps generating while the trainer steps, which requires separate GPUs.

### 4. Fitting it in memory

Parallelism decides how the model is divided; this decides what is allowed to sit in HBM at
all. Four things compete for it: parameters, gradients, optimizer state, and activations.
On bf16 training the optimizer state is the heavy one, at 12 bytes per parameter for the
fp32 master copy plus the two Adam moments, against 2 bytes for a bf16 parameter. Data
parallelism divides that state, so a run with GPUs to spare may need none of what follows,
and a run at DP=1 may need all of it.

Two of the knobs apply to any layout, and the rest exist only because a colocated engine
wants the GPU back.

#### Either layout

**Activations** are the first thing to trade, because recompute is cheap and predictable.
Every recipe passes some form of:

```bash theme={null}
--recompute-granularity full --recompute-method uniform --recompute-num-layers 1
```

**The optimizer step can run on the CPU.** `--optimizer-cpu-offload` keeps the master
weights and moments in host memory and runs Adam there, and
`--overlap-cpu-optimizer-d2h-h2d` hides the copies behind compute. Recipes that use it
usually add `--use-precision-aware-optimizer`, which lets Megatron hold narrower optimizer
state. Note the interaction with rematerialization below: precision-aware on the GPU stores
masters as int16 remainders inside TE FusedAdam, so there is nothing standalone left to
rebuild from.

#### Colocated only

Everything from here down hangs off `--offload-train`, which is on precisely because the
engine needs the same HBM during generation. It is what `sleep` / `wake_up` do, it is turned
on for you by `--colocate`, and in a disaggregated run there is nothing to make room for, so
none of it applies.

| Flag | Effect |
| - | - |
| `--offload-train` / `--offload-rollout` | Which side is offloaded during the other's phase. Both implied by `--colocate`. |
| `--offload-train-target cpu` | Default: the paused actor is backed up in pinned host memory. |
| `--offload-train-target disk` | For when host RAM cannot hold that copy either: stream it to node-local NVMe instead, through a bounded pinned buffer (`--offload-train-disk-dir`, `--offload-train-disk-chunk-mb`). Megatron backend only. |
| `--rematerialize-param-from-master-weight` | Drop the actor's parameter backup during rollout and rebuild it from the optimizer's master weights on the next step. Saves 2 bytes per parameter per rank of host memory on bf16 training. Asserts `--colocate` plus the `cpu` target. |

**If the optimizer state does not fit while the step itself runs**, offloading the actor
cannot help, because pause and resume happen at phase boundaries and everything is resident
again by the time Adam launches. That case is what streaming addresses:

```bash theme={null}
--offload-train --offload-train-target disk \
--stream-optimizer-state-to-disk \
--offload-train-disk-dir /scratch/miles_offload
```

The fp32 masters and Adam moments live in per-bucket files on NVMe, and the step brings in
one bucket at a time, so peak residency is one bucket instead of the whole state. At the
default `fp32` moment dtype it matches keeping the state on the GPU up to the grad norm's
rounding and costs disk traffic every step; `--stream-optimizer-state-moment-dtype bf16` cuts
the volume by a third. It requires the `disk` target and excludes `--optimizer-cpu-offload`.

[Disk Offload](/docs/advanced/disk-offload) has the full picture for both, including the
same-topology resume limit, what checkpointing costs, and measured sleep / wake numbers.

### 5. Getting weights in and out

Megatron trains from its own `torch_dist` format: `.distcp` files that are
parallelism-agnostic, so you can change TP / PP / EP later without re-converting. Convert
once, up front:

```bash theme={null}
MODEL_ARGS_LINE="$(python3 miles/utils/external_utils/model_args_utils.py <family>)" || 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/<model> \
   --save          /root/<model>_torch_dist
```

For models larger than a single node, drive the converter with
`torchrun --nnodes=<N> --nproc-per-node=8 ...`. Each recipe page lists the exact command.

What the run then writes looks like this:

```text theme={null}
/ckpt/
├── latest_checkpointed_iteration.txt
├── iter_0000100/
│   ├── _0_0.distcp
│   └── ...
├── iter_0000200/
└── ...
```

Always pass the **parent** directory to `--load`, never a specific `iter_*`. The loader
reads `latest_checkpointed_iteration.txt` to pick the step.

**Saving on demand.** `--save-trigger-sentinel <path>` forces a save from outside the
process, independent of `--save-interval`:

```bash theme={null}
# trigger a save and wait until the checkpoint is on disk
touch /path/to/save_now && until [ ! -e /path/to/save_now ]; do sleep 5; done
```

A request fired at any moment during an iteration is consumed at that iteration's save
point. The checkpoint is written with `force_sync=True` (so async saves finalize first), and
only then is the sentinel deleted, which is why "file gone" means "checkpoint durable on
disk". If the job crashes mid-save the sentinel survives, so the request stays pending for
the next run. Requires `--save`.

### 6. Hooking into the loop

Three extension points override Megatron behavior without forking it:

| Flag | Runs |
| - | - |
| `--custom-megatron-init-path` | After Megatron initialization |
| `--custom-megatron-before-log-prob-hook-path` | Before every log-probability computation |
| `--custom-megatron-before-train-step-hook-path` | Before every training step |

Typical uses: mixing in an auxiliary loss, instrumenting per-step metrics, clipping weights
surgically. See [Customization](/docs/user-guide/customization#megatron-hooks).

### Going deeper: bringing in a new architecture

Post-training runs on released checkpoints, so this is rarely your problem. When a model does
need a custom module, Miles embeds the model's official HuggingFace module inside Megatron's
scheduling rather than patching Megatron: a spec function under `miles_plugins/models/` is
selected with `--spec <module> <function>`, a bridge under `miles_plugins/mbridge/`
reconciles the parameter layouts, and parameters that must stay fp32 through Megatron's bf16
cast are tagged with `mark_param_dtype` from
`miles/backends/megatron_utils/fp32_param_utils.py`. The model configs in `scripts/models/`
that pass `--spec` are the worked examples.

***

## FSDP

The FSDP backend lives at `miles/backends/fsdp_utils/`. One idea explains the whole thing:
**nothing about the model is re-expressed for the trainer.** Architecture comes from the
HuggingFace `config.json`, weights load through `AutoModelForCausalLM.from_pretrained()`,
and sharding, the distributed optimizer and mixed precision all come from PyTorch FSDP2
rather than from Miles.

So there is no conversion step, no architecture flag block, and no spec to write for a model that
`transformers` already implements. The bill comes due on parallelism, which is why large
models and complex layouts belong on [Megatron-LM](#megatron-lm).

### 1. Pointing it at a model

```bash theme={null}
--train-backend fsdp \
--hf-checkpoint /root/models/<model>
```

`--hf-checkpoint` is the whole model input: tokenizer, config and weights. Layer count is
read from the HF config, so Megatron's architecture flags (`--num-layers`, `--hidden-size`,
`--spec`, and the rest of `scripts/models/`) simply do not apply here.

### 2. Sharding it

This backend is pure data parallel. `miles/backends/fsdp_utils/parallel.py` builds a single
device mesh with two dimensions:

| Dimension | How you set it | What it does |
| - | - | - |
| `dp_replicate` | `--dp-replicate-size` | Replica count for FSDP2 hybrid sharding. Parameters are replicated across replicas, sharded within one. |
| `dp_shard` | derived | Whatever is left: `world_size / dp_replicate`. This is the dimension FSDP2 actually shards parameters, gradients and optimizer state over. |

The default, `dp_replicate=1`, means one flat shard group over every training rank. Tensor,
pipeline, context, expert and expert-tensor parallelism are all fixed at size 1 in the FSDP
`ParallelState`, so the model has to fit within those two dimensions.

> [!NOTE]
>
> Context parallelism is not available here. `--context-parallel-size` above 1 is rejected in
> argument validation (`miles/utils/arguments.py`); the mesh has no CP dimension to build.

The mesh is checked before anything is built: `world_size` must divide by
`--dp-replicate-size`, otherwise the run fails in argument validation instead of deep inside
mesh construction.

Memory, once the layout is set:

| Flag | Effect |
| - | - |
| `--gradient-checkpointing` | Recompute activations. This backend's `--recompute-*`. |
| `--fsdp-cpu-offload` | Offload parameters, gradients and optimizer state to CPU. The optimizer step runs there. |
| `--fsdp-cpu-backend gloo` | CPU process-group backend used by the offload path. |

Under `--colocate` this backend also implements `sleep` / `wake_up` by moving the model and
the optimizer to host memory and back, gated on `--offload-train`. The deeper offload
targets are Megatron-only: both `--offload-train-target disk` and
`--stream-optimizer-state-to-disk` assert the Megatron backend.

### 3. Precision

* bf16 by default; `--fp16` switches the compute dtype.
* An fp32 master copy of the weights is kept by default, which is what makes the
  trainer to rollout weight sync bit-exact. `--no-keep-fp32-master` trades it for memory when
  you do not need that guarantee.
* `--attn-implementation` picks the `transformers` attention backend: `flash_attention_2` by
  default, with `flash_attention_3`, `sdpa` and `eager` passed straight through.
* An architecture with fussier numerics can register its own policy, see
  [when an HF model needs help](#going-deeper-when-an-hf-model-needs-help).

### 4. Checkpoints

`--save` writes PyTorch Distributed Checkpoint directories, one each for model, optimizer
and LR scheduler, plus a `latest_checkpointed_iteration.txt` tracker. So `--load` takes the
**parent** directory exactly like the Megatron backend does. These are FSDP-backend
checkpoints, not `torch_dist` ones, and the two formats are not interchangeable.

### 5. Compute kernels

By default (`--kernel-backend native`) every fused kernel comes from the wheels baked into the
image: `flash_attn`, `causal_conv1d`, `flash-linear-attention` and the rest, built out of band and
installed by `docker/Dockerfile`.

`--kernel-backend hub` additionally resolves a few of them from
[Hugging Face Hub kernel repos](https://huggingface.co/docs/kernels/en/index) instead. A kernel
repo ships one prebuilt variant per `(torch, CUDA, C++ ABI, arch, OS)`, so the client picks a
variant at load time and imports it from the HF cache — **no compiler on the training node, and no
image rebuild to change a kernel.**

| Flag | Effect |
| - | - |
| `--kernel-backend {native,hub}` | `native` (default) never imports `kernels`. `hub` loads the mapping below. |
| `--kernel-mapping-path` | Dotted path to your own `(args) -> dict[str, HubKernelSpec]`, replacing the shipped mapping entirely. |
| `--kernel-strict` | Raise when a repo or a function will not resolve, instead of falling back to the native kernel. |

What miles ships, in `miles/backends/fsdp_utils/plugins/hf_kernels/presets.py` (the whole
feature lives in that plugin package; the main FSDP path only calls `HubKernels.prepare()` once
and `bind()` per model):

| Slot | Repo | Feeds |
| - | - | - |
| `gated_delta_rule` | `kernels-community/fla` v1 | GatedDeltaNet's linear-attention recurrence (Qwen3-Next, Qwen3.5, Qwen3.6) |
| `causal_conv1d` | `kernels-community/causal-conv1d` v1 | GatedDeltaNet's short causal convolution, same architectures |
| `flash_attn_varlen` | `kernels-community/flash-attn2` v2 | The NemotronH attention mixer's varlen path |

Slots resolve independently. Before either model is bound, all ranks agree on each slot's
availability, repository/ref, and kernel build identity. If any rank cannot load a slot or the
identities differ, every rank keeps its native implementation for that slot. `--kernel-strict`
turns that collective fallback into an initialization error on every rank.

Custom mappings may omit entire slots, but each included slot must declare all the functions
listed by `REQUIRED_SLOT_FUNCTIONS` in `presets.py`; unknown slots and incomplete declarations
are rejected before anything downloads, regardless of `--kernel-strict`. Use repositories from
trusted Hub kernel publishers. Local kernel overrides are unsupported because their Hub
provenance cannot be verified across ranks.

These are **module-level** kernels: `kernels.get_kernel()` returns a module and miles rebinds the
free functions HF modeling code already looks up per forward (for GatedDeltaNet, the kernel names in
the model's `transformers` modeling module, process-wide). Nothing rebinds an `nn.Module.forward`,
so `state_dict`, `_no_split_modules` and the DTensor gather are untouched — which is why the binding
runs before `apply_fsdp2` and the ref model takes the same call.

> [!NOTE]
>
> This is worth turning on for the packing path specifically. Every slot feeds a kernel that
> [sequence packing](#going-deeper-when-an-hf-model-needs-help) needs in order to reset state per
> packed document, and each one fails differently when its wheel is missing:
>
> * no `flash-linear-attention` — `transformers` binds `torch_chunk_gated_delta_rule`, whose signature
>   ends in `**kwargs`, so the injected `cu_seqlens` is *accepted and ignored*;
> * no `causal_conv1d` — `transformers` binds its torch `causal_conv1d_fn`, which takes no `seq_idx`;
> * no `flash_attn` — the NemotronH attention patch returns the unpatched dense forward.
>
> In all three the per-document reset stops happening, nothing raises, and the only symptom is a wider
> train/rollout logprob gap. Successfully loading the Hub kernels restores those boundary resets without a wheel build.
> Use `--kernel-strict` when those native wheels are absent: a non-strict fallback does not
> make a native implementation that ignores boundaries safe for packed training.

Hub kernels are rejected together with `--true-on-policy-mode` and `--deterministic-mode`: those
modes require the training kernel to match SGLang's build exactly, and that equivalence has not been
established per kernel yet.

This loader requires access to Hub metadata even when the kernel binaries are already cached.
It does **not** consume `kernels.lock`, and `kernels lock . && kernels download .` does not make
`--kernel-backend hub` work under `HF_HUB_OFFLINE=1`. Offline provisioning remains a follow-up
in [RFC #2207](https://github.com/radixark/miles/issues/2207).

For repeatable online runs, provide a custom mapping with `HubKernelSpec(revision="<commit SHA>",
version=None, ...)` for each selected repository, retaining the slot's required functions. Major
versions such as `version=1` follow moving `v1` branches. The collective check ensures agreement
within a run; immutable revisions also pin the source across runs. A revision pin still requires
online metadata access with this loader.

> [!TIP]
>
> Attention needs none of this. `--attn-implementation` is passed straight to `from_pretrained`, and
> `transformers` resolves a Hub repo ID there on its own:
> `--attn-implementation kernels-community/flash-attn2@v2` works with `--kernel-backend native`.

### Limits

> [!WARNING]
>
> **No TP / PP / CP / EP.** The model must fit under `dp_replicate` × `dp_shard`.
>
> **No LoRA.** [LoRA](/docs/advanced/lora) is Megatron-only.

For large models, multi-rack jobs, or any recipe whose fit depends on tensor, pipeline or
expert parallelism, use [Megatron-LM](#megatron-lm).

### Going deeper: when an HF model needs help

Any HuggingFace causal LM loads. Some need small corrections around the edges: a weight
layout SGLang does not expect, a stateful layer that must be reset per document, a class
that needs patching before construction. Each such architecture has one `ArchAdapter` in
`miles/backends/fsdp_utils/adaptations/specs/<arch>.py` that overrides only the hooks it
needs; an architecture without one runs on the stock HF path.

| Hook | What it fixes |
| - | - |
| `resolve_precision` | Model-specific FSDP compute / autocast policy |
| `patch_classes` | Patch a `transformers` class before construction |
| `patch_model` | Patch the constructed model; runs for both the actor and the ref model |
| `packing_kwargs` | Per-document boundaries under sequence packing, taken from the micro-batch's `cu_seqlens` and passed as HF's own padding-free kwargs (`cu_seq_lens_*`, `seq_idx`) for stateful layers such as Gated-Delta-Net |
| `param_transform` | Train to rollout parameter rename / reshape at weight sync, for example unfusing batched experts into the per-expert names SGLang expects |
| `routing_replay` | Where the MoE router selects experts, for `--use-rollout-routing-replay` (R3) |

Adapters ship today for `qwen3`, `qwen3_moe`, `qwen3_vl`, `qwen3_5` (plus Qwen3.5-MoE and
Qwen3-Next), `glm4_moe_lite` (GLM-4.7-Flash) and `nemotron_h`; `_ADAPTERS` in
`adaptations/specs/__init__.py` is the source of truth for that list.

MoE is part of this backend rather than an exception to it: expert layers use the fused
Triton kernels in `miles/kernels/moe/`, the weight bridge unfuses batched experts at sync
time, and `--use-rollout-routing-replay` (R3) works through per-architecture routing
adapters.

### Try it

```bash theme={null}
export WANDB_API_KEY=<key>

git clone https://github.com/radixark/miles.git && cd miles
pip install -e . --no-deps

# downloads model + datasets itself, no conversion step
python3 scripts/run_qwen3_0_6b_fsdp.py
```

Launchers with the same recipe shape: `scripts/run_qwen3_0_6b_fsdp.py`,
`scripts/run_qwen3_30b_a3b_fsdp.py`, `scripts/run_nemotron_3_nano_4b_fsdp.py`. To compare
the two backends on one model, `scripts/run_mcore_fsdp.py` takes `--train-backend` as a flag.

For profiling: `--use-pytorch-profiler` with `--profile-step-start` / `--profile-step-end`,
`--record-memory-history` with `--memory-snapshot-path`, and `--tensorboard-dir`. See
[Monitoring & Logging](/docs/user-guide/monitoring).

***

## torchtitan

The torchtitan backend lives at `miles/backends/torchtitan_utils/`. One idea explains the
whole thing: **torchtitan's `Trainer` is the black box.** Miles does not assemble torchtitan
parts — it builds one `Trainer.Config` from your flags, hands it to torchtitan, and from then
on owns only three things: feeding it microbatches, taking the optimizer step, and streaming
the resulting weights to SGLang. Model construction, every parallelism, the pipeline
schedule, the HF checkpoint load, the optimizer and the LR schedule are all torchtitan's.

That is why the flags below are torchtitan's own field names: they are copied verbatim into
the config a torchtitan user would write by hand.

### 1. Picking a model

```bash theme={null}
--train-backend torchtitan \
--titan-model-name qwen3 \
--titan-model-flavor 30B-A3B \
--hf-checkpoint /root/models/Qwen3-30B-A3B \
--seq-length 16384
```

`--titan-model-name` names a model package and `--titan-model-flavor` one of the sizes it
registers. The name is looked up first under `miles/backends/torchtitan_utils/models/`, then
under `torchtitan/models/`, so a model torchtitan does not ship can be added on the miles side
without touching torchtitan (see [below](#going-deeper-a-model-torchtitan-does-not-ship)). `--hf-checkpoint` supplies the weights and tokenizer;
torchtitan's own state-dict adapter converts them, so there is no offline conversion step.

`--seq-length` sizes the rotary tables and the buffers pipeline stages exchange, so it has
to be at least as long as the longest sequence you will train on — prompt plus response.
Miles rejects a value that leaves no room for a prompt rather than letting it fail inside a
kernel later.

### 2. Choosing the parallelism

```bash theme={null}
--tensor-model-parallel-size 2 \
--pipeline-model-parallel-size 2 \
--context-parallel-size 2 \
--expert-model-parallel-size 2 \
--dp-replicate-size 1
```

These are the same flags Megatron takes, so a recipe moves between the two backends without
renaming its parallelism. The FSDP shard degree is not a flag: torchtitan infers it from what
the other degrees leave over, so these five settings plus the GPU count fully determine the
layout. All of them compose — tensor, pipeline, context, expert and FSDP have been run together
on one job. `--recompute-granularity full` and `--bf16` are accepted with Megatron's meaning as
well (`--gradient-checkpointing` stays as the FSDP-side spelling).

Two notes on how they behave here:

* **Context parallelism is internal.** The trainer shards the sequence for attention and
  gathers the logits back before the loss sees them, so the RL loss and its metrics behave
  exactly as at `cp=1` and cost the same memory.
* **Pipeline parallelism needs one shape for the whole run**, so every microbatch is padded
  to `--seq-length`. Weight-tied flavors (qwen3 0.6B / 1.7B / 4B) cannot be pipelined —
  torchtitan refuses, by design.

### 3. Fitting it in memory

```bash theme={null}
--gradient-checkpointing \
--micro-batch-size 1
```

`--gradient-checkpointing` selects torchtitan's full activation checkpointing. Per-rank
memory is set by the shard, pipeline and tensor degrees together; MoE dispatch buffers do
*not* shrink with sharding, so adding nodes only helps if it raises the shard degree —
replication alone does not.

### 4. Getting weights in and out

Nothing to configure: weights are streamed to SGLang under HF names through the model's own
state-dict adapter, over IPC when the engines are colocated and over NCCL broadcast when they
are not. `--fully-async` works, and requires the disaggregated layout, since generation
continues while training runs.

Checkpoints are torchtitan's, written under `--save`; a resumed run picks up the trainer's
own step counter. Resuming needs `--ci-disable-weight-update-checker` if `--ci-test` is on,
because that check compares the engine against the original HF checkpoint.

### Limits

> [!WARNING]
>
> **The architecture must be one torchtitan implements or one registered under
> `miles/backends/torchtitan_utils/models/`.** There is no per-architecture spec to write; a
> new model is a Python package that assembles torchtitan's own blocks.
>
> **No LoRA, no optimizer CPU offload, no disk offload, no on-policy distillation, and
> `--ref-update-interval` is rejected** rather than silently ignored.

MoE works, including expert parallelism and `--use-rollout-routing-replay` (R3). R3 matters
more than it looks on MoE: without it, training and rollout routing drift apart as the policy
moves, and the disagreement compounds over rollouts rather than staying flat.

### Going deeper: a model torchtitan does not ship

torchtitan finds a model by importing `torchtitan.models.<name>` and calling its
`model_registry(flavor, attn_backend)`, which returns a `ModelSpec`: the model config, the
parallelize and pipelining functions, and the state-dict adapter that maps HF names to
torchtitan's. Miles looks in `miles/backends/torchtitan_utils/models/<name>/` first and expects
exactly the same function, so a miles-side package is a drop-in peer of a torchtitan one.

Most new architectures are a recombination of blocks torchtitan already has — MLA or GQA
attention, the sigmoid or softmax token-choice router, grouped experts, shared experts — and
then the package is a flavor table plus a `model_registry` that points at an existing
`parallelize_fn` and adapter. `models/glm4_moe_lite/` is the worked example: GLM-4.7-Flash is
DeepSeek-V3's layer with different dimensions and the same HF parameter names, so its package
builds the layer list with torchtitan's DeepSeek-V3 helpers, sets the GLM sizes, and reuses
`DeepSeekV3StateDictAdapter` unchanged.

```bash theme={null}
--train-backend torchtitan \
--titan-model-name glm4_moe_lite \
--titan-model-flavor 30B-A3B \
--hf-checkpoint /root/models/GLM-4.7-Flash
```

When the architecture genuinely differs, the package grows in this order: a `model.py` with
the new block (subclassing torchtitan's `TransformerBlock` / `BaseAttention`), then a
`state_dict_adapter.py` if the HF names differ, and only then a `parallelize.py` if the
existing sharding plans do not apply. Two checks belong with every package: a fast test that pins the flavor's dimensions to the values recorded from the model's `config.json` and checks the adapter's map covers every key pattern recorded from `model.safetensors.index.json`, and an e2e case under `tests/e2e/torchtitan/` — SGLang has to implement the architecture too, since it is what consumes the streamed weights.

### Try it

```bash theme={null}
export WANDB_API_KEY=<key>

# downloads model + datasets itself, no conversion step
python3 scripts/run_qwen3_0_6b_torchtitan.py

# same recipe with tensor parallelism
python3 scripts/run_qwen3_0_6b_torchtitan.py --tp-size 2
```

That launcher mirrors `scripts/run_qwen3_0_6b_fsdp.py` flag for flag, so the two backends'
training curves can be read against each other on one model. The cases in
`tests/e2e/torchtitan/` are the other reference: one per topology, each naming the mechanism
it exercises.

***

## The other half: SGLang

SGLang is the inference engine no matter which training backend you picked. Three pieces of
configuration matter.

**HuggingFace pointer.** SGLang boots from `--hf-checkpoint`. Miles syncs the actor's
weights from the trainer before the first training step, so the checkpoint at that path does
**not** need to be current. The tokenizer and the `config.json`-derived context length are
all SGLang reads at init.

**Context length override.** SGLang takes max context from `config.json`. To serve beyond it
during training, set `--sglang-context-length`.

**Colocation memory.** Under `--colocate` the trainer reserves VRAM during init before
handing off to SGLang, so drop `--sglang-mem-fraction-static` to **0.8** or lower to let both
fit.

### Passthrough convention

Any flag `python -m sglang.launch_server` accepts, Miles accepts with a `--sglang-` prefix:

```bash theme={null}
--sglang-ep-size 8
--sglang-enable-dp-attention
--sglang-dp-size 8
--sglang-mem-fraction-static 0.7
--sglang-log-level INFO
```

Two flags are **set by Miles** rather than by you:

* `--tp-size` from `--rollout-num-gpus-per-engine`
* `--model-path` from `--hf-checkpoint`

The integration lives at
[`miles/backends/sglang_utils/arguments.py`](https://github.com/radixark/miles/blob/main/miles/backends/sglang_utils/arguments.py).

### Router

A router sits in front of the SGLang workers. Router-side flags take a `--router-` prefix:

```bash theme={null}
--router-balance-abs-threshold 0   # force uniform distribution (lowers prefix-cache hit rate)
```

***

## Further reading

* [Core concepts](/docs/user-guide/concepts): the four objects that make up any Miles job.
* [Launch script](/docs/user-guide/launch-script): the launch script,
  argument group by argument group.
* [Fully Async RL](/docs/user-guide/fully-async): keep generation running continuously so rollout
  never waits on a training step.
* [Configuration](/docs/user-guide/cli-reference): the flag taxonomy and defaults.


# Verifiers (Prime Intellect)
Source: https://miles.radixark.com/docs/user-guide/verifiers

Train on Prime Intellect Verifiers environments with Miles.

Miles can train on a Verifiers environment in place of a prompt dataset. The
integration requires Python 3.11 or newer and Verifiers 0.2.0. Verifiers 0.2.1
requires OpenAI 2.9 or newer, while SGLang 0.5.15 pins OpenAI 2.6.1.

## Install

Install the adapter's dependencies and the Prime CLI:

```bash theme={null}
pip install -r examples/experimental/verifiers/requirements.txt
uv tool install prime
```

The recommended workspace keeps local environment packages under `./environments`:

```text theme={null}
workspace/
  environments/
    my-environment/
```

From the workspace root, install a local environment by name. For an environment from
the Environments Hub, authenticate and use its `user/environment` ID:

```bash theme={null}
# Local: ./environments/my-environment
prime env install my-environment

# Environments Hub
prime login
prime env install user/my-environment
```

## Configure

Create a Verifiers `EnvConfig` TOML file. A minimal config selects a taskset:

```toml theme={null}
[taskset]
id = "gsm8k-v1"
```

The config may also define the harness, runtime, judges, retries, and environment
limits supported by Verifiers. Verifiers applies per-rollout and group rewards before
the completed traces are returned to Miles.

The integration implements Verifiers' V1 environment contract. Legacy V0 environment
configs are rejected during startup.

## Run

The integration is a rollout function under
[`examples/experimental/verifiers`](https://github.com/radixark/miles/tree/main/examples/experimental/verifiers);
its launcher wires everything up:

```bash theme={null}
python examples/experimental/verifiers/run.py --verifiers-config /path/to/verifiers.toml
```

The launcher selects the adapter with `--rollout-function-path`, turns off Miles
prompt-data loading with `--disable-rollout-global-dataset`, and points
`VERIFIERS_CONFIG` at the file. This uses the configured taskset instead of Miles
prompt data. Environment behavior comes from the Verifiers config, while Miles
continues to own the model, sampling, batching, concurrency, reward hooks, and
optimizer settings. The Renderers library formats environment messages with Miles'
model and tokenizer settings.

The standard Miles rollout options keep their existing meaning:

| Miles option | Verifiers behavior |
| - | - |
| `--rollout-batch-size` | Number of task groups returned by each training rollout |
| `--n-samples-per-prompt` | Rollouts per training task |
| `--n-samples-per-eval-prompt` | Rollouts per evaluation task |
| `--rollout-shuffle` / `--rollout-seed` | Finite taskset order and sampling seeds |
| `--rollout-*` / `--eval-*` sampling options | Sampling and context limits |
| `--apply-chat-template-kwargs` | Typed template options passed to renderers |
| `--sglang-server-concurrency` | Physical engine capacity |
| Miles reward and filtering options | Applied after Verifiers scoring using the standard Miles hooks |

Evaluation covers every task in the taskset. Training cycles the taskset and advances
from the current Miles rollout ID when a run resumes.

## Environment Support

The adapter supports V1 environments that use the Chat Completions dialect with
text-only Renderers inputs. Tools require a model-specific renderer; use a registered
model identity in `--hf-checkpoint` or the existing `--sglang-tokenizer-path` option.
User simulators, multi-turn episodes, environment runtimes, per-rollout rewards, and
group rewards run through the standard Verifiers environment lifecycle.

Verifiers group rewards apply during both training and evaluation. Miles
`--group-rm` hooks remain training-only, matching the standard Miles rollout path.

## Limitations

`--eval-interval` works through the launcher, which evaluates the whole taskset at
that interval. Miles asserts that eval datasets are configured whenever the flag is
set, so the launcher passes a placeholder `--eval-prompt-data` naming the taskset and
pointing at its EnvConfig; the adapter serves evaluation, so the built-in loader never
opens that path.

`--partial-rollout` is not supported. A Verifiers episode owns live harness and
environment state and has no contract for resuming a partially executed episode. The
adapter rejects this combination when it is constructed, before any episode runs.

`--chat-template-path` is also rejected because Renderers owns message formatting for
Verifiers environments. Use the checkpoint's native template and
`--apply-chat-template-kwargs` instead.

Streaming model requests, Responses and Anthropic dialects, multimodal inputs, OPD,
routing replay, and indexer replay are not supported by the transport. The adapter
rejects the corresponding CLI options at startup.

Traces with multiple graph branches, including compaction, are rejected. Miles does
not currently preserve a trace's rollout-group boundary when it flattens multiple
training samples, which would make group-relative advantages incorrect.