An issue regarding multi-node training with TF code about ferminet HOT 4 CLOSED

We have not tried MultiWorkerMirroredStrategy -- getting everything to work well with MirroredStrategy was very involved but sufficient for all published results. We have successfully done multi-working training using TF-Replicator, so in principle it should be achievable.

Are you trying with ADAM or K-FAC? For any porting work, I would suggest first getting ferminet working with ADAM and then investigate K-FAC. I added support for MirroredStrategy to K-FAC but don't know if MultiWorkerMirroredStrategy will work out of the box. I certainly hit many problems with tensor placement and naming getting MirroredStrategy to work.

Re: JAX multi-node support, please see google/jax#2731. We're not intending to continue developing the TF-version of ferminet.

from ferminet.

Got it, thanks a lot for the info and suggestion @jsspencer ! Let me try harder then.

I was indeed trying KFAC, but it seems to me the error was thrown way before optimization starts.

From the github thread you mentioned about JAX, it seems its multi-node support is still under development. Let me ask them if they have anything to share now.

Thanks again!

from ferminet.

These kinds of errors are normally triggered at graph-construction time rather than (later) at training time.

from ferminet.

These kinds of errors are normally triggered at graph-construction time rather than (later) at training time.

You are absolutely right. I've managed to fix this issue and successfully run the TF code with MultiWorkerMirroredStrategy on multiple nodes (modulo the fact that using kfac would cause core-dump).

from ferminet.