Scale JAX models to multi-GPU systems // TRAIN BRAIN

Scale JAX models to multi-GPU systems

Join the Google Cloud & NVIDIA community → https://g.dev/cloud/google-nvidia-community
Scaling deep learning models across multiple GPUs used to mean rewriting hundreds of lines of complex device communication code.
This is part 3 of JAX on NVIDIA GPUs Crash Course.
Watch along and learn about JAX's modern, compiler-driven sharding model, demonstrating how to distribute workloads automatically across physical device meshes.
* *Master sharding concepts:* Understand how Mesh, PartitionSpec, and NamedSharding declare array layouts across multiple devices. * Implement automatic scaling: Write clean training code that allows the compiler to automatically manage multi-GPU gradient synchronization.
* *Incorporate Flax NNX & Orbax:* See how to manage state and serialize model checkpoints in a distributed training run.
Watch more JAX on NVIDIA GPUs Crash Course → https://g.dev/cloud/jax-nvidia-gpu
? Subscribe to Google Cloud Tech → https://goo.gle/GoogleCloudTech
Speakers: Ivan Nardini, Ekaterina Sirazitdinova
Products Mentioned: Google Cloud, JAX

Google Cloud Tech

Helping you build what's next with secure infrastructure, developer tools, APIs, data analytics and machine learning....

Graph Engineering 101

How to monitor & optimize Google Gemini agents with Grafana Cloud

End-to-end agentic development with Google Gemini, Antigravity, and GitLab

He built a voice controlled drone to avoid using a remote! ?✈️

Don't wait to graduate: A student’s AI ancestry storyteller

AI-powered observability & remediation with Google Gemini CLI and Dynatrace

?‍? Martin Omander takes over the Google Cloud Tech channel

Stop letting AI write code until you fix this one thing (with Antigravity) | The Agent Factory

Automate Google Cloud with Cloud CLI Remote MCP Server

Launch your portfolio website with AI

Launch your to-do web app with AI

Launch your web-based video game with AI

7 AI agent patterns to improve your coding workflow

What "real-time voice AI" actually means

Build a real-time voice AI agent with Gemini Live API

From notebook to production: Serving JAX at scale

Scale JAX models to multi-GPU systems

Build and optimize JAX training loops

Getting started with JAX on NVIDIA GPUs

4 ways loop engineering fails (and how to fix them)

Building apps on Google Cloud

This or that? Imaxxing edition

Real-time voice AI agents, explained (before you build one)

Preview: building a viral movie app with Temporal & Cloud Run

Building a viral Odyssey movie app with Temporal & Cloud Run

This AI agent is raising millions for startups | The Agent Factory

Make your website agent ready with WebMCP

Data agent kit: Your coding agent can now query your data

What is Spanner Omni?

Let /goal take the wheel ? ?️ ?

Architecting production-ready AI apps with Google Cloud & Firebase

Build next-gen AI agents in the All things agentic hackathon

What is tokenmaxxing?

How to vibe code securely (without getting hacked)

?‍? Annie takes over the Google Cloud Tech channel

From tokenmaxxing to tokenomics for your AI agents

Skills vs MCP: How AI tools have evolved

Reviewing code from AI coding agents

Senior Engineer's best practices for using coding agents

11 tips for Claude code on Google Cloud

How to design a game with Claude code on Gemini Agent Platform

AlloyDB: The ultimate hybrid search engine

Build a BigQuery AI Agent with ADK & Cloud Run

AI builder essentials: Tokens, Context Windows, RAG 101

Learn Google's strategy to ship AI prototypes faster

How to find & fix code vulnerabilities autonomously with Google CodeMender

Network Security Integration with Check Point - Demo

Emergent: How Google engineers actually ship AI (official trailer)

Give your app search superpowers: Agent Retrieval (Vector Search 2.0)

What is an Agentic Harness?

Context engineering explained: What every AI developer should know

Is your Kubernetes certification still worth it? | The Agent Factory

Automate your entire SDLC with Google Antigravity

BigQuery Agent Analytics Looker

What is the biggest AI agent controversy today? ?

Agile developers: shifting down with agentic skills

20 minutes to 2 minutes: How Home Depot automated incident responses with OpenTelemetry

Is it possible to fix K8s degradation in one click with Gemini Cloud Assist?

Architecture for running agents in production

Building an MCP-powered autonomous incident response ecosystem

How BBVA manages 1,000+ GCP projects

Inside Google’s new AI hardware

How to deploy an AI agent to Cloud Run (step by step)

Debugging an AI agent like a football match ⚽

Intent-driven development with Claude Code & Fable 5

Make AI agent execute 80% faster ⏱️?

MCP vs API: The protocol every developer needs to know

Breaking silos with agentic AI: How Google Cloud transforms enterprise data into real-time action

Build your first Android app in AI Studio in 5 minutes

Build a multi-agent system using ADK & MCP

Build a full stack app with Antigravity voice prompts & sub-agents

Building a dynamic personal site with Antigravity and Gemma 4

Build a multi-agent system: A2A & Agent Registry

Automating workflows and file organization with Antigravity

5 things every AI agent harness needs

What's new with data agents

Power intelligent agents with AI-native databases

Boost AI context with hybrid search in Spanner

Agent development and AgentOps with BigQuery, ADK, and MCP

Build connected AI: Orchestrate tools and agents with registries and ADK