Implement nas.convert() api for the compress algorithm by danielkorzekwa · Pull Request #482 · NVIDIA/Model-Optimizer

danielkorzekwa · 2025-10-29T18:19:50Z

What does this PR do?

Implement nas.convert() api for the compress algorithm

using MIP-based NAS search algorithm. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

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

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

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

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

…ntal/ folder to not be run by CICD yet. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

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

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

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

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

…o_decilm_convertion

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

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

…as_convert

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

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

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

…as_convert

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

codecov · 2025-10-29T18:32:53Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.40%. Comparing base (cef3655) to head (936556f).
⚠️ Report is 1 commits behind head on feature/compress.

Additional details and impacted files

@@                Coverage Diff                @@
##           feature/compress     #482   +/-   ##
=================================================
  Coverage             73.40%   73.40%           
=================================================
  Files                   180      180           
  Lines                 18127    18127           
=================================================
  Hits                  13306    13306           
  Misses                 4821     4821

☔ 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/nas/plugins/compress_nas_plugin.py

kevalmorabia97 · 2025-10-29T18:37:47Z

tests/experimental/torch/_compress/nas/plugins/test_nas_convert.py

+    hydra_config_dir = project_root_path / "tests/experimental/torch/_compress/resources/configs"
+    hydra_config_name = "Llama-3_1-8B"


Any reason for passing hydra_config_dir and hydra_config_name instead of just 1 argument hydra_config_path?

there could be multiple models supported, not just Llama-3_1-8B, currently, by design for each model there is a dedicated hydra file,

for the user facing interface we could simplify it, but we need to think how.

tests/experimental/torch/_compress/nas/plugins/test_nas_convert.py

kevalmorabia97 · 2025-10-29T18:46:54Z

tests/experimental/torch/_compress/nas/plugins/test_nas_convert.py

+            setup_puzzle_dir(puzzle_dir)
+            save_dummy_dataset(dataset_path)
+
+            # Create a small Llama model
+            tokenizer = create_tokenizer(project_root_path)
+            create_and_save_small_llama_model(
+                llama_checkpoint_path, vocab_size=tokenizer.vocab_size, tokenizer=tokenizer
+            )


We can move this to test_utils.py and reuse in test_compress.py

already did as part of the next nas_search() MR, let's wait till you see it there.

tests/experimental/torch/_compress/compress_test_utils.py

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

kevalmorabia97 · 2025-10-30T18:23:12Z

tests/experimental/torch/_compress/nas/plugins/test_nas_convert.py

+        #
+        # Run the mnt.convert() step
+        #
+        input_model = CompressModel()


Is the CompressModel just a placeholder so it can pass the API requirements? Lets discuss in our meeting tomorrow how to better handle this

it is there because it is needed by mtn api, but I agree it is not really used. let's discuss.

kevalmorabia97 · 2025-10-30T20:10:43Z

modelopt/torch/_compress/hydra.py

+from hydra import compose, initialize, initialize_config_dir
+from omegaconf import DictConfig, OmegaConf


Please add hydra and omegaconf to setup.py puzzle dependencies

done, do we have some integration test that will set setup.py? for:

# Dependedencies for modelopt.torch._compress subpackage "compress": [ "fire", "hydra-core==1.3.2", "omegaconf==2.3.0", ],

I will work on the integration test next week

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

danielkorzekwa and others added 30 commits October 27, 2025 11:50

The main compression function for a model

c758ad5

using MIP-based NAS search algorithm. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Code formatting

8af9903

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

Model search space configuration used by test_compress.py test.

5ba6c27

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

Tokenizer used by test_compress.py test.

0bc5d84

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

Tokenizer utility used by test_compress.py test

87d4fa5

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

e2e tests for compress.py

ced1e99

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

Add convert_llama3_config_to_decilm_config + unit test

5de0bdc

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

Remove unused bypass distillation config files.

800414c

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

Moving integration tests to tests/experimental to not trigger CICD

16abcc9

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

update docs

a5ba1c7

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

Replace mprint with print and replace osp.join with path1 / path2 not…

1bda391

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

Refactor file checking assertions to use .is_file() and .exists()

bb38401

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

Add a new dependency section to setyp.py for the modelopt.torch._comp…

8415548

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

Move test_convert_llama3_config_to_decilm_config.py to tests/experime…

b1b1833

…ntal/ folder to not be run by CICD yet. Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>

Merge branch 'feature/compress' into dkorzekwa/e2e_compression_test

d4ffc91

Fix: Add missing LICENSE headers

6f28e4a

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Use spawn_multiprocess_job for test_compress test (to be able to use …

016fb63

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

Add comments.

0ccf1c4

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

Add _save_dummy_dataset to the test_compress.py

58439ca

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

Refactoring: Move torch distributed env variables to dist_utils.py

2e5f776

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

Refactoring: move torch distributed variables to dist_utils

6274db5

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

Move os.environ["WANDB_DISABLED"] = "true" to dist_utils.py

d942e0a

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

Implement integration test for mnt.convert() for the _compress algori…

f765921

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

Implement mtn.convert() for compress() algorithm.

de876d6

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

Merge branch 'dkorzekwa/e2e_compression_test' into dkorzekwa/llama3_t…

72bdc7a

…o_decilm_convertion

Merge branch 'dkorzekwa/llama3_to_decilm_convertion' into dkorzekwa/n…

40d28af

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

Fix broken test - incorrect package names.

f7fe23c

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

Merge branch 'dkorzekwa/llama3_to_decilm_convertion' into dkorzekwa/n…

3d1d286

…as_convert

Implementing nas.convert for compress algorithm.

a210483

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

Improve docs

739f868

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

danielkorzekwa added 15 commits October 28, 2025 21:17

Merge branch 'feature/compress' into dkorzekwa/llama3_to_decilm_conve…

18cb88b

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

Fix import

1033c81

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

simplify code

0680c45

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

implementing compress_nas_plugin

2d9da30

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

code clean up.

febab44

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

code clean up

86bf394

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

create conftest.py with shared test logic for compress tests.

86e04a0

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

code cleanup

ae61644

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

Merge branch 'dkorzekwa/llama3_to_decilm_convertion' into dkorzekwa/n…

2998cdb

…as_convert

code refactoring

3778ec2

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

refactoring

d940000

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

move test utilities from conftest.py to test_utils.py

0bf9a92

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

Improve comments

b56df9a

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

Merge branch 'feature/compress' into dkorzekwa/nas_convert

fd63130

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

Added TODO.

9bfcc21

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

danielkorzekwa requested a review from a team as a code owner October 29, 2025 18:19

kevalmorabia97 reviewed Oct 29, 2025

View reviewed changes

danielkorzekwa added 2 commits October 30, 2025 18:23

Utilitities for hydra initialization

6504c44

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

Code refactoring

d0fb8f9

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

kevalmorabia97 reviewed Oct 30, 2025

View reviewed changes

danielkorzekwa added 2 commits October 30, 2025 21:32

code refactoring

40f18b2

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

Add compress dependencies to setup.py.

936556f

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

danielkorzekwa requested a review from a team as a code owner October 30, 2025 20:39

danielkorzekwa requested review from kevalmorabia97 and removed request for a team October 30, 2025 20:39

kevalmorabia97 approved these changes Oct 31, 2025

View reviewed changes

danielkorzekwa merged commit 002b8b5 into feature/compress Oct 31, 2025
21 checks passed

danielkorzekwa deleted the dkorzekwa/nas_convert branch October 31, 2025 18:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement nas.convert() api for the compress algorithm#482

Implement nas.convert() api for the compress algorithm#482
danielkorzekwa merged 52 commits intofeature/compressfrom
dkorzekwa/nas_convert

danielkorzekwa commented Oct 29, 2025

Uh oh!

codecov bot commented Oct 29, 2025 •

edited

Loading

Uh oh!

Uh oh!

kevalmorabia97 Oct 29, 2025

Uh oh!

danielkorzekwa Oct 30, 2025

Uh oh!

Uh oh!

kevalmorabia97 Oct 29, 2025

Uh oh!

danielkorzekwa Oct 30, 2025

Uh oh!

Uh oh!

kevalmorabia97 Oct 30, 2025

Uh oh!

danielkorzekwa Oct 30, 2025

Uh oh!

kevalmorabia97 Oct 30, 2025

Uh oh!

danielkorzekwa Oct 30, 2025

Uh oh!

kevalmorabia97 Oct 31, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		hydra_config_dir = project_root_path / "tests/experimental/torch/_compress/resources/configs"
		hydra_config_name = "Llama-3_1-8B"

		from hydra import compose, initialize, initialize_config_dir
		from omegaconf import DictConfig, OmegaConf

Conversation

danielkorzekwa commented Oct 29, 2025

What does this PR do?

Uh oh!

codecov bot commented Oct 29, 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

2 participants

codecov bot commented Oct 29, 2025 •

edited

Loading