<p dir="auto">Hi there. First of all, thanks a lot for making JAX such a great framework. It's awesome.</p>
<p dir="auto">Actually it's only awesome when we only do training on a single node. To my knowledge, the multi-node support is still missing for now. 10 months ago, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/hawkinsp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/hawkinsp">@hawkinsp</a> and <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gnecula/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gnecula">@gnecula</a> mentioned that there's already some experimental support for multi-GPU training across multiple devices in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="600654895" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2731" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2731/hovercard" href="https://github.com/google/jax/issues/2731">#2731</a> and <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="606736362" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2837" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2837/hovercard" href="https://github.com/google/jax/issues/2837">#2837</a> respectively. However, there's no description how we can achieve that, in a hacky way or not. I understand that experimental support could be brittle and/or rough. Just feeling adventurous and want to try out. Could someone share some idea on the multi-node training recipe with JAX?</p>
<p dir="auto">Also what is the priority and roadmap of fully supporting multi-node training with JAX? From the fact that no one's responding to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="600654895" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2731" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2731/hovercard" href="https://github.com/google/jax/issues/2731">#2731</a> any more, I assume it's not of high priority and not likely coming in the near future? Thanks again!</p>
<p dir="auto">I don't have a node with 8 gpus. I have two nodes each with 4 gpus. So is it possible to train a model on multiple nodes?</p>