Technical Summary: SPOTLIGHT

Problem Statement

Reinforcement learning (RL) post-training for Diffusion Transformers (DiTs) is prohibitively expensive, often requiring thousands of high-end GPUs. The standard GRPO-style workflow involves two phases: a rollout phase (generating samples with random seeds) and a training phase (computing gradients and updating weights).

Existing cost-reduction strategies face two primary bottlenecks when applied to DiTs:

1. Seed Exploration Overhead: To accelerate convergence, methods like DanceGRPO and TreeGRPO explore multiple random seeds to find high-contrast samples (high and low rewards). However, this exploration phase sits on the critical path, adding compute time to every iteration and negating potential cost savings.
2. Spot GPU Idleness and Fragmentation: While spot GPUs offer 69–77% lower costs, they are preemptible and have limited inter-node bandwidth.
   • Idleness: Unlike Large Language Models (LLMs) where rollout latencies vary (allowing pipelining with training), DiT rollouts complete in nearly uniform time. This creates a synchronization barrier where spot GPUs sit idle during the entire training phase.
   • Fragmentation: Frequent spot preemptions break Sequence Parallelism (SP) groups, which are essential for running large DiT models. Reconfiguring these groups typically requires minutes to reload models and reinitialize engines, wasting significant resources.

Methodology: SPOTLIGHT

The authors present SPOTLIGHT, the first system designed to harvest spot GPUs for DiT RL post-training by synergizing seed exploration with spot resources. The system is built on two key insights and three core techniques.

Key Insights

1. Exploration Tolerates Stale Weights: The authors demonstrate that the relative ranking of random seeds (which seeds produce high vs. low rewards) remains highly consistent even when using model weights from the previous iteration rather than the current updated weights. This allows seed exploration to run on spot GPUs using stale weights while the reserved cluster performs the training phase, effectively offloading exploration from the critical path.
2. SP Reconfiguration Can Reuse State: The high cost of reconfiguring SP groups stems from CPU scheduler initialization and remote model loading. By decoupling the scheduler and utilizing intra-node weight copying (via NVLink) from co-located peers, SP reconfiguration time can be reduced from minutes to sub-seconds.

System Architecture & Techniques

SPOTLIGHT is implemented on the open-source RL platform ROLL and employs three main techniques:

1. Bandit-Based Exploration Planner:

   • The system formulates exploration planning as a multi-armed bandit problem.
   • It dynamically selects the number of sequences and effective denoising steps (using caching strategies like Teacache) to maximize reward variance within the time budget provided by the training phase.
   • It uses an Upper Confidence Bound (UCB) strategy to balance exploring new configurations and exploiting known high-variance settings.

2. Elastic Sequence Parallelism (SP):

   • Persistent Scheduler: The CPU scheduler is decoupled from GPU workers and kept resident on the node across SP changes, eliminating repeated initialization costs.
   • Intra-Node Weight Loading: When a new GPU worker joins an SP group on the same node, it copies weights from a co-located peer over NVLink rather than pulling from a remote node, drastically reducing reconfiguration latency.

3. Preemption-Aware Pull-Based Request Scheduler:

   • Rollout workers pull requests from a centralized queue rather than being statically assigned tasks, adapting to heterogeneous worker speeds.
   • State Commitment: Intermediate rollout states (latents and metadata) are committed to a persistent Tensor Store on reserved nodes. If a spot GPU is preempted, the system can recover the in-flight state and resume the request on another worker, minimizing progress loss.

Results

The system was evaluated on Qwen-Image (20B) post-training using the DeepSeek-OCR and Geneval datasets at resolutions of 512×512 and 1280×1280.

• Cost Reduction: SPOTLIGHT reduces total training cost by 1.4× to 6.4× compared to state-of-the-art baselines (including systems using spot GPUs without these optimizations and reserved-only systems).
• Convergence Speed: It reaches the same target validation score 4× faster than baselines. This is achieved by reducing the number of iterations required (due to effective seed exploration) and overlapping exploration with training.
• Quality: SPOTLIGHT achieves superior validation scores on all datasets compared to baselines, attributed to the dynamic selection of high-contrast seeds.
• Efficiency:
  • Exploration overhead is kept low (2–3% of iteration time) by fitting exploration within the training window.
  • Elastic SP allows the system to track spot capacity changes in near real-time, whereas naive approaches suffer multi-minute delays during reconfiguration.
  • The system demonstrates near-linear scalability in throughput as spot GPU counts increase from 8 to 64.

Significance and Claims

The paper claims that SPOTLIGHT is the first system to successfully harvest spot GPUs for DiT RL post-training. Its significance lies in:

1. Breaking the Critical Path Dependency: By proving that seed exploration can tolerate stale weights, the system eliminates the trade-off between faster convergence (via exploration) and iteration latency.
2. Resolving Spot Idleness: It solves the specific idleness problem of DiT rollouts by utilizing the training phase's idle spot capacity for exploration, a synergy overlooked by previous LLM-focused systems.
3. Robustness to Preemption: Through elastic SP and state persistence, it makes DiT training resilient to the frequent preemptions inherent in spot markets, turning a fragmented resource pool into a viable, cost-effective training infrastructure.

The authors conclude that these contributions enable a cost-efficient DiT RL post-training system that does not sacrifice image quality, making high-quality diffusion model training more accessible.