[Transform] Better dispatch support for offloaded and multi-gpu

[kylesayrs]

Purpose

• In order for vLLM to properly load transform weights, it must be capable of constructing sharing tensors in a way that is independent of which device map was used during compression by LLM Compressor.
  • Right now, the construction of shared tensors is dependent on the device of the parent module, which is determined by the device map used by LLM Compressor
  • If LLM Compressor decides to generate transforms on a model which is split on two GPUs, the transform will be generated on one GPU and then moved to the other GPU at runtime
• Construct weights based on the precision of the scheme, not the precision of the weight
  • This also simplifies the key used by vLLM loading
  • The dtype will still work due to how we how upcast dtypes when applying weights
• Support creating transforms for both offloaded (sequential) and multi-gpu (basic) dispatches

Changes

• Use scheme.precision rather than module.dtype when constructing parameters
• Do not support different devices when constructing transform weights. Instead use first device seen (typically CPU) and ensure future devices match that device
• Fix get_offloaded_device in the case that the module is not offloaded (such as attention)
• Add a TQDM for transforms

Testing

• Transform precision and shared tensors tests pass
• Ran quip example to completion with sequential and basic pipelines
• Quip correctness tests pass

[brian-dellabetta]

approving with question

[brian-dellabetta]

thanks!

[dsikka]

Instead use first device seen (typically CPU) and ensure future devices match that device

How is this ensured?

[kylesayrs]

@dsikka This is ensured by moving the weight to the device of the value. https://github.com/neuralmagic/compressed-tensors/pull/423/files#diff-be313d6f55c99277b8d747f1d5470f9cf0f08d99ffdb1fc45ac8df80f8784c59R114

This to call only has an effect when the model is dispatched to multiple GPUs. This means that multi-gpu transforms will be supported, but will have a runtime cost. Our examples only focus on single-gpu/ offloaded transforms for now.