Pipeline Extraction

[kylesayrs]

Purpose

• Extract data pipelines from modifiers to enable multiple modifiers to be active at the same time

TODO

• Extract pipeline from quantization modifier
• Compress any modules that didn't receive samples on finalize (accounts for cases like moes where some experts might receive zero samples)

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
• 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
  • Originally, these compression events were triggered by reaching the end of each module’s initialize function. Now a separate event is required
• Implement session.initialize_recipe
  • This is required because, without this, modifiers cannot be created until initialize is called, and initialize cannot be called until the data pipeline is decided
  • This change aligns with the LLM Compressor design where modifiers are created when the recipe is passed
• 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
  • Replace uses of _hf_hook. pre_forward with align_module_device for clarity
  • Specify resolved_mappings_ type hint for clarity
  • 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

[brian-dellabetta]

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

[brian-dellabetta]

If we make pipelines a class satisfying an abstract Pipeline base class with

@abstractmethod
def run_pipeline(model:PreTrainedModel, dataloader:DataLoader):...

would it avoid needing maps like this or types like

PipelineFn = Callable[[PreTrainedModel, torch.utils.data.DataLoader], None]

in the typing.py file?

[kylesayrs]

I purposefully avoided adding a base class, since I think it adds more infrastructure than is required. Such a class would only have one method, which imho doesn't justify a class definition.

[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?