Pathways makes it easy to share weights between trainer and inference mesh over TPU ICI domain. This should improve weight sync performance.
Relevant docs: https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro
I'm willing to take this on but raising this issue first to see if there would be any concerns with adding a jax pathways backend.
Pathways makes it easy to share weights between trainer and inference mesh over TPU ICI domain. This should improve weight sync performance.
Relevant docs: https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro
I'm willing to take this on but raising this issue first to see if there would be any concerns with adding a jax pathways backend.