Catalog Details

  • CATEGORY

    deployment

  • CREATED BY

  • UPDATED AT

    May 03, 2024

  • VERSION

    1.0

Pattern Snapshot

Related Patterns

deployment
Fault-tolerant batch workloads on GKE
MESHERY4b55

What this pattern does:

JAX is a rapidly growing Python library for high-performance numerical computing and machine learning (ML) research. With applications in large language models, drug discovery, physics ML, reinforcement learning, and neural graphics, JAX has seen incredible adoption in the past few years. JAX offers numerous benefits for developers and researchers, including an easy-to-use NumPy API, auto differentiation and optimization. JAX also includes support for distributed processing across multi-node and multi-GPU systems in a few lines of code, with accelerated performance through XLA-optimized kernels on NVIDIA GPUs. We show how to run JAX multi-GPU-multi-node applications on GKE (Google Kubernetes Engine) using the A2 ultra machine series, powered by NVIDIA A100 80GB Tensor Core GPUs. It runs a simple Hello World application on 4 nodes with 8 processes and 8 GPUs each.

Caveats and Consideration:

Ensure networking is setup properly and correct annotation are applied to each resource

Compatibility:



Recent Discussions with "meshery" Tag