Skip to content

fix(prune)!: honour Schedule start_pct/end_pct in Pruner and PruneCallback - #98

Merged
nathanhubens merged 1 commit into
masterfrom
fix/prune-schedule-start-pct
Sep 24, 2026
Merged

nathanhubens merged 1 commit into
masterfrom
fix/prune-schedule-start-pct

Conversation

@nathanhubens

Copy link
Copy Markdown
Collaborator

Breaking (targets 0.5.0). PruneCallback and Pruner(schedule=Schedule) now follow the schedule's start_pct and end_pct.

The bug

Two identical _build_pruning_schedule functions, one in Pruner and one in PruneCallback, evaluated sched_func(start_val, end_val, i/steps) from the first step and never read start_pct/end_pct. As a result:

  • one_shot declares start_pct=0.5 but pruned at the first step of the fit;
  • agp and iterative ramped from step 1 instead of from 20 %.

The fix

  • One owner: Pruner._to_tp_scheduler(schedule) builds torch-pruning's levels from copy(schedule).progress(i/steps) * ratio. The shared module-level schedule objects are never mutated, and a plain callable scheduler passes through unchanged.
  • Both _build_pruning_schedule copies are deleted.
  • PruneCallback passes its Schedule to the Pruner and reads its logged levels from the pruner. The per-layer dict path keeps sparsity_levels == [].
  • A cumulative level of 0 prunes nothing: torch-pruning skips a layer when it has no channel to remove.

Effect

  • one_shot prunes at the first step at or past 50 % of training.
  • agp and iterative begin at their start_pct.
  • end_pct ends the ramp.
  • Schedules with start_pct=0, end_pct=1 (one_cycle, cos, lin, dsd) are unchanged, level for level (golden test). No public notebook uses a non-zero start_pct, so no stored output changes.

Evidence

  • New or changed tests fail on master and pass here:
    • test_pruner, _levels(one_shot) gives [0.5]*9 on master vs [0]*4+[0.5]*5 expected;
    • test_prune_callback, the channel trace gives [8]*6 on master vs [16,16,8,8,8,8] expected.
  • Also covered:
    • agp and windowed-end_pct levels;
    • golden identity for one_cycle, cos, lin and Schedule(sched_agp);
    • the per-layer dict path, and the reg > 0 path getting the same levels;
    • callable passthrough;
    • shared singletons untouched.
  • Full suite nbdev-test passes in a clean clone with CI's nbdev 3.3.21, slow prune tests included. Clean and export leave the tree unchanged.

Not included

PruneCallback.after_epoch prints the level of the previous step: it reads levels[completed_steps - 1], while torch-pruning applies levels[completed_steps]. That one-index fix comes in a small follow-up PR stacked on this one, because it requires re-executing the prune_callback tutorial.

…back

The levels handed to torch-pruning mapped step i of n straight into the
schedule's shape function, so start_pct and end_pct were never read:
one_shot pruned at the first step, agp and iterative ramped from the first
step. Pruner._to_tp_scheduler now builds the levels from
Schedule.progress(i / n) for both entry points; the two
_build_pruning_schedule copies are removed. Schedules with start_pct=0 and
end_pct=1 produce the same levels as before.
@nathanhubens
nathanhubens merged commit 0820785 into master Sep 24, 2026
2 checks passed
@nathanhubens
nathanhubens deleted the fix/prune-schedule-start-pct branch September 24, 2026 10:51
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