Closed kristian-georgiev closed 1 year ago
Projecting and storing in float16 by default https://github.com/MadryLab/trak/commit/ec068c375e116304294e8180ce3dbf5431ac88c4
Would be nice to also add mixed precision by default. I'm not sure how to combine torch.func.functional_call
with torch.autocast
tbh. Maybe the easiest way would be to add an @autocast()
decorator to get_output()
.
Another thing from the AMP docs (e.g. https://pytorch.org/docs/stable/notes/amp_examples.html):
Backward passes under autocast are not recommended.
As far as I can tell, otherwise training becomes unstable, because there's a lot of underflow if you store gradients in float16
. So the usual AMP training accumulates float32
gradients. But here we're projecting with float16
gradients anyways, so I think we could actually wrap both the forward and backward passes within autocast()
?
On a second thought, if we're doing everything in float16
, we might as well skip autocast()
altogether and just .half()
the model. My main concern is that this could be more numerically unstable than calling .half()
on the gradients after we compute them in float32
. Given that compute is dominated by the JL projection, I'll leave things as they are rn, but we can re-open if we think it's worth it.
Added a small script trying out a few options https://github.com/MadryLab/trak/commit/bc627e0e84e415d949a1b36c4dbca41d4be9b77a.
Rough idea: use mixed-precision for gradients by default; only instantiate
float16
gradients and pass them to the projectors; for BasicProjector, do everything infloat16
, for CudaProjector, conver the output tofloat16
; save scores infloat16
to disk.Finally, we should have a flag in
TRAKer
along the lines ofuse_float32: bool = False
.