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

[RyanJDick]

Summary

This PR adds utilities to support partial loading of models from CPU to GPU. The new utilities are not yet being used by the ModelCache, so there should be no functional behavior changes in this PR.

Detailed changes:

• Add autocast modules that are designed to wrap common torch.nn.Modules and enable them to run with automatic device casting. E.g. a linear layer on the CPU can be executed with an input tensor on the GPU by streaming the weights to the GPU at runtime.
• Add unit tests for the aforementioned autocast modules to verify that they work for all supported quantization formats (GGUF, BnB NF4, BnB LLM.int8()).
• Add CachedModelWithPartialLoad and CachedModelOnlyFullLoad classes to manage partial loading at the model level.

Alternative Implementations

Several options were explored for supporting inference on partially-loaded models. The pros/cons of the explored options are summarized here for reference. In the end, wrapper modules were selected as the best overall solution for our use case.

Option 1: Re-implement the .forward() methods of modules to add support for device conversions

• This is the option implemented in this PR.
• This approach is the most manual of the three, but as a result offers the broadest compatibility with unusual model types. It is manual in that we have to explicitly add support for all module types that we wish to support. Fortunately, the list of foundational module types is relatively small (e.g. the current set of implemented layers covers all but 0.04 MB of the full FLUX model.).

Option 2: Implement a custom Tensor type that casts tensors to a target_device each time the tensor is used

• This approach has the nice property that it is injected at the tensor level, and the model does not need to be modified in any way.
• One challenge with this approach is handling interactions with other custom tensor types (e.g. GGMLTensor). This problem is solvable, but definitely introduces a layer of complexity. (There are likely to also be some similar issues with interactions with the BnB quantization, but I didn't get as far as testing BnB.)

Option 3: Override the __torch_function__ dispatch calls globally and cast all params to the execution device.

• This approach is nice and simple: just apply a global context manager and all operations will happen on the compute device regardless of the device of the participating tensors.
• Challenges:
  • Overriding the __torch_function__ dispatch calls introduces some overhead even if the tensors are already on the correct device.
  • It is difficult to manage the autocasting context manager. E.g. it is tempting to apply it to the model's .forward(...) method, but we use some models with non-standard entrypoints. And we don't want to end up with nested autocasting context managers.
  • BnB applies quantization side effects when a param is moved to the GPU - this interacts in unexpected ways with a global context manager.
  • CPU tensors that get used in more than one operation would be cast to the device multiple times rather than casting once and re-using the on-device tensor.

QA Instructions

Most of the changes in this PR should not impact active code, and thus should not cause any changes to behavior. The main risks come from bumping the bitsandbytes dependency and some minor modifications to the bitsandbytes quantization code.

• Regression test bitsandbytes NF4 quantization
• Regression test bitsandbytes LLM.int8() quantization
• Regression test on MacOS (to ensure that there are no lingering bitsandbytes import errors)

I also tested the new utilities for inference on full models in another branch to validate that there were not major issues. This functionality will be tested more thoroughly in a future PR.

Merge Plan

• Partial Loading PR1: Tidy ModelCache #7492 should be merged first so that the target branch can be updated to main.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)