Thoughts on Role Confusion
The other day, I came across "Prompt Injection as Role Confusion"
(via Simon Willison). It's a really
interesting blog-style version of a paper by Charles Ye, Jasmine Cui and Dylan Hadfield-Menell,
where they find that LLMs seem to almost ignore 'role' tags like <system>, <user> or <think>, and
instead use the tone of text to infer roles. This seems to explain a lot of jailbreaks.
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.
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.
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.
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.
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.
On first looking into JAX
Much have I travell'd in the realms of gold,
And many goodly states and kingdoms seen;
Round many western islands have I been
Which bards in fealty to Apollo hold.
Oft of one wide expanse had I been told
That deep-brow'd Homer ruled as his demesne;
Yet did I never breathe its pure serene
Till I heard Chapman speak out loud and bold:
Then felt I like some watcher of the skies
When a new planet swims into his ken;
Or like stout Cortez when with eagle eyes
He star'd at the Pacific -- and all his men
Look'd at each other with a wild surmise --
Silent, upon a peak in Darien.
John Keats, On First Looking into Chapman's Homer
I've been working with PyTorch quite a lot for the last couple of years, and feel like I've come to a reasonably solid understanding of how it all fits together. Working through Sebastian Raschka's book "Build a Large Language Model (from Scratch)", training my own LLMs locally and in the cloud, rebuilding Andrej Karpathy's 2015-vintage RNNs -- over time, it all adds up!
But, of course, there are other frameworks, and one I kept hearing about was JAX. While it's less dominant than PyTorch, it has a reputation for a certain cleanliness, a certain purity. And having spent time over the last couple of weeks working through the tutorials, and translating small PyTorch examples into it, I've been really impressed.
In this post I want to give an overview -- to report back to beginners like me, still living in PyTorch-land, on my new discovery. Less like Herschel discovering Uranus, and more like a 16th-century European coming back after having discovered something that the people who lived there were perfectly well aware of. What is this JAX thing, and how does it differ from PyTorch?
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.
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?
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.