-
Notifications
You must be signed in to change notification settings - Fork 22k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
svd_backward: more memory and computationally efficient. #50109
Conversation
💊 CI failures summary and remediationsAs of commit 8b34f0b (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 This comment was automatically generated by Dr. CI (expand for details).Follow this link to opt-out of these comments for your Pull Requests.Please report bugs/suggestions to the (internal) Dr. CI Users group. This comment has been revised 11 times. |
L = L - L.conj(); | ||
Tensor imag_term = 0.5 * at::matmul(u * (L * sigma_inv).unsqueeze(-2), vh); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: Could it be better to replace L
with just 2 * L.imag()
and use the float-float mul
in L * sigma_inv
instead of the float-complex one, then multiply by the imaginary unit in the most efficient way?..
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's a simple change that might give better performance, so I think it's worth adding it to this PR.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also you don't need to do the 2*
if you remove the 0.5*
right? :D
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, so one more reason to follow the nit.
Resolves #38353. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It looks much better now! For completeness could you do some benchmarks to know how much performance and memory saving do we gain here?
L = L - L.conj(); | ||
Tensor imag_term = 0.5 * at::matmul(u * (L * sigma_inv).unsqueeze(-2), vh); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That's a simple change that might give better performance, so I think it's worth adding it to this PR.
@@ -1947,9 +1944,12 @@ Tensor svd_backward(const std::vector<torch::autograd::Variable> &grads, const T | |||
|
|||
if (gu.defined()) { | |||
auto guh = gu.conj().transpose(-2, -1); | |||
u_term = at::matmul(u, at::matmul(F.mul(at::matmul(uh, gu) - at::matmul(guh, u)), sigma_mat)); | |||
u_term = at::matmul(u, (F.mul(at::matmul(uh, gu) - at::matmul(guh, u)) * sigma.unsqueeze(-2))); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit there is an extra parenthesis here that is not needed for the second argument of the outer most matmul:
u_term = at::matmul(u, F.mul(at::matmul(uh, gu) - at::matmul(guh, u)) * sigma.unsqueeze(-2));
L = L - L.conj(); | ||
Tensor imag_term = 0.5 * at::matmul(u * (L * sigma_inv).unsqueeze(-2), vh); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also you don't need to do the 2*
if you remove the 0.5*
right? :D
…ed/svd_complex_backward_minor_fix
…ed/svd_complex_backward_minor_fix
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM! Thanks for the update!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
…ed/svd_complex_backward_minor_fix
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
As per title.
CC @IvanYashchuk (unfortunately I cannot add you as a reviewer for some reason).