Dkorzekwa/puzzletron use importance hooks from prune by danielkorzekwa · Pull Request #1115 · NVIDIA/Model-Optimizer

danielkorzekwa · 2026-03-24T20:14:46Z

What does this PR do?

Use pruning importance hooks from prune/importance_hooks/ in modelopt.torch.puzzletron

…om_prune Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

codecov · 2026-03-25T10:48:56Z

Codecov Report

❌ Patch coverage is 0% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.19%. Comparing base (f460d16) to head (94cb2fe).
⚠️ Report is 1 commits behind head on feature/puzzletron.

Files with missing lines	Patch %	Lines
modelopt/torch/utils/robust_json.py	0.00%	4 Missing ⚠️

Additional details and impacted files

@@                  Coverage Diff                   @@
##           feature/puzzletron    #1115      +/-   ##
======================================================
- Coverage               70.22%   70.19%   -0.03%     
======================================================
  Files                     228      228              
  Lines                   25952    25956       +4     
======================================================
- Hits                    18224    18221       -3     
- Misses                   7728     7735       +7

☔ View full report in Codecov by Sentry.
📢 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.

kevalmorabia97 · 2026-03-25T12:44:45Z

Suggestion to keep model-specific hooks inside the model descriptor so someone wanting to add a new model support can see all model-specific stuff in single place

I thought the same, but let's rethink it, added to TODOs

kevalmorabia97 · 2026-03-25T12:46:15Z

Wasnt this change already merged to feature/puzzletron in your previous PR?

maybe it auto-resolved wrongly during your merge from main to feature/puzzletron? I saw the same when merging from main to dkorzekwa/puzzletron_use_importance_hooks_from_prune and I think I noticed it.

### What does this PR do? Implement puzzletron compression algorithm based on Puzzle paper (https://arxiv.org/abs/2411.19146) <details> <summary> Th list of reviewed and merged MRs that resulted in the feature/puzzletron branch</summary> Merging dkorzekwa/any_model to feature/puzzletron [Add anymodel directories to feature/puzzletron by danielkorzekwa · Pull Request #974 · NVIDIA/Model-Optimizer](#974) - merged [Draft: anymodel activation scoring by danielkorzekwa · Pull Request #989 · NVIDIA/Model-Optimizer](#989) - merged [Draft: Merge anymodel pruning by danielkorzekwa · Pull Request #990 · NVIDIA/Model-Optimizer](#990) - merged [Draft: Merging anymodel:build_library_and_stats by danielkorzekwa · Pull Request #993 · NVIDIA/Model-Optimizer](#993) - merged [Dkorzekwa/any model calc one block scores by danielkorzekwa · Pull Request #994 · NVIDIA/Model-Optimizer](#994) - merged [Draft: merge any_model: mip_and_realize_models by danielkorzekwa · Pull Request #995 · NVIDIA/Model-Optimizer](#995) - merged [Dkorzekwa/any model other modeqls by danielkorztiekwa · Pull Request #1007 · NVIDIA/Model-Optimizer](#1007) - merged PR to 1007: #1039 - merged [Dkorzekwa/anymodel gptoss by danielkorzekwa · Pull Request #1020 · NVIDIA/Model-Optimizer](#1020) - merged [Merge any_model tutorial by danielkorzekwa · Pull Request #1035 · NVIDIA/Model-Optimizer](#1035) - merged [Merge mbridge distillation for any_model by danielkorzekwa · Pull Request #1036 · NVIDIA/Model-Optimizer](#1036) - merged [MR branch for the remaining difference between dkorzekwa/any_model an… by danielkorzekwa · Pull Request #1047 · NVIDIA/Model-Optimizer](#1047) - merged [Dkorzekwa/decilm hf code cleanup by danielkorzekwa · Pull Request #1071 · NVIDIA/Model-Optimizer](#1071) - merged [Dkorzekwa/decilm hf code cleanup 2 by danielkorzekwa · Pull Request #1073 · NVIDIA/Model-Optimizer](#1073) - merged [Dkorzekwa/anymodel subblock stats by danielkorzekwa · Pull Request #1085 · NVIDIA/Model-Optimizer](#1085) - merged [Dkorzekwa/anymodel subblock stats nodecilm by danielkorzekwa · Pull Request #1102 · NVIDIA/Model-Optimizer](#1102) - merged [Dkorzekwa/decilm cleanup post subblockstats by danielkorzekwa · Pull Request #1103 · NVIDIA/Model-Optimizer](#1103) - merged [code clean up by danielkorzekwa · Pull Request #1110 · NVIDIA/Model-Optimizer](#1110) - merged Merging into main: [Activation hooks redesign (reuse hooks component across both minitron and puzzletron) by danielkorzekwa · Pull Request #1022 · NVIDIA/Model-Optimizer](#1022) - merged [Dkorzekwa/puzzletron use importance hooks from prune by danielkorzekwa · Pull Request #1115 · NVIDIA/Model-Optimizer](#1115) - merged </details>  ### Usage Puzzletron tutorial: https://github.com/NVIDIA/Model-Optimizer/tree/feature/puzzletron/examples/puzzletron ### Testing The main e2e test for compressing 9 models with Puzzletron: https://github.com/NVIDIA/Model-Optimizer/blob/feature/puzzletron/tests/gpu/torch/puzzletron/test_puzzletron.py 2-gpu nightly tests: - https://github.com/NVIDIA/Model-Optimizer/actions/runs/24468209205/job/71501061203 - https://github.com/NVIDIA/Model-Optimizer/actions/runs/24470214159/job/71508152952 ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅  ## Summary by CodeRabbit * **New Features** * Added Puzzletron: end-to-end heterogeneous pruning & NAS workflow with AnyModel support, example pipelines, deployment and evaluation utilities, and tools for converting/pruning and exporting compressed checkpoints. * **Documentation** * Comprehensive Puzzletron tutorials, model-specific guides, evaluator instructions, example configs, and changelog entry. * **Chores** * CI/workflow updates (extras installation, longer GPU test timeout), pre-commit hook exclusion updated, and CODEOWNERS entries added.  --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com> Signed-off-by: Liana Mikaelyan <lmikaelyan@nvidia.com> Signed-off-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com> Signed-off-by: Daniel Korzekwa <daniel.korzekwa@gmail.com> Signed-off-by: jrausch <jrausch@nvidia.com> Signed-off-by: root <root@pool0-00848.cm.cluster> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Liana Mikaelyan <lmikaelyan@nvidia.com> Co-authored-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com> Co-authored-by: J Rausch <38429553+j-rausch@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

danielkorzekwa added 3 commits March 24, 2026 11:46

Merge branch 'main' into dkorzekwa/puzzletron_use_importance_hooks_fr…

b26d37a

…om_prune Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Restore to main version

558c9e7

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Refactoring of activation hooks

39ec6b6

Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

danielkorzekwa requested review from a team as code owners March 24, 2026 20:14

danielkorzekwa requested review from ChenhanYu, ajrasane, jenchen13, rohansjoshi and vishalpandya1990 March 24, 2026 20:14

kevalmorabia97 requested review from kevalmorabia97 and removed request for a team, ChenhanYu, ajrasane, cjluo-nv, jenchen13, realAsma, rohansjoshi and vishalpandya1990 March 25, 2026 10:45

kevalmorabia97 reviewed Mar 25, 2026

View reviewed changes

kevalmorabia97 approved these changes Mar 25, 2026

View reviewed changes

kevalmorabia97 merged commit 2f55c73 into feature/puzzletron Mar 25, 2026
33 checks passed

kevalmorabia97 deleted the dkorzekwa/puzzletron_use_importance_hooks_from_prune branch March 25, 2026 16:15

danielkorzekwa mentioned this pull request Mar 25, 2026

Merge puzzletron compression algorithm #1121

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Dkorzekwa/puzzletron use importance hooks from prune#1115

Dkorzekwa/puzzletron use importance hooks from prune#1115
kevalmorabia97 merged 4 commits into
feature/puzzletronfrom
dkorzekwa/puzzletron_use_importance_hooks_from_prune

danielkorzekwa commented Mar 24, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Mar 25, 2026 •

edited

Loading

Uh oh!

kevalmorabia97 Mar 25, 2026

Uh oh!

danielkorzekwa Mar 25, 2026

Uh oh!

kevalmorabia97 Mar 25, 2026

Uh oh!

danielkorzekwa Mar 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

danielkorzekwa commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

codecov Bot commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

kevalmorabia97 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

danielkorzekwa Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

kevalmorabia97 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

danielkorzekwa Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

danielkorzekwa commented Mar 24, 2026 •

edited

Loading

codecov Bot commented Mar 25, 2026 •

edited

Loading