[fsdp] Trim dense padding per final microbatch

[EazyReal]

Summary

Dense FSDP forwards currently inherit the sequence width chosen when the controller rectangularizes the whole rollout batch. Final sample- or token-based microbatch membership is decided later in the worker, so a short microbatch can still execute at the longest sequence and response widths of an unrelated sample elsewhere in the global batch.

This PR projects each finalized dense-FSDP microbatch to its own local maxima immediately before model execution:

• sequence-aligned tensors are sliced to the longest real sequence in that microbatch;
• response-aligned tensors are sliced to its longest real response;
• policy, reference, and critic outputs are left-restored to the controller's canonical response width before microbatch outputs are recombined.

The controller batch remains the single full-width representation used by advantage calculation, dispatch, and replay. Packed inputs, intentional DP-equalization padding microbatches, SFT batches without an explicit response boundary, remove_microbatch_padding=True, and all Megatron paths remain unchanged.

Why this boundary

The worker iterator is the first place that knows the final sample/token bin membership. Re-collating earlier in the controller cannot choose the correct per-rank token bins, while trimming inside the model would mix scheduling metadata with model semantics. The existing left-padding contract makes the projection and output restoration unambiguous: real sequence and response values are right-aligned, and semantic zeros in loss_mask are not used to infer response length.

Performance trade-off

Local trimming reduces dense token work and memory but increases sequence- and response-shape variability across final microbatches, which may reduce kernel or compilation-cache reuse. This PR has no SkyRL GPU evidence of repeated execution-plan builds, so it does not assume the net throughput direction or add a user-facing mode for a theoretical cost. A fixed-width opt-out should be earned by an observed SkyRL workload; the current PR claims only the verified layout reduction.

Verification

• Focused sample/token microbatch and output-restoration regressions: 25 passed
• Related batch, VLM-plumbing, worker-utils, and preprocessing CPU suite: 68 passed, 1 skipped
• Megatron packing/THD regression suite: 14 passed
• Broad CPU train/backend suite: 1,467 passed, 9 skipped
• Changed-file Ruff, Black, gitleaks, and git diff --check: passed

Two optional-test failures require unavailable vllm, and two collection errors require unavailable torchvision; they are environment setup gaps rather than failures in the changed paths. I did not run distributed FSDP/VLM GPU parity or claim measured peak-memory/throughput improvements. The CPU tests establish the tensor-layout and output-contract invariants.

An independent exact-head review found no merge-blocking issues and additionally exercised all schema fields, VLM TensorList selection, token/sample batching, restoration/reordering, and packed/dummy no-op behavior in a CPU adversarial probe.

Scope

This addresses premature dense-FSDP padding only. SkyRL's Megatron packing keeps token padding inside a fixed [1, T] packed transport and does not exhibit the separate dummy-BSHD-row issue.

---

Note

Medium Risk
Changes FSDP dense forward tensor layouts and restores logprob/value widths, which can affect memory, packing, and output contracts if trim/restore is wrong. Scope is gated off packed/Megatron paths and covered by CPU layout tests, not GPU parity.

Overview
Dense FSDP microbatches no longer inherit the controller's global sequence/response widths. After sample- or token-based binning, each microbatch is sliced to its own longest real sequence and response so short bins do not run at an unrelated sample's pad width.

Policy, critic, and ref inference outputs are left-padded back to the controller response width before recombination, so advantage/dispatch/replay still see one rectangular batch. Trimming is gated to FSDP with remove_microbatch_padding=False; packed rows, DP dummy microbatches, SFT without response_mask, and Megatron are unchanged. Forward passes now also receive response_mask so the response boundary can be recovered.

^{Reviewed by Cursor Bugbot for commit 0253d94. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request implements microbatch padding trimming to optimize memory usage during FSDP training. It introduces new utility functions, trim_microbatch_padding and restore_microbatch_response_padding, and updates the batch iterators and worker methods to support this functionality. The review feedback highlights the need for additional defensive checks: one to handle empty batches in trim_microbatch_padding to avoid potential runtime errors, and another to validate tensor dimensions in restore_microbatch_response_padding to prevent silent bugs with higher-dimensional tensors.

[EazyReal]

Hosted CI note: all CPU/code-quality suites passed, including the 58-minute repository-wide skyrl_tests job. The GPU workflow exited before submitting tests because the fork run had no ANYSCALE_CLI_TOKEN; Vercel likewise requires deployment authorization. Those two red statuses need maintainer-side credentials/authorization rather than a code change in this PR.