axlearn icon indicating copy to clipboard operation
axlearn copied to clipboard

Adding support for Pathways proxy

Open jesus-orozco opened this issue 1 year ago • 0 comments

Enable JAX+Pathways single-controller architecture for coordination of accelerators.

  • Set jax_backend to "proxy"
  • Additional logic to handle new jax_backend type
  • Import previewutilities library to identify pathways-enabled workloads

jesus-orozco avatar Sep 09 '24 17:09 jesus-orozco