fix(train): reject a train_batch_size that truncates the prompt batch to zero

[pranavraja99]

Closes #1609.

RayPPOTrainer._remove_tail_data keeps the largest multiple of stride = lcm_dp_size // gcd(lcm_dp_size, n_samples_per_prompt) prompts so that the resulting samples shard evenly across the data-parallel sizes of the enabled models. When train_batch_size is below that stride, (len(entries) // stride) * stride is 0: the generator is called with zero prompts and the run dies several frames later in get_rollout_metrics on np.min of an empty array (ValueError: zero-size array to reduction operation minimum which has no identity), with nothing pointing at the batch size.

validate_batch_sizes did check a related condition, train_batch_size * n_samples_per_prompt >= lcm_dp_size, but that is necessary rather than sufficient. For the config in the issue (lcm_dp_size=6, n_samples_per_prompt=4, train_batch_size=2) it passes — 8 >= 6 — while the stride is 3, so the batch still truncates to zero.

Since tbs * n % lcm == 0 is equivalent to tbs % (lcm // gcd(lcm, n)) == 0, the stride is the quantity the truncation actually depends on, so this checks it directly:

• Assert train_batch_size >= prompt_stride at startup. This is exactly the condition under which at least one prompt survives truncation, and it subsumes the previous assert. The message reports policy_dp_size / ref_dp_size / lcm_dp_size / n_samples_per_prompt and the stride to set train_batch_size to.
• Warn when train_batch_size is at or above the stride but not a multiple of it. That case is not fatal, but it silently drops the same prompts from every training batch, so the effective batch size is smaller than configured.
• Raise in _remove_tail_data as well. validate_batch_sizes only sees train_batch_size, while the truncation also runs on batches sized by the caller (examples/train/async/async_trainer.py and skyrl-agent's trainer), so the guard has to sit where the stride is applied.

Only the fatal case is a hard failure; the lossy-but-nonempty case stays a warning so that existing configs keep running.

Testing

• test_remove_tail_data_rejects_batch_smaller_than_prompt_stride covers the raise and that a batch at or above the stride still truncates to the largest multiple.
• Test cases 17 and 18 in test_validate_batch_sizes cover the startup assert (using the placement from the issue) and the non-multiple warning.
• uv run --isolated --extra skyrl-train --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm" — 1451 passed, 7 skipped.

The pre-existing test_validate_batch_sizes_lcm_dp_requirement passes unchanged: it uses n_samples_per_prompt=1, where the stride equals lcm_dp_size and the new assert reduces to the old one.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request improves batch size validation and handling in SkyRL. It introduces a check to ensure that the prompt batch size is not smaller than the prompt stride, preventing all prompts from being dropped during truncation. It updates validate_batch_sizes to assert this constraint and warn when the batch size is not a multiple of the stride, and adds corresponding unit tests. The reviewer suggested refining the check in _remove_tail_data to ensure that a ValueError is not misleadingly raised when the input batch is already empty.

[gemini-code-assist]

If entries is empty (i.e., len(entries) == 0), kept_prompts will be 0. In this case, raising a ValueError stating that the batch is smaller than the prompt stride is misleading because the batch was already empty, not truncated to zero. Adding a check for len(entries) > 0 ensures we only raise this error when actual prompts are being truncated to zero.

Suggested change

	if kept_prompts == 0:
	if len(entries) > 0 and kept_prompts == 0: