The goal, stated honestly
The endpoint worked. The model was open weights, we had it serving privately, and the complaints about it were consistent and specific: it thinks a lot and does little, and its design taste in generated interfaces was mediocre. Both are trainable complaints. So the objective was a LoRA fine-tune of the exact approved checkpoint on a sealed internal curriculum, followed by adapter selection and evaluation.
The economics were stated up front, in writing, before any pod was provisioned: a $500 total ceiling for the whole pilot, including model loading while the GPU was billing, canary steps, checkpointing and reload tests, evaluation, storage and export. Not a $500 canary followed by a bigger ask. If the full two-arm experiment could not fit under the ceiling, the decision was no run, and we would say so.
Writing the budget rule down as a hard constraint before you start changes your behaviour in a way that a mental budget does not. It makes "stop" a pre-authorised outcome rather than an admission of failure.
Frameworks: most of them cannot load this model
The architecture is not exotic, but it is not vanilla either: hybrid attention, hash routing, hyper-connections, packed MXFP4 experts. The practical consequence is that generic toolchains — plain Transformers, PEFT, Axolotl — do not handle it. We used NVIDIA's NeMo AutoModel at a pinned commit, which ships a native implementation and is the source of the official LoRA recipe.
Then the trap: the official "LoRA" recipe is not QLoRA. It dequantizes the base checkpoint into the training path rather than training on the quantized weights in place. That one line in a config file is the difference between "a LoRA needs modest extra memory on top of the base model" and "the training path materializes a much larger representation of a 304B model." Everything downstream — your GPU count, your budget, your timeline — follows from which one you are actually doing.
Every gate passed except the real one
This is the part I would want to read if I were starting the project again. The validation work was genuinely thorough, and it was aimed at the wrong risk.
- The frozen-base evaluation completed: 618 questions, no crashes.
- The dataset was repacked and reduced training gates passed on a clean container.
- Representative-layer validation ran through versions v4 to v10 — gradient contracts, layer selection, numerics — to establish that the LoRA path was wired correctly.
- None of that touches the question that killed the run: does the full model fit in memory once aggregation begins?
We validated components. The failure was in the whole. Reduced gates are cheap to run and satisfyingly green, and a green gate is a sedative. If you take one thing from this postmortem: the last gate before an expensive, irreversible step should be the expensive part at full size, not a faithful miniature of it.
The failure
The first complete training-path load failed during BF16 expert aggregation. Each GPU reached roughly 263 GiB allocated. The pod had four B300s. It ran out anyway, because the memory peak is not the model — it is the model plus the dequantized training representation plus transient aggregation buffers, and the last term scales with experts per layer and the batch it is being asked to gather.
The pod was stopped with its persistent workspace preserved.
The same checkpoint loaded for inference on two B300s. Nothing was misconfigured in the serving stack. Inference holds the quantized weights and runs one forward pass; training holds a dequantized representation plus gradients plus optimizer state plus communication buffers, and needs every expert group resident and synchronized. The gap between those two numbers is not a percentage. It is a different order of magnitude, and it is the single most useful thing to model before you rent anything.
The three-session record
The work spanned three long agent sessions. Reconstructing them rather than trusting memory was worth it: across roughly 63,000 recorded events and about 1.66 GB of source data, the pattern was visible in a way it never is in the moment.
- The 2×B300 phase: initial benchmark, repeated inference-path integration, the successful 618-row base evaluation, and a no-load training handoff.
- The 4×B300 migration: dataset repacking, runtime construction, a clean-container run, and most pre-load training gates passing.
- The final phase: the LoRA gradient contract, representative-layer validation v4 through v10, and the full-load authorisation work that ended in the out-of-memory failure.
Note the shape. Session one ends optimistic and hands off. Session two scales hardware by 2× and passes more gates. Session three is entirely preparation for a step that never executes. Effort accumulated monotonically right up to the wall.
What I would do differently
- Model the memory peak first, in writing. Weights, training representation, gradients, optimizer state, aggregation buffers, activation peak. If the total exceeds the hardware by 20%, the project is a no before any pod is rented.
- Verify what "LoRA" means in the recipe you are using. Read the dequantization flag. Do not infer it from the word LoRA.
- Run the full-load test before the expensive phase, not after. A load that succeeds and is immediately torn down costs minutes. Discovering it cannot load costs the whole pilot.
- Distinguish harness failures from model failures in the record. Half the debugging time in a project like this goes to control-plane and validation-harness problems that have nothing to do with the model. If your postmortem does not separate those categories, it will teach you the wrong lesson.
- Keep the negative result. The pilot produced a precise hardware floor for this architecture, a corrected cost model, and a restart plan. That is an asset. Discarding it because the run did not finish is how the same money gets spent twice.
The honest conclusion
The pilot failed, and the failure was predictable from information we had before we started. That is a worse outcome than an unpredictable failure and a more useful one, because predictable failures are the only kind you can design out.
What survived is worth more than the adapter would have been: a measured memory model for a 304B sparse MoE training path, a documented framework trap that would have burned the next person just as badly, and a habit of writing the budget rule down before the meter starts.