Skip to content

Conversation

@CedricHwong
Copy link

@CedricHwong CedricHwong commented Dec 30, 2025

What does this PR do?

Overview:
This PR extends GradNAS gradient scoring with
(1) optional score averaging
(2) early stopping when gradient scores converge. New search config keys are added with defaults: average_scores, score_convergence_tol, score_convergence_patience, and score_convergence_min_updates. The gradient estimation loop now exits early once all hparam score trackers converge. Unit tests cover the averaging math, convergence boundaries, and disabled-convergence behavior.

Usage

  import modelopt.torch.prune as mtp

  # Example: enable/adjust score averaging + convergence early stop
  pruned_model, search_history = mtp.prune(
      model=model,
      mode="gradnas",
      constraints={"flops": "90%"},
      dummy_input=dummy_input,
      config={
          "data_loader": train_loader,
          "loss_func": loss_func,
          # New options:
          "average_scores": True,
          "score_convergence_tol": 1e-3,
          "score_convergence_patience": 5,
          "score_convergence_min_updates": 10,
      },
  )

Testing

  • pytest tests/unit/torch/prune/test_gradnas.py
  • pre-commit run --all-files

Before your PR is "Ready for review"

  • Make sure you read and follow Contributor guidelines (https://github.com/NVIDIA/Model-
    Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed.
  • Is this change backward compatible?: No Default GradNAS behavior now averages gradient scores and may stop early on convergence, which can change score scale and number of batches processed vs. previous behavior.
  • Did you write any new necessary tests?: Yes
  • Did you add or update any necessary documentation?: No
  • Did you update Changelog (https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: No

Additional Information

N/A

@CedricHwong CedricHwong requested a review from a team as a code owner December 30, 2025 10:30
@copy-pr-bot
Copy link

copy-pr-bot bot commented Dec 30, 2025

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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