[Transform] Factory classes with shared memory and offloading

[kylesayrs]

Purpose

• Support applying transforms to models while maintaining maximal shared memory across transform applications

Prerequisites

• [Transform] Accelerate Utilities #328
• [Transform] Hadamard and Matrix Transform Utils #330

Transform Factory and Submodules

• Transform weights are cached using the ParameterizedDefaultDict, typically keyed by matrix size, dtype, and device
• Transforms are defined as submodules which apply their weights to a module
  • Transforms are applied to modules either through hooks (activations) or updating the module’s weight
• Transforms are added as submodules to the modules they are applied to
  • This enables easy serialization with a lora-style checkpoint structure
  • This allows easy utilization of pytorch's parameterization feature for learning transform during training
  • If the parent module is offloaded, then the transform submodule also becomes offloaded

Transform Types

TransformFactory

• Base transform factory class
• Each transform factory corresponds to one transform scheme in the transform config
• Subclasses must implement create_transform and are responsible for caching/sharing memory
• Implements apply_to_model, which applies transform created via create_transform to target modules in the model

HadamardFactory

• Creates hadamard transforms
• Weights are cached according to matrix size, dtype, and device
• In the future, will also implement permutations, see [Transform] Hadamard Permutations #329

RandomMatrixFactory

• Creates random matrix transforms
• Weights are cached according to matrix size, dtype, and device
• In order to save runtime, inverse are also cached

RandomHadamardFactory

• Subclasses HadamardFactory, but uses random_hadamard_matrix to create matrix weights

Testing

• Add tests for correctness when applying transforms across models
• Add tests for transform weight memory sharing (both with and without offloading)