Repository navigation
Conversation
With --train-iters and --step-batch-size-schedule, get_optimizer_param_scheduler sized every LR and weight decay phase as iterations times cfg.train.global_batch_size. Under a schedule that value is micro_batch_size * batch_parallel_size, while every step advances the scheduler by the full scheduled global batch size. Warmup, cosine and WSD decay, and the weight decay ramp all ended far too early. With batch 8 then 16 (mbs 1, dp 1) and 100 iterations, warmup ended at iteration 2 (configured: 10) and the LR reached min_lr at iteration 13. Convert the iteration counts to samples by walking the batch size schedule, as update_train_iters already does for sample-based training. The WSD anneal length is the samples consumed over the last lr_wsd_decay_iters iterations of the decay, and lr_warmup_fraction applies to lr_decay_iters. A one-entry schedule gives the same horizons as the constant batch size formula, and the horizon no longer depends on the data parallel size. The constant batch size path is unchanged. Runs that already combine --train-iters with a schedule stored the old horizons and stop on resume; --override-opt-param-scheduler adopts the corrected ones. Also drop --global-batch-size from examples/gpt3/train_gpt3_175b_distributed.sh. validate_args rejects it next to --step-batch-size-schedule, so the example failed at startup. Without it the eval batch falls back to micro_batch_size * data_parallel_size, so pin --eval-global-batch-size 1536, as NVIDIA#5022 did for the GPT 15B GB200 release configs. With this fix the example (DP 1) warms up for 430 iterations and reaches min_lr at iteration 430000. Without the fix, and with --global-batch-size removed, those were iterations 27 and 26875. Signed-off-by: Elia Liu <Elialiu760317@outlook.com>
Signed-off-by: Elia Liu <Elialiu760317@outlook.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fix LR warmup, cosine/WSD decay and weight decay horizons when
--train-itersis combined with--step-batch-size-schedule. In a 100-iteration CPU reproduction with batches of 8 then 16, main finishes warmup at iteration 2 and decay at iteration 13; the configured boundaries are 10 and 100. This PR reaches the configured boundaries.Issue tracking
The regression follows the removal of the iteration-mode batch-ramp guard in #4411. The same PR changed the shipped GPT-3 175B example to specify both a schedule and
--global-batch-size, which argument validation rejects. No separate issue has been filed.Changes
PretrainConfigContainer._resolve_scheduler_step_fields. WSD covers the final configured anneal iterations. A fractional warmup iteration uses the next iteration's batch.Merged current main at b291235 and resolved the conflict caused by upstream moving scheduler derivation into configuration finalization. The fix now lives in
megatron/training/config/container.py.Validation
CPU, Python 3.12 and torch 2.14.1+cpu, on the merged head:
The original negative control produced 15 failures for scheduled-batch cases and 7 passing constant-batch cases. Ruff and isort passed after adapting the change to the upstream configuration refactor. GPU, functional training, and the full 175B example have not been run.
Compatibility and review questions
--override-opt-param-schedulerto adopt corrected values.--use-checkpoint-opt-param-schedulerretains their previous horizons.train_iters * global_batch_size; correcting that can change dataset indices and sample order and needs a separate change.Contribution process
AI assistance was used for analysis, implementation, regression tests, and PR preparation. The configuration-refactor adaptation and its test reruns were completed after the author's review.