kingoflolz / mesh-transformer-jax

Model parallel transformers in JAX and Haiku

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

TPU Instance Creation

zzj0402 opened this issue · comments

Please elaborate on the TPU creation process. Which TPU to use? Which version of the software to use? I am getting not enough space error.

Is the error you are getting something like this:
"jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 32.00MiB (33554432B) on device ordinal 0: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well)."
If so, I have the same problem. It doesn't seem to matter what size dataset I use...

Is the error you are getting something like this: "jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 32.00MiB (33554432B) on device ordinal 0: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well)." If so, I have the same problem. It doesn't seem to matter what size dataset I use...

No. I get a dependency installation error. The Google TPU VM cannot install the requirements. I haven't gone that far just yet.