← All transcripts

From notebook to production: Serving JAX at scale Transcript, AI Summary & Key Points

Google Cloud Tech · 2 hours ago · Science & Technology · 03:05 · EN-US

AI Summary

Serving a model from a notebook can cause a long compilation delay on the first request. Four serving patterns are presented: in-process jax.jit for Python services, ahead-of-time compilation for predictable startup behavior, jax.export for portable JAX artifacts, and jax2tf for TensorFlow-based serving. The deployment target determines which serving wrapper to use, while the checkpoint remains the same.

Key Points

  • A model that works in a notebook can underperform when served because the first request may trigger compilation.
  • jax.jit is the simplest in-process option for a Python service, but its first call may compile; warming it up before real traffic can avoid the initial latency spike.
  • Ahead-of-time compilation separates StableHLO lowering from executable compilation so server calls do not incur the first-call compilation spike.
  • StableHLO is the lowest-level program seen by the compiler and can be inspected for debugging or to confirm what was staged.
  • Different input shapes may compile separately, while larger batches can improve tokens per second by amortizing overhead.
  • jax.export serializes a StableHLO-based artifact that can be written, read back, executed, and checked for matching predictions.
  • jax2tf converts JAX graphs into standard TensorFlow SavedModels for TensorFlow-based serving infrastructure.
  • The same checkpoint can be used across deployment targets; the serving wrapper changes.

🔒 3 more in the full analysis

AI in practice

Used for

Model inference serving
What
Deploy a JAX model on a GPU with predictable latency and suitable runtime integration.
Serve a rebuilt transformer model in a Python service after loading its checkpoint parameters onto the serving device.
With
JAX
How
Rebuild the model structure, load the checkpoint parameters, place them on the serving device, and use JAX jit for in-process execution.
Outcome
The model can serve predictions, but the first call may incur compilation.
Compile a JAX model ahead of serving so requests execute without a first-call compilation spike.
With
ahead-of-time compilation
How
Lower the model to stable HLO, compile the resulting representation into an executable, and have the server call the executable.
Outcome
Server calls execute without the first-call compile spike.

🔒 3 more in the full analysis

From this video

6 products

JAX jax2tf jax.export StableHLO TensorFlow XLA

Links mentioned

🔒 Full analysis locked

Unlock more videos and the full analysis

Buy credits to process more videos. Each run includes the full analysis, not just the summary — and you get access to the locked analysis across the library.

Inquire for details →

Transcript

Searchable transcript of From notebook to production: Serving JAX at scale — Google Cloud Tech (03:05). Search for a phrase, then click its timestamp to jump straight to that moment in the video.

Captions sourced from the original video on YouTube, published by Google Cloud Tech. The video, its captions and all related intellectual property remain the property of their respective owners; AINotes claims no ownership. Provided for research, accessibility and search — see the Transcript Notice and Copyright Policy.

00:00 A model that works in a notebook can still underperform when you serve it. The first user asks for a prediction and the server goes into a long compilation who pays the compile cost. >> In the last video, we saved a train transformer checkpoint. Now we rebuild the model structure, load the orbox parameters and place them on the serving device. >> There are different patterns you can consider.

00:25 Jax G for a Python service. Hi head of time compile for startup time compilation. Jax export for portable Jax artifacts and JAX 2TF when the serving stack is TensorFlow based. >> First Jaxjit is the simplest in process pass. It's great for Python service but the first call may compile. Consider warming it up before real traffic. Second option is ahead of time compilation which makes the stage explicit.

00:58 Lower produce stable HLO and compile produces the executable. Then server calls execute without the first call compile spike. >> Stable HLO is not a new model. It's the lowest program the compiler sees. You inspect it when you need to debug or confirm what got staged. >> [music] >> Batch size is still a thing to consider. New shapes may compile separately, but larger batches can improve tokens per second by a mortisizing overhead.

01:29 Same shape lesson we share before, but now at a second time. Third option is Jax export which serializes a stable HLO based artifact. The notebook writes it, reads it back, calls it and checks that predictions match. If your platform speaks Jax runtime, this is the portable route. >> Fourth and last option is Jax 2TF for TensorFlow infrastructure. If the serving stack expects a save model, convert with native serialization and save the standard TF artifact.

02:09 To wrap up, choose by deployment target, Python service, AOTJ runtime, portable JAX artifact or TensorFlow serving stack. The checkpoint is the same but the serving wrapper changes. >> So to conclude, in the whole course we had one mental model J traces, XLA compiles and the Nvidia stack runs. When something surprises you, keep in mind these four layers and start your investigation.

02:40 >> Don't forget to verify the GPU stabilized [music] shapes. Time is blocking profile before guessing and avoid accidental host transfers. >> And that's it. This is the path from your first GPU chuck to serving a Jax model on GPU. >> Now go make your GPU busy.