TRC로 월 몇만원에 GPU 수십개급의.. TPU 사용 가능

TL;DR : TRC를 신청하시면 월 몇만원에 3090x4 이상급의 TPUv3-8을 5대 사용하실 수 있습니다. (사실 더 쓸 수 있습니다)

TRC를 신청하시면 TPUv3-8 (16*8GB HBM), TPUv2-8 (8*GB HBM)을 5대씩 (+TPUv2-8 24시간제한 100대) 사용할 수 있습니다.

30일만 무료로 준다고 하지만.. 연구하면서 필요할 때마다 요청하면, 제공해주었습니다. 저는 구글의 관대함을 TRC를 통해 알게 되었습니다..

pytorch에서 TPU 사용시

• pytorch_xla를 써야하지만, Pytorch Lightning, fairseq, HF Trainer API, HF Accelerate 등, TPU 사용을 돕는 library가 꽤 있습니다. (Accelerate가 제일 유연한 것 같아요.)
• pytorch는 경험상 'TPU VM (물리적 연결)'이 TPU node + VM (네트워크 연결)보다 8배 빨랐습니다. 다만 'TPU VM'은 내부 스토리지가 shm 합쳐서 200G정도밖에 안됩니다.
• TRC를 신청하더라도 그 외의 비용은 부과해야합니다. (ex. 클라우드 스토리지 : 월 3만/1TB 등)
• 단점 : 코딩&디버깅 overhead
  • TPU에서 돌아가는 pytorch는 아직은 개발되는 중이라, 예상치 못한 곳이 병목이 될 수 있고, 커뮤니티가 작아, 오류가 났을 때 도움을 얻기 어려울 수 있습니다.
  • 그러다보니 유명 library 내부 코드가 애초에 병목이 되기도 합니다 (ex. transformers의 'generate')
  • 대표적인 병목은
    • 구현되지 않은 operation : 예전에는 floordiv (//) 가 병목이 되기도 했답니다..
    • tensor shape의 변화 (ex. dynamic padding)
    • computational graph의 변화 (ex. huggingface의 layer drop 구현)
    • .item() 등의, Tensor 값을 요청하는 코드

tensorflow에서 TPU 사용시

Tensorflow에서 TPU 쓰는건 documentation도 비교적 잘 되어있어 줄입니다.

• 경험상 TPU node + VM을 써도 'TPU VM'에 비해 속도차이가 크지 않았던 것 같습니다.
• 여기서도 코딩&디버깅 overhead가 조금 있습니다
  • 예를 들면... x가 float tensor일 때, 1 / x -> error -> 1.0 / x로 고쳐야.. truediv는 1을 float64로 바꾸는데, float64는 TPU에서 지원 안됨