ringattention icon indicating copy to clipboard operation
ringattention copied to clipboard

Published 20 hours ago • verified publisher icon forhaoliu

vmem OOM on TPU

Open hxssgaa opened this issue 1 year ago • 1 comments

Hi,

I tried to run your script on Cloud TPU v4-64, but failed with following error:

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space vmem. Used 59.79M of 16.00M vmem. Exceeded vmem capacity by 43.79M.

I tried mesh dim of 1,1,1,32 and 1,1,4,8 all failed.

Any suggestion what caused the error? Thanks.

hxssgaa avatar Dec 11 '23 09:12 hxssgaa