Add a configuration for max tokens per microbatch

[justinvyu]

Summary

Introduces token-based binpacking for microbatches, enforcing a max_tokens_per_microbatch limit on the total number of real tokens in a microbatch.

Replaces the even sample-based chunking BatchIterator with a BalancedBatchIterator that respects the microbatch token limits.

Problem

Previously, microbatches created by chunking sequences evenly so that each microbatch was a certain batch size (micro_train_batch_size_per_gpu), which could lead to:

• Sequence packed microbatches can exceed GPU memory limits when sequences vary in length, so you need to conservatively set the micro batch size.
• Inefficient GPU utilization due to uneven token distribution across microbatches

Solution

Introduce balanced_binpacking and BalancedBatchIterator utilities that:

• Enforce max_tokens_per_microbatch — no microbatch exceeds the token limit
• Roughly balance the microbatches so that we avoid straggler microbatches.
• Ensure that every worker still receives the same number of microbatches to satisfy Deepspeed/FSDP requirements (every DP worker must do the same number of forward passes).

API changes

• Adds a BalancedBatchIterator that follows the BatchIterator interface, which is just an iterator over Experiences.
• Unifies Worker.forward to also use the BatchIterator interface, rather than call TrainingBatch.chunk directly.
• Replaces the BatchIterator with BalancedBatchIterator when the max_tokens_per_microbatch is set, for the forward, Critic.ppo_train and Policy.ppo_train.

Example

Example ppo_train call from the unit test:

minibatch_sequences=tensor([[88, 96,  3, 23, 44, 65, 55,  0, 66, 49],
        [25, 85, 10, 86, 31, 55, 31, 51,  4, 56],
        [28, 38, 76,  8,  0,  0,  0,  0,  0,  0],
        [21, 97, 46, 75, 73,  0,  0,  0,  0,  0]])
minibatch_attention_mask=tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0]])
Policy Train epoch [1/1]:   0%|          | 0/2 [00:00<?, ?it/s]
(DeepSpeedPolicyWorkerBase pid=182679) micro_batches_per_mini_batch=2
(DeepSpeedPolicyWorkerBase pid=182679) experience.sequences=tensor([[88, 96,  3, 23, 44, 65, 55,  0, 66, 49],
(DeepSpeedPolicyWorkerBase pid=182679)         [28, 38, 76,  8,  0,  0,  0,  0,  0,  0]])
Policy Train epoch [1/1]:  50%|█████     | 1/2 [00:01<00:01,  1.28s/it, pg=-0.203, glen=4, policy_lr=1e-6, ent=3.08]
(DeepSpeedPolicyWorkerBase pid=182679) experience.sequences=tensor([[25, 85, 10, 86, 31, 55, 31, 51,  4, 56],
(DeepSpeedPolicyWorkerBase pid=182679)         [21, 97, 46, 75, 73,  0,  0,  0,  0,  0]])

Follow-up

This only does balanced microbatching at the individual worker level. We can further balance the minibatches sent to each worker at a global level to ensure every DP worker has roughly the same amount of work.

There's also different algorithms that we could plug into the microbatch chunking step.

[justinvyu]

TODO: test a multi-worker case where you have padding microbatches

[SumanthRH]

Suggested change

	from skyrl_train.workers.worker_utils import BatchIterator, TokenBasedBatchIterator, reduce_metrics
	from skyrl_train.workers.worker_utils import BaseBatchIterator, SampleBasedBatchIterator, TokenBasedBatchIterator, reduce_metrics

[SumanthRH]

QQ: Should the dynamic batch size related packing happen at the batch level or the mini-batch level?

The most efficient one would be at the batch level, but then we won't have a good way to get mini batches afterwards. It seems like we should just do this at the mini batch level first. (you'd need to introduce an explicit loop over mini batches)

[SumanthRH]

This also allows us to naturally define the accumulation steps

[justinvyu]

Should the dynamic batch size related packing happen at the batch level or the mini-batch level?

Yeah it should also be done at the global batch level, but I'm leaving that for a follow-up optimization, see the end of the PR description. The global batch level batch balancing is useful to assign each worker a similar number of tokens.

This PR's main goal is the max tokens per microbatch parameter. The balancing of microbatches is kind of a convenient feature add-on that is a slight optimization.

Here's why the balancing at the microbatch level is useful:

• Worker 0 mini-batch contains 4 sequences of length: [10, 10, 5, 5]
• Worker 0 has a max microbatch token length of 15.
• If you just do a single front-to-back pass, you'll end up with 3 microbatches: [10], [10, 5], [5]
• However, the better way to chunk this would be [10, 5], [10, 5], which is what the balanced microbatching method does here.
• Basically, we want to minimize number of microbatches while satisfying max_tokens_per_microbatch.

[justinvyu]

This also allows us to naturally define the accumulation steps

Could you elaborate on this a bit more?

[SumanthRH]

Oh sorry here's what I meant;

We have "mini batches" per worker that is set based on policy_mini_batch_size. This is separate from the per-worker shard you're referring to.

train_batch_size / policy_mini_batch_size dictates the number of optimization steps with the current training batch.

So for example, with train_batch_size=16, num_workers=8, policy_mini_batch_size=8 , each worker will receive a shard of 2 samples, and this wil be split further into train_batch_size/ policy_mini_batch_size=2 mini batches. Each mini batch will correspond to one optimization step.

For each mini batch per worker, we further split it into micro batches.

My point was that the dynamic batch size related binpacking should happen at the worker-level mini batches not the worker-level overall batch.

[justinvyu]

Closing this PR since it's a bit outdated. Will redo this now that #817 has landed.