Add subblock stats to the compress algorithm by danielkorzekwa · Pull Request #623 · NVIDIA/Model-Optimizer

danielkorzekwa · 2025-12-01T15:52:47Z

What does this PR do?

Add subblock stats to the compress algorithm.

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

codecov · 2025-12-01T16:17:09Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.37%. Comparing base (97fe7f0) to head (6cb644a).
⚠️ Report is 1 commits behind head on feature/compress.

Additional details and impacted files

@@                Coverage Diff                @@
##           feature/compress     #623   +/-   ##
=================================================
  Coverage             74.37%   74.37%           
=================================================
  Files                   182      182           
  Lines                 18219    18219           
=================================================
  Hits                  13550    13550           
  Misses                 4669     4669

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

modelopt/torch/_compress/subblock_stats/calc_subblock_params_and_memory.py

kevalmorabia97 · 2025-12-01T17:35:13Z

modelopt/torch/_compress/subblock_stats/calc_subblock_params_and_memory.py

+    raise_unknown_subblock_config_error(subblock_config)
+
+
+def calculate_subblock_params(


I think param count is a bit hacky and can be simplified by just running 1 forward pass on sample input and calculate params. We have this simple utility in modelopt already (param_num_from_forward): https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/modelopt/torch/utils/network.py#L129 which is generic and works for any HF model (moe or dense).

We can run both functions and compare the numbers as well

Good candidate for a shared component. Added as a high priority internal Nvidia issue: issues/74.

kevalmorabia97 · 2025-12-01T17:39:30Z

modelopt/torch/_compress/subblock_stats/calc_subblock_params_and_memory.py

+)
+
+
+def calculate_subblock_memory(


Quick question - is memory = (active params * param dtype memory) + (kv cache params * kv cache dtype memory) ? Or is a bit more complicated than that? I see a lot of logic for different types of layers but there isnt any docstring so I'm not sure whats the main reason for custom logic per layer

Sepehr raised similar questions and concerns. Added to: issues/74

kevalmorabia97 · 2025-12-01T17:43:44Z

modelopt/torch/_compress/subblock_stats/calc_subblock_stats.py

+        from puzzle_tools.subblock_stats.runtime_stats.calc_runtime_stats import (
+            calc_runtime_ms_for_subblocks,
+        )


Is this for TRT-LLM stats?

yes, this is used only if runtime_stats is enabled via a config param.

kevalmorabia97 · 2025-12-01T17:45:23Z

modelopt/torch/_compress/subblock_stats/calc_subblock_stats.py

+        # TODO: fix
+        # from puzzle_tools.calc_subblock_runtime import measure_non_block_runtime_ms
+        # non_block_runtime_ms, embedding_runtime_ms, lm_head_runtime_ms = \
+        #     measure_non_block_runtime_ms(batch_size, prefill_seq_len, generation_seq_len, n_embd, vocab_size,
+        #                                  benchmark_iterations, use_cuda_graph)


will this be added in follow-up PR?

Once scoring/mip are in, we can prioritize what to do next. I added an internal issue for subblock runtime stats: issues/75

kevalmorabia97 · 2025-12-01T17:46:15Z

modelopt/torch/_compress/subblock_stats/calc_subblock_stats.py

+    # ==== START === Setup for attach-helper ====
+    # import sys
+    # import os
+    # sys.path.insert(0, os.environ["ATTACH_HELPER_INSTALLATION_PATH"])
+    # from attach_helper import debugging_setup
+    # debugging_setup()  # You can optionally pass a name to identify the job (e.g. `debugging_setup(name="my_script")`)
+    # ==== END === Setup for attach-helper ====


what is this for?

for debugging likely, I removed

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

danielkorzekwa added 30 commits November 3, 2025 21:05

Add decilm modelling code

694c317

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

Add decilm modelling code.

991659f

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

Add transformers codebase

8489cee

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

Add transformers code

f0afefe

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

Add decilm modelling code

b3ed5bc

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

Add decilm modelling code

a700da5

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

Correct licence headers

b59b679

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

Correct licence headers

1abdf3e

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

Add decilm code

66609b1

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

Add decilm code

7da0a8a

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

Add decilm code

6e09a81

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

Add decilm code

2e3f5da

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

Add decilm code

418890e

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

Make llama3 converter self-contained (no deps on internal Nvidia code)

01f4fc1

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

Add common module

c57eed4

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

module refactoring

3dc37b3

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

refactoring

10ffdfe

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

add shared_checkpointing_utils

27a4456

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

Add json tools

b0e22b7

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

add logger

52e7827

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

import refactoring

f5c1c87

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

add post_init_sparse module

0aa6320

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

Add post_init_sparse

35d0dbc

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

merginy hydra.py and hydra_utils.py

e39a1ad

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

Add integrationt test for attention pruning

1bd0c67

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

add score_pruning_activations

0ecd52b

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

import refactoring

278c6b7

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

add dist_utils

7a0af16

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

Add validate_model

0f0cbbd

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

Add activation scoring hooks for pruning

cb5cf25

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

danielkorzekwa added 16 commits November 6, 2025 11:56

add pruning_ckpts

845d453

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

import refactoring

4fd921b

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

refactor imports

3641847

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

import refactoring

8d6333b

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

Add build_replacement_library

dcb86e2

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

import refactoring

dfd3adc

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

add replacement_library

daf94d3

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

refactor imports

ab6e9e3

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

refactor imports

01a6aee

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

add subblock_stats

79cf2cf

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

add subblock_stats

46d2ea2

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

import refactoring

561e04a

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

import refactoring

22da6b0

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

import refactoring

d4dd0d7

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

Merge branch 'feature/compress' into dkorzekwa/subblock_stats

2796a3e

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

replace frozendict with immutabledict

ebb1338

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

danielkorzekwa requested a review from a team as a code owner December 1, 2025 15:52

danielkorzekwa added 5 commits December 1, 2025 16:53

Remove not needed tokenizer

4824928

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

Delete empty module

0069ee2

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

Add doc strings

6533895

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

code clean up

e1ec39c

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

Add doc strings

8b4ea16

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

AAnoosheh approved these changes Dec 1, 2025

View reviewed changes

kevalmorabia97 reviewed Dec 1, 2025

View reviewed changes

rename calc_subblock_memory.py to calc_subblock_params_and_memory.py

6cb644a

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

kevalmorabia97 approved these changes Dec 1, 2025

View reviewed changes

danielkorzekwa merged commit 954103e into feature/compress Dec 1, 2025
21 checks passed

danielkorzekwa deleted the dkorzekwa/subblock_stats branch December 1, 2025 20:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add subblock stats to the compress algorithm#623

Add subblock stats to the compress algorithm#623
danielkorzekwa merged 60 commits intofeature/compressfrom
dkorzekwa/subblock_stats

danielkorzekwa commented Dec 1, 2025

Uh oh!

codecov bot commented Dec 1, 2025 •

edited

Loading

Uh oh!

Uh oh!

kevalmorabia97 Dec 1, 2025

Uh oh!

danielkorzekwa Dec 1, 2025

Uh oh!

kevalmorabia97 Dec 1, 2025 •

edited

Loading

Uh oh!

danielkorzekwa Dec 1, 2025

Uh oh!

kevalmorabia97 Dec 1, 2025

Uh oh!

danielkorzekwa Dec 1, 2025

Uh oh!

kevalmorabia97 Dec 1, 2025

Uh oh!

danielkorzekwa Dec 1, 2025

Uh oh!

kevalmorabia97 Dec 1, 2025

Uh oh!

danielkorzekwa Dec 1, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		raise_unknown_subblock_config_error(subblock_config)


		def calculate_subblock_params(

Conversation

danielkorzekwa commented Dec 1, 2025

What does this PR do?

Uh oh!

codecov bot commented Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevalmorabia97 Dec 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov bot commented Dec 1, 2025 •

edited

Loading

kevalmorabia97 Dec 1, 2025 •

edited

Loading