feat(autogram): Add ModuleBasedGramianComputer.

[PierreQuinton]

No description provided.

[ValerianRey]

This is big for architectures with very big linear layers, like AlexNet.
For AlexNet, on cuda, with batch_dim=0, this leads to:

• Double max batch size (from batch_size=19 to batch_size=38 on my gpu). Sadly, this is still very far from the max batch size of 1268 of SGD with autograd). Do you think there might be a theoretical way to bridge this gap even more? EDIT: I have no idea why, but re-running the same tests (or maybe I did a mistake before) yields completely different results. Max batch size is now 18 for main and 468 for this PR - much, much closer to the 1268 of autograd). EDIT2: the big memory improvement (and a small speed improvement) comes from just installing opt_einsum (without even changing the code). uv pip install opt_einsum.
• x2 to x4 speed (depending on the batch size) of the whole autogram_forward_backward function (so this includes not only the gramian computation of the linear layers, but also of all other layers + the forward passes).

For other architectures, the differences are not very noticeable though. But this is very promising. Let's fully focus on this direction IMO.

[ValerianRey]

This can be replaced by:

G_W = oe.contract("ik,il,jl,jk->ij", dY1, X, X, dY2, optimize="optimal", backend="torch")

with import opt_einsum as oe
but it seems to be the exact same runtime and memory usage.

[ValerianRey]

Actually whenever opt_einsum is installed, the contraction is already done even without changing the line:

G_W = torch.einsum("ik,il,jl,jk->ij", dY1, X, X, dY2)

We could still add the line just to make it explicit maybe.

[PierreQuinton]

whatever you prefer. I prefer not having to give the two additional parameters, for me what is important here is

1. It is an einsum
2. It is fast

But the second criteria is more of an "how" than a "what" so we don't really need to know. For this reason I would vouch slightly for torch.einsum. The negative part is that a user could set the global settings of opt_einsum to non-optimized thereby making it slow, but I guess that is the user's responsability.

[ValerianRey]

There's actually no guarantee that X, dY1 and dY2 are matrices.

From the documentation of nn.Linear:

In particular, when there is no batch dim, I think the * dimension could be empty, and in transformers, the * dimension is (batch_size, seq_length), which is why transformers fail with this PR.

[ValerianRey]

I made it work on Transformers with that:

if dY1.ndim == 1:
    G_b = torch.einsum("k,k->", dY1, dY2)
    G_W = torch.einsum("k,l,l,k->", dY1, X, X, dY2)
elif dY1.ndim == 2:
    G_b = torch.einsum("ak,ik->ai", dY1, dY2)
    G_W = torch.einsum("ak,al,il,ik->ai", dY1, X, X, dY2)
elif dY1.ndim == 3:  # Typical in transformers
    G_b = torch.einsum("abk,ijk->ai", dY1, dY2)
    G_W = torch.einsum("abk,abl,ijl,ijk->ai", dY1, X, X, dY2)
else:
    raise ValueError("Higher dimensions not supported. Open an issue if needed.")

Not elegant at all but it seems to work. Maybe there's a clean way to write this that works for any number of dimensions without having ifs. Also, please review the equations. I did them basically with trial and error until the tests passed.

[PierreQuinton]

Well it needs to be at least matrices (2<=ndim) as we know it's a batched scenario. However, We could in principle add the no batched dimension scenario, but I'm not sure it would be faster than the classical Jacobian based GramianComputer.

I did them basically with trial and error until the tests passed.

I which I could have done that ^^

[ValerianRey]

Also pretty big for Transformers (with the change i suggested to handle higher order tensors).

Times for forward + backward on WithTransformerLarge with BS=256, A=Mean on cuda:0.
Reduced from 3.13 sec (main) to 2.20 (this PR)

Memory is however increased. Max batch size went from 273 (main) to 256 (this PR).

[ValerianRey]

This seems to break NoFreeParam (tiny errors) and ModuleReuse (large errors). Need to investigate that.

For ModuleReuse, my guess is that it simply doesn't consider cross terms anymore, so it's normal that it fails the test.

[PierreQuinton]

Of interest: https://optimized-einsum.readthedocs.io/en/stable/reusing_paths.html
If we can try to explore what contraction is the optimal and if it is essentially always the same, then we may want to use a self forged contraction. It would be very helpful to know the optimal contraction order.

[ValerianRey]

@PierreQuinton I found a way to compute the gramian with autograd with no cross terms from module reuse / inter-module param reuse: 30fdc00. Basically, the idea is to have a module pre-hook that clones each parameter before using them, and a module post-hook that restores to the original params. This way, each module usage corresponds to a different clone, and you can compute a gradient wrt each clone. The implementation is with a context manager so it's quite clean IMO. Current limitations:

• Does not work on WithMultiheadAttention, WithTransformer and WithFreeParam, because they all involve some indirect parameters for the hooked module. Need to investigate and fix that (it's probably doable).
• Still counts cross-terms from intra-module parameter reuse: we'd need a node-based algo (rather than module-based) to fix that. But since autogram is still module based, it doesn't matter yet.