PaperlayerView exact revision on arXiv

Source-linked AI summary

Rollplex: Cross-Phase GPU Spatial Sharing for Vision Language Model Post-Training

Hanfeng Lu, Tianyu Feng, Suyi Li, Yuheng Zhao, Wei Gao, Shaopan Xiong, Ju Huang, Siran Yang, Jiamang Wang, Lin Qu, Wei Wang

arXiv:2608.14498v1cs.LGcs.DC

TL;DR

VLM on-policy RL runtimes serialize phases even though prompt-prefix processing is response-independent and dominates visual workloads. Rollplex overlaps prefix computation with rollout decoding using phase-aware memory management and parallelism-aware weight sharing, achieving 1.23×–1.30× speedup over serial colocation and 1.57×–2.24× over disaggregation on 32 H800 GPUs while preserving synchronous updates.

  • Problem

    VLM RL workloads contain substantial response-independent prompt-prefix computation, but on-policy runtimes serialize rollout, scoring, and training phases.

  • Method

    Rollplex overlaps prefix computation with rollout decoding through phase-aware memory management and parallelism-aware weight sharing.

  • Results

    1.23×–1.30× speedup over serial colocation and 1.57×–2.24× over disaggregation was achieved on 32 H800 GPUs under the same GPU budget.

  • Takeaways & Limitations

    Rollplex enables overlapping VLM RL phases while preserving strict synchronous on-policy ordering under the evaluated GPU budget.

  • Takeaways & Limitations

    Rollplex benefits diminish when prefix computation is small or a clean rollout-training split is affordable.

Abstract

from arXiv · show

Vision-language models (VLMs) enable embodied agents to reason and act from visual observations and language instructions. Reinforcement learning (RL) post-training enhances these capabilities using task feedback, but current on-policy RL runtimes execute rollout, reference scoring, and actor training in strict serial phases. While effective for text-only RL, this phase-granular execution is wasteful for VLMs, where processing dense video inputs and prompt prefixes occupies a large fraction of each phase. Because prefix processing is independent of the generated response, it can be run alongside rollout decoding, which leaves GPU compute capacity underutilized, without breaking synchronous on-policy semantics. We present Rollplex, a runtime that decomposes the reference and training phase and moves the prefix computation into the rollout decode window. Realizing this schedule requires more than concurrent kernel launches: naive colocation of Qwen2.5-VL-32\,B requires roughly 165\,GiB per GPU, while rollout and training prefer different tensor-parallel (TP) degrees and weight layouts. Rollplex addresses these constraints with two mechanisms. Phase-aware memory management controls HBM residency according to producer--consumer lifetimes. Parallelism-aware weight sharing uses the same physical storage for layout-compatible tensors across distinct TP degrees and reconstructs only incompatible tensors, avoiding a complete second actor copy. On 32 H800 GPUs, Rollplex achieves $1.23\times$--$1.30\times$ speedup over serial colocation and $1.57\times$--$2.24\times$ over disaggregation under the same GPU budget, while preserving the synchronous RL update.

1 Introduction

Rollplex targets the input-heavy execution profile of VLM reinforcement learning by overlapping rollout-independent prefix computation with rollout decoding while preserving synchronous update semantics. It enables this schedule through phase-aware memory management and parallelism-aware weight sharing, achieving substantial speedups over serial colocation and disaggregation.

  • Motivation: VLM RL is input-heavy: video-workload prompt tokens account for 79%–98% of each sampled sequence, versus 11%–19% for MATH and GSM8K.This makes phase-level serialization unnecessarily coarse because prompt-side computation rivals autoregressive decoding.
  • Execution schedule: Rollplex overlaps rollout-independent prefixes with rollout decoding, then resumes response-dependent suffixes after generation while preserving synchronous RL semantics.The runtime decomposes later phases at the first response position and treats decode as the critical path.
  • System challenges: Naive concurrent execution requires roughly 165 GiB per GPU for Qwen2.5-VL-32 B, exceeding an H800’s 80 GB capacity.The overlap also risks decode interference from prefix kernels, which can extend the critical path and erase its benefit.
  • Design: Rollplex uses phase-aware memory management to control HBM residency by producer–consumer lifetimes and parallelism-aware weight sharing to alias compatible layouts while copying only incompatible tensors.A CUDA-VMM-backed allocator preserves virtual addresses as physical residency changes, avoiding engine-level tensor reconstruction.
  • Evaluation: 1.23×–1.30× speedup over serial colocation and 1.57×–2.24× over disaggregation is achieved on 32 H800 GPUs under the same GPU budget.The evaluation uses Qwen2.5-VL-32 B across four video-reasoning workloads.

2 Motivation and Challenges

VLM RL is input-heavy because long visual prompts make prefix processing a substantial cost, while rollout decoding leaves compute capacity underused. Overlapping prefix work with decoding preserves synchronous on-policy ordering but creates HBM and tensor-parallel layout challenges.

  • VLM RL motivation: Video prompts contain 470–5,900 median tokens, making prompt processing 79%–98% of each video sample.Text baselines contain 60–90 tokens, with prompt shares of only 11%–19%.
  • VLM RL motivation: Prefix computation combines vision encoding and prompt prefill, while rollout decoding maintains below 19% active-SM utilization in the H800 profile.Their complementary compute profiles motivate running prefix work during the decode window.
  • Cross-phase overlap: The overlap schedule preserves strict on-policy ordering by using the same model weights and inputs as serial execution.Its maximum benefit is bounded by min(T_decode,T_prefix) prefix work hidden within the decode window.
  • Feasibility challenges: Cross-phase state can exceed HBM because rollout KV, boundary KV, and training state must coexist instead of residing only during their active phases.Sequential execution avoids this simultaneous residency requirement.
  • Feasibility challenges: Training prefers TP=8 while rollout prefers TP=4; matched TP=4 cannot fit training state, whereas matched TP=8 slows rollout by up to 1.31×.Different shard boundaries and parameter ordering further complicate sharing one physical actor-weight copy.

3 Rollplex Design

Rollplex overlaps prefix computation with rollout decoding while preserving synchronous phase dependencies through barriers. It makes this schedule feasible with phase-aware HBM residency and parallelism-aware sharing of actor weights across different TP layouts.

  • Cross-phase execution: Three barriers ensure suffixes begin after response tokens and prefix boundaries are ready, while updates wait until all readers of θ_k stop.Physical mappings are released or rebound only when no active kernel can dereference the old view.
  • Cross-phase execution: Rollplex runs reference and training prefix work alongside rollout decode, while response-dependent suffixes and backward follow generation.An orchestrator treats rollout decode as the critical path and enforces phase barriers around shared actor snapshot θ_k.
  • Phase-aware memory management: Rollplex keeps each object in HBM only from its first latency-critical use through its last use, classifying state as permanent, regenerable, or boundary state.VMM reservations preserve tensor addresses across eviction while physical pages are unmapped and remapped as needed.
  • Phase-aware memory management: Peak update memory is independent of the full FP32 optimizer footprint because adjacent chunks load, update, and offload optimizer state through a pipelined stream schedule.Chunk sizing accommodates the pipeline, partially regenerated BF16 snapshot, and current gradient slice simultaneously.
  • Parallelism-aware weight sharing: Rollplex shares one VMM-backed physical placement for compatible actor tensors across different TP degrees and reconstructs tensors when mappings cannot safely alias.For q=TP_train/TP_rollout∈Z>0, grouped training ranks populate the rollout-sized actor snapshot through exchange and predetermined offsets.

4 Evaluation

Rollplex reduces end-to-end step time across datasets by overlapping prefix computation with rollout decoding while preserving training quality and synchronous on-policy semantics. Its phase-aware memory management and parallelism-aware weight sharing keep the overlapped execution within HBM and enable preferred rollout parallelism.

  • 4.2 End-to-end performance: 1.23×–1.30× lower step time than Colocate and 1.57×–2.24× lower than Disagg makes Rollplex fastest on every dataset.The comparison uses the same 32-GPU budget and matched model, data order, rewards, batch size, and sequence lengths.
  • 4.3 Training quality: Near-zero reward differences across all four datasets show that Rollplex improves time-to-train without affecting training quality.The overlapped prefixes use the same on-policy snapshot and preserve the prefix KV state that serial execution would recompute.
  • 4.4 Phase-aware memory management: All four memory-management ablations run out of memory, whereas full Rollplex runs continuously below the 80 GB per-GPU capacity.Phase-scoped residency releases rollout KV, reference weights, and actor snapshots at producer–consumer boundaries, while optimizer state is streamed in chunks.
  • 4.5 Parallelism-aware weight sharing: 1.06×–1.17× lower end-to-end step time on longer datasets results from rollout TP=4 instead of the slower matched TP=8 configuration.Parallelism-aware weight sharing lets rollout use TP=4 against TP=8 training while both read one physical actor-weight allocation.
  • 4.6 GPU execution sharing: 1.09×–1.41× lower total step time is achieved with MPS-Default, which remains within 2.7% of the best policy without tuning.Rigid 70%/30% limits increase total time by 1.30×–3.44% relative to the corresponding baseline.
  • 4.7 Extensions: 1.08×–1.21× and 1.04×–1.23× per-step reductions extend Rollplex’s benefit to partial rollout and oversampling, respectively.Each method is compared only with and without Rollplex because the methods change the amount of rollout work per step.

5 Discussion

Rollplex applies when training-side prefix computation depends on the prompt and snapshot θ_k but not the rollout response, as in VLM video encoding and visual-token prefill. Its benefits are largest when prefix work rivals decoding and separate preferred-TP rollout and training pools are infeasible.

  • Applicability Envelope: Rollplex targets response-agnostic prefix work dependent on the prompt and snapshot θ_k, including VLM video encoding and visual-token prefill.These computations can materialize prefix KV before the response is known.
  • Applicability Envelope: Benefits are largest when prefix work rivals decode and the GPU budget cannot support separate rollout and training pools at their preferred TP degrees.The passage states that benefits diminish when the prefix is not substantial relative to decoding.

6 Related Work

Rollplex distinguishes itself from prior RL, GPU multiplexing, memory-management, and serving systems by spatially overlapping response-independent prefix work across rollout and training engines within one synchronous iteration. It coordinates phase dependencies, shared state, and cross-TP weight views while preserving on-policy semantics.

  • Cross-phase overlap: Rollplex spatially overlaps independent rollout and training work from the same synchronous iteration, crossing engines rather than pipelining within training.Resident engines read one snapshot at different TP degrees without converting layouts or transferring snapshots.
  • GPU multiplexing: Rollplex derives co-runners from the on-policy dependence graph instead of multiplexing independent jobs, coordinating autograd state, KV caches, optimizer residency, and shared weights.It uses MPS but overlaps models and engines across RL phases, unlike systems that overlap operations inside one serving engine.
  • Memory management: Rollplex uses VMM for phase-scoped memory residency and cross-process weight views, rather than stable KV addresses or reduced fragmentation.BEEMS smooths vision-inference memory demand, while vAttention and GMLake target KV-address stability and fragmentation reduction.
  • Decode-window utilization: Rollplex fills decode windows with response-independent reference and training prefixes from the same iteration, without changing sampled responses or PPO/GRPO data dependence.This differs from APRIL and RollPacker, which add rollout work, and serving systems, which use other requests for prefill–decode co-scheduling.
  • Prefix sharing: Rollplex’s speedups come from moving already-deduplicated prefix computation across phases, not from eliminating redundant recomputation in GRPO baselines.Each role computes its GRPO prefix once per prompt group and reuses it across sampled responses.

7 Conclusion

Rollplex overlaps VLM prefix processing with rollout decoding using phase-aware memory management and parallelism-aware weight sharing. On 32×H800 GPUs, it improves throughput over serial colocation and disaggregation under the same GPU budget.

  • 7 Conclusion: Rollplex overlaps VLM prefixes with rollout decode through phase-aware memory management and parallelism-aware weight sharing.The runtime is designed to share GPU resources across phases while managing memory and weights across differing parallelism configurations.
  • 7 Conclusion: 1.23×–1.30× speedup over serial colocation is achieved on 32×H800 GPUs under the same GPU budget.This is the reported improvement against the serial-colocation baseline.
  • 7 Conclusion: 1.57×–2.24× speedup over disaggregation is achieved on 32×H800 GPUs under the same GPU budget.This is the reported improvement against the disaggregation baseline.

A Appendix

The appendix provides detailed implementation contracts for reproducing Rollplex, while noting that these details are not necessary to understand the main execution schedule.

  • The appendix documents implementation contracts useful for reproducing Rollplex but not required to follow the main execution schedule.

A.1 Intermediate-data lifecycle

Rollplex fits the overlap window by retaining only latency-critical boundary data across phases while evicting bulky training state from HBM. Its per-object lifecycle contract limits the Phase 1 peak to rollout, boundary, weight, and residual state, excluding suffix training state.

  • Intermediate-data lifecycle: The lifecycle contract tracks each tensor’s production, consumption, autograd attachment, and memory residency across Rollplex’s three phases.Table 3 defines these fields per object and ties their lifetimes to the P1/P2/P3 schedule.
  • Intermediate-data lifecycle: Rollout KV remains HBM-resident only during generation, while boundary prefix KVs and prefix logprobs persist across the phase boundary for Phase 2.Only the training actor’s boundary-KV copy remains attached to the backward graph, because gradients re-enter through that boundary.
  • Intermediate-data lifecycle: Bulky training state is removed from HBM after production and later offloaded, discarded for recomputation, or rematerialized for Phase 2 backward.This includes training-prefix activations and checkpoint state, gradient buffers, FP32 master weights, and Adam moments.
  • Intermediate-data lifecycle: The Phase 1 peak includes rollout KV, boundary KVs, the shared actor snapshot, transient reference weights, and small residuals, but excludes suffix activations, gradients, and optimizer state.This per-object accounting explains why the combined working set fits despite naive colocation exceeding available capacity.
Loading 2608.14498v1…