Full Width [alt+shift+f] Shortcuts [alt+shift+k]
Sign Up [alt+shift+s] Log In [alt+shift+l]
2
This post is the capstone of the most long-running series on my blog. In December 2024 (!), I started reading Sebastian Raschka's book "Build a Large Language Model (from Scratch)", and worked through it carefully. Being who I am, despite trying to apply a strict "no side quests" policy, I found myself zooming off and digging into all kinds of things. It's time to wrap it up. I had decided that the endpoint would be to build and train an LLM from scratch just using my notes -- no reference to the book, no reference to the model code I'd written when following the book. After an X/Twitter poll, I decided to use JAX for that, just to make sure that I really was building it from scratch and not regurgitating bits of PyTorch code like a bad coding LLM spitting out half-digested lumps of Stack Overflow. In my last post, I showed how I built a JAX training script that mirrored what I had built for the original PyTorch version of the model. To test it as I went along, I used it to train a really dumb "LLM", which instead of trying to predict the next token for every token in an input sequence, instead predicted the input -- that is, if you fed it The fat cat sat on the mat It would return the same thing. I called that an A-to-A model. In this post, I'll show you how I turned it into a GPT-2 model, and then trained it from scratch on my RTX 3090 (using the parameter counts for the original paper's "small" size). What turned out really well with this is that I found a route that meant that almost every component I added made the model better! That's not guaranteed -- sometimes different aspects of an AI model depend on each other, so adding A without also adding B makes things worse. But (admittedly with a bit of backtracking in places) I was able to find a route that shows a nice clear progression. The final training run took 37 hours 15 minutes -- compared to 40 hours, 38 minutes for an equivalent PyTorch model. That is despite it being full-fat 32-bit -- the...
8th Jul 2026

Stay updated

Get a weekly newsletter with the top 5 articles worth reading every week.

More from Giles' blog

Putting my JAX-trained models on the Hugging Face Hub

I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because Transformers has been PyTorch-only since version 5 (though they say they're working to add interoperability with JAX in the future), so it would have been tough to get them working natively with AutoModelForCausalLM and the like. But then it dawned on me that I'd already written a conversion script that could take my JAX safetensors files and convert them into ones compatible with my PyTorch code. It's actually those converted models that I use for my evals -- so I could use my existing PyTorch script to upload them. So, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models: "Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)" gpjt/jax-no-mha-bias-no-dropout -- the first full LLM trained in the post, in the "Adding LayerNorm" section. gpjt/jax-no-mha-bias-with-dropout -- the second full LLM trained in the post, in the "Dropout" section. gpjt/jax-with-mha-bias-no-dropout -- the third full LLM trained in the post, in the "Adding bias to the MHA output projections" section. "Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining" gpjt/jax-with-mha-bias-no-dropout-extended -- the single-epoch, double-Chinchilla-tokens model. gpjt/jax-with-mha-bias-no-dropout-2-epoch -- the model trained on two epochs over the Chinchilla-optimal number of tokens. "A quick(ish) Chinchilla check" gpjt/jax-with-mha-bias-larger-chinchilla-1 -- the slightly-larger model. gpjt/jax-with-mha-bias-larger-chinchilla-2 -- the slightly-smaller model. I've also added links to the posts in question.

a week ago 1 votes
A quick(ish) Chinchilla check

I recently overtrained a couple of GPT-2 style models, training them both on 40 tokens per parameter rather than the 20 per parameter that is generally regarded as "Chinchilla-optimal". The normal heuristic is that instead of doing that, you should scale up the number of tokens and the number of parameters equally -- so I would have been better off scaling up the model by 2 and the token count by the same amount. By doing that, I should expect to get a better model in terms of loss on my held-back test set than I did with my 40-tokens-per-parameter models. My training machine poppy wasn't doing anything, so I decided to give that a go. Would the Chinchilla rule-of-thumb hold up? As you might expect, it did. But it was a surprisingly close-run thing, and could conceivably have been in the noise. Let's take a look. The Chinchilla heuristic If you already know all about the Chinchilla paper -- regular readers in particular must be sick and tired of it by now :-) -- then click here to skip this section. In "Training Compute-Optimal Large Language Models", which is always called the Chinchilla paper after the name of the model they trained at the end, the authors tried to work out the optimal number of tokens to train an LLM on based on its number of parameters. In particular, they were pushing back on a trend they were seeing at the time, where people were making models ever-larger, but not increasing the amount of data they were training on. The authors were all at Google DeepMind, and this was the kind of project that only a large lab could do: they trained "over 400 language models ranging from 70 million to over 16 billion parameters on 5 to 500 billion tokens". Their conclusion was "for compute-optimal training, the model size and the number of training tokens should be scaled equally: for every doubling of model size the number of training tokens should also be doubled". They don't actually state an overall optimal number of tokens to train on in the paper, but in table 3 they provide an estimate of the optimal training FLOPs and tokens for models of various sizes, and it's approximately 20 tokens per parameter. That number has become a heuristic, and people talk about a model as being trained for the Chinchilla-optimal number of tokens. Models that were trained on fewer tokens per parameter are referred to as "undertrained", and models that were trained on more as "overtrained". It's worth noting that overtraining a model is not, in itself, a bad thing. If you have a model of a particular size and you continue training it past the Chinchilla-optimal number of tokens, it will -- in general -- get better. The point of the heuristic is that doing that is not the best way to spend whatever budget you have in terms of compute time. You'll get better results, as they say, by scaling the number of tokens and the number of parameters equally. But let's say you're creating a model for specific target hardware -- say, a mobile device. You have a hard restriction on how large the model can be -- the device has only so much RAM to hold it. So it might make sense to overtrain to get a better model. 1 But if you're not so limited in how many parameters you can use, then you should indeed scale the model up, and that's what I wanted to try. How would that work? Scaling the model A week or two back, I was investigating whether I could make my GPT-2 style models better at a specific instruction-following task by overtraining them. The details of that experiment aren't important here, but what it meant was that I had three GPT-2-style models, each of exactly the same size, roughly 163M parameters A Chinchilla-optimal one, which I'll call jax-gpt2-chinchilla here. One trained on twice the Chinchilla-optimal tokens, jax-gpt2-2x-chinchilla. One trained on the Chinchilla-optimal tokens, with two epochs (so that it was trained for as long as #2): jax-gpt2-2-epoch-chinchilla When I tested them against a held-back test set of sequences -- stuff that they'd never seen before -- they got results rather like you might expect: Test loss jax-gpt2-2x-chinchilla 3.324953 jax-gpt2-2-epoch-chinchilla 3.326482 jax-gpt2-chinchilla 3.418784 A lower loss is better, and you can see that the longer-trained models were noticeably better than the Chinchilla-optimal one. The difference between them was tiny; they were trained starting with the same initial weights, and the training runs themselves were deterministic, but a difference of 0.05% in loss doesn't seem like it could be meaningful -- an extra batch for one or one fewer for the other could easily swap them around, you'd think. Now, these models each had 163,009,536 parameters -- they were the small-size model from the GPT-2 paper, modified to not have QKV bias or weight-tying. jax-gpt2-chinchilla had been trained on 3,260,190,720 tokens (rounded up to fit into a round number of full batches), and the other two on 6,520,381,440 tokens each -- double the amount (rounded up too). What I needed to do for my Chinchilla check was to try training a model that used the same amount of compute, scaling the parameters and the number of training tokens equally. Because training compute increases roughly linearly with both parameters and tokens, that would mean scaling both up by 2, giving us: 163,009,536*2≈230,530,296parameters ...and thus 4,610,605,920 tokens. How to scale the model up? In the GPT-2 paper, they train four models: Name Parameters 2 Layers d_emb MHA heads 3 small 124M 12 768 12 medium 345M 24 1024 16 large 762M 36 1280 20 xl 1542M 48 1600 25 I wanted to scale my own model up from 163M parameters to about 231M. Which of those numbers would I want to increase, and by how much? The first thing that stands out is that the number of heads is always 1/64th of the number of embedding dimensions. So that sorted that one out. I just needed to adjust the number of layers, and the number of embedding dimensions, but ensure that the latter was a multiple of 64. I decided to see if I could fit some kind of curve to the relationship between the number of parameters and the GPT-2 authors' choices. This was made a bit more complicated by one thing: they were using weight-tying, and I was not. That meant that they re-used the embedding matrix at the start of the LLM as an output head at the end -- which is why they had 38M fewer parameters. Embeddings and the output head make up a surprisingly large percentage of the parameters for small models like this -- about 47% without weight-tying, 23% with. I couldn't work out a solid way to scale things up and wound up doing some rather messy hacking around in a spreadsheet. I came up with two proposed model sizes that were within a couple of percentage points of the right size: Name Layers d_emb MHA heads Parameters % diff slightly-larger 15 896 14 235,621,120 +2.21% slightly-smaller 14 896 14 225,978,368 -1.97% Interestingly, I found that because d_emb could only change in increments/decrements of 64, it was a pretty coarse control -- my first attempt at making a slightly-smaller model changed it to the next step down, 832, but that led to a model that was 9.25% too small. That was an interesting first lesson. I'd previously been thinking of the Chinchilla rule as being something like "don't double the tokens, just scale the model and the tokens equally". But that "just" was wrong. Scaling a model is hard -- even with just two dials to fiddle with, like in this case, it was tricky to get something right -- and I can't say for sure that my choices were the right ones. Anyway, the next step was to double-check that these models would use the right amount of compute to train. Training FLOPs As I said earlier, the compute time scales roughly linearly with the number of parameters. Let's dig into that "roughly". Different kinds of parameters take different amounts of FLOPs to train, and scale differently with things like the embedding dimensions, sequence length, and so on. Now, for very large models, a lot of that comes out in the wash, but with tiny models like these where the embeddings make up such a large proportion of the parameters, it might matter. Conveniently, in appendix F of the Chinchilla paper, they provide a set of formulae for estimating the number of training FLOPs for a normal dense LLM like these ones. I coded that up into a script that, given the JSON configuration files I was using for my models and training runs, would work out the number of FLOPs for a single epoch of training. It didn't take account of the fact that my real training runs round the number of tokens up so that we do a round number of full batches, but I felt that so long as the results weren't very close that wouldn't matter. I got these results (multiplying the two-epoch numbers by two): Est. training FLOPs jax-gpt2-chinchilla 3,544,967,596,946,227,200 jax-gpt2-2x-chinchilla 7,089,935,193,892,454,400 jax-gpt2-2-epoch-chinchilla 7,089,935,193,892,454,400 slightly-larger 7,419,664,885,127,577,600 slightly-smaller 6,804,429,215,367,168,000 The numbers were indeed different enough that I wasn't worried about the batch-rounding. And the good news was that slightly-larger and slightly-smaller would indeed use slightly more and slightly less compute to train than the overtrained models -- about 4.6% more and 4% less respectively. A true Chinchilla-equivalent model would lie somewhere between them. It was time to train some models! Training I kicked off the run for the slightly-larger model first. Because it was bigger than the 163M models I'd been training, I couldn't fit such large batches into my VRAM; previously I'd been running with a batch size of 6, and now I could only fit in a batch of 4. Luckily, though, I was using gradient accumulation, so by bumping that up from 16 steps to 24 steps I could keep the same overall batch size and keep the training runs comparable. Even despite that, the training run ran out of VRAM about 60 hours in -- I'm guessing due to VRAM fragmentation, as I did not have TF_GPU_ALLOCATOR set to cuda_malloc_async -- but I was able to restart from the most recent checkpoint and complete the run. After just less than four days total training time, it completed. When it was done, I copied the last checkpoint 4 over to my dev box, perry, and ran my standard smoke test against it, asking it to complete "Every effort moves you" with 20 tokens, using greedy sampling. I got something reasonably coherent: Every effort moves you. I’m not sure what you’re thinking. I’m Next, I converted the safetensors file -- which had been saved by my JAX code -- into a format compatible with my PyTorch code, because that's what I use for evals. I ran another smoke test (this one with temperature 1): Every effort moves you through the motions for your life, your soul, your body, and your soul’s happiness Very spiritual. Next, it was time to work out the loss on my held-back test set: giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla/model.json ../jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla/checkpoints/latest/pytorch-model.safetensors Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3485.09it/s] 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [07:11<00:00, 7.42it/s] Loss against our test dataset: 3.280028 Well, it was certainly better than the 3.324953 that the best of the overtrained models got -- but only by a bit over 1% better. Interesting! I decided to train the second model, slightly-smaller. This one crashed mid-way through with an error that I've seen before: jax.errors.JaxRuntimeError: INTERNAL: CUDA error: Failed to end stream capture: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED: operation failed due to a previous error during capture [executable_name='jit_train_step'] I'm going to have to investigate that more in future, but for now, I just restarted from the checkpoint, and again after a bit less than four days, I had a model. The JAX smoke test was solid: Every effort moves you forward. The best way to get started is to start with a free trial. You can ...and so was the PyTorch one: Every effort moves you forward in love with our products. I love the way it’s easy to use. Both quite commercial this time! It was time for the proper test loss eval: giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run test_loss.py datasets/ ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla-2/model.json ~/Dev/jax-gpt2-from-scratch/runs/full-llm-full-train-with-mha-output-bias-larger-chinchilla-2/checkpoints/latest/pytorch-model.safetensors Fetching 4 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 3151.24it/s] 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3200/3200 [06:45<00:00, 7.90it/s] Loss against our test dataset: 3.292937 So, slightly worse than the 3.280028 from the larger model, better than the 3.324953 from the best overtrained one. Time to put this all together. Results Here's an updated version of the table from the start of this post; I've added in the two new models, and the improvement they each had over jax-gpt2-2x-chinchilla in both absolute terms and as a percentage rounded to 3sf. Test loss Improvement Improvement % slightly-larger 3.280028 0.044925 1.35% slightly-smaller 3.292937 0.032016 0.962% jax-gpt2-2x-chinchilla 3.324953 - - jax-gpt2-2-epoch-chinchilla 3.326482 - - jax-gpt2-chinchilla 3.418784 - - Now, unlike the overtrained models, prior to training these two new ones started with different initial weights to the jax-gpt2-chinchilla one -- after all, they had to, because they had more of them! A while back, I did a bit of analysis of how random variation in weight initialisation can change the resulting test loss. It wasn't anything in-depth, but I trained three models with different explicit seeds set prior to the model initialisation, but with the same seed set before the training run started 5. Those three models wound up with test losses of 3.681356, 3.673943, and 3.664345. Doing statistics with three data points is a bit flaky, but the cost of training models is so high that I'll leave the Proper Science to the likes of Google DeepMind and wing it :-) Mean: ~3.673215 Sample variance: ~0.000073 Standard deviation (SD): ~0.008529 Now, piling statistical flakiness on statistical flakiness, we'll compare these. You'd normally expect about two thirds of results to be within one SD of the mean, 95.4% to be within two SDs, and 99.7% to be within three. Three SDs on that (yes, different, I know) distribution is 0.025587. That's smaller than both of the improvements that our Chinchilla-optimal runs had over the overtrained ones. So what does that tell us? Well, perhaps not much given the statistical flakiness. But I think it is useful directionally. It suggests that we might be able to take these results seriously as an improvement, and that Chinchilla held: scaling up the model and the number of tokens evenly did give us a better model than just scaling up the number of tokens. In particular, the fact that the loss for slightly-smaller was lower -- even though it had 4% less compute spent on it than the overtrained models -- was encouraging. But it's certainly far from a slam-dunk. A larger test, training lots of overtrained models and lots of Chinchilla-optimal ones, all with different random seeds, would give actual real serious data. Not worth it for me, and perhaps not for anyone. Conclusion I wanted to do a quick sanity check of the Chinchilla heuristic of 20 tokens per parameter. I came up with results that were certainly in line with it -- perfectly so in terms of the ordering of the models I trained. But the effect was small enough that I could imagine that it was in the noise, especially given the small numbers of models I'm able to train. I'll chalk it up as a tentative success. In addition, I learned one useful thing: when talking about scaling up a model to more parameters, you actually have to think quite hard about where you want to put those parameters. I wound up doing a rough curve-fit to the models in the GPT-2 paper, but I have no idea if that was optimal. At some point I should try to dig up some research into optimising embedding dimensions, numbers of layers, and so on. But not now, as I've a bunch of other stuff I want to investigate first. Anyway, I hope you found this experiment interesting, and as ever, comments and questions welcome below. Thanks for reading! I'm less familiar with arguments for under-training -- that is, for fewer than 20 tokens per parameter. I've heard that these days, modern LLMs get a lot more reinforcement learning than they do pre-training, and perhaps that might mean that some very big ones are undertrained prior to RL? I'm uncertain. It's unlikely to be raw lack of data; even for those of us outside the big labs, FineWeb has 18.5T tokens. On its own, that would be enough to train a 0.925T-parameter model, and given that you can apparently do four epochs over the same data before you start getting diminishing returns, that takes us up to 3.7T. That's frontier-lab size, and I'm sure they have better datasets than FineWeb. ↩ Parameter counts are from the paper, apart from the "small" model, which is known to be wrong -- I used my own calculation, and the result is in line with what I've seen elsewhere. ↩ The paper doesn't mention the number of heads; these numbers are from "Build a Large Language Model (from Scratch)", and match up with the ones on this Hugging Face page. ↩ Regular readers might have noticed that I'm ignoring what I've been calling the "best" checkpoint. I've come to the conclusion that because for my training script, "best" means best in terms of training loss, and the training loss changes based on what training data the model has seen recently, it's actually not a very useful metric and just confuses things. At some point I'll probably re-introduce pre-checkpoint evals and use that for "best", which would be the right way to do it. ↩ At the time I was using dropout, so training runs were not deterministic without a known seed. ↩

7th Aug 2026 2 votes
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. If you look at the Safetensors docs, you'll see that it doesn't mention a JAX implementation -- indeed, searching for "safetensors jax" at the time I'm writing this gives you a link to this GitHub repo by Alvaro Bartolome -- which was last updated in 2023. However, if you look more closely at the docs, they do have a link to the Flax API. I feel this is somewhat misnamed, as it is actually a JAX API. There's no reference (again, as of the time of writing) to Flax in the source -- it's all just JAX code. And in fact Bartolome's library uses it under the hood. There is one problem, though. The API works with simple single-level dictionaries, with strings mapping directly to JAX arrays. For example, the save_file function has this signature: def save_file( tensors: Dict[str, Array], filename: Union[str, os.PathLike], metadata: Optional[Dict[str, str]] = None, ) -> None This can cause problems if you're not careful. If you look at the Flax documentation on checkpointing, it suggests that you use Orbax 1, which has its own API and file format, but then goes on to say: When interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the nnx.State.to_pure_dict and nnx.State.replace_by_pure_dict API to convert an nnx.State to and from pure nested dictionaries. I initially put two and two together -- that and the dictionary-based API for Safetensors -- and got five, and tried feeding one of those "pure" dicts into Safetensors. I got a very confusing error: SafetensorError: dtype object is not covered It's worth digging in to why that happens. The problem is that although Safetensors is expecting a dict of strings mapping to tensors, it doesn't check that that is what it actually gets. And while the dictionaries from nnx.State.to_pure_dict are "pure", they are also nested (as the docs say!). Even for the simple model I was working with, I got a structure like this: { 'output_head': { 'kernel': Array([...], dtype=float32) }, 'token_embedding': { 'embedding': Array([...], dtype=float32) } } So, we had strings mapping to dicts, and those dicts mapped from strings to the JAX arrays. More complex models would have had deeper dict structures. Now, internally inside Safetensors, the Flax/JAX API is a simple wrapper. It iterates over the keys in the dictionary it's been provided with, and tries to convert their respective values into NumPy arrays. It does that by passing them into NumPy's asarray function, which accepts things like lists, tuples, and NumPy arrays, and converts them into arrays. JAX's own Array class exposes an interface that it recognises, so they're converted without trouble. Once it's done that, it passes the result to a lower-level Rust implementation that actually converts everything to Safetensors format. But because Safetensors didn't check types, in my case it was iterating over the top level of the dict, trying to convert the values to NumPy arrays, and got something like this: { 'output_head': numpy.array({'kernel': Array([...], dtype=float32)}, dtype=object), 'token_embedding': numpy.array({'embedding': Array([...], dtype=float32)}, dtype=object) } That is -- because it assumed that the values in the top-level dict were JAX Arrays, it blindly tried to convert them to NumPy arrays. But they were dicts (that happened to map from strings to arrays) -- and if you ask asarray to create an array based on a random object, it happily does so and wraps that object in a NumPy array, with a dtype of object. When that is then fed into the lower-level Rust code that is trying to write the file, it encounters NumPy arrays that have a dtype it can't handle, object -- hence that error: SafetensorError: dtype object is not covered It all makes sense when you read through the code, but I was a bit perplexed for a while! I think all this might be the reason why Bartolome created his GitHub repo. In the README, he says that: There are no plans from HuggingFace to extend safetensors to support anything more than tensors e.g. FrozenDicts, see their response at huggingface/safetensors/discussions/138. So the motivation to create safejax is to easily provide a way to serialize FrozenDicts using safetensors as the tensor storage format However, you don't need to use that library to serialise simple Flax models. Consider how PyTorch models get serialised to Safetensors; my LLMs have keys with names like out_head.weight, pos_emb.weight, and trf_blocks.0.att.out_proj.weight. They're "flat" dictionaries mapping strings to PyTorch Tensors, similar to what Safetensors wants for these Flax ones, but they use dots to separate different levels, with integers for list items and strings for field names. Looking at the pure-dict structure I had for my model: { 'output_head': { 'kernel': Array([...], dtype=float32) }, 'token_embedding': { 'embedding': Array([...], dtype=float32) } } ...you can see that you could walk the dictionary structure to generate keys like output_head.kernel and token_embedding.embedding. That would be easy enough to code up. But -- as Adithya Dsilva points out on GitHub -- you can get there even faster by using nnx.to_flat_state. That returns a (non-dict) structure like this: FlatState([ (('output_head', 'kernel'), Param( # 786,432 (3.1 MB) value=Array([[ 2.3581974e-02, 3.0957451e-02, -3.5088759e-02, ..., -4.5880198e-02, 5.3717274e-02, -2.6590331e-02], ..., [-9.6302675e-03, -3.3276502e-02, 5.7173111e-02, ..., -7.9063717e-03, 2.0532632e-02, 5.4753982e-02]], dtype=float32) )), (('token_embedding', 'embedding'), Param( # 786,432 (3.1 MB) value=Array([[ 0.00273973, -0.01754938, 0.04656043, ..., -0.04276522, -0.03986642, -0.00781331], ..., [ 0.01421758, -0.0219186 , -0.01701825, ..., -0.00793659, 0.00500103, 0.03839901]], dtype=float32) )) ]) If you iterate over that FlatState, you get tuples where the first element is that tuple of strings, like ('output_head', 'kernel'), and the second is a Param object wrapping the JAX Array. The tuples mirror the dot-separated string format in the PyTorch-style Safetensors files. Param objects also implement an interface that asarray can understand, so you can quickly and easily convert the FlatState to a regular dict for Safetensors: from safetensors.flax import save_file ... model_state = nnx.state(model) flat_state = nnx.to_flat_state(model_state) simple_dict = {} for tuple_key, param in flat_state: key = ".".join(str(key) for key in tuple_key) simple_dict[key] = param save_file(simple_dict, "model.safetensors") (You need to wrap key in a str because if you have a nnx.Sequential in your model, the item in the tuple will get an integer index rather than a string). You can go the other way pretty easily too; given a model, you can load the saved checkpoint into it like this (because from_flat_state accepts raw JAX Arrays in place of explicit Params): from safetensors.flax import load_file ... simple_dict = load_file("model.safetensors") dict_flat_state = {} for key, array in simple_dict.items(): elements = key.split(".") list_key = [] for element in elements: try: list_key.append(int(element)) except ValueError: list_key.append(element) dict_flat_state[tuple(list_key)] = array new_flat_state = nnx.from_flat_state(dict_flat_state) nnx.update(model, new_flat_state) A little more work than I'd ideally like, but given that it can be tucked away in general save_checkpoint/load_checkpoint functions, not too big a deal. Hope that's of use for other people coming across this problem! I'm beginning to feel a bit swamped with all of these libraries with names ending in -ax. It reminds me of the names of the characters in Asterix's village... ↩

4th Jun 2026 1 votes
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. I bought a 40-piece set of heatsinks made by the improbably-named VooGenzek on Amazon for €8, and attached two of them like this -- see the bottom module, with the yellow cable: That was 24 hours ago, and here's a chart of temperatures from that module showing the 24 hours before and after: You can see the big drop-off in the middle of the chart; it even overshot a bit (I'm guessing because the heatsinks absorbed a bunch of heat initially when I put them on). The difference looks more dramatic than it is! See where the Y-axis starts. But given that the weather has been pretty much the same today as it was yesterday, that looks like a 3.5°C improvement. Not great, but not nothing either. In the copious discussion about the last post on Hacker News, one of the most popular comments -- from xxpor -- was that there are two generations of SFP+ modules for this kind of thing; an older one, using a Marvell chip, and the newer one using one from Broadcom. blunden on the ServeTheHome forums made the same point. They both mentioned that a good indicator of which type a module is using is that the older ones tend to be rated up to 30 metres, while the newer ones are rated up to 100. This one is a MikroTik S+RJ10, which definitely is one of the older ones -- the specific chip is mentioned in the docs. I'm not sure which chip the Protectli modules in my router reggie are -- they're these modules -- but they say they're rated up to 30 metres, so I guess they're probably the older type too. Looking into switching those out might be a good next step! I probably won't do that in the short term, though, unless I start getting issues as we move into summer.

18th May 2026 1 votes

More in AI

AI and Existential Dread

There's a lot of polarising discourse right now about the threat AI poses to humanity. Some think it's a farce and others think we face extinction. Here are my thoughts.

3 days ago 1 votes
The AI-as-Normal-Technology view of loss-of-control incidents

A middle ground between the cybersecurity and AI safety communities

3 days ago 2 votes
AI Skills and Job Market, Q3 2026

An overview of the current state of the engineering market and the AI skills that are in demand

5 days ago 1 votes
Why Is AI Bad at Writing?

And Why It’ll Likely Stay That Way—Em Dashes or Not!

a week ago 2 votes
Wendell Berry and the Promise of the Deep Life

Wendell Berry died last week at 92 at his home in Port Royal, Kentucky, where he farmed his land using traditional techniques and wrote with ... Read more The post Wendell Berry and the Promise of the Deep Life appeared first on Cal Newport.

a week ago 1 votes
📚 BoredReading

You seem to be enjoying this.

Join free to unlock everything.

Create free account

Already have an account? Sign in