Rollplex: Cross-Phase GPU Spatial Sharing for Vision Language Model Post-Training
Organizations: HKUST Hong Kong SAR · Alibaba Inc China
Abstract
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 -- speedup over serial colocation and -- over disaggregation under the same GPU budget, while preserving the synchronous RL update.
Figures & tables
| Role | Component | GiB |
| Training actor | bfloat16 (BF16) weights | 7.6 |
| (TP=8) | FP32 gradient buffer | 15.1 |
| FP32 master weights | 15.1 | |
| FP32 Adam moments | 30.3 | |
| Activations (one microbatch) | 28.3 | |
| Preserved prefix KV | 16.0 |
| Policy | CLEVRER | LLaVA | Percep. | STAR |
|---|---|---|---|---|
| Sequential (no overlap) † | 44.5 15.0 | 216.1 133.7 | 163.4 79.7 | 56.2 23.9 |
| MPS-Default | 39.2 16.7 | 197.8 134.9 | 115.5 81.4 | 41.1 23.8 |
| ThreadCap 80%/20% | 40.0 19.0 | 199.6 135.4 | 120.8 84.5 | 40.7 23.6 |
| ThreadCap 70%/30% | 126.6 73.8 | 259.5 177.0 | 222.6 150.1 | 91.6 65.0 |
| SM-Affinity 80%/20% | 38.2 17.3 | 200.2 137.3 | 115.4 76.3 | 42.2 24.3 |
| SM-Affinity 70%/30% | 131.2 75.5 | 259.5 175.7 | 213.5 140.4 | 93.7 65.5 |
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
| Object | Produced | Consumed | Autograd | Residency policy | In P1 peak |
|---|---|---|---|---|---|
| Rollout KV | P1 prefill | P1 decode | n/a (inference) | HBM; released after generation | yes |
| Boundary prefix KV (training actor) | P1 prefix fwd | P2 suffix fwd/bwd | attached | HBM; kept until suffix consumes it | yes |
| Boundary prefix KV (reference) | P1 prefix fwd | P2 scoring | detached | HBM; kept until scoring consumes it | yes |
| Reference-model parameters | permanent (frozen) | P1 ref prefix / P2 ref scoring | n/a (frozen) | PinnedPool; loaded to HBM around reference compute, offloaded when it completes | yes (during ref prefix) |
| Prefix logprobs | P1 prefix fwd | P2 scoring/training | detached or attached by role | HBM or host; retained until combined with suffix logprobs | yes (small) |
| Training-prefix activation/checkpoint state | P1 prefix fwd | P2 backward (reload or recompute) | attached | Offloaded immediately to PinnedPool or dropped for recompute; rematerialized during backward | yes (handles) |