Scale and Serve JAX Models on NVIDIA GPUs

1. What is the best first check before debugging JAX GPU performance?

   Increase the batch size

   Run jax.devices() and confirm JAX sees a GPU

   Convert all JAX arrays to NumPy

   Disable jax.jit

2. If jax.default_backend() prints cpu and jax.devices() shows only a CpuDevice, what is the most likely cause?

   The model is too large for the GPU

   A CPU-only jaxlib, a driver mismatch, or a container started without GPU access

   The batch size is set too high

   jax.jit has not been called yet

3. Why can the first call to a jax.jit function be much slower than later calls?

   JAX always runs the first call on CPU

   The first call may include tracing and compilation

   JAX disables GPU memory on the first call

   Gradients are computed only on the first call

4. What is a common cause of accidental recompilation in JAX?

   Reusing the same input shape and dtype

   Keeping batches on the GPU

   Changing input shapes between calls

   Calling block_until_ready() before timing

5. Inside a @jax.jit function, why does if x > 0: on a traced array raise an error?

   JAX does not support the > operator

   During tracing x is an abstract value with known shape and dtype but no concrete value, so Python cannot branch on it

   The comparison must use jnp.greater instead of >

   jit disables all Python control flow permanently

6. Why must you call block_until_ready() before stopping a timer around GPU work?

   It frees GPU memory before the next step

   JAX dispatches GPU work asynchronously, so without blocking you may only measure how long Python took to enqueue the work

   It recompiles the function for accurate timing

   It moves the result back to the CPU

7. Calling float(loss) every step inside a training loop can slow it down because it:

   Recompiles the loss function each step

   Forces a host-device transfer and can synchronize the device every step

   Changes the dtype of the loss to float64

   Disables the JAX compilation cache

8. Why does the training loop reshape data into fixed-size batches and drop the partial final batch?

   To shuffle the data more effectively

   To keep the input shape constant so the jitted step is not recompiled for a smaller last batch

   Because JAX cannot process batches larger than 512

   To reduce the model's parameter count

9. In a JAX training step, how are parameters updated?

   The weights are mutated in place inside the model

   A pure function takes params, state, and a batch in and returns new params and state out — no in-place mutation

   jax.jit updates the weights automatically

   Gradients are applied directly by the GPU driver

10. What does setting implementation="cudnn" in jax.nn.dot_product_attention require?

    A CPU fallback and float64 inputs

    bfloat16 or float16 inputs and compute capability >= 8.0

    At least two GPUs

    Disabling causal masking

11. Why does naive attention use more memory than a fused cuDNN kernel on long sequences?

    It stores the model weights twice

    It materializes the full (B, N, T, S) attention matrix, which grows with the square of the sequence length

    It runs in float64 by default

    It keeps the data on the CPU

12. In data-parallel training with JAX, what actually creates the parallelism?

    Rewriting the training step with explicit loops over GPUs

    How the arrays are placed — sharded data and replicated parameters — while the training step code stays the same

    Calling jax.jit twice

    Increasing the learning rate per GPU

13. What does PartitionSpec('data', None) mean for a 2-D array on a mesh with a 'data' axis?

    Replicate the whole array on every GPU

    Shard dimension 0 across the 'data' axis and replicate dimension 1

    Shard both dimensions across the 'data' axis

    Move the array to the CPU

14. How does the Lesson 7 transformer get causal (backward-only) attention?

    It builds a triangular mask matrix by hand each step

    It passes causal_sdpa (which calls dot_product_attention with is_causal=True) as the attention_fn of nnx.MultiHeadAttention

    It sorts the tokens in reverse order

    It uses a separate decoder library

15. When saving an NNX model with Orbax StandardCheckpointer, what is saved and what is not?

    The full Python class is pickled, including methods

    The parameters (nnx.state(model, nnx.Param)) are saved; the graph structure is not — it comes from the Python class on restore

    Only the optimizer state is saved

    Nothing is saved until training finishes

16. What problem does AOT compilation (lower() → compile()) solve for serving?

    It reduces the model's parameter count

    It separates compilation from execution so you compile once at startup and serve requests with no cold-start latency

    It converts the model to float16 automatically

    It shards the model across multiple GPUs

17. Which serving path produces a portable artifact that runs in any JAX runtime without the original model code?

    jax.jit in-process

    jax.export — serialized StableHLO bytes

    Calling the model inside a Python loop

    AOT compilation with lower().compile()

Submit answers