[Transforms] Enable shared memory and introduce permutations

[dsikka]

Summary

1. Makes a series of updates to the registry by supporting matrix caching of the transform parameter

• With these changes, instead of generating hadamards for a particular size repeatedly, cache the matrix using its size as the key. This significantly speeds up transform set-up as by leveraging the property of hadamards being deterministic
• For randomized hadamards, splits up the math so that the underlying hadamard can be cached and the randomness is introduced as a separate permutation matrix [will need to add a few tests to make sure my math is correct]
• For the matrix-multiply transform, caches based on a user provided name

2. Also uses shared memory, so that layers with identical transforms use the same underlying transform data.

• Significantly reduces the memory required by transforms
• NOTE: For training, we may need to make updates depending on how transform updates are expected to be made during training as we are now using shared data

3. Move update/register functionality to be done inside the registry; introduce permutation parameter

• The update and register steps are identical as before but now happen inside the registry for slightly more clarity
• Introduce the permutation parameter. Only being used for the random-hadamard for now but will follow up with the remaining transfoms as well

4. Swap global to be called "shared"

[brian-dellabetta]

i think this makes sense 👍

[brian-dellabetta]

it seems like we have to do some sort of name matching, but i'm wondering if some name collision down the road is going to cause this to run when we don't want it? if we came up with a more unique name or something to prevent false positives

[dsikka]

we have this problem with any parameter we introduce (e.g weight_scale, weight_g_idx, etc) but yeah, we can work on making them more unique

[brian-dellabetta]

do we need update_parameter_data here too in case of offloading?

[dsikka]

this is in the case if the parameter isn't registered to a module for whatever reason. update_parameter_data handles module_params. I'm not sure if this case is totally necessary but yeah, we would have to add offloading/onloading around it if we decide to keep it

[brian-dellabetta]

this comment might need to go too if we are not normalizing?

[brian-dellabetta]

so that all matrices live in a single global key-value store, right?

[dsikka]

yeah, we can expand this but it seems like there will be a lot of repetition across decoder layers for example
I think if this goes too big in scope, we may have to consider other data stores to handle it

[kylesayrs]

Is a new class singleton necessary? Since each matrix Transform requires own registry, shouldn't this be implemented on the class itself?

class Hadamard(Transforms):
    registry: Dict[int, torch.Tensor] = {}

    def __new__(cls, size, empty, transform_name, *args, **kwargs):
        if empty:
            matrix = ...
        else:
            matrix = cls.registry.get(size, torch.Tensor(deterministic_hadamard_matrix(size=self.size)))
        return super().__new__(transform=matrix, transform_name=transform_name)

[kylesayrs]

For randomized hadamards, splits up the math so that the underlying hadamard can be cached and the randomness is introduced as a separate permutation matrix

Since permutations are specific to RandomHadamards, shouldn't we be implementing this logic on the RandomHadamard class, not the general Transforms class?