[Performance] Sequential onloading

[kylesayrs]

Purpose

• Reduce hardware requirements when calibrating large models by only onloading one layer at a time when calibrating using the sequential pipeline
• Updating the examples can be done after pipeline extraction lands. Examples which only use the basic pipeline should dispatch to "auto", while examples which use GPTQ should dispatch to the cpu and set oneshot_device.

Usage

When using the sequential pipeline, a few behaviors change

• If your model is dispatched to the gpu (has parameters which execute on a gpu), then a warning is raised

logger.warning(
    "Calibrating a model dispatched to the gpu can potentially lead to OOM "
    "errors. Consider loading the model without a `device_map` and instead "
    "executing with `cuda:0` (set `oneshot_device` to override this default)"
)

• Otherwise (if you model is dispatched to the cpu), then the oneshot_device argument is used to determine the onload device (this defaults to cuda if a cuda device is available)

elif oneshot_device is None:
    has_cuda = torch.cuda.is_available()
    oneshot_device = torch.device("cuda:0") if has_cuda else torch.device("cpu")
    logger.info(f"No oneshot_device passed, using {oneshot_device}")

This policy encourages users to dispatch to the CPU when using the sequential pipeline, and to dispatch to "auto" when using the basic pipeline

Changes

• Keep layer parameters onloaded during the entire sequential calibration + compression step

Testing

• Calibrated and GPTQ-compressed one layer of Deepseek-V3 with a single H100 in 50 seconds
  • 4.5x Improvement over original 236 seconds
  • Peak memory of ~40 GB, which can be further reduced by increasing the granularity of sequential targets
• Not offloading activations did not result in a performance improvement

[github-actions]

👋 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.

[dsikka]

Why not keep this in compressed tensors with the other cpu offloading tools?

[kylesayrs]

Yep! Implementing here before the next CT release

[kylesayrs]

neuralmagic/compressed-tensors#282

[dsikka]

Seems like all we're doing is wrapping the forward passes in this context manager, if I'm reading this correctly?

[kylesayrs]

Yes. Rather than onloading then discarding for each of the 512 forward passes, we onload once for the layer and keep it onloaded through compression and propagation.

[dsikka]

Do you mind explaining what we're changing in our graph partition here?

[kylesayrs]

The graph partition doesn't change, this change just collects all the modules used by this subgraph for use in onloading/offloading by the sequential pipeline.

[brian-dellabetta]

sorry, i approved this thinking it was the one-liner removing clear-ml, will have to take a closer look

sorry, i approved this thinking it was the one-liner removing clear-ml, will have to take a closer look

[brian-dellabetta]

I am understanding this for the most part -- very cool!