google cloud中使用TPU

首先,应该先尝试google colab —— 因为直接注册google cloud中的VM和TPU,只能免费用3个月,注册早了实在是亏了,等colab里试得差不多了,再注册cloud不迟。

正式使用的步骤:

  1. 创建VM,要注意选区,目前只有几个区支持TPU,所以最好同时把创建TPU的界面打开对照着看,不同区之间,要想访问只能通过公网了,这对于机器学习的密集运算显然不可行
  2. VM的核越多,内网带宽越大,所以有钱的情况下,尽量选多核的
  3. VM的抢占式虽然便宜,但一旦被抢占,所有状态就丢失了,慎选
  4. VM的磁盘可以放大一些,因为需要保存训练中的snapshot,大模型建议300GB

创建TPU

  1. v3比v2贵,但性能也高,所谓v3-8, v2-8,都是一块TPU(内置8个单元,是最小单位了),慎选v3-128 v3-512这种配置,看上去价格没变,实际上标的是单价,最终收费会成倍上涨
  2. TPU可以选择抢占式,因为即使被抢占,下次还可以从快照继续,而且价格便宜好多
  3. 总体来说,TPU比VM贵多了,有数量级的差异

?VM和TPU怎么结合呢?
貌似要在VM里创建一个指向TPU的配置:

export TPU_IP_ADDRESS=10.2.3.2;  # 这个ip是tpu的ip,通过gcloud compute tpus list --zone=europe-west4-a命令可以看到
export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470";

这里还还有待进一步学习探索。

如果图省事的话,创建TPU的之后,直接选择TPU虚拟机(而不是TPU节点)即可,这时会得到一个VM+TPU一体的机器,可以支持ssh登录访问,连TPU版本的tensorflow或者pytorch(看选择哪个版本的操作系统)都是预安装好的,直接运行自己的代码即可。

你可能感兴趣的