Partial Loading PR2: Add utils to support partial loading of models from CPU to GPU#7494

Merged

RyanJDick merged 13 commits intomainfrom

ryan/model-offload-2-partial-load-utils

Dec 27, 2024

Commits on Dec 24, 2024

Bump bitsandbytes. The new verson contains improvements to state_dict loading/saving for LLM.int8 and promises improved speed on some HW.
RyanJDick
committed
Add torch module autocast utilities.
RyanJDick
committed
Add torch module autocast unit test for GGUF-quantized models.
RyanJDick
committed
Simplify the state management in InvokeLinear8bitLt and add unit tests. This is in preparation for wrapping it to support streaming of weights from cpu to gpu.
RyanJDick
committed
Add CustomInvokeLinear8bitLt layer for device streaming with InvokeLinear8bitLt layers.
RyanJDick
committed
Add CustomInvokeLinearNF4 to enable CPU -> GPU streaming for InvokeLinearNF4 layers.
RyanJDick
committed
Add CachedModelWithPartialLoad to manage partially-loaded models using the new autocast modules.
RyanJDick
committed
Make CachedModelWithPartialLoad work with models that have non-persistent buffers.
RyanJDick
committed
Add CachedModelOnlyFullLoad to mirror the CachedModelWithPartialLoad for models that cannot or should not be partially loaded.
RyanJDick
committed
Fix bitsandbytes imports to avoid ImportErrors on MacOS.
RyanJDick
committed
Reduce peak memory used for unit tests.
RyanJDick
committed
Workaround a weird quirk of QuantState.to() and add a unit test to exercise it.
RyanJDick
committed
Skip flaky test when running on Github Actions, and further reduce peak unit test memory.
RyanJDick
committed