[jax 오류] INTERNAL: nvlink exited with non-zero error code 65280, output: nvlink error

d3pm code를 뜯어보던 과정에서 jax 관련 오류가 발생하였다.

INTERNAL: nvlink exited with non-zero error code 65280, output: nvlink error

대충 위와 같은 오류였는데, 디버깅을 해봐도 non-zero(?)와 관련된 부분은 찾을 수 없었다.

 

그러던 중 jax의 공식 레포(https://github.com/google/jax)에서 jax가 TPU에서 제공 된다는 것을 알게되었고, GPU 설정이 아닌 TPU 환경에서 실행해야 할 것 같다는 생각이 들었다.

 

참고로 개발 환경이 colab pro + vscode였는데,

colab - 런타임 - 런타임 유형 변경 - 하드웨어 가속기 TPU 선택 을 통해 문제를 해결할 수 있었다.

참고로 tpu는 구글이 자체적으로 개발한 칩으로, tensorflow(및 jax)에 특화되어 있다.
jax 사용 시 참고해야겠다 ..