einsum problem with linear grad sample #242
Conversation
facebook-github-bot
added
the
CLA Signed
opacus/grad_sample/linear.py
Outdated
| # With activations with 3 dimensions (transformers) the original formula does not aggregate correctly | ||
| if activations.ndim == 2: | ||
| gs = torch.einsum("n...i,n...j->nij", backprops, activations) | ||
| else: |
There was a problem hiding this comment.
Noob question: I thought ellipses in einsum handles all dimensions, and by extension even 3D. Do you know why it didn't work?
There was a problem hiding this comment.
After careful investigation today, it appears that its linked to a previous version fo PyTorch that we used in our environment (specifically 1.7.1). The bug appears to be solved for 1.9 since we could not reproduce for this version.
My suggestion would then be either (1) not to merge the PR but update the requirements for PyTorch or (2) merge the PR so that we have backward compatibility with older PyTorch versions.
There was a problem hiding this comment.
Updating PyTorch requirements past 1.8 is problematic, because we depend on csprng, which doesn't support 1.9 and beyond.
Is the bug still reproducable in 1.8?
There was a problem hiding this comment.
Good news, bug disappears for 1.8.0. Let's update the requirements then?
There was a problem hiding this comment.
Let's update the requirements then?
+1
There was a problem hiding this comment.
Updated requirements.txt with torch>=1.8. Sorry for the multiple commits to revert grad_sample/linear.py.
There was a problem hiding this comment.
@ffuuugor Are you sure csprng doesn't support torch-1.9 (in it's current state). Did you see my comment? Just checking. I didn't test it deeply.
There was a problem hiding this comment.
You're probably right - but I haven't looked too deeply either.
I was only relying on the fact that PyPI release depend on 1.8 and their documentation explicitly lists all pyTorch versions they support, notably lacking 1.9 and beyond.
Is the fix as simple as installing from main branch instead of the latest release?
Make linear einsum compatible in the case of 3D activations (transformers use case). Maybe there is a way of getting rid of the
if/elsestatement.