Skip to content

[DRAFT] Make distributed Puzzletron campaigns reliable - #2430

Open
chochowski wants to merge 1 commit into
feature/puzzletron_v2from
fix/puzzletron-campaign-reliability
Open

chochowski wants to merge 1 commit into
feature/puzzletron_v2from
fix/puzzletron-campaign-reliability

Conversation

@chochowski

@chochowski chochowski commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

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

    • Added depth-importance finalization for validating completed runs and generating final manifests.
    • Added support for configuring worker-group size independently in distributed evaluation workflows.
    • Improved execution planning for embedding-pruning replacement scoring across configured embedding widths.
    • Updated default setup suggestions to cover 75% of available teacher-value options.
  • Bug Fixes

    • Improved distributed distillation support for sharded language-model heads.
    • Preserved worker log paths when active stages fail.
    • Added automatic depth-importance finalization after successful depth evaluation.

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>
@coderabbitai

coderabbitai Bot commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

Review Change StackReview Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 96cd6aec-fed0-46ec-a20f-1a3f36a70e7a

📥 Commits

Reviewing files that changed from the base of the PR and between 60df4db and 9b212fa.

📒 Files selected for processing (16)
  • examples/puzzletron/distributed_eval/run_coordinator.sh
  • examples/puzzletron/distributed_eval/run_depth_coordinator.sh
  • examples/puzzletron/distributed_eval/run_depth_pool.sh
  • examples/puzzletron/distributed_eval/run_replacement_pool.sh
  • examples/puzzletron/finalize_depth_importance.py
  • modelopt/torch/puzzletron/distillation/global_kd_recipe.py
  • modelopt/torch/puzzletron/orchestration/adapters/pool.py
  • modelopt/torch/puzzletron/orchestration/controller.py
  • puzzletron_setup/v2/bundle.py
  • puzzletron_setup/v2/wizard_common.py
  • tests/unit/torch/puzzletron/test_depth_finalizer.py
  • tests/unit/torch/puzzletron/test_orchestration_executors.py
  • tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py
  • tests/unit/torch/puzzletron/test_setup_v2_quick.py
  • tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py
  • tests/unit/torch/puzzletron/test_training_flash_kld.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.


📝 Walkthrough

Walkthrough

The 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.

Changes

Puzzletron execution updates

Layer / File(s) Summary
Distributed evaluation sizing and finalization
examples/puzzletron/distributed_eval/*, examples/puzzletron/finalize_depth_importance.py, tests/unit/torch/puzzletron/test_depth_finalizer.py, tests/unit/torch/puzzletron/test_orchestration_executors.py
Evaluation scripts resolve PUZZLETRON_WORKER_WORLD_SIZE with WORLD_SIZE fallback. Depth evaluation runs finalization after successful coordinator completion. The finalizer validates trajectories and writes the terminal manifest. Tests cover successful and incomplete trajectories and worker environment variables.
Distributed teacher projection handling
modelopt/torch/puzzletron/distillation/global_kd_recipe.py, tests/unit/torch/puzzletron/test_training_flash_kld.py
DTensor alignment handles non-TP heads and rejects mixed meshes. Teacher projection invokes the module call path for sharded heads. Distributed tests cover rank-local hidden states and hook behavior.
Pool environment and stage-view reporting
modelopt/torch/puzzletron/orchestration/adapters/pool.py, modelopt/torch/puzzletron/orchestration/controller.py, tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py
Pool execution paths export worker-group sizing. Active-stage log paths are collected before status selection, so failed stage views retain worker log paths.
Execution instance and axis planning
puzzletron_setup/v2/bundle.py, puzzletron_setup/v2/wizard_common.py, tests/unit/torch/puzzletron/test_setup_v2_*
Stage instance calculation accounts for smoke runs and embedding widths. Default axis selection targets three-quarters of the teacher value. Tests cover both planning changes.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~30 minutes

Change: Bug fix

Suggested reviewers: j-rausch

Merge Risk: ⚪ Minimal · up to 9b212

No actionable regression remains identified; the change is ready to merge with normal checks.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 30 functions across 16 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: improving reliability for distributed Puzzletron campaigns. It is concise and directly related to the changeset.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns ✅ Passed PASS. The pull-request diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comm…
  • Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/puzzletron-campaign-reliability

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2430/

Built to branch gh-pages at 2026-09-14 14:40 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@codecov

codecov Bot commented Sep 14, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.61538% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 51.03%. Comparing base (60df4db) to head (9b212fa).

Files with missing lines Patch % Lines
.../torch/puzzletron/distillation/global_kd_recipe.py 91.66% 1 Missing ⚠️
...elopt/torch/puzzletron/orchestration/controller.py 0.00% 1 Missing ⚠️
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     
Flag Coverage Δ
puzzletron 29.56% <84.61%> (+<0.01%) ⬆️
regression 8.71% <0.00%> (?)
unit 28.66% <0.00%> (-0.01%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant