Improve JAX for pretraining notebook. by f33ac · Pull Request #293 · jax-ml/jax-ai-stack

f33ac · 2026-01-06T22:13:12Z

Description

This PR extends the notebook by replacing all placeholder "REVAMP" sections with actual content. The notebook ready for review initial review.

Note:
The vLLM section was intentionally not fully implemented, as running it inside the notebook would require more code and some hacky setup. Instead, it is briefly mentioned with guidance.

Changes

Replaced all "REVAMP" placeholders with content

Testing

Ran the notebook E2E on a T4 GPU & v5e-1 TPU on Google Colab with no issues.

f33ac · 2026-01-06T22:14:07Z

@emilyfertig

emilyfertig · 2026-01-14T20:16:32Z

Hi @f33ac thanks for the PR -- I was going to start reviewing it but it looks like the file size is now 41 MB when it used to be 110 kb, do you know why that is and could you mitigate it?

emilyfertig

Thanks, overall I think this looks great! I left some comments (I might take another pass later), but in general I think your additions tie it all together really nicely.

A couple high-level things:

Some of the cells are tagged nbval-skip or nbval-ignore-output. Some of these tags got deleted, so we'll need to reinstate them so the presubmits pass (some of the cells you added might need these tags too, if they're too long-running for the tests or the cell output is different on the CPU test runners e.g.)
The sphinx build is failing and I'm not sure why: https://app.readthedocs.org/api/v2/build/31206675.txt It looks like this might be the relevant line:

/home/docs/checkouts/readthedocs.org/user_builds/jax-ai-stack/checkouts/293/docs/source/JAX_for_LLM_pretraining.ipynb:650002: WARNING: skipping unknown output mime type: application/vnd.google.colaboratory.intrinsic+json [mystnb.unknown_mime_type] [mystnb.unknown_mime_type]

If you could take a look, that would be great.

emilyfertig · 2026-01-30T20:25:12Z

+    "\n",
+    "By the end of this tutorial, you will be able to:\n",
+    "\n",
+    "- **Understand JAX's parallelism capabilities**: Learn how to leverage data and tensor parallelism to distribute training across multiple TPU devices\n",


say "TPU or GPU" to emphasize JAX is multi-platform

emilyfertig · 2026-01-30T20:30:53Z

   },
-   "outputs": [],
+   "outputs": [


Revert the outputs?

emilyfertig · 2026-01-30T20:34:55Z

-    "tags": [
-     "nbval-ignore-output"
-    ]
+    "outputId": "29c2e1dc-e49c-484a-96db-3b6c0df901a5"


We probably want to keep the nbval-ignore-output tags. The notebook integration tests compare the output of the run cell with the saved output and error if they're different, and the tag bypasses the error.

Also, could you clear the outputs of this cell?

emilyfertig · 2026-01-30T20:36:00Z

+   "outputs": [],
+   "source": [
+    "# Verify that an accelerator (GPU/TPU) is available\n",
+    "# This tutorial requires hardware acceleration and cannot run on CPU\n",


Is there a reason not to allow it to be run on CPU, besides that it's slow? I'm wondering if this should be a warning instead.

emilyfertig · 2026-01-30T20:38:40Z

+    "- **Automatic parallelization**: JAX can automatically distribute computations across multiple devices using SPMD (Single Program, Multiple Data) parallelism\n",
+    "- **Functional programming**: JAX encourages pure functions, which enables reliable transformations and better performance\n",
+    "\n",
+    "For training large language models like our miniGPT, JAX's automatic parallelization capabilities are particularly valuable. They allow us to efficiently utilize multiple TPU cores without manually managing device placement and communication."


Let's mention GPUs here too, again to emphasize multi-platform capabilities and not imply we're TPU-only

emilyfertig · 2026-01-31T00:21:19Z

   "metadata": {
-    "id": "Mtzb0NXb8TVY",
-    "tags": [
-     "nbval-ignore-output"


Reinstate nbval-ignore-output?

emilyfertig · 2026-01-31T00:24:32Z

-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m242.4/242.4 kB\u001b[0m \u001b[31m24.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m221.6/221.6 kB\u001b[0m \u001b[31m22.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m377.3/377.3 kB\u001b[0m \u001b[31m32.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.4/21.4 MB\u001b[0m \u001b[31m104.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",


Clear cell output?

emilyfertig · 2026-01-31T00:30:01Z

I don't think we need tensorflow anymore -- could you take a look at https://docs.jax.dev/en/latest/profiling.html and check that this is up to date?

Hey! According to my colleague, using Perfetto here is not really feasible.

It's blocking execution, analizing single trace rather than a batch as would be the case in training (multiple training steps executions).

We can dive deeper into that if you think it's really important for this tutorial.

emilyfertig

Hi @f33ac could you please sign the Contributor License Agreement? (See the link in the cla/google presubmit). Thanks!

emilyfertig · 2026-02-03T22:45:22Z

     ]
    },
    {
     "data": {
+      "application/vnd.google.colaboratory.intrinsic+json": {


Could you delete these three lines? (2543-2545) That should fix the readthedocs build.

f33ac · 2026-02-10T17:42:10Z

Hey @emilyfertig!
I've released a new version with the tags fixed. Here are all the validatoin that I've run on my side:

>>>> jupytext --sync docs/source/JAX_for_LLM_pretraining.ipynb
[jupytext] Reading docs/source/JAX_for_LLM_pretraining.ipynb in format ipynb
[jupytext] Loading docs/source/JAX_for_LLM_pretraining.md
[jupytext] Updating docs/source/JAX_for_LLM_pretraining.md
---
>>>>> git add docs/source/JAX_for_LLM_pretraining.*
---
>>>>> pre-commit run         
check python ast.....................................(no files to check)Skipped
check for merge conflicts................................................Passed
check toml...........................................(no files to check)Skipped
check yaml...........................................(no files to check)Skipped
fix end of files.........................................................Passed
trim trailing whitespace.................................................Passed
debug statements (python)............................(no files to check)Skipped
pyink................................................(no files to check)Skipped
ruff.....................................................................Passed
jupytext.................................................................Passed

---
>>>> pytest --nbval docs/source/JAX_for_LLM_pretraining.ipynb 
====================================== test session starts ======================================
platform darwin -- Python 3.13.5, pytest-9.0.2, pluggy-1.6.0
rootdir: /Users/adam/Projects/JAX/jax-ai-stack
configfile: pyproject.toml
plugins: xdist-3.8.0, jaxtyping-0.3.7, nbval-0.11.0
collected 36 items                                                                              

docs/source/JAX_for_LLM_pretraining.ipynb s.ss...........s.ssssssss.ssssss..ss            [100%]

================================ 16 passed, 20 skipped in 6.15s =================================

Improve JAX for LLM pretraining notebook.

78c083d

f33ac added 4 commits January 29, 2026 22:31

Update the notebook. Fix size.

1741da0

Remove widgets

c42601a

Sync notebook with Markdown format.

f32942b

Validate outputs with pre-commit.

82223fd

emilyfertig requested changes Jan 31, 2026

View reviewed changes

readd nbval tags; update markdown

40a764d

emilyfertig reviewed Feb 3, 2026

View reviewed changes

f33ac added 5 commits February 4, 2026 14:39

Validate outputs with pre-commit.

35727cc

Clean metadata.

91306ee

Remove Google Colab metadata.

eed8c50

Fix minor issues.

0d20f18

Add proper nbval tags.

cee9e4e

emilyfertig approved these changes Mar 10, 2026

View reviewed changes

emilyfertig merged commit 75cafc8 into jax-ml:revamp-2025 Mar 10, 2026
2 of 3 checks passed

Uh oh!

Conversation

f33ac commented Jan 6, 2026

Description

Changes

Testing

Uh oh!

f33ac commented Jan 6, 2026

Uh oh!

emilyfertig commented Jan 14, 2026

Uh oh!

emilyfertig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

emilyfertig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

f33ac commented Feb 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants