Closed nickk124 closed 10 months ago
EDIT 2: after that modification, now the OOM error is happening simply when trying to initialize the empty score tensor on GPU -- I'll try it on CPU, but otherwise I think this is just a limitation of my dataset size being so large. Marking as resolved :)
Hey @nickk124. Thanks for opening this issue. We are in fact in the process of fixing this! Check out PR #43 (it's still a bit rough around the edges; it'll be merged in the next version of TRAK in the coming weeks). Essentially, we won't initialize the scores
array on GPU and instead only compute blocks of it (on GPU). I'll close this issue once we merge.
@kristian-georgiev, thanks, great to hear!
Resolved by https://github.com/MadryLab/trak/commit/16e9d4627c41292a4b81a0d28962dbc42803239c. Currently in the v0.3.0 branch.
Hi,
First of all, thank you for an awesome library/paper and for the high quality documentation and support!
I'm using a very large attribution training set (> 1 M) and test set (> 50K), so I followed your tutorial at https://trak.readthedocs.io/en/latest/slurm.html to parallelize the featurization. Now that I'm at the step of score finalization, I'm getting an OOM error requiring >200 GB of GPU memory to compute this product:
I will likely try to divide up this computation into "blocks" to do separately to make it tractable, so I figured I'd raise an issue for it here too. I'm doing this using 20 models, by the way.
EDIT: it looks like this is similar to the PR here https://github.com/MadryLab/trak/pull/43. I'm going to attempt to resolve it by lowering
CUDA_MAX_DIM_SIZE
intrak.score_computers.BasicScoreComputer
.thanks!