[DRAFT] Make distributed Puzzletron campaigns reliable - #2430
chochowski wants to merge 1 commit into
Conversation
Preserve worker topology, stage completion, and FSDP projection semantics so heterogeneous pruning campaigns can progress reliably from scoring through distillation. Signed-off-by: mchochowski <mchochowski@nvidia.com>
|
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 (16)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe changes separate worker-group sizing from scheduler sizing, add depth-importance finalization, support additional DTensor teacher-head layouts, preserve failed-stage log paths, and adjust execution instance and axis-value planning. ChangesPuzzletron execution updates
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~30 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to No actionable regression remains identified; the change is ready to merge with normal checks. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## feature/puzzletron_v2 #2430 +/- ##
=========================================================
+ Coverage 50.63% 51.03% +0.39%
=========================================================
Files 713 713
Lines 94313 94323 +10
=========================================================
+ Hits 47756 48135 +379
+ Misses 46557 46188 -369
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Preserve worker topology, stage completion, and FSDP projection semantics so heterogeneous pruning campaigns can progress reliably from scoring through distillation.
What does this PR do?
Key fixes:
Correct distributed worker world-size handling.
Publish completed depth-stage manifests.
Allocate replacement workers per embedding width.
Preserve failed-stage log paths.
Support FSDP dp_shard_cp LM-head projection during KD.
Use safer 75% initial pruning targets.
Added regression tests for each behavior.
Verification:
Full Nemotron3 production campaign succeeded.
All pre-commit hooks passed.
Shell syntax checks passed.
Generated checkpoints/ data was excluded.### Usage
Summary by CodeRabbit
New Features
Bug Fixes