Skip to content

forward_backward queue drain can create oversized backend batches #1873

Description

@j316chuck

Summary

A client can submit many pending FORWARD_BACKWARD requests with uneven sequence lengths. The engine currently drains all batchable pending backward requests before the next destructive barrier into one backend call, so a long queue can become one very large training batch. With large Megatron + LoRA models, this can surface as CUDA OOM during backward.

Fix PR: #1874

Public Real-Service Repro

The PR includes a public repro at:

python examples/repro_forward_backward_queue_drain.py summarize
python examples/repro_forward_backward_queue_drain.py run-forward-backward \
  --base-url http://$SERVICE_HOST:8000 \
  --future-timeout-s 7200 \
  --skip-leading-blocker \
  --max-pending-requests 17

The repro uses synthetic token IDs to avoid private data, but it sends actual forward_backward requests through a SkyRL/Tinker-compatible service URL and waits for real service results.

Full queued shape:

Client shape: many pending FORWARD_BACKWARD requests with uneven sequence lengths.
Optional first single request: sequence_length=10,000.
single queue drain before limiting: requests=193, examples=193, max_sequence_length=35,000, prepared_input_tokens=816,247
pressure symptom: process_batch_requests(forward_backward) can coalesce the pending requests into one large train call
padding note: sample microbatching can pad all rows in the coalesced batch to the longest sequence before backend microbatching
expected unbounded padded batch shape: rows=193, sequence_slots=6,755,000
model config note: text hidden_size is 5,120; 2,304 is vision_config.num_position_embeddings, not a text activation width.

The live OOM proof uses the first 17 pending requests: eight short requests, one 35,000-token request, then eight more short requests.

Endpoint shape:

model: Qwen/Qwen3.6-27B
backend: SkyRL-Train Megatron + LoRA rank 8
tensor parallelism: 4
micro_train_batch_size_per_gpu: 16
micro_forward_batch_size_per_gpu: 16
remove_microbatch_padding: false
max_tokens_per_microbatch: -1

Observed Failure

Case Result
No request-count cap OOMs on a real endpoint. Warmup n=1 completes, then process_batch_requests(forward_backward, n=16) fails with torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 9.08 GiB while running backward through the LoRA MLP path.
Request-count cap only Still OOMs. The engine splits the queue, but SkyRL-Train logs Padded batch from 1 to 16 (alignment=16), so the one 35,000-token request is expanded back to the configured microbatch size before backward.
Cap + cap-aware padding Passes. The engine logs forward_backward request batch split 16, processes sixteen n=1 backend calls, the 35,000-token request returns metrics, and there is no OOM.

Expected Behavior

The engine should preserve request ordering and barrier semantics while avoiding unbounded backward work per backend call. When request-count limiting creates small chunks, SkyRL-Train should not pad those chunks back up to the full configured microbatch size before backward.

Proposed Fix

Add optional forward_backward_max_request_count engine config. When set, the engine splits pending FORWARD_BACKWARD requests into ordered chunks before calling the backend. Also make SkyRL-Train padding cap-aware by using min(actual_batch_size, configured_micro_batch_size) for the padding alignment.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions