Closed izmailovpavel closed 9 months ago
Try by reducing your Batch size since your workaround of using a data disk and mounting it on the TPU VM seems to have alleviated the issue by reducing memory usage.
Only saw this now. Knowing Pavel, I think it's not relevant to him anymore, but for reference, here are two more options, at the slight expense of a little speed:
cache_raw
to False in the config: https://github.com/google-research/big_vision/blob/c01707f710170baf162dc217db7f1f044fc5be9c/big_vision/configs/vit_s16_i1k.py#L48cache_final
to False: https://github.com/google-research/big_vision/blob/c01707f710170baf162dc217db7f1f044fc5be9c/big_vision/evaluators/classification.py#L59
Hi! I am trying to run the vit_s16_i1k.py script on a TPU-v3-8 machine. I put the data in a google-storage bucket and I am running the following command
The training runs for a few iterations, and then fails with the
killed
message. When I look athtop
outputs, the memory used by the process grows all the way up to 335G available before the process crashes.I have been able to work around this issue by creating a data disk, mounting it on the TPU VM and putting the data there. In that case the same process only uses 205G of RAM and runs normally.