VLM/LLM Infra Updating

[CjhHa1]

Background

UniRL now supports mainly on diffusion models, but AR models and unified models are highly required. We plan to quickly improve our infrasturcture based on opensource AR performance enhancement technoligies.

Implement Table

Priority	#	TECH	Rating	Modalities	Benefit	Effort
P0	1	AR rollout prefix-cache hits (RadixAttention / prefix cache)	⚠️	AR/VLM	High (long prompts + multi-sampling)	Small to medium (expose config; real hits need fanout/routing changes)
	2	FSDP2 forward_prefetch communication/compute overlap	✅	All	Medium (needs benchmark)	Small (config switch + adjacent-block prefetch)
	3	Enable and benchmark torch.compile	⚠️	All	Medium (already supported, needs measurement)	Very small (config)
	4	Rollout engine tuning cleanup (SGLang arg correction + recommended defaults)	⚠️	AR	Low-medium	Small (allowlist/recipes/docs)
P1	5	Sequence-length load balancing (Karmarkar-Karp bucketing)	✅	AR/VLM	High (reduces DP stragglers, especially variable-length GRPO)	Medium (dispatch layer)
	6	Dynamic batching by token budget (max_token_len_per_gpu)	✅	AR/VLM	High (memory + throughput)	Medium (train stack)
P2	7	Sequence packing / remove-padding (FlashAttention varlen)	⚠️	AR/VLM	High (removes padding waste)	Large (per-model replay changes)
	8	Activation offload (FSDP)	⚠️	All	Medium (marginal gain after checkpointing)	Medium
	9	LigerKernel / fused kernels (RMSNorm/SwiGLU/RoPE)	⚠️	AR/VLM	Medium (10-20% AR model speedup)	Medium (model side)
P3	10	Ulysses sequence parallelism (long context)	⚠️	AR	High (only when >32k context)	Large
	11	Rollout/train pipelining + one-step off-policy	⚠️	All	High (GPU utilization)	Large
	12	FP8 training (TransformerEngine)	⚠️	AR	Medium	Large

[leviking98z-rgb]

Performance parity with verl: per-technique inventory (measurement-driven)

Setup: 16x H20, Qwen3-4B-Base + DAPO, DRPO, 512 samples/step, max_new_tokens 8192, 4 mini-updates, fp32 logprob. Phase times are perf/<phase>_time_s means over a matched step window (steps 6-15, light regime, response len ~1k); every run keeps ratio_mean ~= 1.000 throughout. wandb project unirl-grpo: baseline perfphase16_unirl_base, optimized perfphase16_unirl, verl perfphase16_verl (patched to mirror the same perf/* keys, so all runs overlay in one panel); ablations abl_v11a_balance / abl_v11b_defer / abl_v11c_conc16. Note: these numbers were measured while the packing budget counted response tokens only (a cost-model accessor bug, since fixed on the #42 branch); they are real for the configs as run — a re-validation on the fixed branch is pending.

tech	verl	unirl (before)	unirl (after)	benefit for unirl (phase before -> after)	index (PR)
token-budget packed micro-batching (dynamic bsz)	✅ use_dynamic_bsz + Karmarkar-Karp + same_micro_num_in_dp	❌ micro_batch_size=1, one sequence per fwd/bwd	✅	train 75.2s -> 54.3s	#42
prompt trim to the micro's true max width	— (varlen removes padding entirely; problem does not exist)	❌ (padded to the worker-shard global max)	✅	train 54.3s -> 46.8s	#43
packed varlen replay (padding removal)	✅ use_remove_padding (flash-attn varlen)	❌ (dense rectangular replay)	✅	train 46.8s -> 27.9s (joint with flex_attention; varlen alone over SDPA regresses to 63.3s)	#43 + #44
flex_attention (fused block-causal kernel)	✅ (flash-attn-2 serves the same purpose)	❌ (SDPA cannot fuse packed block-causal masks: dense O(L^2) mask, no cross-sequence block skipping)	✅	included in the joint number above; single-GPU micro-bench fwd+bwd 7.6s -> 2.63s (5 seqs / 10.6k tokens)	#44
DP seqlen load balancing	✅ trainer.balance_batch (on by default)	❌	✅ (recipe default on)	train 28.4s -> 15.3s (-46%), whole step 78.1 -> 64.0s; rank token spread -> 0.6-0.9%. Effect is large because GRPO groups dispatch as consecutive siblings: one rank draws all 8 long answers of a hard prompt	#45
one-wave decode (concurrency / cudagraph bs 32)	— (vllm async server; different mechanism)	❌ (defaults 16)	knob rejected	structurally ineffective (review): at the documented 32-GPU config each engine receives 512/32 = 16 requests, so the concurrency semaphore (16) never throttles; and with mem_fraction_static: 0.3 the KV pool admits only ~12-16 concurrent 8k requests, so graphs for bs 17-32 sit unused. Clean ablation confirmed no gen delta (48.9 vs 48.3s)	#47 (closed)
defer_grad_sync (one reduce-scatter per optimizer step)	❌ (reduce-scatters every micro backward)	✅ in code, off by default	unchanged (kept off)	no measurable gain (train 29.4 vs 28.4, within noise) — the per-micro reduce-scatter is already overlapped; also requires reshard_after_forward=false (memory cost). Not enabled	(no PR)
rollout engine TP>1 colocate	✅ (vllm tp=2)	❌	✅ capability (default stays tp=1)	negative: gen 47.7 -> 70.0s (len ~1k), 110.5 -> 177.7s (len ~3.4k). With custom all-reduce disabled (container CUDA-graph IPC), per-layer NCCL allreduce latency dominates — the vllm-tp2 advantage does not transfer to sglang-NCCL. Capability + 3 correctness fixes still land	#46

Whole-step totals (light regime, matched window): baseline 126.7s -> 64.0s with all effective techniques (verl 62.2s in the same window) — gap 2.2x -> ~1.03x, with train at parity (15.3s vs 16.6s) and the residual in generate tail latency (46.3 vs 42.7s; the rare long sequence decodes on one GPU vs verl's tp2).

Heavy regime — controlled test (ignore_eos forces every response to exactly 4096 tokens on both stacks, removing length-distribution confounds; runs heavy4k_unirl n=35 steps / heavy4k_verl n=4 steps, both ended by an external SIGKILL that does not affect completed-step timings): unirl 112.0s (gen 57.0 / train 52.9) vs verl 113.8s (gen 56.4 / train 53.5) — parity within 2% on every phase. An earlier estimate of a ~1.9x heavy-regime train gap was measured BEFORE DP balancing and was dominated by the GRPO-group straggler effect, not by kernels; the controlled test shows flex_attention+packed matches flash-attn-2+remove_padding per token at 4k. This also deprioritizes installing flash-attn as a follow-up.

Correctness evidence: packed-vs-dense replay equivalence in fp32, mean |dlogp| = 1e-5; packed-micro loss weighting mathematically exact (unit-tested to 1e-12); ratio_mean ~= 1.000 on every run; balanced run reward/loss curves match unbalanced.

End-to-end phase breakdown (all effective optimizations enabled)

Light load (natural lengths, response len ~1k; matched window steps 6-15)

runs: perfphase16_unirl_base / abl_v11a_balance / perfphase16_verl

	step total	generate	train	reward	weight_sync	wake_up	sleep
unirl (before)	126.7s	49.9	75.2	0.7	0.8	0.0	0.1
unirl (after)	64.0s	46.3	15.3	1.0	1.2	0.0	0.2
verl	62.2s	42.7	16.6	0.0*	2.8	—	—

Heavy load (uniform fixed-length 4096 via ignore_eos on both stacks; steps 2-5)

runs: heavy4k_unirl_base / heavy4k_unirl / heavy4k_verl

	step total	generate	train	reward	weight_sync	wake_up	sleep
unirl (before)	149.3s	64.3	81.7	1.9	1.1	0.0	0.2
unirl (after)	111.8s	56.8	52.6	0.7	0.7	0.0	0.3
verl	113.7s	56.5	53.6	0.0*	3.6	—	—

* verl computes rewards inside the async agent loop (overlapped with generation), so its timing_s/reward reads ~0; wake_up/sleep do not exist on the vllm async path. Phase sums cross-check against the step totals (e.g. unirl-after light: 46.3+15.3+1.0+1.2+0.2 = 64.0). The heavy benchmark forces every sequence to decode exactly 4096 tokens (UNIRL_FORCE_IGNORE_EOS=1; verl needed a one-line patch for its async-vllm path to honor rollout.ignore_eos).

Bottom line: with #42-#45 applied, UniRL is at parity with verl on this workload — light regime 64.0s vs 62.2s (1.03x, residual = tp1 tail-latency in generate) and heavy regime 111.8s vs 113.7s (0.98x, unirl ahead; generate and train both at parity).

[leviking98z-rgb]

VLM addendum: do the LLM optimizations transfer? (measured on Qwen2.5-VL-7B + geo3k_mc, 16x H20)

Setup: qwen_vl_grpo_geo3k_mc_sglang_4x8 recipe at num_devices=16 (32 prompts x 16 samples = 512/step, max_new 512, mbs=2, grad-ckpt on); phase times = perf/<phase>_time_s means over steps 2-5; runs vl_base / vl_bal (wandb project unirl-grpo). This workload is prompt/vision-dominated: responses average ~2 tokens (multiple-choice answers), the opposite shape of the LLM DAPO workload — which drives every verdict below.

tech	verl	unirl (before)	unirl (after)	benefit for unirl (measured)	index (PR)
token-budget packed micro-batching	✅ (generic)	❌ (mbs=2)	knob available (#42), cannot engage	none available on this workload: 7B fp32 master + optimizer shards + SRT 0.25 reservation leave activation room for only ~2 vision-heavy rows even with grad-ckpt — packed micros OOM at budget 10240 AND 6144 (158 MiB alloc vs 60 MiB free, deterministic). Packing is memory-capped to the existing mbs=2	#42 (knob)
DP seqlen load balancing	✅	❌	knob available (#45), keep off here	negative on this workload: responses are uniform (~2 tokens, rank spread already 0.0%) so there is nothing to balance, while the driver-side hydrate+reorder costs real time — step 28.2s -> 37.8s. Balance pays only on variable-length workloads (it was -46% train on LLM DAPO)	#45
prompt trim to the micro's true max	— (varlen)	✅ already present (VL replay strips right-padding per micro)	✅	n/a (pre-existing in the VL stage)	—
packed varlen replay + flex_attention	✅	❌	❌ not ported	blocked: Qwen2.5-VL M-RoPE uses 3D position ids; transformers' packed-sequence detection (1D position diffs) cannot express it — needs an explicit BlockMask + VL-aware rope indexing	—
SGLANG_DISABLE_CUDNN_CHECK=1	n/a	required	required	the VL vision tower trips sglang's cudnn version check; without the env the SRT server dies at startup (LLM runs do not trigger it)	—

Phase breakdown (VL, natural load)

	step total	generate	train	reward	weight_sync
unirl vl_base	28.2s	7.3	16.1	3.4	0.9
unirl vl_bal (balance on)	37.8s	6.9	20.7	3.3	1.0

Takeaways

1. The LLM wins (packing/balancing) target response-length variance and response-dominated compute; geo3k_mc has neither. On VL workloads with long CoT responses they should transfer (memory permitting), but that is untested here — left blank rather than claimed.
2. The actual VL-specific levers, in measured order: train 16.1s is prompt/vision replay (the M-RoPE-aware prompt-side varlen port would attack this), and reward 3.4s is ~12% of the step (largest relative reward share seen in this series).
3. PR perf(recipe): enable token-budget packing + DP balancing for Qwen-VL recipes #48 is closed accordingly: nothing in the current VL recipe should flip; the perf(train): token-budget packed micro-batches (verl dynamic-bsz parity) (issue #40) #42/perf(trainer): driver-side DP seqlen balancing (verl balance_batch parity) (issue #40) #45 knobs remain available for variable-length VL workloads.