Limit forward_backward queue drain size

[j316chuck]

Summary

Closes #1873.

A client can submit many pending FORWARD_BACKWARD requests with uneven sequence lengths. Today the engine can drain every batchable pending backward request before a destructive barrier into one backend call, so a long pending queue can become one oversized training step.

This PR bounds that path in two places:

• forward_backward_max_request_count splits pending backward requests into ordered chunks before backend processing.
• SkyRL-Train padding now uses min(actual_batch_size, configured_micro_batch_size) so a capped one-request chunk is not padded back up to the full configured microbatch size.

The default remains unlimited unless operators configure the cap.

Public Repro

This PR includes a public real-service repro. The payloads are synthetic token IDs to avoid any private data, but the script sends actual forward_backward calls through a SkyRL/Tinker-compatible service URL and waits for real service results.

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

Shape summary:

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 repro uses the first 17 pending requests from that shape: eight short requests, one 35,000-token request, then eight more short requests.

Endpoint shape used for before/after validation:

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

Before / Cap-Only / After

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.
Current PR	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.

flowchart TD
  A["Client queues uneven FORWARD_BACKWARD requests"] --> B["Engine drains pending backward work"]
  B --> C1["Before: one backend call with 16 pending requests"]
  C1 --> D1["Sample microbatch padding sees max seq len 35,000"]
  D1 --> E1["Backward OOM: attempted 9.08 GiB allocation"]
  B --> C2["Cap-only: split into one-request chunks"]
  C2 --> D2["Backend pads each 1-row chunk back to 16 rows"]
  D2 --> E2["Single long request still OOMs"]
  B --> C3["Current PR: split + cap-aware padding"]
  C3 --> D3["1-row chunk keeps padding alignment at 1"]
  D3 --> E3["Long request completes; full client succeeds"]

Loading

Code Changes

• Adds EngineConfig.forward_backward_max_request_count.
• Splits pending FORWARD_BACKWARD requests into ordered request-count chunks before calling process_batch_requests.
• Preserves ordering/barrier behavior: all eligible backward chunks run before forward/sample batches and single requests.
• Adds effective_padding_micro_batch_size(...) and applies it before _pad_batch(...) for forward and forward-backward SkyRL-Train paths.
• Adds a public real-service repro script in examples/.
• Adds focused unit tests for default behavior, chunking, ordered processing, repro shape, and cap-aware padding.

Validation

uv run --active --no-sync ruff format examples/repro_forward_backward_queue_drain.py skyrl/backends/microbatch_padding.py skyrl/backends/skyrl_train_backend.py tests/tinker/test_forward_backward_queue_repro.py tests/tinker/skyrl_train/test_backend_microbatch_padding.py
uv run --active --no-sync ruff check examples/repro_forward_backward_queue_drain.py skyrl/backends/microbatch_padding.py skyrl/backends/skyrl_train_backend.py tests/tinker/test_forward_backward_queue_repro.py tests/tinker/skyrl_train/test_backend_microbatch_padding.py
uv run --isolated --no-project --with pytest python -m pytest tests/tinker/test_forward_backward_queue_repro.py -q
uv run --isolated --no-project --with pytest --with pydantic python -m pytest tests/tinker/skyrl_train/test_backend_microbatch_padding.py -q

Focused test results: 5 passed for the queue repro tests and 3 passed for the padding tests.

Real endpoint validation:

Before/no cap: process_batch_requests(forward_backward, n=16) -> CUDA OOM, attempted 9.08 GiB allocation.
Cap only: forward_backward request batch split 16, but Padded batch from 1 to 16 -> CUDA OOM on the 35,000-token request.
Current PR: forward_backward request batch split 16; all sixteen n=1 chunks complete; request 9 returns num_tokens:sum=35000.0; client exits successfully.

[gemini-code-assist]

Code Review

This pull request introduces a new configuration option forward_backward_max_request_count to limit the number of pending forward_backward requests coalesced into a single backend call, preventing potential CUDA OOM issues. It includes a synthetic reproduction script, engine changes to chunk the requests, and corresponding unit tests. The reviewer suggests adding Pydantic validation (gt=0) to the new configuration field to prevent non-positive values from being silently accepted.

[gemini-code-assist]

To prevent silent configuration errors where a user might accidentally configure a non-positive value (e.g., 0 or a negative integer) and experience unexpected CUDA OOMs due to the engine falling back to unlimited batching, we should enforce that forward_backward_max_request_count is strictly greater than zero using Pydantic's gt=0 validation.

    forward_backward_max_request_count: int | None = Field(
        default=None,
        gt=0,
        description=(
            "Optional cap on how many pending forward_backward requests the engine "
            "coalesces into one backend call. Default 'None' preserves the current "
            "unlimited batching behavior."
        ),
        json_schema_extra={"argparse_type": lambda v: None if v == "None" else int(v)},
    )