Pipeline Extraction

[kylesayrs]

Purpose

• Extract data pipelines from modifiers to enable multiple modifiers to be active at the same time
  • This enables faster compression of larger models
  • This enables more memory efficient compression of larger models (not limited to just GPTQ/SGPT)

Prerequisites

• Implement QuantizationMixin #1351
• Use align_module_device util #1298

Callback Changes

• Implement calibration_epoch_start
  • This callback should be called at the start of every calibration pipeline
  • This callback causes modifiers to attach hooks
• Implement sequential_epoch_end
  • This callback should be called after one sequential layer has been calibrated with one epoch
  • This callback triggers compression and replaces passing a callback_modifier
• Implement calibration_epoch_end
  • This callback triggers at the end of a calibration epoch, and is used to trigger compression in between pipelines composed using the independent pipeline and remove hooks in between independent pipelines

Lifecycle Changes

• Oneshot modifiers implement on_end, which removes hooks when calibration finishes
  • In the future, calibration_epoch_start is treated like batch_start, where it is an opportunity for modifiers to start
  • In the future, calibration_epoch_end is treated like batch_end, where it is an opportunity for modifiers to end
  • In the future, finalize is treated like batch_end, where it is an opportunity for modifiers to end
  • Right now, these opportunities are implemented manually on each oneshot modifier, rather than being a lifecycle rule

Data Pipeline Changes

• Implement data pipeline registry
  • Inferred pipeline is selected using modifiers and can be overridden by user
• Implement independent pipeline
  • This pipeline treats each modifier as a separate stage and assigns a pipeline to each modifier
  • Meant to replicate current LC behavior
  • Originally, these compression events were triggered by reaching the end of each module’s initialize function. Now a separate event is required
• Implement session.get_modifiers
  • In order to perform data pipeline inference and other sequential pipeline inference, these functions must get the list of active modifiers before they initialize
  • This function gets all the active modifiers across all ModifierStages
• Prepare smoothquant for pipeline extraction
  • Trigger _apply_smoothing on the sequential_epoch_end and calibration_epoch_end
  • Add a guard which allows the _apply_smoothing function to be called multiple times per session (as is required by sequential pipeline)

Testing

• Quantized llama3-8b using both the independent (basic + sequential) and sequential pipelines
  • There was no accuracy regression from using a shared pipeline, although we keep the independent pipeline as the default for now
• Transformers tests pass
• https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14622080074

[brian-dellabetta]

Definitely looks cleaner this way! Leaving comments rather than approving, as I am still getting up to speed with pipelines

[brian-dellabetta]

I know you're looking for feedback on this, but I'm not sure I understand it enough to approve. I do like the removal of all the try/catch code in GPTQ. Maybe we can have a deep dive session on this next week?

[kylesayrs]

Looks like there's one perplexity failure, although I wasn't able to replicate locally https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14622080074/job/41024772318#step:13:31981

[brian-dellabetta]

really cool! excited to try this out. Should we run the e2e/lmeval tests before merging this in? Lots of moving pieces, they might catch something

[kylesayrs]

Reran job: https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14839620945

[kylesayrs]

I've validated the previously failing awq e2e test passes locally

[brian-dellabetta]

👍 the refactor from pipelines as functions to classes looks really good to me!

[rahul-tuli]

Looks good overall, left some minor comments.
One change/resolution/explanation requested for independent pipelines.

Generally I see a lot of similar TODO's trickled across multiple file(s); would like to address/delete/link out to ticket or issues before merge.

Great job on this!

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14870766869

[brian-dellabetta]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14870766869

Looks like tests passed, but there was a issue reporting timings, possibly expected for this manual run and unrelated to these changes. So I think we're good to go on this!