Flax debugging: making a hash of things

I was debugging an issue with a JAX/Flax NNX training loop the other day, and found a neat little trick to help debug it. Specifically, I wanted to see if the issue was with my model, my loss function, my optimiser settings, or the "plumbing" of the training loop itself -- were gradients actually coming through and being applied to the parameters?

I could print out the loss and the gradients, but printing out the parameters to see if they were changing was unhelpful -- any given update might only change a small number of parameters, or might change them such a small amount that I'd not notice -- especially given that the model had 77 million of them!

Let's take a look.

[ Read more ]

10Gb/s Ethernet: switching to a Broadcom SFP+ module

Back in April, I upgraded my home LAN to 10Gb/s. The in-wall cabling is CAT-6 or similar, so I had to use 10GBASE-T. Now, the router I'm using, and the switch in my study, provide 10Gb/s through SFP+ cages; that meant that they needed 10GBASE-T SFP+ modules in order to connect.

That kind of module is known to run hot -- sometimes too hot to actually work. The modules in reggie, the router, appeared to be running OK (see the linked post above for charts), but the one in nigel, the study switch, was a worrying 93C. I tried sticking some mini-heatsinks on it, which seemed to help a bit. But the weather got warmer, and eventually the module overheated. I lost access to the Internet from the study, and checking the metrics showed me this:

You can see that it's "flapping": the temperature gets up to a level where the module shuts itself down for its own protection -- about 95C, I think -- and then when it has recovered, it switches on again, the temperature rises, and the process repeats.

I was able to work around the problem by switching on the air conditioning in the study. But normally I only have it on when I'm in there, and keeping aircon on 24/7 just to keep the network working felt like the wrong solution.

It was time to switch to a more power-efficient SFP+ module.

[ Read more ]

JAX: commitment issues

Imagine you have JAX code like this, and run it on a machine with CUDA set up:

    key = jax.random.key(42)

    cpu0 = jax.devices("cpu")[0]
    with jax.default_device(cpu0):
        array = jax.random.randint(
            key,
            (530640, 6, 1024),
            0, 50_000,
            dtype=jax.numpy.uint16
        )
        array.block_until_ready()

    item = array[0]
    item.block_until_ready()

We're creating a big array, blocking until it's ready (JAX is asynchronous, so this makes sure that it's actually finished creating it), then getting the first item, and as a belt-and-braces thing making sure that that is ready too. How long do you think those last two lines -- a simple retrieval of a 6 x 1024 array from a larger one -- will take? Some tiny fraction of a second would seem reasonable.

But running it on my machine just now, the answer is a bit of a surprise: just over 5 seconds. And if you try to get array[1] immediately afterwards, it still takes about 1.2s. Further lookups into array consistently take more than a second -- so while the larger initial number might be something to do with setup -- maybe internal stuff being JITted -- that's clearly not the whole story. Something is making these seemingly-simple array lookups take much longer than you'd expect them to.

Let's dig into that.

[ Read more ]

JAX backends and devices

There's nothing like writing your own code with a framework to clarify how things fit together! Continuing with my port of my PyTorch LLM code to JAX, I wanted to load up a large dataset: the 10,248,871,837 16-bit unsigned integers in the train split of gpjt/fineweb-gpt2-tokens. That's just over 19GiB of data.

from safetensors.flax import load_file
...
full_dataset = load_file(dataset_dir / f"train.safetensors")["tokens"]

When I ran that, I got a CUDA out-of-memory error:

jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 19.09GiB.

That makes sense! The allocation it was trying to do is exactly the size of the data I was trying to load. I have an RTX 3090 with 24 GiB, but some is already used up by the OS, various apps, and a model that the code creates earlier on.

But in PyTorch land, I was used to things being loaded into RAM by default, and only moved over to the GPU when I asked it to do that. JAX was clearly loading to the GPU by default. How could I stop it from doing that for this case? The load into the GPU was happening inside Safetensors, in code I couldn't directly control.

Understanding how to do it helped me understand a little bit more about JAX.

[ Read more ]

Using Safetensors with Flax

I'm porting my PyTorch LLM code to JAX, using Flax as the neural network layer. For various reasons I wanted to use Safetensors to store checkpoints of the model. It took a little while to get it working; here's the trick I learned.

[ Read more ]

10Gb/s Ethernet: using mini-heatsinks with a 10GBASE-T SFP+ module

In my last post I showed the somewhat-scary temperatures I was getting on the MikroTik 10GBASE-T SFP+ module I have plugged into nigel, the 10Gb/s switch I have in my study. As I mentioned then, the plan was to try using some of the mini-heatsinks that people use on Raspberry Pis, to see if that would help.

Here's how it went.

[ Read more ]

10Gb/s Ethernet: what I actually did to get it working in my home

Having learned enough about 10Gb/s Ethernet to be comfortable about setting it up in my house, it was time to bite the bullet: order it from the ISP, buy some kit, and get started.

I already had 2.5Gb/s working. The apartment has structured cabling -- each room has one or more RJ45 sockets in the wall, and there's a patch panel downstairs by our front door that has a matching patch socket for each wall socket. So when we moved in, I simply set things up so that there was a 2.5Gb/s switch down by the patch panel, and wired everything together there. Most of our stuff works over WiFi, of course, but I needed a wired backbone to connect the excessive number of computers in my study both to each other, and to the outside world.

What did I need to do?

[ Read more ]

10Gb/s Ethernet: what I had to (re)learn

My ISP recently started offering a 10Gb option, and my "shiny new thing!" Pavlovian response immediately kicked in. So of course, I had to upgrade the wired networking in my home -- which meant I had to learn a few things to get it all working, and relearn a bunch of stuff I'd forgotten over the years.

Wired networking for home and small offices hasn't really moved forward that much in the last 20-odd years. Back in 2006, gigabit Ethernet was standard for businesses, and most home users moved to it not long after. Perhaps due to the rise of WiFi for most "last few metres" connections, it's pretty much stagnated there, perhaps with a bit of a push towards 2.5Gb/s more recently.

But with faster ISP connections arriving, I think things are starting to become a bit more interesting. Even the fastest WiFi 7 connections are only able to get up to around 6Gb/s to a single device -- and that's in an ideal "super-fast machine sitting right next to the AP in a shielded lab" setup.

Here's what I had to drag up from my memory, and the new stuff I had to learn, in order to get this all working. I'll write about the background in this post, and then tomorrow I'll post about what I actually put in place.

[ Read more ]

Getting MathML to render properly in Chrome, Chromium and Brave

The other day I posted about adding mathematical typesetting to this blog using markdown2, LaTeX and MathML. One problem that remained at the end of that was that it looked a bit rubbish; in particular, the brackets surrounding matrices were just one line high, albeit centred, like this:

...rather than stretched to the height of the matrix, like this example from KaTex:

After posting that, I discovered that the problem only existed in Chromium-based browsers. I saw it in Chromium, Chrome and Brave on Android and Linux, but in Firefox on Linux, and on Safari on an iPhone, it rendered perfectly well.

Guided by the answers to this inexplicably-quiet Stack Overflow question, I discovered that the prolem is the math fonts available on Chromium-based browsers. Mathematical notation, understandably, needs specialised fonts. Firefox and Safari either have these pre-installed, or do something clever to adapt the fonts you are using (I suspect the former, but Firefox developer tools told me that it was using my default body text font for <math> elements). Chromium-based browsers do not, so you need to provide one in your CSS.

Using Frédéric Wang's MathML font test page, I decided I wanted to use the STIX font. It was a bit tricky to find a downloadable OTF file (you specifically need the "math" variant of the font -- in the same way as you might find -italic and -bold files to download, you can find -math ones) but I eventually found a link on this MDN page.

I put the .otf file in my font assets directory, then added the appropriate stuff to my CSS -- a font face definition:

@font-face {
    font-family: 'STIX-Two-Math';
    src: url('/fonts/STIXTwoMath-Regular.otf') format('opentype');
}

...and a clause saying it should be used for <math> tags:

math {
    font-family: STIX-Two-Math;
    font-size: larger;
}

The larger font size is because by default it was rendering about one third of the height of my body text -- not completely happy about that, as it feels like an ad-hoc hack, but it will do for now.

Anyway, mathemetical stuff now renders pretty well! Here's the matrix from above, using my new styling:

I hope that's useful for anyone else hitting the same problem.

[Update: because RSS readers don't load the CSS, the bad rendering still shows up in NewsBlur's Android app, which I imagine must be using Chrome under the hood for its rendering. Other RSS readers are probably the same :-(]

Adding mathematical typesetting to the blog

I've spent a little time over the weekend adding the ability to post stuff in mathematical notation on this blog. For example:

It should render OK in any browser released after early 2023; I suspect that many RSS readers won't be able to handle it right now, but that will hopefully change over time. [Update: my own favourite, NewsBlur, handles it perfectly!]

Here's why I wanted to do that, and how I did it.

[ Read more ]