[skyrl-train] Refactor training loop structure to explicitly batch at two levels (minibatch -> microbatch)

[justinvyu]

Summary

This PR refactors the training loop structure in both PolicyWorkerBase and CriticWorkerBase to use a consistent two-level batching strategy (minibatch → microbatch) with a microbatch_weight used for gradient accumulation.

Refactored ppo_train to use a two-level loop: iterate over minibatches, then subdivide each minibatch into microbatches

• Changed forward_backward signature from accumulation_steps: int to microbatch_weight: float
• Loss is now scaled by microbatch_weight (i.e., 1.0 / num_microbatches) instead of dividing by accumulation_steps
• Optimizer step is now called once per minibatch (after all microbatches are processed), rather than conditionally based on step count

Motivation

The motivation for this PR was for batch balancing which I tried adding in this PR: #640

The problem was that minibatch boundaries are not explicitly defined right now. We have 2 configurations policy_mini_batch_size and micro_train_batch_size_per_gpu, and the minibatch is implicitly constructed by doing gradient accumulation over (policy_mini_batch_size // micro_train_batch_size_per_gpu) identical microbatches. #640 breaks the "same sized microbatches" assumption, which is why we need to partition explicitly at the minibatch level first, before chunking into possibly uneven microbatches. That way, it's simpler to tell when to stop accumulating the gradient.

Also, the introduction of a more general microbatch_weight is also motivated by dynamic microbatch sizes introduced by #640. Each microbatch should contribute N_i / sum(N_j) * loss_i to the accumulated gradient. In the default case, this is just 1/accumulation_steps.

Testing

• ✅ FSDP metrics are equivalent to master
• ✅ Megatron codepath still works

[gemini-code-assist]

Code Review

This pull request refactors the training loop to use a two-level batching strategy (minibatch → microbatch), which is a good architectural improvement for supporting uneven microbatches in the future. The core logic changes are sound, but I've identified a few critical issues. The most significant one is that in both policy and critic training loops, metrics are not being aggregated correctly across microbatches, leading to inaccurate logging. Only the status from the last microbatch of a minibatch is being recorded. Additionally, the optimizer step has been unintentionally removed from the critic's training_step method, which will affect tests relying on it. I've provided detailed comments and suggestions to address these points.

[justinvyu]

The next step is basically just to change this microbatch iterator from one that's doing sample-based chunking to one that's token-based.

[erictang000]

this change also affects the use of BatchIterator for the megatron backend, which implements ppo_train differently FSDP/Deepspeed.

SkyRL/skyrl-train/skyrl_train/workers/megatron/megatron_worker.py

Line 514 in 3725635

dataloader = BatchIterator(

could you make sure the conversion to experience is also handled correctly for the megatron code path? Making sure one of these tests pass:

SkyRL/skyrl-train/tests/gpu/test_megatron_worker.py

Line 440 in 3725635

async def test_megatron_train(

is probably a good way to check this.

[justinvyu]

Sounds good. I added a followup TODO to update the megatron worker's ppo loop as well. Made a minimal change for now to prevent this PR from getting too large.

[erictang000]

left a comment about megatron + i think the gemini comments are worth a look - essentially we want to make sure that the metrics are aggregated in the same way before and after this PR. Maybe it would be nice to show with running a test before and after this PR (maybe this guy:

SkyRL/skyrl-train/tests/gpu/gpu_ci/test_ppo_train.py

Line 33 in 3725635

def test_ppo_train_basic_execution(ray_init_fixture, cfg, use_entropy_loss, use_kl_loss):

) that this PR doesn't change the metrics

[justinvyu]

To sanity check that this PR didn't introduce any regressions for metrics, I printed the output status from this test: pytest tests/gpu/gpu_ci/test_ppo_train.py::test_gradient_accumulation_scenarios[accumulation_calculation] -s

Printed status on master:

train_status={'final_loss': 0.0018931262311525643, 'policy_loss': -0.0024274957249872386, 'ppo_clip_ratio': 0.0, 'policy_entropy': 5.9921875, 'policy_kl': 10.0, 'policy_lr': 9.999999974752427e-07, 'raw_grad_norm': 0.33137789368629456, 'policy_update_steps': 1.0}, actual_optimizer_steps=1.0

Printed status with this PR:

train_status={'final_loss': 0.0018931262311525643, 'policy_loss': -0.0024274957249872386, 'ppo_clip_ratio': 0.0, 'policy_entropy': 5.9921875, 'policy_kl': 10.0, 'policy_lr': 9.999999974752427e-07, 'raw_grad_norm': 0.33137789368629456, 'policy_update_steps': 1}, actual_optimizer_steps=1

[justinvyu]

Tested Megatron code path as well:
uv run --isolated --extra dev --extra mcore -- pytest tests/gpu/gpu_ci/test_megatron_worker.py::test_megatron_train[tp2_pp2_policy_seq_packing]

megatron results:  {'final_loss': -0.021974159637466073, 'policy_loss': -0.030495601160509977, 'policy_lr': 9.999999974752427e-07, 'ppo_clip_ratio': 0.0, 'policy_entropy': 8.893153190612793, 'policy_kl': 8.521441221237183, 'raw_grad_norm': 4.0175676345825195, 'policy_update_steps': 1}

fsdp results:  {'final_loss': -0.021487861638888717, 'policy_loss': -0.030021082626679796, 'ppo_clip_ratio': 0.0, 'policy_entropy': 0.5429687462747097, 'policy_kl': 8.533220887184143, 'raw_grad_norm': 3.5146021320670116, 'policy_lr': 9.999999974752427e-07, 'policy_update_steps': 4}

[erictang000]

/gemini review

[gemini-code-assist]

Code Review

This PR refactors the training loops in PolicyWorkerBase and CriticWorkerBase to use a two-level batching structure (minibatch -> microbatch), which is a great improvement for supporting batch balancing and dynamic microbatch sizes. The changes are well-motivated and mostly well-executed.

My feedback focuses on a few key areas:

• An inconsistency in megatron_worker.py where the new training loop structure has not been applied.
• A potential bug in the calculation of critic_update_steps.
• Several opportunities for minor refactoring to improve code clarity and reduce duplication, such as extracting helper methods for status recording and memory snapshotting.
• Identifying potentially unused code marked with TODO comments that should be cleaned up.

Overall, this is a solid refactoring. Addressing these points will improve the consistency and maintainability of the codebase.

[erictang000]

looking good! I think just one super minor gemini comment, and we need to fix the failing cpu test and we should be good to merge this

cpu test:

=========================== short test summary info ============================
FAILED tests/cpu/test_trainer.py::test_ppo_train_batch_calculations - TypeError: test_ppo_train_batch_calculations.<locals>.mock_policy_forward_backward() got an unexpected keyword argument 'microbatch_weight'
============ 1 failed, 353 passed, 83 warnings in 101.86s (0:01:41) ============

[erictang000]

good catch...

[erictang000]

🚀🚀🚀