Skip to main content

We fine-tuned a 235B model on four GPUs. Then we didn't need it.

· 8 min read
Savitha Suresh
Software Engineer

We had been using fine-tuned OpenAI and Gemini models to classify entities in our customers' documents. Accuracy was fine. The monthly bill kept growing, so we wanted to know whether we could fine-tune an open model ourselves, on hardware we control, and do at least as well.

We could. One epoch of fine-tuning on about 3,600 examples improved accuracy for every customer we measured, by 7.6 F1 points on average, at about a third of the per-token cost of the API models. The run took 26 hours on four H100s.

We never shipped it. That is the second half of this story. The first half is the four training setups we went through, three of which failed before a single training step ran, because the reason surprised us and it applies to anyone trying to fine-tune a model this size on one machine.

Why a 235-billion-parameter model​

We did not pick the biggest model for its own sake. We benchmarked several open models on our task before fine-tuning anything, and Qwen3-235B-A22B-Thinking scored highest, so we went with it.

It is a Mixture of Experts model. Of its 235 billion parameters, only about 22 billion are used for any one token: each layer holds 128 small "expert" networks and routes each token through a few of them. That keeps inference cheap. It also, as we found out, makes the model awkward to load.

It is also a "thinking" model, meaning it writes out reasoning before it answers. We left that on during training and evaluation, and our training examples included the reasoning traces.

The task and the result​

Given a long, structured document, the model has to find the relevant entities and give each one a label from a fixed set. Every customer (we call them tenants) has its own document formats, so we scored each tenant separately with F1, which balances precision and recall.

TenantBase F1Fine-tuned F1Change
A79.9788.48+8.51
B78.2883.74+5.46
C78.1983.86+5.67
D72.5379.11+6.58
E81.9887.14+5.16
F72.8884.84+11.96
G76.2487.14+10.90
H80.9885.79+4.81
I58.1767.77+9.60
All75.4783.10+7.63

Every tenant improved. The fine-tuned model also beat the fine-tuned API models we had been paying for, on accuracy and on cost per token.

Why it was hard to even start​

The model is 470 GB in its normal 16-bit form. Compressed to 4-bit, which is the usual way to make a big model cheap to fine-tune, it is about 120 GB. Our machine had four H100 GPUs with 80 GB each (320 GB of GPU memory in total) and 960 GB of ordinary system RAM.

120 GB into 320 GB looked comfortable. The plan was the standard one: compress the model to 4-bit, split it across the four GPUs, and train small adapter layers on top. That method is called QLoRA.

The trouble is that compressing and splitting are done by different libraries, and they want things in a different order. The compression library needs to see each weight matrix whole before it compresses it. The splitting library wants to cut everything into pieces as early as possible so that no single process ever has to hold the full model. In between, something has to hold a lot of the model at once. With a 7-billion-parameter model that something is 14 GB and nobody notices. With this model it is 470 GB, and whether a run starts at all depends on where that 470 GB lands.

Each failure below is the same thing happening in a different place: one process ended up holding the whole model somewhere it did not fit.

Four setups​

Axolotl: out of memory, both ways​

Axolotl has a setting that makes one process load the model into system RAM while the others wait for their share. With it on, loading worked (470 GB fits in 960 GB) but the run died while moving the model to the GPUs. With it off, each process loaded straight onto its own GPU, where 120 GB does not fit in 80 GB.

In every one of Axolotl's loading paths, each process holds the whole model at some point, either in RAM or on its GPU. There is one path that loads, compresses and splits one piece at a time, but it does not recognise how this model stores its experts, so it skips them and crashes later. Someone else hit that exact failure with the same model (axolotl #3445). The only reported way to run this model in Axolotl keeps the weights in system RAM and streams them to the GPUs as needed (axolotl #3443), which is the same idea that eventually worked for us.

Answer.AI's FSDP+QLoRA: could not start​

Answer.AI published the reference implementation for combining 4-bit compression with PyTorch's model splitting (FSDP). It was written for our exact problem, big models on few GPUs, and it errored the moment we pointed it at our model.

It was built for Llama and Mistral and last updated in June 2024, before Qwen3 existed. It imports a class that Transformers removed in version 4.48, and Qwen3-MoE support only arrived in version 4.51. So any Transformers version that can load our model cannot run this tool.

LlamaFactory with FSDP: thirty minutes, zero steps, RAM climbing​

We expected this one to work. LlamaFactory sets everything the FSDP and QLoRA combination needs, and its docs describe fine-tuning a 70B model on two consumer GPUs. We turned on the same "only one process loads the model" setting as in Axolotl. After more than thirty minutes no training step had finished, and system RAM rose steadily until the process was killed.

This one took us the longest to understand. Transformers, the library underneath LlamaFactory, turns that memory-saving setting off whenever the model is being compressed to 4-bit. It does this without a warning or an error. So all four processes loaded the full 470 GB model into RAM at the same time, about 1.9 TB requested against 960 GB available, and the steady climb was four copies being read in parallel. The guard is still in Transformers today (in modeling_utils.py, the check is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized), and a related LlamaFactory issue has been open since July 2024 (#4785).

What worked: DeepSpeed ZeRO-3 with CPU offload​

The setup that ran was LlamaFactory with DeepSpeed's ZeRO-3 and CPU offload turned on. It worked because it gives up on keeping the model on the GPUs.

ZeRO-3 sharded the adapter weights and the optimizer state across the four processes. It did not shard the frozen 4-bit base, for the reason from earlier: the compression library needs each weight matrix whole. Each of the four processes held its own complete 120 GB copy in system RAM, and CPU offload copied each layer to the GPUs as it was needed, then dropped it again. The GPUs only ever held the current layer, the adapters being trained, and the working memory for the batch. Our 960 GB of RAM is what made the run possible.

The price is speed. Every step moves weights from RAM to the GPUs over a bus that is much slower than GPU memory. One epoch on 3,600 examples took 26 hours. A setup that kept the model on the GPUs would have been several times faster. We accepted that because the run only had to happen once.

Saving the result​

With the trained parameters spread across four processes, saving a normal checkpoint means gathering everything into one process, and we did not have enough RAM left for that. We saved each process's piece separately and rebuilt just the adapter weights afterwards with DeepSpeed's zero_to_fp32.py --exclude_frozen_parameters. The adapters are small, so this takes a few minutes.

What came next​

We never put it in production. Getting a 235B thinking model to answer fast enough, at a cost that kept the savings we had just won, was shaping up to be a bigger project than the fine-tune. While we were working on that, the cost problem got solved from a different direction: new open-source models came out that solved a lot of the classification workloads without any fine-tuning.

Weeks of fighting loaders and memory limits, a result that beat everything we had, and then a model release made it unnecessary. What is hard this month is easy next month, and the work you put into the hard version is mostly sunk cost.

The lesson is not "don't fine-tune." It is to be picky about where you spend this kind of effort. Fine-tuning earns its place when what you are teaching the model cannot come from anywhere else: your customers' document formats, your internal conventions, data no public model will ever see. If the gap is just "the base model isn't quite good enough yet," the next release will probably close it for free. Wait a month and check again.