fix(rollout): validate aligned prompt-group chunk boundaries

[0z5a]

Incremental review: Files changed against main.

Summary

Replace whole-group shape bucketing with an optional validate_group_boundaries check (default false). GRPO siblings can run in separate forwards: results are regrouped before advantages are computed. The check accepts whole groups or aligned equal subgroups and never changes forward_batch_size, row order, or memory limits.

For example, group size 16 with FBS 4 keeps four 4-row forwards. Unsupported heterogeneous-shape scheduling is deferred until per-row geometry and ragged trajectories exist.

Related Issue

Refs #483. This is boundary validation only, not completion of shape bucketing or continuous scheduling.

Test Plan

RTX 5090, torch 2.14.0+cu130. External harness: PYTHONPATH=. python ../validate_unirl.py, using the real engine, Sample/Part, trajectory and GRPO advantage code with a deterministic tiny BF16 linear module; no checkpoint or dataset downloads.

Check	Result
CPU/GPU validation on/off parity	14 cases passed; identical calls, row/noise IDs, latents, log-probs, rewards and advantages
Invalid boundaries	3 helper cases rejected; engine rejects before forward and restores training mode
Built-in trainside recipes	61/61 boundary checks passed
Group/FBS 16/4, 16/8, 16/1, 24/1	All accepted with the original chunk sizes

ruff check unirl/rollout/engine/trainside/, ruff format --check unirl/rollout/engine/trainside/, python lint/check_docstring_lines.py, and python lint/check_core_dependencies.py pass.

Full pretrained training was not run; this change only validates existing chunk boundaries. No throughput or reward-improvement claim.

Compatibility / Risk

Existing recipes retain their exact schedule by default. The earlier proposed shape-bucket options are removed from this unmerged PR; the only new option is validate_group_boundaries.

Reviewer Notes

Checked open overlapping work, including #445 and #465; this adds no reward-overlap or replay-backward machinery. Heterogeneous execution scheduling remains separate.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated documentation and validated the change with a repo-external harness.

[CjhHa1]

Thanks for working on this. Before discussing the implementation details, I am not yet convinced that keeping every prompt group in one model forward is an algorithmic requirement.

I checked the 61 built-in trainside recipes at this PR's merge base:

• 48 split one prompt group across multiple pipeline.generate calls.
• 13 avoid splitting through aligned/unset forward_batch_size or group size 1.
• The existing split configurations are deliberate aligned sub-chunks (group_size % forward_batch_size == 0); I found no built-in recipe with an accidental non-divisor boundary.

The split does not currently break GRPO grouping:

• chunks are concatenated in original sample_ids order before reward scoring;
• group advantages are computed on the complete reconstructed frontier;
• driver-authored x_T is keyed by sample/group identity, so init_same_noise remains consistent across calls;
• rollout trajectories and log-probs remain row-aligned.

Changing chunk geometry can change global SDE RNG assignment and BF16 numerics, but that is a reproducibility/batch-geometry effect, not by itself evidence that GRPO siblings must share one forward.

Conversely, many recipes intentionally use a smaller forward batch for memory:

• WAN21: group 16, FBS 4
• Z-Image: group 16, FBS 8
• WAN22: group 16, FBS 1
• VideoAlign: group 24, FBS 1

shape_bucket rejects these unless one whole group fits the declared capacity, so it cannot replace most current trainside configurations without increasing peak memory, potentially to OOM.

Could you first establish the algorithmic motivation with a concrete invariant or experiment?

1. What becomes incorrect when siblings are generated in separate forwards?
2. Is there a measured reward/optimization regression, rather than only small batch-shape numerical differences?
3. Which real recipe can enable this mode while fitting a whole group in memory?
4. What throughput or tail-latency improvement does that recipe obtain?

Also, heterogeneous shape bucketing is not reachable today: Part.sampling_params and init_noise_latent_shape are shared, _execution_keys() repeats one key for every row, and dense trajectories cannot merge different geometries. No existing recipe enables this mode.

If there is no demonstrated same-forward algorithm invariant, I would prefer a much smaller change: validate or plan group-aware boundaries while preserving memory-safe aligned sub-group execution. The full execution-key/shape-bucket machinery should wait until per-row geometry and ragged trajectory support exist.

[0z5a]

Agreed. Replaced the shape-bucket planner with a small optional boundary check, preserving aligned subgroups and the existing forward-size cap. No claim that GRPO siblings must share a forward; heterogeneous scheduling is deferred.

Validation	Result
CPU/GPU output and advantage parity	14 cases passed
Built-in recipe boundary checks	61/61 passed
16/4, 16/8, 16/1, 24/1	Accepted, unchanged chunk sizes
Lint and core guards	Passed

These use the real engine with a tiny deterministic BF16 module; full pretrained training and throughput were not measured.