Spotlight: Synergizing Seed Exploration and Spot GPUs for DiT RL Post-Training

2026-06-17

AI Summary

This work addresses the high computational cost of Diffusion Transformer (DiT) reinforcement learning (RL) post-training, which hinders effective utilization of low-cost Spot GPU resources. The authors propose an efficient training framework that synergistically combines seed exploration with Spot GPU scheduling, enabling the first effective deployment of DiT RL post-training on Spot instances. Key innovations include a Bandit-driven exploration mechanism tolerant to weight staleness, sub-second elastic sequence parallelism reconfiguration, and a preemption-aware pull-based request scheduler—all integrated into the ROLL platform. Experiments demonstrate that the proposed approach achieves the target validation score 4× faster than baseline methods, reduces total training cost by 1.4–6.4×, and yields superior image generation quality on high-resolution tasks.

Abstract

Reinforcement learning (RL) post-training of Diffusion Transformers (DiTs) is prohibitively expensive, requiring thousands of high-end GPUs. Existing works explore two directions to reduce cost: seed exploration improves training convergence by selecting high-contrast samples, yet adds compute to the critical path; spot GPUs offer 69--77\% lower cost, yet sit idle during training because DiT rollouts finish nearly simultaneously, which prevents LLM-style pipelining of rollout with training. Spot preemptions further break Sequence Parallelism (SP) groups, fragmenting GPU topology. We present Spotlight, the first system that harvests spot GPUs for DiT RL post-training. Spotlight rests on two key insights we devise: (1)~we show that exploration can tolerate stale model weights because exploration that uses the model weights from the previous iteration preserves the relative ranking of random seeds, allowing exploration to run on idle spot GPUs during training. (2)~SP reconfiguration can reuse on-node state, reducing group recovery from minutes to sub-second launches. Built on these insights, Spotlight introduces three techniques: a bandit-based exploration planner that maximizes reward variance within the training time budget, elastic sequence parallelism that reconfigures SP groups on the fly via persistent schedulers and intra-node weight copying, and a preemption-aware pull-based request scheduler that balances load and commits in-flight state upon preemption. We implement Spotlight on the open-source RL platform ROLL and evaluate it on Qwen-Image post-training. Spotlight reaches the same target validation score $4\times$ faster than baselines, reducing total cost by $1.4$-$6.4\times$ while achieving superior image quality on DeepSeek-OCR and Geneval datasets with resolution $512\times512$ and $1280\times1280$.

Problem

• Diffusion Transformers
• Reinforcement Learning Post-Training
• Spot GPUs
• Sequence Parallelism
• Seed Exploration

Innovation

• Spot GPUs
• Diffusion Transformers
• Reinforcement Learning Post-Training
• Elastic Sequence Parallelism
• Seed Exploration