# radixark/miles_diffusion architecture

> Architecture diagram of the GitHub repository radixark/miles_diffusion, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/radixark/miles_diffusion
- Repository: https://github.com/radixark/miles_diffusion

## Overview

Miles Diffusion trains diffusion models through SFT, DiffusionNFT, or Flow-GRPO. A Ray-managed rollout generates samples with SGLang-diffusion; reward and filtering produce scored samples, which are converted into training pairs for an FSDP2 actor. The actor updates model weights and synchronizes them to rollout engines for subsequent rounds. Model-family configs and parallel plans support the documented image, video, and omni families. The README documents LoRA sync, deterministic controls, and multiple parallelisms; sampled excerpts show portions of rollout, actor, engine, router, and conversion behavior, so uncertain wiring is omitted.

Main entry points: [`miles/rollout/sglang_diffusion_rollout.py`](https://github.com/radixark/miles_diffusion/blob/main/miles/rollout/sglang_diffusion_rollout.py) (Diffusion rollout), [`miles/backends/fsdp_utils/actor.py`](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/actor.py) (FSDP2 actor), [`miles/ray/rollout.py`](https://github.com/radixark/miles_diffusion/blob/main/miles/ray/rollout.py) (Ray rollout manager), [`miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py`](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py) (SGLang engine)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_orchestration["Ray Orchestration"]
  node_trainer["Training entry<br/>[train_diffusion.py]"]
  node_ray["Ray rollout manager<br/>[rollout.py]"]
  node_convert["Pair conversion<br/>[flow_grpo.py]"]
end

subgraph group_rollout["Generation and Rewards"]
  node_rollout_fn["Diffusion rollout"]
  node_sources["Prompt data<br/>[data_source.py]"]
  node_engine["SGLang engine"]
  node_reward["Reward models"]
  node_filters["Sampling filters"]
end

subgraph group_training["Training Backend"]
  node_actor["FSDP2 actor<br/>[actor.py]"]
  node_loss["Training losses"]
  node_model["Model backend<br/>[model_backend.py]"]
  node_config["Family configs"]
  node_parallel["Parallel plans<br/>[plan.py]"]
  node_weights["Weight synchronization"]
end

subgraph group_operations["Routing and Observability"]
  node_router["Miles router<br/>[router.py]"]
  node_dashboard["Rollout dashboard"]
end

node_user(("Training operator"))
node_sglang{{"SGLang-diffusion"}}

node_user -->|"launches"| node_trainer
node_trainer -->|"starts workers"| node_ray
node_ray -->|"dispatches rollout"| node_rollout_fn
node_rollout_fn -->|"loads prompts"| node_sources
node_rollout_fn -->|"requests generation"| node_engine
node_engine -->|"serves generation"| node_sglang
node_engine -.->|"registers worker"| node_router
node_rollout_fn -->|"scores samples"| node_reward
node_rollout_fn -->|"filters samples"| node_filters
node_ray -->|"converts samples"| node_convert
node_ray -->|"provides rollout data"| node_actor
node_actor -->|"uses conversion"| node_convert
node_actor -->|"computes objective"| node_loss
node_actor -->|"trains model"| node_model
node_config -->|"selects backend"| node_model
node_config -->|"configures parallelism"| node_parallel
node_actor -->|"updates rollout weights"| node_weights
node_weights -->|"synchronizes weights"| node_engine
node_rollout_fn -->|"records rollout telemetry"| node_dashboard

click node_trainer "https://github.com/radixark/miles_diffusion/blob/main/train_diffusion.py"
click node_ray "https://github.com/radixark/miles_diffusion/blob/main/miles/ray/rollout.py"
click node_rollout_fn "https://github.com/radixark/miles_diffusion/blob/main/miles/rollout/sglang_diffusion_rollout.py"
click node_sources "https://github.com/radixark/miles_diffusion/blob/main/miles/rollout/data_source.py"
click node_engine "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py"
click node_reward "https://github.com/radixark/miles_diffusion/tree/main/miles/rollout/rm_hub"
click node_filters "https://github.com/radixark/miles_diffusion/tree/main/miles/rollout/filter_hub"
click node_convert "https://github.com/radixark/miles_diffusion/blob/main/miles/ray/data_conversion_hub/flow_grpo.py"
click node_actor "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/actor.py"
click node_loss "https://github.com/radixark/miles_diffusion/tree/main/miles/backends/fsdp_utils/loss_hub"
click node_model "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/model_backend.py"
click node_config "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/configs/train_pipeline_config.py"
click node_parallel "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/sequence_parallel/plan.py"
click node_router "https://github.com/radixark/miles_diffusion/blob/main/miles/router/router.py"
click node_dashboard "https://github.com/radixark/miles_diffusion/tree/main/miles/dashboard"
click node_weights "https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/diffusion_update_weight_utils.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_trainer,node_ray,node_convert toneBlue
class node_rollout_fn,node_sources,node_engine,node_reward,node_filters toneAmber
class node_actor,node_loss,node_model,node_config,node_parallel,node_weights toneMint
class node_router,node_dashboard toneRose
class node_user,node_sglang toneIndigo
```

## Components

### Ray Orchestration

- [Training entry](https://github.com/radixark/miles_diffusion/blob/main/train_diffusion.py): `train_diffusion.py`
- [Ray rollout manager](https://github.com/radixark/miles_diffusion/blob/main/miles/ray/rollout.py): `miles/ray/rollout.py`
- [Pair conversion](https://github.com/radixark/miles_diffusion/blob/main/miles/ray/data_conversion_hub/flow_grpo.py): `miles/ray/data_conversion_hub/flow_grpo.py`

### Generation and Rewards

- [Diffusion rollout](https://github.com/radixark/miles_diffusion/blob/main/miles/rollout/sglang_diffusion_rollout.py): `miles/rollout/sglang_diffusion_rollout.py`
- [Prompt data](https://github.com/radixark/miles_diffusion/blob/main/miles/rollout/data_source.py): `miles/rollout/data_source.py`
- [SGLang engine](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py): `miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py`
- [Reward models](https://github.com/radixark/miles_diffusion/tree/main/miles/rollout/rm_hub): `miles/rollout/rm_hub`
- [Sampling filters](https://github.com/radixark/miles_diffusion/tree/main/miles/rollout/filter_hub): `miles/rollout/filter_hub`

### Training Backend

- [FSDP2 actor](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/actor.py): `miles/backends/fsdp_utils/actor.py`
- [Training losses](https://github.com/radixark/miles_diffusion/tree/main/miles/backends/fsdp_utils/loss_hub): `miles/backends/fsdp_utils/loss_hub`
- [Model backend](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/model_backend.py): `miles/backends/fsdp_utils/model_backend.py`
- [Family configs](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/configs/train_pipeline_config.py): `miles/backends/fsdp_utils/configs/train_pipeline_config.py`
- [Parallel plans](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/sequence_parallel/plan.py): `miles/backends/fsdp_utils/sequence_parallel/plan.py`
- [Weight synchronization](https://github.com/radixark/miles_diffusion/blob/main/miles/backends/fsdp_utils/diffusion_update_weight_utils.py): `miles/backends/fsdp_utils/diffusion_update_weight_utils.py`

### Routing and Observability

- [Miles router](https://github.com/radixark/miles_diffusion/blob/main/miles/router/router.py): `miles/router/router.py`
- [Rollout dashboard](https://github.com/radixark/miles_diffusion/tree/main/miles/dashboard): `miles/dashboard`

### Other components

- **Training operator**
- **SGLang-diffusion**

## Connections

- Training operator → Training entry: launches
- Training entry → Ray rollout manager: starts workers
- Ray rollout manager → Diffusion rollout: dispatches rollout
- Diffusion rollout → Prompt data: loads prompts
- Diffusion rollout → SGLang engine: requests generation
- SGLang engine → SGLang-diffusion: serves generation
- SGLang engine → Miles router: registers worker
- Diffusion rollout → Reward models: scores samples
- Diffusion rollout → Sampling filters: filters samples
- Ray rollout manager → Pair conversion: converts samples
- Ray rollout manager → FSDP2 actor: provides rollout data
- FSDP2 actor → Pair conversion: uses conversion
- FSDP2 actor → Training losses: computes objective
- FSDP2 actor → Model backend: trains model
- Family configs → Model backend: selects backend
- Family configs → Parallel plans: configures parallelism
- FSDP2 actor → Weight synchronization: updates rollout weights
- Weight synchronization → SGLang engine: synchronizes weights
- Diffusion rollout → Rollout dashboard: records rollout telemetry