fix: change default batch sizes for dp workloads - #528
Conversation
binaryaaron
changed the title
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
📜 Recent review details⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)
WalkthroughUpdated DP training batch sizes and added/retained training.gradient_accumulation_steps=16 across SLURM job configs and corresponding E2E required test fixtures (Mistral: batch_size 8→2; TinyLlama & SmolLM3: batch_size 8→4). ChangesDP Training Batch Size Adjustments
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Suggested labels
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Greptile SummaryThis PR reduces the physical microbatch size for all three DP SLURM configs (Mistral-7B, SmolLM3-3B, TinyLlama) and adds explicit
Confidence Score: 4/5Safe to merge for smollm3 and tinyllama; the Mistral config silently halves the effective batch size from 64 to 32, contradicting the PR's stated invariant. Two of the three model configs correctly preserve effective batch size at 64. The Mistral config uses batch_size: 2 instead of 4, yielding 2x16=32, not 64 as stated. This affects both the SLURM run config and the e2e required config and needs resolution before merging. script/slurm/configs/mistral-dp.yaml and tests/e2e/required_configs/mistral-dp.yaml — both use batch_size: 2, giving an effective batch size of 32 instead of 64. Important Files Changed
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 193f104c-9a00-4f3a-a58a-0bbc68157478
📒 Files selected for processing (6)
script/slurm/configs/mistral-dp.yamlscript/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yaml
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)
- GitHub Check: Smoke Tests
- GitHub Check: Unit Tests (3.11)
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Analyze (Python)
- GitHub Check: Typecheck
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
🧰 Additional context used
📓 Path-based instructions (6)
**/*.yaml
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.yaml: Use 2-space indentation in YAML files.
Use colon-space (:) for key-value pairs in YAML.
Include SPDX copyright headers at the top of YAML files.
Use unquoted values in YAML unless special characters require them.
Include a newline at the end of YAML files.
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
⚙️ CodeRabbit configuration file
Review YAML for 2-space indentation, SPDX headers when required, unquoted values unless needed, and newline at EOF.
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
**/*.{py,sh,yaml,yml}
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Include SPDX copyright header at the top:
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.and# SPDX-License-Identifier: Apache-2.0. Themake formatcommand handles this automatically.
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced bypre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured inruff.toml).
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
⚙️ CodeRabbit configuration file
**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.
- Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.- Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.
- Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.- Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.- Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.- Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
**/*.{yaml,yml}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All YAML files require SPDX copyright headers, automatically added by
make format
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yamltests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
script/**
⚙️ CodeRabbit configuration file
Review standalone scripts for reproducibility and operational safety. Check argument validation, quoting, repo-root detection, environment variables, generated artifacts, external commands, GPU/cluster assumptions, and whether the script should be wired through Makefile or documented in README/docs.
Files:
script/slurm/configs/smollm3-dp.yamlscript/slurm/configs/tinyllama-dp.yamlscript/slurm/configs/mistral-dp.yaml
tests/**
📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)
Tests should mirror the
src/directory structure intests/
Files:
tests/e2e/required_configs/smollm3-dp.yamltests/e2e/required_configs/tinyllama-dp.yamltests/e2e/required_configs/mistral-dp.yaml
🔇 Additional comments (3)
script/slurm/configs/mistral-dp.yaml (1)
5-5: LGTM!script/slurm/configs/smollm3-dp.yaml (1)
5-5: LGTM!script/slurm/configs/tinyllama-dp.yaml (1)
5-5: LGTM!
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
e53acab to
54b6543
Compare
coderabbitai
Bot
added
bug
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
kendrickb-nvidia
left a comment
kendrickb-nvidia
left a comment
There was a problem hiding this comment.
Approval pending positive outcome from slurm runs
Apply the reduced DP microbatch shape to SmolLM3 and remove the stale TinyLlama comment so Slurm configs consistently use 4 x 16 batching. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Signed-off-by: Matthew Kornfield <mkornfield@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 053b75c7-101c-49e5-b1fd-7f940040c04d
📒 Files selected for processing (1)
script/slurm/configs/mistral-dp.yaml
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)
- GitHub Check: Smoke Tests
- GitHub Check: Unit Tests (3.12)
- GitHub Check: Unit Tests (3.11)
- GitHub Check: Unit Tests (3.13)
- GitHub Check: Typecheck
- GitHub Check: Analyze (Python)
- GitHub Check: Analyze (python)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.yaml
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*.yaml: Use 2-space indentation in YAML files.
Use colon-space (:) for key-value pairs in YAML.
Include SPDX copyright headers at the top of YAML files.
Use unquoted values in YAML unless special characters require them.
Include a newline at the end of YAML files.
Files:
script/slurm/configs/mistral-dp.yaml
⚙️ CodeRabbit configuration file
Review YAML for 2-space indentation, SPDX headers when required, unquoted values unless needed, and newline at EOF.
Files:
script/slurm/configs/mistral-dp.yaml
**/*.{py,sh,yaml,yml}
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
Include SPDX copyright header at the top:
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.and# SPDX-License-Identifier: Apache-2.0. Themake formatcommand handles this automatically.
Files:
script/slurm/configs/mistral-dp.yaml
**/*
📄 CodeRabbit inference engine (STYLE_GUIDE.md)
**/*: Include a newline at the end of all files, never trailing whitespace. This is enforced bypre-commit.
Use line length of 120 characters for code, comments, and docstrings (configured inruff.toml).
Files:
script/slurm/configs/mistral-dp.yaml
⚙️ CodeRabbit configuration file
**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.
- Refactor suggestion: use for local maintainability problems introduced
by the diff when they have clear future cost, such as duplicated setup,
unclear boundaries, over-mocking, avoidable complexity, or opaque test
helpers.- Nitpick: avoid in chill mode. Do not emit formatting, import-order,
wording, or style-only comments unless automated tools cannot catch the
issue and it affects maintainability.Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.
- Major: incorrect generation/training/evaluation behavior, broken
CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
cleanup and process-isolation bugs likely to fail CI or production
runs.- Minor: localized bugs, missing focused tests for changed behavior, or
bad test patterns that weaken regression coverage.- Trivial: small cleanup with no behavior impact. Usually suppress in
chill mode.- Info: context only. Avoid unless it helps reviewers understand risk.
Safe-Synthesizer-specific review focus: - Data ...
Files:
script/slurm/configs/mistral-dp.yaml
**/*.{yaml,yml}
📄 CodeRabbit inference engine (CONTRIBUTING.md)
All YAML files require SPDX copyright headers, automatically added by
make format
Files:
script/slurm/configs/mistral-dp.yaml
script/**
⚙️ CodeRabbit configuration file
Review standalone scripts for reproducibility and operational safety. Check argument validation, quoting, repo-root detection, environment variables, generated artifacts, external commands, GPU/cluster assumptions, and whether the script should be wired through Makefile or documented in README/docs.
Files:
script/slurm/configs/mistral-dp.yaml
🔇 Additional comments (1)
script/slurm/configs/mistral-dp.yaml (1)
1-13: LGTM!
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
|
Actionable comments posted: 0 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
3 participants
followup from #483 - there's a regression in memory performance in transformers v5. this is the minimal fix for our specific workloads; I'll put up a more comprehensive change that should help with observability and memory pressure during training after this.
Summary
This PR lowers the physical microbatch size used by the DP run configs from
8to4and raisesgradient_accumulation_stepsto16.The effective batch size stays at
64, but each forward/backward pass uses fewer examples. This is the minimal fix for the recent larger-model DP OOMs without changing the shared training defaults for non-DP runs.Why
DP training is not inherently constrained to
batch_size: 1, but larger physical batches increase activation and per-sample-gradient memory. The prior DP configs usedbatch_size: 8with the global default accumulation of8. Moving to4 x 16preserves optimizer-step batch semantics while reducing peak memory pressure.Changes
training.batch_size: 4training.gradient_accumulation_steps: 16Test Plan
make checkSummary by CodeRabbit