alibaba / alibaba/ROLL

LR scheduler progress can be inconsistent with dynamic batching in Megatron actor training

Open
#442 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.4k
Forks
312
Avg merge
1h 2m
Merged PRs (30d)
2

Description

## Summary

There seems to be a potential mismatch between LR scheduler budgeting and LR scheduler consumption in the Megatron actor training path when dynamic batching is enabled.

The scheduler budget appears to be computed from nominal/static batch semantics, for example based on `rollout_batch_size`, `ppo_epochs`, `per_device_train_batch_size`, and `gradient_accumulation_steps`.

However, with `use_dynamic_batching_in_train: true`, the actor train batch is first split into token-bounded dynamic micro-batches. These micro-batches are then grouped by `gradient_accumulation_steps`, so the number of scheduler-consuming optimizer updates can scale roughly as `ceil(num_dynamic_micro_batches / gradient_accumulation_steps)`, rather than being determined only by the nominal sample-level batch formula.

As a result, changing dynamic batching behavior may change how fast the LR scheduler consumes its step budget, even when the intended logical training batch semantics are unchanged.

## Potential issue

The code seems to mix two notions of training progress:

- scheduler budget: nominal/static sample-level batch progress
- scheduler consumption: actual optimizer update count after token-level dynamic batching and gradient-accumulation grouping

This can cause the LR schedule to progress faster than intended when dynamic batching produces more actual optimizer updates than the nominal static-batch formula expects.

In one long-context agentic RL run, I observed the saved scheduler state progressing far beyond the configured decay horizon: `lr_decay_steps = 2400`, while `num_steps = 8429`.

This example is only meant as a symptom. The underlying concern is more general: it may be unexpected for LR scheduler progress to depend on token-level dynamic batching decisions if the logical training batch configuration is unchanged.

## Expected behavior

Changing `max_tokens_per_microbatch_in_train` should affect memory usage and internal batching, but it should not silently accelerate LR decay relative to the nominal training budget.

The LR scheduler should either explicitly follow the actual optimizer update count after dynamic batching, or follow the nominal/static training progress implied by the user configuration.

If the second behavior is intended, then scheduler progress may need to be normalized when dynamic batching changes the actual number of optimizer updates.

## Possible direction

One possible fix direction is to compute, for each actor train batch:

- `nominal_steps`: expected optimizer update count under static batch semantics
- `actual_steps`: actual scheduler-consuming optimizer update count after dynamic batching
- `scheduler_step_increment = nominal_steps / actual_steps`

Then the scheduler can advance by normalized progress rather than one full step per actual optimizer update.

The exact implementation does not have to use this API literally. The main point is to decouple LR scheduler progress from token-level dynamic batching when the logical training progress is unchanged.

## Questions

Could you clarify the intended LR scheduler semantics for Megatron actor training with dynamic batching?

1. Should the scheduler follow the actual optimizer update count after dynamic batching?
2. Or should it follow the nominal/static training progress from the configured batch semantics?
3. If the latter is intended, would normalizing scheduler progress under dynamic batching be an acceptable fix direction?

Contributor guide

No contributing guide indexed for this repository

Research direction

The issue names no files or tests. Trace the Megatron actor-training scheduler budget and the dynamic-batching optimizer-update path, compare nominal and actual step counts, and establish which progress semantics are intended. Done means the chosen behavior is implemented consistently and covered by a regression test for dynamic batching.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.