How radixark/miles_diffusion works
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: sglang_diffusion_rollout.py Diffusion rollout · actor.py FSDP2 actor · rollout.py Ray rollout manager · sglang_diffusion_engine.py SGLang engine
Components 18
Ray Orchestration
- Training entry
train_diffusion.py - Ray rollout manager
miles/ray/rollout.py - Pair conversion
miles/ray/data_conversion_hub/flow_grpo.py
Generation and Rewards
- Diffusion rollout
miles/rollout/sglang_diffusion_rollout.py - Prompt data
miles/rollout/data_source.py - SGLang engine
miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py - Reward models
miles/rollout/rm_hub - Sampling filters
miles/rollout/filter_hub
Training Backend
- FSDP2 actor
miles/backends/fsdp_utils/actor.py - Training losses
miles/backends/fsdp_utils/loss_hub - Model backend
miles/backends/fsdp_utils/model_backend.py - Family configs
miles/backends/fsdp_utils/configs/train_pipeline_config.py - Parallel plans
miles/backends/fsdp_utils/sequence_parallel/plan.py - Weight synchronization
miles/backends/fsdp_utils/diffusion_update_weight_utils.py
Routing and Observability
- Miles router
miles/router/router.py - Rollout dashboard
miles/dashboard
Other components
- Training operator
- SGLang-diffusion
Connections 19
- 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