[Bugfix] Fix polynomial_decay pruning scheduler for even exponents by winklemad · Pull Request #3096 · vllm-project/llm-compressor

winklemad · 2026-08-25T17:49:13Z

Summary

polynomial_decay_scheduler computes the AGP (Zhu & Gupta, 2017) sparsity decay factor as:

scaled_complete = pow(per_complete - 1, exponent) + 1

This only equals the intended 1 - (1 - per_complete)**exponent when exponent is odd. For even exponents — including the default of 2 — (per_complete - 1)**exponent == (1 - per_complete)**exponent, so the decay term keeps the wrong sign: sparsity starts above final_sparsity and decreases over training instead of rising from init_sparsity.

cubic_scheduler sets exponent=3 (odd), so it happens to be correct — which is likely why this went unnoticed — but selecting "polynomial_decay" (default exponent=2) or any even exponent produces a broken schedule.

Reproduce

from llmcompressor.core import Event
from llmcompressor.modifiers.pruning.helpers import (
    polynomial_decay_scheduler, PruningCreateSettings,
)

s = PruningCreateSettings(start=0, end=10, update=1,
                          init_sparsity=0.0, final_sparsity=0.9, args={})
sched = polynomial_decay_scheduler(s)   # default exponent=2

def at(i):
    e = Event(); e.current_index = i
    return sched(e, None)

print(at(0), at(5), at(10))   # 1.8  1.125  0.9   -> starts at 1.8 (>1.0) and falls
# expected a rise 0.0 -> ~0.675 -> 0.9

Fix

Use 1 - pow(1 - per_complete, exponent), which is correct for all exponents and identical to the previous behaviour for odd exponents, so cubic_scheduler (exponent=3) is byte-for-byte unchanged.

After the fix the same schedule rises 0.0, 0.171, 0.324, …, 0.9.

Testing

Adds tests/llmcompressor/modifiers/pruning/test_helpers.py:

polynomial_decay is monotonic and bounded in [init, final] for exponents 2, 3, and 4;
the default (even) exponent starts at init_sparsity;
cubic_scheduler output is unchanged.

The even-exponent tests fail on main and pass with the fix. ruff check / ruff format --check clean.

polynomial_decay_scheduler computed the AGP decay factor as (t - 1)**exponent + 1, which only equals the intended 1 - (1 - t)**exponent for odd exponents. For even exponents -- including the default of 2 -- the sign of the decay term flips, so sparsity starts above final_sparsity and decreases over training instead of rising from init_sparsity. For example with init=0.0, final=0.9 the schedule started at 1.8 and fell to 0.9. Use 1 - (1 - t)**exponent, which is correct for all exponents and identical to the previous behaviour for odd exponents, so the cubic scheduler (exponent=3) is unchanged. Add regression tests for the schedulers. Signed-off-by: Madan Kumar <[email protected]>

github-actions · 2026-08-25T17:49:20Z

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

coderabbitai · 2026-08-25T17:49:20Z

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: c21caec8-8119-4cb9-a259-aaa3969fcc14

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

mergify · 2026-08-25T17:49:49Z

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 Require one maintainer review

All PRs must have at least one approving review from a maintainer before merging.

gemini-code-assist

Code Review

This pull request corrects the AGP decay calculation in the polynomial decay scheduler to properly handle even exponents, ensuring sparsity increases monotonically, and adds corresponding unit tests. The review feedback highlights a potential issue where the completion ratio could exceed 1.0, leading to complex numbers with fractional exponents or non-monotonic behavior with even exponents. It is recommended to clamp the completion ratio to the [0.0, 1.0] range and expand the test suite to cover fractional exponents and out-of-bounds indices.

kylesayrs

You are indeed correct, thanks for the fix!

winklemad · 2026-08-25T23:38:53Z

Thanks for the quick review and approve, @kylesayrs!

dsikka · 2026-08-26T12:25:43Z

@winklemad do you mind looking into the failing tests?

The new tests/llmcompressor/modifiers/pruning/test_helpers.py had no __init__.py, so under pytest's default prepend import mode it was imported under the bare module name 'test_helpers', colliding with tests/llmcompressor/pipelines/ sequential/test_helpers.py and aborting collection with an import-file-mismatch error. Add the package marker (matching sibling dirs like modifiers/reap and modifiers/utils) so it imports as a unique dotted module. Signed-off-by: Madan Kumar <[email protected]>

winklemad · 2026-08-26T16:26:55Z

Looked into it — the red build isn't the scheduler change, it's a pytest collection error that aborts the whole suite before any test runs:

ERROR collecting tests/llmcompressor/pipelines/sequential/test_helpers.py
import file mismatch:
imported module 'test_helpers' has this __file__ attribute:
  .../tests/llmcompressor/modifiers/pruning/test_helpers.py
which is not the same as the test file we want to collect:
  .../tests/llmcompressor/pipelines/sequential/test_helpers.py
HINT: ... use a unique basename for your test file modules

Root cause is on me: I added tests/llmcompressor/modifiers/pruning/test_helpers.py but the new pruning/ dir had no __init__.py. Under the default prepend import mode a test module in a dir without __init__.py is imported by its bare basename (test_helpers), which then collides with the pre-existing bare tests/llmcompressor/pipelines/sequential/test_helpers.py. main is green only because pruning/ was the second bare test_helpers in the tree.

Fixed in 9af1153 by adding an empty tests/llmcompressor/modifiers/pruning/__init__.py, so it imports as the unique dotted tests.llmcompressor.modifiers.pruning.test_helpers — matching the sibling dirs that already have one (modifiers/reap, modifiers/utils, observers). Locally the full tests/llmcompressor collection now completes with 0 errors, and the pruning tests themselves pass (the helpers.py math change is unchanged). Should be green on the next run.

winklemad requested review from HDCharles, brian-dellabetta, dsikka and kylesayrs as code owners August 25, 2026 17:49

gemini-code-assist Bot reviewed Aug 25, 2026

View reviewed changes

Comment thread src/llmcompressor/modifiers/pruning/helpers.py Resolved

Comment thread tests/llmcompressor/modifiers/pruning/test_helpers.py Resolved

kylesayrs approved these changes Aug 25, 2026

View reviewed changes

kylesayrs added the ready When a PR is ready for full CI testing before merge label Aug 25, 2026

kylesayrs enabled auto-merge (squash) August 25, 2026 18:03

Merge branch 'main' into fix/polynomial-decay-even-exponent

50b2997

auto-merge was automatically disabled August 26, 2026 16:26
Head branch was pushed to by a user without write access

Merge branch 'main' into fix/polynomial-decay-even-exponent

1eaa75e

dsikka enabled auto-merge (squash) August 27, 2026 23:54

dsikka merged commit 6b48754 into vllm-project:main Aug 28, 2026
8 checks passed

Navigation Menu

Repository navigation

[Bugfix] Fix polynomial_decay pruning scheduler for even exponents - #3096

[Bugfix] Fix polynomial_decay pruning scheduler for even exponents#3096
dsikka merged 4 commits into
vllm-project:mainfrom
winklemad:fix/polynomial-decay-even-exponent

winklemad commented Aug 25, 2026

Uh oh!

github-actions Bot commented Aug 25, 2026

Uh oh!

coderabbitai Bot commented Aug 25, 2026 •

edited

Loading

Review skipped

Uh oh!

mergify Bot commented Aug 25, 2026 •

edited

Loading

🟢 Require one maintainer review

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment •

edited

Loading

Uh oh!

winklemad commented Aug 25, 2026

Uh oh!

dsikka commented Aug 26, 2026

Uh oh!

winklemad commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

winklemad commented Aug 25, 2026

Summary

Reproduce

Fix

Testing

Uh oh!

github-actions Bot commented Aug 25, 2026

Uh oh!

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

Uh oh!

Review skipped

Uh oh!

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

Uh oh!

Merge Protections

🟢 Require one maintainer review

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment • 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!

winklemad commented Aug 25, 2026

Uh oh!

dsikka commented Aug 26, 2026

Uh oh!

winklemad commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai Bot commented Aug 25, 2026 •

edited

Loading

mergify Bot commented Aug 25, 2026 •

edited

Loading

kylesayrs left a comment •

edited

Loading