TPUでDeepfloyd IF動かしたい

TRCっていうのがあるらしいからTPUで動かしてみたい（VRAM難民）あわよくば追加学習もしたい

https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py

torchのdiffusersのやつ

とりあえずtokenizerとencoderは

https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_flax_t5.py
Flaxのやつ使えば良さそう

unet, schedularはstable diffusion用のやつからいじってくればよさそう
IFWatermarkerっていうのがあるのか

TPUならVRAMは無制限みたいなものだからCPUオフロードは考えなくていいかな

https://cloud.google.com/tpu/docs/system-architecture-tpu-vm?hl=ja

TPUではGPUでのVRAMのかわりにHBMがあるみたいで、TPU のバージョンによってこれのサイズがちょっと違います。

https://cloud.google.com/tpu/docs/supported-tpu-versions?hl=ja

多分 HBM2 の容量 というのが VRAM 的なやつだと思います。

使える TPU 構成は バージョン-コア数 の表記になっていて、v4-64 なら v4 の 64 コア TPU になります。

ただし、(少なくとも私の場合) TRC で与えられた無料枠の TPU 構成は v2-8 と v3-8 だけで、しかも謎の問題が発生(空きがない？)して v3-8 を使うことができなかったです。

もしかしたら申請するときに高性能なものを要求する必要があるのかもしれないですが、一応 TPU でも物理的な上限があるのでちょっと注意です。