We've moved up in the world! These days we'd train on one of our TFRC TPU pods. Colab is a lot of trouble and better avoided...

(We actually can create a TPU-2048 pod. For a short while, anyway, before it preempts. But we wouldn't need to if we used lucidrain's efficient attention implementations I mention in my other comment.)