[Transform] Accelerate Utilities

[kylesayrs]

Purpose

• Prerequisite to [Transform] Factory classes with shared memory and offloading #316

Changes

• Implement register_offload_module
  • Used to register transform modules during apply
• Implement delete_offload_module
  • Used to delete transform modules during serialization
• Implement force_cpu_offload
  • Used for testing
• Misc: Add missing fallbacks for disable_offload and align_modules

Testing

• Add tests for register_offload_module, delete_offload_module, and force_cpu_offload

[brian-dellabetta]

I'm assuming force name means that takes precedent over any other context, e.g. force_cpu_offload overrules disable_offload, and the module/tensor will be moved to CPU?

[kylesayrs]

@brian-dellabetta Not necessarily. I use the word "force" because dispatch_model does not always cause tensors to go to the meta device (become offloaded). For example, passing a device map with just one device will not trigger offloading for small models.

This method "forces" the tensors to go onto the meta device by nature of specifying a very granular device map (dispatch_model is a weird method :/ )