This repository was archived by the owner on Apr 23, 2025. It is now read-only.
Transformer model requires more parameters than supported on TPU #638
Description
It has been pointed out by Wojtek Czarnowski that in specific cases the Transformer model (or components used within it) can trigger a compilation error in X10 on TPU:
2020-07-16 22:51:03.077357: F tensorflow/compiler/xla/xla_client/xla_util.cc:90] Invalid argument: From /job:tpu_worker/replica:0/task:0:
Computation requires more parameters (333) than supported (limit 237).
[[{{node XRTCompile}}]]
Current stack trace:
frame #17: 0x00007f6da8c0ceb2 $__lldb_expr102`partial apply for closure #1 in update(model:using:for:) at <Cell 14>:12:9
frame #23: 0x00007f6da8c0c268 $__lldb_expr102`update(model=<unavailable>, optimizer=<unavailable>, batch=<unavailable>) at <Cell 14>:4:18
frame #24: 0x00007f6d5000a483 $__lldb_expr132`closure #1 in at <Cell 19>:20:31
frame #25: 0x00007f6da48245b7 libjupyterInstalledPackages.so`time(repeating=1, f=0x00007f6d50009230 $__lldb_expr132`closure #1 () -> () in __lldb_expr_131 at <Cell 19>:4) at timing.swift:15:9 [opt]
frame #26: 0x00007f6d5000914b $__lldb_expr132`main at <Cell 19>:4:1
He provided a reproducer notebook which can be opened and run in Colab. Choosing a GPU-backed instance lets this succeed, but running this notebook with a TPU-backed instance triggers the above crash.