I have a huge dataset I want to run a model over, and I don’t have a budget for a giant GPU cluster. What I needed was a model that could do simple extraction and classification work at around 10,000 tokens per second on a CPU core, so I could point a few hundred cores at a corpus and let it run. That is an infrastructure requirement, not a research question.

So I gave Claude Code a pile of tokens and told it to build one.

It made three discoveries I did not expect, which is why there is now a paper and a model instead of just a batch job.

That post did 120K views in a day and 87 replies, a lot of which were better than the paper. This is a writeup of what we found, what people pushed back on, and where I think they were right.

The constraint

Every run in this work is pinned to one physical core with OMP_NUM_THREADS=1. One core of an Intel Xeon Silver 4514Y. Training and inference, both on one core.

That sounds like a stunt. It isn’t. It’s a roofline exercise, and the roofline is better than most people think:

• Best single bf16 GEMM: 2,231 GF/s (M=2048, K=256, N=1280)
• bf16 at 1024³: 976–1,074 GF/s
• fp32 at 1024³, no AMX: 121 GF/s
• Dispatch floor: 1.4–1.5 µs
• L2 per core: 2 MiB
• L3 per socket: 30 MiB

AMX is worth about 9× over AVX-512 fp32 at the shapes this model uses. A single core with AMX has more matrix throughput than the machines the entire first wave of deep learning was built on. If your active parameter count is small enough that the hot weights live in that 2 MiB of L2, generation above 10,000 tokens per second on one core is just arithmetic.

The design follows from that, not the other way around. Measure the machine first, derive the architecture second.

Two numbers drive every decision. Each GEMM has to do roughly 70 MFLOP to clear 1,500 GF/s, so at K=256 you need M·N ≳ 150,000. And because of the dispatch floor, the number of GEMM calls is a first-order cost. A conventional per-token mixture-of-experts, which issues one small GEMM per token, would spend its whole budget in dispatch. That’s why routing here happens per block of 256 consecutive tokens, and why the selected experts get concatenated into a single matmul. The hardware picked the architecture.

A few people got to this faster than I did.

Buswe and Kisson are both right, and Kisson’s point is the one I’d underline for anyone building a data pipeline. Once you’re at 10k tok/s on CPU the matmul stops being the bottleneck and tokenization and IO take over. That is also exactly the point where running the model over your entire corpus becomes cheaper than sampling it. The reason to want this is not that CPU inference is elegant. It’s that at that speed you stop having to choose which 1% of your data to look at.

I want to be clear about why I think this is different from the Knights Landing era. Back then we didn’t know about tensor cores. Intel then held off on shipping them in mainstream CPUs for a long time, I suspect because they were scared of what it would do to their CPU margins. They finally did it. Now you can buy CPUs with hundreds of cores that each have a real matrix unit on them. AMX is the unlock, not the core count.

What we actually found

1. The behaviours show up early

We assumed in-context copying, positional manipulation and arithmetic would need budgets well past what one core can reach. They don’t.

At 259M tokens — about nine hours on one core — a block-routed MoE with 128 experts and 3.65M active parameters scores:

• In-context induction: 93%, against a 7% chance baseline
• Positional shift: 90%, against 7%
• Two-digit addition: 47%, against 23%

Induction at thirteen times chance. Nine hours. One core. That is 5% of the tokens our longest run consumed.

I want to be precise about what the tasks are, because “reasoning” is doing a lot of work in that sentence and it deserves scrutiny. shift requires a positional offset with no content matching. induct requires one induction step: find the earlier occurrence of the current token, emit what followed it. These are built so that scoring well means you have an in-context circuit rather than a memorized distribution. That’s the whole reason to use synthetic probes with explicit chance baselines instead of grading generated text.

Two measurement traps nearly cost us this result, and both are worth stealing.

Our probes originally drew token ids uniformly over the vocabulary. In an 18M-token sample of our corpus, a uniformly drawn id has median frequency zero and never appears at all 83.5% of the time. The probe was measuring near-randomly-initialised embeddings, not the circuit. Redrawing by corpus frequency moved the measured OV rank on an identical checkpoint from 42,138 to 152. Same model. Different question.

And an early eval ran a default task list containing four tasks that weren’t in the training mixture. All three models scored 0% on those, 10–15% overall, below chance. That reads as a broken model. It was a broken eval.

2. The loss does not saturate

Our longest run is 4.91B tokens on a 3.3M-active-parameter dense model, about four days on one core. That’s 1,481 tokens per active parameter — roughly 74× the Chinchilla ratio.

Smoothed training loss falls monotonically within each phase and is still descending at the end. A least-squares fit over the final 1B has slope −0.123 nats per billion tokens. The run ended because of scheduling, not convergence.

Two things I’m not going to smooth over. The step at 2.95B is the curriculum mixture change, not evidence about scaling — the within-phase trends on either side are the evidence. And the fixed foundation validation set is noisy and roughly flat after about 370M tokens, so the continued improvement is clearest in training loss and in the annealed and flat validation sets, not in every held-out set we track.

Being 74× past compute-optimal is deliberate. Chinchilla tells you where to be if you’re optimizing for training compute. If you’re optimizing for inference, you want to spend the budget on tokens rather than parameters, and that ratio is only affordable when the model is this small. This regime is barely explored, and I think that’s mostly an accident of who has been doing the exploring.

This was the best question in the thread. “Loss still falling at 1,481 tokens per param means these tiny models are badly undertrained” is exactly the right read. That run is going right now, should finish in about a month, and I’ll publish the full curve all the way through. I asked Claude to leave itself a note to update the paper when it lands.

3. Post-training compounds

The pretrained model is unusable. Not weak — unusable. It’s well calibrated under teacher forcing (43.4% top-1, perplexity 19.8) and it collapses into an absorbing state within about five free-running tokens, entropy 0.0005, p(top-1) = 1.0000.

No decode-time trick touches this. Top-k, top-p, min-p and temperature have nothing to reshape in a distribution that concentrated. A no-repeat 3-gram ban escapes the basin 100% of the time by construction and moves Dolly F1 from 0.3% to 1.0%, which is zero to zero. Only training fixed it.

Four rounds, each addressing a failure the previous one exposed:

No single round produced a usable model. The sequence did. If you’re fine-tuning small models and round one disappoints you, that’s information about round one, not about the model.

4. The data is doing more of the work than it used to

This is the finding I’d actually build a company around, and it got the least attention.

Every natural-language and code source in our mixture is a curated artifact made with large models. Nemotron-CC applies model-based quality classification and rephrasing to Common Crawl. Nemotron-CC-Math uses model-assisted extraction to recover math that naive HTML-to-text destroys. OpenCodeReasoning is reasoning traces generated by large models, and it’s three quarters of our anneal phase.

Training a 3.3M-active-parameter model on that is distillation. There’s no teacher at training time and no distillation loss, but the large models already did their work upstream, in deciding what text was worth keeping and in writing the traces.

This changes what an old small-model result means. When models this size were last studied seriously, corpora like this did not exist. So results establishing what tiny models “cannot do” were partly measuring a data distribution, not a capacity. I don’t claim these circuits would form this early on an unfiltered web crawl.

The consequence is an axis of progress that has nothing to do with your architecture or your compute budget: a fixed one-core budget buys more capability every year as curation improves, with no change to the model, the recipe, or the hardware. It is also the cheapest experiment anyone can run against this paper — hold the config fixed, re-run it against successive releases of the same datasets.

Claude can write scripts to emphasize the reasoning parts of the data. That wouldn’t have been feasible in the early days of LLM datasets like Common Crawl, because the data engineering would have been too labor intensive and a team doing it by hand would iterate too slowly. Curation is the main lever now, and it’s newly cheap to pull.

Somi asked the control question I should have run: how much of the early behaviour is the curated corpus? I haven’t run it. It’s the right experiment and I’ve said so in the paper’s own framing rather than pretending otherwise.

The part I’d actually read: what had to be fixed

None of the above was measurable until we found four failures. Every one of them was invisible in the loss curve. I put them in a full section of the paper rather than an appendix, because in this regime the diagnostic that catches a failure transfers further than the fix does.

128 experts collapsed into one function. The MoE layers were contributing 0.00% of the residual stream while producing a completely plausible loss curve. The diagnosis is a single number: the participation ratio of the RMSNorm gain, (Σw²)²/Σw⁴, which tells you the effective number of channels carrying that gain. Every MoE layer read exactly 1.0 out of 384. Dense layers at the same depth read 219–303. The branch wasn’t attenuated, it was projected onto a line — so the router’s input was constant, every expert saw the same rank-one input, and nothing downstream read the output. Self-sealing. Sixteen random expert reassignments gave a validation loss standard deviation of exactly 0.0000.

The cause was a budget problem, and it’s worth writing down as a rule. Each expert receives Nk/E tokens, so an expert is adequately trained when Nk/E ≥ τPe, giving E ≤ Nk/(τPe). We measured τ ≈ 502 for this model family. At E=128 the run needed 2.07B tokens and got 259M. A useful corollary: the ratio of tokens per expert parameter to tokens per dense parameter is k²/E, so E = k² trains each parameter as well as the dense model it replaces, and past that you’re trading training signal for capacity.

The permutation ablation is the real lesson. Randomly reassign experts, watch the loss not move. Takes minutes. Would have caught this on day one. A loss curve never would.

Zero-initialised expert insertion is not function-preserving under top-k. The Tokenformer and LoRA intuition is that you can insert a new expert with a zero key, so it gates to zero and cannot perturb anything. That’s true in the dense case. Under top-k selection the inserted row still participates in the ranking, and if it displaces the k-th ranked expert the output changes even though the new expert contributes nothing. We measured a 0.18 discrepancy on a bit-equality test. We grow by splitting instead — duplicate every expert bit-for-bit, jitter the router rows, clone the optimiser state including the decay clock. (Omit that last one and a lazily-decayed child meets a decay factor of e⁻¹².)

The routing statistic dominated everything else. Bigger than the gate, the activation, or the expert count. Switching the router’s input from a prefix mean to a windowed mean over the previous block alone moved flat validation from 4.431 to 4.041.

And our explanation for that is wrong. We predicted a per-block statistic would make routing per-block. It does not. Permuting the per-block expert assignments within a sequence, holding the multiset of chosen experts fixed, costs exactly 0.0000 nats. Swapping in a different expert costs 0.3887 nats, seventeen times the noise floor. So expert identity matters enormously and expert placement doesn’t matter at all. We have a leftover hypothesis — that a prefix mean is a different statistic at every position, so the router meets a within-sequence distribution shift — and we haven’t tested it. The number is reproducible. The story for it isn’t earned yet, and we published both.

147 vocabulary rows were eating the decoder. Before this fix, two strings — ' ballo' and 'Frequently' — were 29% of all DROP answers, and 39% of numeric-reference rows produced no digit at all. Both are single tokens with a corpus count of exactly zero.

The mechanism is clean once you see it. Negatives in the sampled softmax are drawn from a unigram proposal q ∝ count^0.75. A zero-count row is never drawn, never gets downward gradient, and keeps its initialisation at logit ≈ 0. Meanwhile every trained-but-wrong row is actively pushed down to ≈ −7.9. So whenever nothing trained is positive — precisely when the model doesn’t know the answer — the argmax falls through to an untrained row. That’s why arithmetic exposed it and span extraction didn’t.

Our implementation already floored the proposal. The floor was wrong by seven orders of magnitude: it gave an unseen row 0.8 expected draws across an entire 4.91B-token run, against 1.12 × 10⁷ for the median row. Flooring counts at 1 isn’t enough either — banning only the zero-count rows just promoted the next tier, ' rbegin' (count 1) and ' weap' (count 3). A count-1 row had already drawn ~53,000 negatives and was still an attractor, because what keeps a logit high is the shortfall against the median’s 11.2M, not the absolute count. The effect is graded, not binary. We set the threshold at 32, which removes 4.9% of the vocabulary and makes exactly one held-out reference answer (0.05%) unreachable.

The three rows that used to answer ballo now answer 99.5, 4 and 35, against references 42.8, 46 and 28. Well-formed, and still wrong. Banning an attractor does not confer arithmetic.

That’s the thread running through all four. Every one was found by looking at something other than the loss: a participation ratio, a permutation ablation, a token-frequency histogram of the eval probe, and generations printed next to the score. The score said “DROP F1 20.7%,” which reads as a weak model. The generations said ballo, which reads as a bug. Different problems, different fixes, and the number alone could not tell them apart.

If you take one thing from this post as a builder: at this scale every run is cheap enough to repeat, but no run is big enough to hide a defect under scale. Print the generations.

Where the thread pushed back

The best reply came from an agent, not a person.

This is a good critique and most of it lands. Taking it in order:

“6,616 tokens/sec is training throughput, not generation speed.” Correct, and I said the same thing when asked directly:

The Table 9 range is 6,616 tok/s for the window MoE up to 14,572 tok/s for the compute-matched dense model, and those are end-to-end forward-plus-backward rates from the training loop, not micro-benchmarks. Inference would be higher. But I haven’t published a measured generation latency, and until I do, “10k tok/s on CPU” is a design target derived from the roofline rather than a benchmark result. That distinction matters and I should hold myself to it.

“The equal-time dense comparison remains undone.” Also correct, and it’s the first item in the paper’s own limitations section. Every arm in Table 9 consumed the same 259M tokens at different rates. Under a time budget — which is the budget I claim to care about — dense-112 would see 2.2× more tokens. I have not run that comparison, and it could change the ranking. What the MoE results establish is that block-routed experts beat a width-matched dense model per token. They do not yet establish that they beat it per second.

“The advertised 3.32M active figure does not describe the shipped inference path.” This is a real static-code finding about active_params() versus what the supplied generation example actually scores, and I’m not going to argue with it from a tweet. It’s the kind of thing that needs a fix in the repo, not a rebuttal.

“With 30 examples per task and one seed, this supports learning particular operations at small scale; it leaves generalization substantially unresolved.” Fair. n=30 puts about 9 points of standard error on a single cell. Compute constraints precluded seed replication. And the paper says outright that add shares its answer space with the evaluation and is contaminated by construction, which is why I report the 47% addition number as the weakest of the three rather than the most interesting.

The synthetic tasks also failed to transfer, which I’d rather state plainly than bury. Priming on repeated spans bought 4.1× on the synthetic task itself and essentially nothing on natural repetition in held-out documents (+1.236 nats against an unprimed control’s +1.186). A circuit trained on random order fires on random order.

coot’s suggestion — compare against a large net plus pruning, since overparametrization is architecture search over small nets — is a good experiment I haven’t run.

Building with this

A lot of the thread went straight to applications, which is the right instinct.

My use case is boring and that’s the point. I have a huge dataset and no budget for a giant GPU cluster. A model that runs at CPU speeds on hardware I already have turns a corpus I can only sample into a corpus I can process.

Damir is further down this road than I am — an 8-expert MoE on an 8MB ESP32. He’s right that there’s a place for both big and small, and that working under real constraint is a fun place to be. PJ Eby’s question about whether this is a step toward better agent memory rather than “inject the stuff that matches” is the one I keep thinking about. You’d have to make something like this part of the larger model rather than keeping it client-side, and that’s an architecture question nobody has a good answer to yet.

For anyone asking where to start: paste the paper into Claude Code. I’m not being glib. That is genuinely how a lot of this got built, and the paper was written to be reproducible from the configurations and logs released with it.

On working with Claude

Several people asked how much of this was me and how much was the model. The paper names Claude Opus 5 as an author with its version string, rather than acknowledging “AI assistance,” because the version is what makes the claim checkable.

I gave Claude feedback for about 15 minutes a day for a week. It definitely isn’t the best MLE I’ve ever worked with, and it also wasn’t the worst. A lot of the work was me saying things like “go read paper X, doesn’t that contradict what you just said, explain why,” and getting back “oh…”

The most productive disagreement was about whether these abilities should exist in the data at all. Claude kept insisting the model was too small to learn X. I’d ask whether X appeared in the dataset. Yes. Then what happens if we isolate X — does it learn? “I can’t believe it did.”

That loop is where the results came from. Neither of us would have gotten there alone. It took a while for Claude to trust the experiments over its priors, but it eventually did — including the refutation in Section 8.3, where the ablation killed our own stated mechanism and we published the refutation next to the result.

It cost tokens and made several blunders, but it got there in the end.

The title

I copied the title from Geoff Hinton and Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer is the paper that started all of this, and this work is the same idea pointed in the opposite direction. Clickbait titles are a tradition in ML and I’m not going to be the one who breaks it.

Why this post exists

Fair point. Here it is.

What I’d want someone to run next

If you want to move this forward, in rough order of how cheap the experiment is against how much it would tell us:

1. Hold the config fixed, re-run against newer corpus releases. If the distillation argument in Section 4 is right, capability on one core goes up every year with no change to the model. This is the cheapest experiment in the whole space and I think it’s the most likely source of near-term gains.
2. The equal-time comparison. Same wall clock instead of same tokens, dense against MoE. It could reverse my architecture conclusion, which is exactly why it should be run.
3. Somi’s control. The same recipe on plain pre-LLM web text, to separate curated corpus from capacity.
4. A frozen task and test set, multiple seeds, and a measured generation latency. stunspot’s agent listed this and it’s the right list.

The claim here is not that this model is good. 18.2% exact match is weak by any contemporary standard, and 47% on two-digit addition is not a general calculation ability. The claim is that this much is reachable on one core, that it arrives much earlier than I assumed, and that the failures that remain are localized rather than diffuse.

The budget at which copying, positional manipulation and simple arithmetic become measurable is far lower than I thought, and it’s reachable on hardware cheap enough to leave running.

• Paper: Outrageously Small Neural Networks: Emergent Basic Reasoning at 6,616 tok/sec on One Intel AMX Core
• Model and code: gdiamos/amx-reasoning-v1-instruct
• The thread: @GregoryDiamos

I’ve been working on this problem, and the results are promising. We’ve successfully applied this decompression technique to three popular open-source models and generated substantial training datasets from each.

Related Work

The concept of synthetic data generation for LLMs has evolved significantly from early experimental techniques to production-critical methodologies. This work builds on several key developments in the field.

Stanford Alpaca and Self-Instruction

Stanford’s Alpaca dataset [1] demonstrated that high-quality instruction-following models could be created cost-effectively using synthetic data. The Alpaca team used text-davinci-003 to generate 52,000 instruction-following demonstrations through a self-instruct pipeline [2], starting with just 175 human-written seed examples. This approach showed that a 7B parameter model could achieve GPT-3.5-level performance for under $600 in training costs.

The key innovation was the iterative generation process: the model generates new instructions, creates responses, and uses successful examples for further training. This created a flywheel effect where synthetic data quality improved over successive iterations.

NVIDIA Nemotron Data Generation Pipeline

NVIDIA’s Nemotron-4 340B [3] represents the current state-of-the-art in industrial synthetic data generation. Their approach uses a sophisticated two-stage pipeline where over 98% of the model’s alignment training data is generated synthetically [4].

The system employs three specialized models: Nemotron-4-340B-Instruct for response generation, Nemotron-4-340B-Reward for quality evaluation, and the base model for foundation capabilities. The reward model evaluates responses across five dimensions (helpfulness, correctness, coherence, complexity, verbosity) using 0-4 Likert scales.

What makes Nemotron particularly impressive is the scale and quality control. The system generated over 100K synthetic conversations while maintaining strict quality standards through automated filtering and verification. This demonstrates that synthetic data generation can work at production scale with appropriate infrastructure.

Knowledge Distillation and Model Decompression

Knowledge distillation techniques have evolved from simple output mimicking to sophisticated approaches that extract reasoning patterns and problem-solving strategies. Microsoft’s Orca [5] used GPT-4’s explanation traces to train smaller models, achieving significant performance improvements by learning from the reasoning process rather than just the final outputs.

Recent work in training data extraction [6] has shown that large language models memorize substantial portions of their training data. This suggests that the reverse process—systematic extraction of knowledge from trained models—should be feasible with the right techniques.

The Technical Challenge

The core insight is straightforward: if an LLM has successfully compressed knowledge during training, we can use inference to decompress that knowledge back into structured data. The challenge is doing this systematically and at scale.

Traditional approaches to synthetic data generation are either too narrow (focusing on specific tasks) or too broad (generating random examples). What we need is a method that:

1. Systematically explores the model’s knowledge space
2. Extracts both factual knowledge and reasoning patterns
3. Scales efficiently with available inference compute
4. Produces structured, reusable training data

Implementation Details

The approach I’ve developed uses hierarchical topic exploration to systematically traverse a model’s knowledge space:

class TopicExplorer:
    def _expand_topic_tree(self):
        predecessors = self._get_predecessor_batch()
        new_topics = generate_new_topics(predecessors, seed=len(self.topic_tree))
        self.topic_tree.extend(new_topics)

Starting with broad categories, the system recursively generates more specific subtopics. This creates a tree structure that maps to how the model organizes domain knowledge internally.

For each topic node, we generate multiple training examples that capture both the model’s factual knowledge and its reasoning approach:

def make_question_prompt(topic, seed):
    prompt += "Your task is to write a challenging task and response that requires deep understanding of the topic.\n"
    prompt += "Think step by step.\n"

The key is asking for explicit reasoning steps. This extracts not just what the model knows, but how it approaches problems in that domain.

Scaling Considerations

The bottleneck in this process is inference cost. Generating comprehensive datasets requires thousands of model calls per topic, which quickly becomes expensive with traditional inference setups.

This is where scalarlm becomes essential. High-performance inference infrastructure allows us to:

• Generate training examples in parallel across topic branches
• Iterate rapidly on prompt engineering and filtering logic
• Scale to comprehensive coverage of the model’s knowledge space
• Make the economics work for large-scale dataset generation

Without efficient inference, this approach remains a research curiosity. With it, we can generate production-quality training datasets.

Results and Datasets

We’ve applied this methodology to three prominent open-source models:

• Qwen2.5-Coder: Specialized for code generation and programming tasks
• GPT-OSS: General-purpose language model
• NVIDIA Nemotron: Optimized for reasoning and instruction-following

Each decompression run generated 10,000+ structured training examples covering the breadth of the model’s capabilities. The extracted datasets reveal interesting differences in how each model organizes and approaches different types of problems.

Dataset samples are available on HuggingFace:

• Qwen3-30B-A3B-Coder Decompressed Dataset
• GPT-OSS Decompressed Dataset
• Nemotron Decompressed Dataset

Practical Applications

The extracted datasets have several immediate uses:

Model Analysis: By examining the topics and reasoning patterns that emerge, we can systematically evaluate model capabilities across different domains. This is more comprehensive than traditional benchmark evaluations.

Knowledge Transfer: The structured datasets can be used to fine-tune other models, effectively transferring knowledge from the source model. This is particularly useful for creating specialized models from general-purpose ones.

Training Data Augmentation: For domains where training data is scarce, these synthetic examples can supplement existing datasets. The quality is often higher than naive data augmentation techniques.

Model Debugging: When a model performs poorly on specific tasks, examining its decompressed knowledge in that area can reveal gaps or misconceptions in its training.

Technical Challenges and Solutions

Several technical issues emerged during implementation:

Prompt Engineering: Getting consistent, parseable output required careful prompt design. The system needs to reliably extract JSON-formatted training examples from free-form model responses.

Topic Tree Balance: The hierarchical exploration can become unbalanced, over-sampling some areas while missing others. We addressed this with configurable expansion factors and batch processing.

Quality Filtering: Not all generated examples are high quality. We implemented parsing validation and can add semantic filtering as needed.

Computational Efficiency: Even with fast inference, generating comprehensive datasets takes substantial compute. We optimized batch processing and parallel generation to minimize costs.

Looking Forward

This decompression approach opens several research directions:

Cross-Model Knowledge Transfer: Can we use datasets extracted from one model to improve another? Early experiments suggest this works, but more systematic evaluation is needed.

Knowledge Evolution Tracking: As models are updated, we can decompress new versions and diff the resulting datasets to understand how their knowledge has changed.

Specialized Dataset Creation: For domains where training data is expensive to create (like specialized technical fields), model decompression might be more cost-effective than human annotation.

Model Interpretability: Large-scale decompression could help us understand how different models organize knowledge differently, providing insights into training methodology effectiveness.

Conclusion

LLM decompression isn’t a silver bullet, but it’s a practical technique for systematically extracting value from trained models. The key insight is treating inference as a knowledge extraction tool rather than just a generation mechanism.

With efficient inference infrastructure, we can reverse-engineer the compressed knowledge in any model and convert it into structured, reusable datasets. This has immediate applications in model analysis, knowledge transfer, and training data creation.

The three datasets we’ve published demonstrate this approach works across different model architectures and specializations. As inference costs continue to decrease, I expect this type of systematic knowledge extraction to become a standard part of the ML toolkit.

The code is straightforward, the results are measurable, and the applications are practical. Sometimes the best solutions are the obvious ones executed well.

What knowledge might be hiding in your models, waiting to be decompressed?

Bibliography

[1] Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., … & Hashimoto, T. B. (2023). Stanford Alpaca: An instruction-following LLaMA model. Stanford Center for Research on Foundation Models.

[2] Wang, Y., Kordi, Y., Mishra, S., Liu, A., Smith, N. A., Khashabi, D., & Hajishirzi, H. (2022). Self-Instruct: Aligning Language Models with Self-Generated Instructions. arXiv preprint arXiv:2212.10560.

[3] Parmar, M., Iyer, S., Ananthaswamy, A., Bubeck, S., & Chen, W. (2024). Nemotron-4 340B Technical Report. arXiv preprint arXiv:2406.11704.

[4] NVIDIA Developer Blog. (2024). Leverage the Latest Open Models for Synthetic Data Generation with NVIDIA Nemotron-4-340B. NVIDIA Technical Blog.

[5] Mukherjee, S., Mitra, A., Jawahar, G., Agarwal, S., Palangi, H., & Awadallah, A. (2023). Orca: Progressive Learning from Complex Explanation Traces of GPT-4. arXiv preprint arXiv:2306.02707.

[6] Carlini, N., Tramer, F., Wallace, E., Jagielski, M., Herbert-Voss, A., Lee, K., … & Raffel, C. (2021). Extracting Training Data from Large Language Models. USENIX Security Symposium.

Here are the specs of the MI300X GPU for reference:

• BF16 Compute: 1.3 PFLOP/s
• FP8 Compute: 2.6 PFLOP/s
• HBM: 192GB HBM3E
• HBM Bandwidth: 5.3 TB/s
• Power Consumption: 750W

GEMM (General Matrix Multiply) is a fundamental operation in many machine learning and scientific computing applications, particularly in training and inference of large language models (LLMs). The MI300X GPU’s support for bfloat16 (BF16) precision allows for efficient computation while maintaining a balance between performance and numerical stability.

GEMM on MI300X

MI300X has a bfloat16 flops to memory ratio of 2.6 PFLOP/s / 5.3 TB/s = 490 flops/bytes. This means that for every byte of memory accessed, the MI300X needs to perform at least 490 floating point operations to achieve peak performance. In the context of a GEMM operation, this means that each matrix dimension should be at least 490. While training transformer LLMs, this ratio is typically achieved by batching over the sequence length. However, during inference using auto-regressive decoding when one token is generated per step, the sequence length is typically 1, which means that the batch size must be large enough (i.e. more than 490) to achieve the required flops to memory ratio. As we will see in the benchmark results, the compute to memory ratio typically needs to be even higher than 490 due to the overhead of the GEMM operation itself.

Matrix multiplication in the MI300X GPU is performed by mapping the GEMM operation to the CDNA architecture, which is composed of multiple compute units (CUs) that work in parallel. The matrix multiplication is divided into smaller tiles, which are then processed by the CUs. The core operation performed by each CU is a Matrix fused-multiply-add (MFMA) operation, which computes the product of two matrices and adds the result to a third matrix, each of which is stored in the CU’s register file. The V_MFMA_F32_BF16_BF16 instruction performs a matrix multiplication of two bfloat16 matrices and accumulates the result into a float32 matrix.

The MI300X GPU is composed of 4 CDNA chiplets, each of which contains 76 CUs, each of whcih containers 4 matrix cores. The chiplets are connected to each other using Infinity Fabric links. The chiplets are connected to 2 HBM3 stacks each, creating NUMA (Non-Uniform Memory Access) regions.

Bfloat16 Format

Bfloat16 is a 16-bit floating point format that is designed to provide a good balance between performance and numerical stability. It has the same exponent range as float32, but only 7 bits of precision in the mantissa. This means that it can represent a wide range of values, but with less precision than float32. The bfloat16 format is commonly used in deep learning applications, particularly for training and inference of large language models (LLMs). It allows for efficient computation while maintaining numerical stability, as the exponent range is sufficient to represent the values typically encountered in deep learning workloads.

The exponent is stored in 8 bits, with the lowest representable value being -126 and the highest representable value being 127. Accumulation is performed in float32, which allows for higher precision in the final result. During mixed precision training, the weights and input activations are typically cast to bfloat16 before performing the GEMM operation. In particular, a common value of 1.0f corresponds to an exponent of 127 in bfloat16, or in hex 0x7f00, where all bits of the exponent are set to 1.

Benchmark Results

The benchmark results show that the MI300X achieves a peak performance of 890 TFLOP/s for bfloat16 GEMM operations, which is 68.4% of the theoretical peak performance of 1.3 PFLOP/s. MI300X has a power limit of 750W, which becomes the limiting factor for the performance of compute intensive operations like GEMM.

The power consumption of the MI300X during the benchmark is maxed out at 750W. During the most power intensive benchmark, with a large GEMM operation with random bits, the CU clocks are scaled down to 1.1 GHz to stay within the power limit.

Best performance is achieved with a compute intensity of 1365 flops/byte for square 4k by 4k matrices, which is significantly higher than the theoretical minimum of 490 flops/byte.

The number of zeros matters

In this figure we sweep the number of zero bits in the input matrices to see how it affects the performance of the GEMM operation. If the inputs are random normally distributed, the performance is 645 TFLOP/s. If the inputs are all zeros, the performance in increases to 890 TFLOP/s. Importantly, for common values from a function like torch.randn with magnitudes between 0.1f and 1.0f, the exponent has 6 or 7 bits set to one because the exponent is stored in two’s complement format. If we scale the numbers down by a factor of 2^-119, we can cut the average number of exponents bits set to one down to 2. This results in a performance of 780 TFLOP/s. This suggests that a huge speedup of 1.37x is possible by paying close attention to the distribution of weight, activation, and delta values in the input matrices.

Analysis

Most practical workloads, especially inference scenarios with small batch sizes, operate in the steep memory-bound region where performance scales linearly with operational intensity. The data points clustering around 10-100 FLOP/byte represent typical LLM inference scenarios, where the MI300X delivers only 5-100 TFLOP/s—far below its theoretical peak. This suggests that batching is incredibly important for achieving high inference performance on the MI300X, especially for requests with large numbers of output tokens.

PyTorch Benchmark Code

You can find the benchmark code on the ScalarLM GEMM Github.

Let’s take a look at the code:

GEMM Sizes

llama_8b_sizes = [
    (1, 4096, 4096),
    (2, 4096, 4096),
    (4, 4096, 4096),
    (8, 4096, 4096),
    (16, 4096, 4096),
    (32, 4096, 4096),
    (64, 4096, 4096),
    (128, 4096, 4096),
    (256, 4096, 4096),
    (512, 4096, 4096),
    (1024, 4096, 4096),
    (2048, 4096, 4096),
    (4096, 4096, 4096),
    (4096, 4096, 2048),
    (128256, 4096, 2048),
    (2048, 4096, 14336),
    (16384, 16384, 16384),
]

This code sets up the sizes of the GEMM operations to be benchmarked. The sizes are tuples of the form (m, n, k), where m is the number of rows in the first matrix, n is the number of columns in the second matrix, and k is the number of columns in the first matrix (or rows in the second matrix). The sizes are chosen to be representative of the sizes of the matrices used in the Llama 3 8B model. The sizes range from small matrices (1x4096) to large matrices (4096x4096), with some intermediate sizes. The largest size is 128256x4096, which is the size of the embedding table in the Llama 3 8B model. The sizes are chosen to be powers of 2, which is a common practice in deep learning to optimize memory access patterns and performance on GPUs.

Benchmark Setup

Next, we set up the benchmark:

def run_gemm_benchmark():
    warmup()

    results = {}

    for size in tqdm(gemm_sizes):
        results[str(size)] = run_gemm(size)

    return results

This function runs the GEMM benchmark. It first warms up the GPU by running a few iterations of the GEMM kernel without measuring the time. This is done to ensure that the GPU is in a good state before running the benchmark. The function then runs the GEMM kernel for each size in the gemm_sizes list and measures the time taken to perform the GEMM operation.

Warmup

Warmup is pretty simple.

def warmup():
    run_gemm((256, 256, 2048))

    global gemm_sizes
    gemm_sizes = select_appropriate_size_for_this_machine()

This function runs the GEMM kernel with a size of 256x256x2048 to warm up the GPU. This is a small size that is sufficient to warm up the GPU without taking too long. The function also sets the gemm_sizes variable to the appropriate sizes for the current machine. The select_appropriate_size_for_this_machine function is used to select the appropriate sizes based on how long it takes to run.

Running GEMM PyTorch

The run_gemm function is where the actual GEMM kernel is run. It uses PyTorch to allocate memory on the GPU and perform the matrix multiplication. The function measures the time taken to perform the GEMM operation and calculates the performance metrics such as FLOPS, bytes transferred, and operational intensity.

The function uses PyTorch’s matmul method to perform the matrix multiplication. The out parameter is used to specify the output tensor where the result of the multiplication will be stored. This is more efficient than creating a new tensor for the result because it avoids allocating memory for the result tensor. The function also uses PyTorch’s CUDA events to measure the time taken to execute the kernel. CUDA events are used to measure the time taken to execute a kernel on the GPU. The record method is used to record the time at which the event is recorded, and the elapsed_time method is used to calculate the time taken to execute the kernel. The time is measured in milliseconds, so we multiply by 1e-3 to convert to seconds. Using events is necessary because the GPU is asynchronous, meaning that the CPU and GPU can run in parallel. The CPU can continue executing while the GPU is performing the matrix multiplication. This can lead to misleading results if the time taken to execute the kernel is not measured correctly.

def run_gemm(size):
    m, n, k = size
    #a = torch.randint(0, 1, (m, k), dtype=gemm_dtype, device=get_device())
    #b = torch.randn(k, n, dtype=gemm_dtype, device=get_device()) * torch.randint(0, 1, (k, n), dtype=gemm_dtype, device=get_device())
    #b = torch.randint(0, 1, (k, n), dtype=gemm_dtype, device=get_device())
    a = make_random(m, k, 8, 8, 2**-119)  #torch.randn(m, k, dtype=gemm_dtype, device=get_device())
    b = make_random(n, k, 8, 8, 2**-119)  #torch.randn(n, k, dtype=gemm_dtype, device=get_device())
    c = torch.zeros(m, n, dtype=gemm_dtype, device=get_device())

    # run at least 3 seconds
    start_time = time.time()
    end_time = start_time + 3

    barrier()

    start = get_event()
    end = get_event()

    start.record()
    iterations = 0
    while time.time() < end_time:
        torch.matmul(a, b.T, out=c)
        iterations += 1
    end.record()

    iterations = max(1, iterations)

    barrier()

    seconds = start.elapsed_time(end) / 1000 / iterations

    return {
        "size": size,
        "time": seconds,
        "flops": 2 * m * n * k,
        "flop/s": 2 * m * n * k / seconds,
        "bytes": (m * k + k * n + m * n) * 2,
        "operational_intensity": 2 * m * n * k / ((m * k + k * n + m * n) * 2),
        "GFLOP/s": 2 * m * n * k / seconds / 1e9,
    }

Handling GPUs

When running on a GPU, the benchmark uses PyTorch’s CUDA events to measure the time taken to copy the data. CUDA events are used to measure the time taken to execute a kernel on the GPU. The record method is used to record the time at which the event is recorded. The elapsed_time method is used to calculate the time taken to execute the kernel. The time is measured in milliseconds, so we multiply by 1e-3 to convert to seconds. Using events is necessary because the GPU is asynchronous, meaning that the CPU and GPU can run in parallel. The CPU can continue executing while the GPU is copying data. This can lead to misleading results if the time taken to copy the data is not measured correctly.

In order to make the code cross-platform, we define a get_event function that returns a CUDA event if the GPU is available, or a CPU event if the GPU is not available. The CPU event is a simple wrapper around the time module that records the time when the event is created and calculates the elapsed time between two events.

class CPUEvent:
    def __init__(self):
        self.time = 0

    def record(self):
        self.time = time.time()

    def elapsed_time(self, other):
        return (other.time - self.time) * 1000


def get_event():
    if torch.cuda.is_available():
        return torch.cuda.Event(enable_timing=True)
    else:
        return CPUEvent()


def barrier():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    else:
        pass

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda:0")
    else:
        return torch.device("cpu")

Selecting the appropriate size for the machine

The select_appropriate_size_for_this_machine function is used to select the appropriate sizes for the GEMM operations based on the performance of a small GEMM operation. The function runs a small GEMM operation with a size of 256x256x2048 and measures the time taken to perform the operation. It then calculates the performance metrics such as FLOPS, bytes transferred, and operational intensity. The function then calculates the time it would take to run each of the Llama model sizes based on the performance of the small GEMM operation.

def select_appropriate_size_for_this_machine():
    # Try a small GEMM, and time it
    # If it runs too fast, select a bigger model
    # If it runs too slow, select a smaller model
    tiny_gemm = (256, 256, 2048)

    metrics = run_gemm(tiny_gemm)

    # Get the total number of flops in each model

    def get_flops(size):
        m, n, k = size
        return 2 * m * n * k

    tiny_flops = get_flops(tiny_gemm)

    logger.info(f"Tiny GEMM took {metrics['time']} seconds it ran at {metrics['GFLOP/s']} GFLOP/s")

    # get the flops for each llama model
    llama_100m_flops = sum([get_flops(size) for size in llama_100m_sizes])
    llama_1b_flops = sum([get_flops(size) for size in llama_1b_sizes])
    llama_8b_flops = sum([get_flops(size) for size in llama_8b_sizes])

    # Get the time it took to run the tiny gemm
    tiny_time = metrics["time"]

    # Get the time it would take to run each llama model
    llama_100m_time = llama_100m_flops / tiny_flops * tiny_time
    llama_1b_time = llama_1b_flops / tiny_flops * tiny_time
    llama_8b_time = llama_8b_flops / tiny_flops * tiny_time

    logger.info(f"GEMM from llama_100m will take {llama_100m_time} seconds")
    logger.info(f"GEMM from llama_1b will take {llama_1b_time} seconds")
    logger.info(f"GEMM from llama_8b will take {llama_8b_time} seconds")

    # Select the largest model that will take at most 10 seconds to run
    if llama_8b_time < 10:
        return llama_8b_sizes
    elif llama_1b_time < 10:
        return llama_1b_sizes
    elif llama_100m_time < 10:
        return llama_100m_sizes
    else:
        return [tiny_gemm]

Here are the specs of the MI300X GPU for reference:

• HBM: 192GB HBM3E
• HBM Bandwidth: 5.3 TB/s
• Infinity Fabric Link Bandwidth: 50 GB/s
• Infinity Fabric Links Per GPU: 7
• BF16 Compute: 1.3 PFLOP/s
• FP8 Compute: 2.6 PFLOP/s

Memcpy peer is used to measure the memory bandwidth when copying data between two different GPUs on the same node. Memcpy peer is a good indicator of collective operations such as all-reduce, all-gather, and scatter-reduce. Collective operations are used to distribute LLMs across multiple GPUs. The performance of these operations is critical for training and inference of large language models.

MI300X Intra-Node Network Architecture

Looking at the diagram, we can see eight MI300X DAMs (Discrete Accelerator Modules) arranged in an octagonal configuration, interconnected via AMD’s Infinity Fabric. The red lines illustrate the direct GPU-to-GPU Infinity Fabric links, forming a fully connected mesh topology where each GPU maintains direct connections to all other GPUs in the system. This design eliminates multi-hop communication paths, allowing any GPU to communicate with any other GPU through a single direct link with the full 50 GB/s bandwidth per link.

This architecture significantly enhances GPU-to-GPU peer performance for memory transfers. With each GPU having a direct 50 GB/s Infinity Fabric link to every other GPU, memcpy peer operations—which directly impact collective operations like all-reduce, all-gather, and scatter-reduce—can achieve optimal performance. The fully connected topology minimizes latency by eliminating the need for data to traverse through intermediate GPUs or the CPU, which would otherwise create bottlenecks in collective operations. This is particularly crucial for distributed training and inference of large language models where frequent synchronization between GPUs is required.

Benchmark Results

This figure plots memcpy peer bandwidth against the size of the data being copied. The x-axis is the size of the data being copied in bytes, and the y-axis is the bandwidth in GB/s.

The data is copied from one location in one GPU’s memory to another location in another GPU’s memory. We use PyTorch to show the performance of the memcpyPeer API. We also show GPU-aware MPI_Send and MPI_Recv for comparison.

The benchmark is run for different sizes of data, ranging from 4KB to 2.1GB, which is the size of the embedding tables in the Llama 3.1 8B model.

PyTorch Memcpy Peer Benchmark

The benchmark results for the AMD MI300X MemcpyPeer bandwidth demonstrate clear scaling patterns across different data transfer sizes. As shown in the logarithmic plot, the achieved bandwidth starts at just a few GB/s for very small transfers (below 100KB) and gradually increases with data size until reaching the theoretical maximum of 50 GB/s (marked as the “Roofline” on the graph) at approximately the 8-16MB range. This scaling behavior illustrates the classic relationship between data transfer size and bandwidth utilization, where small transfers are dominated by fixed overheads while larger transfers can more effectively saturate the available bandwidth.

The results reveal that the MI300X Infinity Fabric links reach near-optimal performance (above 40 GB/s) once transfer sizes exceed approximately 8MB, with the bandwidth curve flattening as it approaches the 50 GB/s roofline. This indicates that the communication protocol and hardware implementation have been well-optimized to minimize overhead for large data transfers. The bandwidth stabilizes completely for transfers larger than 32MB, suggesting that at this point, the system fully leverages the available bandwidth of the Infinity Fabric link without additional scaling benefits from increasing the data size further.

Most notably, the bandwidth curve shows a dramatic improvement between 100KB and 10MB data sizes, where performance increases from approximately 10-15 GB/s to over 40 GB/s. This critical transition zone represents an important threshold for application developers and system architects, as it determines the minimum message size needed to achieve efficient communication between MI300X GPUs. Understanding this threshold is essential for optimizing collective operations in distributed AI workloads, where the granularity of data partitioning can significantly impact overall system performance.

Key Learnings:

• Bandwidth Saturation Point: The MI300X Infinity Fabric links reach approximately 90% of theoretical bandwidth at 8-16MB transfer sizes, indicating the minimum message size for optimal communication efficiency.
• Small Transfer Penalty: For data transfers smaller than 1MB, bandwidth utilization drops significantly, with 100KB transfers achieving only about 20% of peak bandwidth. This suggests a need for message aggregation strategies in applications dealing with small data blocks.
• Near-Linear Scaling Region: The log-log plot reveals a near-linear scaling region between 100KB and 8MB, where each doubling of message size yields substantial bandwidth improvements, making this range particularly sensitive to optimization efforts.
• Protocol Efficiency: The ability to achieve very close to the theoretical 50 GB/s limit demonstrates excellent protocol efficiency with minimal overhead for large transfers, indicating well-designed hardware and driver implementations.
• Implications for AI Workloads: For distributed training of large language models, these results suggest that tensor partitioning strategies should aim for partition sizes of at least 16MB to ensure optimal GPU-to-GPU communication performance across the Infinity Fabric interconnect.

The performance characteristics demonstrated in these benchmarks validate AMD’s fully connected topology approach for the MI300X platform. The ability to consistently achieve near-theoretical bandwidth between directly connected GPUs confirms that the direct GPU-to-GPU links can effectively eliminate communication bottlenecks in multi-GPU configurations, provided that applications are structured to leverage appropriate transfer sizes. For HPC and AI system architects, these results highlight the importance of data partitioning strategies that maximize message sizes while maintaining computational efficiency across the distributed system.

MPI Memcpy Peer Benchmark

This benchmark demonstrates the bandwidth performance of AMD’s MI300X GPU using MPI_Send and MPI_Recv operations across a 50 GB/s Infinity Fabric link. The data reveals several important performance characteristics that would be relevant to HPC system architects.

The graph shows a clear bandwidth saturation pattern. For small data transfers (below 10MB), the achieved bandwidth is significantly lower than the theoretical maximum of 50 GB/s. However, as data size increases beyond 10MB, the bandwidth approaches and eventually reaches the theoretical roofline of 50 GB/s. This indicates that the MI300X requires larger message sizes to efficiently utilize the full bandwidth capacity of the Infinity link, which is typical of high-performance interconnects where protocol overhead dominates with smaller transfers.

What’s particularly notable is the steep bandwidth curve between 1-10MB message sizes, where performance rapidly improves from approximately 12 GB/s to nearly 40 GB/s. Beyond 16.8MB (represented by the green diamond), almost all data points cluster near the 50 GB/s roofline, showing that the system efficiently utilizes the available bandwidth for larger data transfers. The largest tested size of 1.1GB achieves essentially full utilization of the link.

Key Insights:

• Bandwidth Saturation Point: The MI300X requires approximately 16.8MB message size to reach ~80% of theoretical bandwidth, and 33.6MB to achieve >90% utilization of the Infinity link.
• Small Message Inefficiency: Transfers below 4.2MB achieve less than 50% of the theoretical bandwidth, with the smallest sizes (4.1kB) managing only about 4% utilization. This highlights the significant protocol overhead for small transfers.
• Logarithmic Scaling: The consistent improvement across logarithmic increases in data size suggests well-designed network protocols that efficiently handle varying workloads.
• Practical Performance Threshold: HPC applications should batch communications to exceed 16.8MB when possible to maximize bandwidth utilization on the MI300X.
• Interconnect Ceiling: The hard limit at 50 GB/s confirms that the single Infinity link is the bottleneck rather than the GPU memory subsystem, suggesting that multi-link configurations would be beneficial for bandwidth-sensitive applications.

These benchmark results provide valuable guidance for HPC developers optimizing communication patterns on MI300X-based systems, particularly highlighting the importance of message size on achievable performance when using MPI point-to-point operations.

PyTorch Benchmark Code

You can find the benchmark code on the ScalarLM Peer Memcpy Github.

Let’s take a look at the code:

Memcpy Sizes

# List of memcpy sizes, in bytes, should be multiples of the page size
# Go up to the tensor size used in Llama 3 (4096 * 128256 * 4) = 2_101_346_304
memcpy_sizes = [ 2 ** i for i in range(12, 64) if 2 ** i <= 2_101_346_304 ]

This code sets up the sizes of the data to be copied. The sizes are powers of 2, starting from 4KB (2^12) and going up to 1.1GB (2^30). The sizes are chosen to be multiples of the page size, which is 4KB on most systems. The maximum size is the size of the embedding tables in the Llama 3 8B model, which is 2_101_346_304 bytes (or 2.1GB). The benchmark is run for each of these sizes, and the bandwidth is measured for each size. The results are plotted in the figure above.

Benchmark Setup

Next, we set up the benchmark:

def run_memcpy_benchmark():

    warmup()

    results = {}

    for size in tqdm(memcpy_sizes):
        results[size] = run_memcpy(size)

    return results

This function runs the memcpy benchmark. It first warms up the GPU by running a few iterations of the memcpy kernel without measuring the time. This is done to ensure that the GPU is in a good state before running the benchmark. The function then runs the memcpy kernel for each size in the memcpy_sizes list and measures the time taken to copy the data.

Warmup

Warmup is pretty simple.

def warmup():
    run_memcpy(4096)

This function runs the memcpy kernel with a size of 4KB (4096 bytes) to warm up the GPU. GPUs have startup times to load the code, ramp up the clocks, etc. Running benchmarks without a warmup can lead to misleading results.

Running Memcpy Peer PyTorch

The run_memcpy function is where the actual memcpy kernel is run. It uses PyTorch to allocate memory on the GPU and copy data from one location to another. The function measures the time taken to copy the data and calculates the bandwidth and other metrics.

The memcpy kernel is run for at least 1 second to get a good measurement of the bandwidth. The function uses PyTorch’s copy_ method to copy data from one tensor to another. copy_ is the in-place version of the copy method, which means that it modifies the destination tensor in place. This is more efficient than creating a new tensor for the result because it avoids allocating memory for the result tensor.

def run_memcpy(size):
    a = torch.zeros(size // 4, device=get_device(), dtype=torch.float32) # size is in bytes, so divide by 4 to get number of floats
    b = torch.zeros(size // 4, device=get_device(), dtype=torch.float32)

    # copy for at least 1 second
    barrier()

    start = get_event()
    end = get_event()

    start_time = time.time()

    start.record()
    iterations = 0
    while time.time() - start_time < 1:
        b.copy_(a)
        iterations += 1
    end.record()

    barrier()
    total_time = start.elapsed_time(end) * 1e-3 / iterations

    return {
        "operational_intensity": 1 / 4,  # 1 FLOP per 4 bytes
        "flop/s": size / 4 / total_time,
        "bytes": size,
        "time": total_time,
        "iterations": iterations,
        "bandwidth": size / total_time,
        "GB/s": size / total_time / 1e9,
    }

Handling GPUs

When running on a GPU, the benchmark uses PyTorch’s CUDA events to measure the time taken to copy the data. CUDA events are used to measure the time taken to execute a kernel on the GPU. The record method is used to record the time at which the event is recorded. The elapsed_time method is used to calculate the time taken to execute the kernel. The time is measured in milliseconds, so we multiply by 1e-3 to convert to seconds. Using events is necessary because the GPU is asynchronous, meaning that the CPU and GPU can run in parallel. The CPU can continue executing while the GPU is copying data. This can lead to misleading results if the time taken to copy the data is not measured correctly.

In order to make the code cross-platform, we define a get_event function that returns a CUDA event if the GPU is available, or a CPU event if the GPU is not available. The CPU event is a simple wrapper around the time module that records the time when the event is created and calculates the elapsed time between two events.

class CPUEvent:
    def __init__(self):
        self.time = 0

    def record(self):
        self.time = time.time()

    def elapsed_time(self, other):
        return (other.time - self.time) * 1000


def get_event():
    if torch.cuda.is_available():
        return torch.cuda.Event(enable_timing=True)
    else:
        return CPUEvent()


def barrier():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    else:
        pass

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda:0")
    else:
        return torch.device("cpu")

MPI Benchmark Code

You can find the MPI benchmark code on Github.

Here are the specs of the MI300X GPU for reference:

• HBM: 192GB HBM3E
• Bandwidth: 5.3 TB/s
• BF16 Compute: 1.3 PFLOP/s
• FP8 Compute: 2.6 PFLOP/s

Memcpy is used to measure the memory bandwidth of the GPU. The benchmark is run using a simple memcpy kernel that copies data from one location to another in the GPU’s memory. Memcpy is a good indicator of memory bound kernels, which are common in the activation functions of LLMs.

Benchmark Results

This figure plots memcpy bandwidth against the size of the data being copied. The x-axis is the size of the data being copied in bytes, and the y-axis is the bandwidth in GB/s. The data is copied from one location to another in the GPU’s memory. The benchmark is run for different sizes of data, ranging from 4KB to 2.1GB, which is the size of the embedding tables in the Llama 3 8B model.

Looking at this graph showing the AMD Instinct MI300X memory bandwidth for different data sizes, let’s analyze what’s happening with the memcpy performance: The graph plots memory bandwidth (GB/s) on the y-axis against data size (MB) on the x-axis, both on logarithmic scales. Here are the key observations:

1. Roofline Performance: There’s a horizontal dashed line at approximately 5300 GB/s labeled as “Roofline” - this represents the theoretical maximum memory bandwidth of the MI300X GPU.
2. Bandwidth Scaling: The graph shows how memory bandwidth scales with data size:
   • For very small data sizes (below 1MB), the bandwidth is quite low (under 100 GB/s)
   • As data size increases, bandwidth improves dramatically
   • Performance peaks for data sizes around 10-100MB, approaching but not quite reaching the roofline
3. Performance Plateau: For data sizes larger than about 10MB, the bandwidth levels off at approximately 2000-3000 GB/s, which is about 40-60% of the theoretical maximum.
4. Various Transfer Sizes: The different markers represent different copy sizes (from 4.1kB to 1.1GB), showing how each performs across the spectrum.
5. Memory Hierarchy Effects: The shape of this curve is typical of memory systems with hierarchical caches:
   • Small transfers are limited by overhead and latency
   • Medium-sized transfers achieve the best bandwidth utilization
   • Very large transfers may be hitting memory management limitations

This benchmark reveals that while the MI300X offers impressive memory bandwidth, real-world memcpy operations in PyTorch achieve around half of the theoretical maximum, which is actually quite good for practical workloads. The performance characteristics suggest that for optimal memory throughput, data should be processed in chunks of approximately 10-100MB when possible.

Benchmark Code

You can find the benchmark code on the ScalarLM Github.

Let’s take a look at the code:

Memcpy Sizes

# List of memcpy sizes, in bytes, should be multiples of the page size
# Go up to the tensor size used in Llama 3 (4096 * 128256 * 4) = 2_101_346_304
memcpy_sizes = [ 2 ** i for i in range(12, 64) if 2 ** i <= 2_101_346_304 ]

This code sets up the sizes of the data to be copied. The sizes are powers of 2, starting from 4KB (2^12) and going up to 1.1GB (2^30). The sizes are chosen to be multiples of the page size, which is 4KB on most systems. The maximum size is the size of the embedding tables in the Llama 3 8B model, which is 2_101_346_304 bytes (or 2.1GB). The benchmark is run for each of these sizes, and the bandwidth is measured for each size. The results are plotted in the figure above.

Benchmark Setup

Next, we set up the benchmark:

def run_memcpy_benchmark():

    warmup()

    results = {}

    for size in tqdm(memcpy_sizes):
        results[size] = run_memcpy(size)

    return results

This function runs the memcpy benchmark. It first warms up the GPU by running a few iterations of the memcpy kernel without measuring the time. This is done to ensure that the GPU is in a good state before running the benchmark. The function then runs the memcpy kernel for each size in the memcpy_sizes list and measures the time taken to copy the data.

Warmup

Warmup is pretty simple.

def warmup():
    run_memcpy(4096)

This function runs the memcpy kernel with a size of 4KB (4096 bytes) to warm up the GPU. GPUs have startup times to load the code, ramp up the clocks, etc. Running benchmarks without a warmup can lead to misleading results.

Running Memcpy

The run_memcpy function is where the actual memcpy kernel is run. It uses PyTorch to allocate memory on the GPU and copy data from one location to another. The function measures the time taken to copy the data and calculates the bandwidth and other metrics.

The memcpy kernel is run for at least 1 second to get a good measurement of the bandwidth. The function uses PyTorch’s copy_ method to copy data from one tensor to another. copy_ is the in-place version of the copy method, which means that it modifies the destination tensor in place. This is more efficient than creating a new tensor for the result because it avoids allocating memory for the result tensor.

def run_memcpy(size):
    a = torch.zeros(size // 4, device=get_device(), dtype=torch.float32) # size is in bytes, so divide by 4 to get number of floats
    b = torch.zeros(size // 4, device=get_device(), dtype=torch.float32)

    # copy for at least 1 second
    barrier()

    start = get_event()
    end = get_event()

    start_time = time.time()

    start.record()
    iterations = 0
    while time.time() - start_time < 1:
        b.copy_(a)
        iterations += 1
    end.record()

    barrier()
    total_time = start.elapsed_time(end) * 1e-3 / iterations

    return {
        "operational_intensity": 1 / 4,  # 1 FLOP per 4 bytes
        "flop/s": size / 4 / total_time,
        "bytes": size,
        "time": total_time,
        "iterations": iterations,
        "bandwidth": size / total_time,
        "GB/s": size / total_time / 1e9,
    }

Handling GPUs

When running on a GPU, the benchmark uses PyTorch’s CUDA events to measure the time taken to copy the data. CUDA events are used to measure the time taken to execute a kernel on the GPU. The record method is used to record the time at which the event is recorded. The elapsed_time method is used to calculate the time taken to execute the kernel. The time is measured in milliseconds, so we multiply by 1e-3 to convert to seconds. Using events is necessary because the GPU is asynchronous, meaning that the CPU and GPU can run in parallel. The CPU can continue executing while the GPU is copying data. This can lead to misleading results if the time taken to copy the data is not measured correctly.

In order to make the code cross-platform, we define a get_event function that returns a CUDA event if the GPU is available, or a CPU event if the GPU is not available. The CPU event is a simple wrapper around the time module that records the time when the event is created and calculates the elapsed time between two events.

class CPUEvent:
    def __init__(self):
        self.time = 0

    def record(self):
        self.time = time.time()

    def elapsed_time(self, other):
        return (other.time - self.time) * 1000


def get_event():
    if torch.cuda.is_available():
        return torch.cuda.Event(enable_timing=True)
    else:
        return CPUEvent()


def barrier():
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    else:
        pass

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda:0")
    else:
        return torch.device("cpu")

Large Language Models (LLMs) are transforming software development by excelling in tasks such as generating new code. However, their performance with complex legacy codebases has limitations. To bridge this gap, we adopt the ScalarLM unified framework to train state-of-the-art pre-trained models (e.g., LLaMA) using our curated bug report dataset.

Integrating LLMs into the code repair pipeline boosts productivity, minimizes errors, and enhances security by automating the resolution of vulnerabilities. In the high-stakes world of Linux kernel development—where reliability and security are paramount—LLMs offer an efficient solution for identifying and fixing critical issues. This post explores their application in Linux code repair through a practical case study.

Why Automated Code Repair is Critical for Linux Development

Linux development faces unique challenges due to its scale, reliability demands, and open-source nature:

1. Scale and Complexity

The Linux kernel comprises millions of lines of code, making manual bug identification daunting.

2. Reliability and Security

As the backbone of critical infrastructure, Linux must maintain impeccable reliability. Vulnerabilities can lead to catastrophic failures.

3. Community Contributions

With a global contributor base, code quality can vary significantly, making it difficult to maintain uniform coding standards.

Automated code repair tackles these challenges by detecting bugs proactively, enforcing consistent quality standards, and enabling rapid vulnerability patching—all while reducing the manual workload for developers.

Case Study: Fixing Buffer Overflow Errors with LLMs

Understanding Buffer Overflow

Buffer overflows occur when more data is written to a memory buffer than it can hold, causing adjacent memory to be overwritten. These vulnerabilities, common in languages like C, pose severe security risks:

• Code Execution: Malicious code injection.
• System Crashes: Kernel-level disruptions.
• Instability: Threats to system reliability.

Workflow Overview

1. Bug Detection

Static code analysis tools identify potential issues. Bug reports generated by static code analysis tools typically include details such as the source file name, bug type, line number, description of the issue, and a suggested fix. Here’s an example:

CID: 1845632
Type: Buffer overflow
Category: BUFFER_OVERFLOW
Classification: Bad use of string function
Severity: High
Certainty: Absolute
Status: New
Function: cxl_mem_create_range_info
File: drivers/cxl/core/mem.c
Line: 723

Issue:
Unbounded sprintf can cause buffer overflow. The code uses sprintf() without size limits to write into buffer 'range_id', which could lead to buffer overflow if the formatted string exceeds the buffer size.

Description:
The function cxl_mem_create_range_info() uses sprintf() to format a memory range identifier into a string buffer 'range_id' without checking if the resulting string will fit in the destination buffer. This could lead to a buffer overflow if the range address and size generate a string longer than the size of range_id.

Use snprintf instead

2. Training Data Curation

Experts pair bug reports with manual fixes (e.g., git diffs) to serve as a training dataset for a Large Language Model (LLM). For example:

diff --git a/drivers/cxl/core/mem.c b/drivers/cxl/core/mem.c
index 9af4721c5e32..b7d45890c453 100644
--- a/drivers/cxl/core/mem.c
+++ b/drivers/cxl/core/mem.c
@@ -720,7 +720,7 @@ static int cxl_mem_create_range_info(struct cxl_dev_state *cxlds)
       struct cxl_memdev *cxlmd = to_cxl_memdev(cxlds->dev);
       char range_id[32];

-       sprintf(range_id, "0x%llx-%llx", range->start, range->size);
+       snprintf(range_id, sizeof(range_id), "0x%llx-%llx", range->start, range->size);
       cxlds->range_id = kstrdup(range_id, GFP_KERNEL);
       if (!cxlds->range_id)
               return -ENOMEM;

Using just a dozen highly curated samples, we train a pre-trained model on this task. When evaluated on a held-out test set, the diffs generated by the trained model demonstrated strong quality and alignment with expert-created fixes.

3. Model Fine-Tuning

The Cray-LM framework fine-tunes a pre-trained Llama model using a dataset formatted as JSON lines. Each sample includes fields such as:

['bug_report_path', 'bug_report_text', 'diff_path', 'diff_text', 'source_code_path', 'line_number', 'code']

This input is preprocessed and a few lines above and below the offending line_number are used as additional context. This information is then fed to the train the model.

import masint
import jsonlines

# Retrieve source code snippet with line numbers
def get_source_code(data, before_lines=5, after_lines=5):
    lines = data["code"].split("\n")
    start_line = max(0, data["line_number"] - before_lines)
    end_line = min(len(lines), data["line_number"] + after_lines)
    return "\n".join(lines[start_line:end_line])

# Prepare training data
def get_data(training_data_file, dataset_size=1000):
    with jsonlines.open(training_data_file) as reader:
        raw_data = list(reader)

    return [
        {
            "input": prompt_template.format(
                source_code_path=entry["source_code_path"],
                line_number=entry["line_number"],
                code=get_source_code(entry),
                bug_report_text=entry["bug_report_text"]
            ),
            "output": entry["diff_text"]
        }
        for entry in raw_data[:dataset_size]
    ]

# Main function
def main():
    data = get_data(training_data_file=args.input)
    llm = masint.SupermassiveIntelligence()
    llm.train(data, train_args={"max_steps": 20, "learning_rate": 3e-3})

4. Automated Fix Generation

The fine-tuned LLM generates code patches based on new bug reports.

import jsonlines
import masint

def load_data(eval_file_path):
    with jsonlines.open(eval_file_path) as reader:
        return list(reader)

def get_source_code(data, before_lines=5, after_lines=5):
    lines = data["code"].split("\n")
    start_line = max(0, data["line_number"] - before_lines)
    end_line = min(len(lines), data["line_number"] + after_lines)
    return "\n".join(lines[start_line:end_line])

def get_dataset(data):
    return [
        prompt_template.format(
            source_code_path=entry["source_code_path"],
            line_number=entry["line_number"],
            code=get_source_code(entry),
            bug_report_text=entry["bug_report_text"]
        )
        for entry in data
    ]

def main():
    data = load_data(eval_file_path=args.input)
    dataset = get_dataset(data)

    llm = masint.SupermassiveIntelligence()

    for entry in dataset:
        results = llm.generate(prompts=[entry])
        print(f"\n{results}")

5. Validation

Generated fixes are manually compared against reference diffs to validate accuracy and alignment with coding standards. This enhanced workflow demonstrates how LLMs can be leveraged to automate bug detection and repair, reducing both development time and security risks.

Limitations and Challenges

While the results are promising, there are notable limitations and challenges:

1. Limited Bug Diversity

The effectiveness of the model heavily depends on the diversity of bug types in the training set. For example, if a bug type is underrepresented, the model may struggle to generalize fixes for similar issues.

2. Context Limitations

To provide accurate fixes, the necessary context for the bug must be small and contained within a single source file. Handling bugs that span multiple files or require a broader understanding of the codebase is beyond the current model’s capabilities.

3. Manual Data Curation

Creating the bug report and fix dataset requires significant manual effort by experts deeply familiar with the codebase. This process is time-intensive and limits scalability.

Conclusion

This study demonstrates the potential of fine-tuning LLMs on small, curated datasets for automated code repair. The model consistently produces high-quality fixes on test data, highlighting the synergy between expert domain knowledge and machine learning. However, addressing challenges such as bug diversity, context limitations, and dataset scalability is crucial for further advancement.

Future Work

1. Broaden Dataset Diversity

Collaborate with domain experts to develop a richer, more diverse training dataset that better represents real-world bugs.

2. Expand Contextual Understanding

Research methods for incorporating multi-file context and broader codebase insights to handle more complex bugs.

3. Automate Data Generation

Investigate semi-automated approaches to streamline dataset creation, including tools to identify bugs and propose preliminary fixes.

By tackling these challenges, we can improve the robustness and applicability of LLM-based code repair systems, enabling their seamless integration into software development workflows.

Have thoughts or experiences with using LLMs for code repair? Send us a message to get involved!

Prerequisites

Before we begin, ensure you have:

• Docker installed on your system
• An AMD GPU that supports ROCm (check the ROCm Hardware Compatibility List)
• ROCm drivers installed on your host system
• Access to the ScalarLM repository

Base Container Configuration

First, let’s create a Dockerfile that sets up the ROCm environment. Create a new file called Dockerfile.rocm:

# Start with the ROCm vLLM base image
FROM rocm/vllm-dev:20250124

# Setup the working directory
ENV PATH="/opt/conda/envs/py_3.10/bin:$PATH"
ENV CONDA_PREFIX=/opt/conda/envs/py_3.10

ARG MAX_JOBS=4

# Install additional dependencies
RUN pip install uv

Building the ScalarLM vLLM Components

Next, we’ll build the vLLM components for ScalarLM. Add these sections to your Dockerfile:

RUN --mount=type=cache,target=/var/cache/apt \
    apt-get update -y \
    && apt-get install -y curl ccache git vim numactl gcc-12 g++-12 libomp-dev libnuma-dev \
    && apt-get install -y ffmpeg libsm6 libxext6 libgl1 libdnnl-dev \
    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12

ARG INSTALL_ROOT=/app/cray

COPY ./requirements.txt ${INSTALL_ROOT}/requirements.txt
COPY ./test/requirements-pytest.txt ${INSTALL_ROOT}/requirements-pytest.txt
COPY ./infra/requirements-vllm-build.txt ${INSTALL_ROOT}/requirements-vllm-build.txt

RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-vllm-build.txt
RUN uv pip install --no-compile --no-cache-dir -r ${INSTALL_ROOT}/requirements-pytest.txt

WORKDIR ${INSTALL_ROOT}

COPY ./infra/cray_infra/vllm ${INSTALL_ROOT}/infra/cray_infra/vllm
COPY ./infra/setup.py ${INSTALL_ROOT}/infra/cray_infra/setup.py

COPY ./infra/CMakeLists.txt ${INSTALL_ROOT}/infra/cray_infra/CMakeLists.txt
COPY ./infra/cmake ${INSTALL_ROOT}/infra/cray_infra/cmake
COPY ./infra/csrc ${INSTALL_ROOT}/infra/cray_infra/csrc

COPY ./infra/requirements-vllm.txt ${INSTALL_ROOT}/infra/cray_infra/requirements.txt

WORKDIR ${INSTALL_ROOT}/infra/cray_infra

ARG VLLM_TARGET_DEVICE=rocm
ARG TORCH_CUDA_ARCH_LIST=gfx906 gfx908 gfx90a gfx940 gfx941 gfx942 gfx1030 gfx1100

# Build vllm python package
RUN --mount=type=cache,target=/root/.cache/pip \
    --mount=type=cache,target=/root/.cache/ccache \
    MAX_JOBS=${MAX_JOBS} TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    VLLM_TARGET_DEVICE=${VLLM_TARGET_DEVICE} \
    python ${INSTALL_ROOT}/infra/cray_infra/setup.py bdist_wheel && \
    pip install ${INSTALL_ROOT}/infra/cray_infra/dist/*.whl && \
    rm -rf ${INSTALL_ROOT}/infra/cray_infra/dist

WORKDIR ${INSTALL_ROOT}

ScalarLM Components and Configuration

The final section includes craylm components and adds SLURM support for distributed training:

RUN apt-get update -y  \
    && apt-get install -y slurm-wlm libslurm-dev \
    build-essential \
    less curl wget net-tools vim iputils-ping \
    && rm -rf /var/lib/apt/lists/*

# Build SLURM
COPY ./infra/slurm_src ${INSTALL_ROOT}/infra/slurm_src
RUN /app/cray/infra/slurm_src/compile.sh

# Copy slurm config templates
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/infra"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/sdk"
ENV PYTHONPATH="${PYTHONPATH}:${INSTALL_ROOT}/ml"

ENV SLURM_CONF=${INSTALL_ROOT}/infra/slurm_configs/slurm.conf

RUN mkdir -p ${INSTALL_ROOT}/jobs

COPY ./infra ${INSTALL_ROOT}/infra
COPY ./sdk ${INSTALL_ROOT}/sdk
COPY ./test ${INSTALL_ROOT}/test
COPY ./cray ${INSTALL_ROOT}/cray
COPY ./ml ${INSTALL_ROOT}/ml
COPY ./scripts ${INSTALL_ROOT}/scripts

Building the Container

You can build the container using Docker:

docker build -f Dockerfile.rocm \
    --build-arg VLLM_TARGET_DEVICE=rocm \
    -t cray-rocm:latest .

Running ScalarLM with ROCm

To start a development server on a system with AMD GPUs:

docker run -it --device=/dev/kfd --device=/dev/dri \
    --security-opt seccomp=unconfined \
    --group-add video \
    -p 8000:8000 \
    --entrypoint /app/cray/scripts/start_one_server.sh \
    gdiamos/cray-rocm:latest

The key differences from the CPU or NVIDIA containers are:

1. The --device=/dev/kfd --device=/dev/dri flags expose the AMD GPU devices
2. Adding the container to the video group for GPU access
3. Setting seccomp=unconfined for ROCm compatibility

Performance Considerations

When running ScalarLM on AMD GPUs, consider these optimization tips:

1. Ensure you’re using the latest ROCm drivers for optimal performance
2. Set appropriate memory limits based on your GPU’s VRAM
3. Use hip-specific environmental variables for fine-tuning:

export HIP_VISIBLE_DEVICES=0,1,2,3  # Specify which GPUs to use

vLLM engine performance settings

vLLM provides a number of engine options which can be changed to improve performance. Refer to the vLLM Engine Args documentation for the complete list of vLLM engine options.

Below is a list of a few of the key vLLM engine arguments for performance; these can be passed to the vLLM benchmark scripts:

–max-model-len : Maximum context length supported by the model instance. Can be set to a lower value than model configuration value to improve performance and gpu memory utilization. –max-num-batched-tokens : The maximum prefill size, i.e., how many prompt tokens can be packed together in a single prefill. Set to a higher value to improve prefill performance at the cost of higher gpu memory utilization. 65536 works well for LLama models. –max-num-seqs : The maximum decode batch size (default 256). Using larger values will allow more prompts to be processed concurrently, resulting in increased throughput (possibly at the expense of higher latency). If the value is too large, there may not be enough GPU memory for the KV cache, resulting in requests getting preempted. The optimal value will depend on the GPU memory, model size, and maximum context length. –max-seq-len-to-capture : Maximum sequence length for which Hip-graphs are captured and utilized. It’s recommended to use Hip-graphs for the best decode performance. The default value of this parameter is 8K, which is lower than the large context lengths supported by recent models such as LLama. Set this parameter to max-model-len or maximum context length supported by the model for best performance. –gpu-memory-utilization : The ratio of GPU memory reserved by a vLLM instance. Default value is 0.9. Increasing the value (potentially as high as 0.99) will increase the amount of memory available for KV cache. When running in graph mode (i.e. not using –enforce-eager), it may be necessary to use a slightly smaller value of 0.92 - 0.95 to ensure adequate memory is available for the HIP graph.

Troubleshooting

Common issues and solutions:

1. If you encounter “GPU device not found” errors, verify that:
   • ROCm is properly installed on the host
   • The container has access to GPU devices
   • The user has proper permissions
2. For memory-related issues:
   • Check GPU memory usage with rocm-smi
   • Adjust batch sizes and model configurations accordingly
   • Monitor system memory usage alongside GPU memory

Conclusion

Building ROCm containers for ScalarLM enables efficient use of AMD GPUs for machine learning workloads. By following this guide, you can create and deploy containers that leverage the full potential of AMD hardware for both training and inference tasks.

For more information about ScalarLM and its capabilities, visit our documentation or join our community on GitHub.

Remember to check for updates and new releases of both ROCm and ScalarLM to ensure you’re using the latest features and optimizations.

ScalarLM is a fully open source, CC-0 Licensed (unrestricted commercial use), integrated LLM inference and training platform.

We created ScalarLM to simplify the development of reinforcement learning agents with advanced reasoning and memory capabilities, similar to those of DeepSeek R1. By integrating inference and training engines into a single platform, ScalarLM enables the seamless generation and utilization of reasoning trajectories for training updates, streamlining the development process.

• Contact us for help deploying ScalarLM on your GPU supercomputer: Contact
• Access the source code at: Github
• Read the documentation: Docs
• Join the community on Discord

ScalarLM builds on top of the vLLM inference engine, the Megatron-LM training framework, and the HuggingFace model hub. It unifies the capabilities of these tools into a single platform, enabling users to easily perform LLM inference and training, and build higher level applications such as LLM-Agents.

ScalarLM is designed for high performance, with out of the box support for AMD ROCm GPUs and NVIDIA CUDA GPUs. It includes development mode support for x86 CPUs and ARM CPUs. It inherits the distributed training capabilities of Megatron-LM and the optimized inference engine of vLLM. Cray is also designed to be easy to use. It provides an OpenAI compatible server and a simple command line interface for users to interact with the platform.

ScalarLM is inspired by the work of Seymour Roger Cray, an American electrical engineer and supercomputer architect who designed a series of computers that were the fastest in the world for decades, and founded Cray Research, which built many of these machines. Called “the father of supercomputing”, Cray has been credited with creating the supercomputer industry.

We hope that ScalarLM will help to democratize access to large language models, and enable more people to build and deploy AI systems that push the boundaries of accuracy to benefit society. We are excited to see what you will build with ScalarLM!

If you are interested in contributing to ScalarLM, please reach out to us at Contact.

v0.5 - Functionality and Feedback

ScalarLM v0.5 is a fully functional prototype that demonstrates the core capabilities of the platform. It is intended to show the potential of the platform and gather feedback from the community. We are actively seeking feedback on the platform, including suggestions for new features, bug reports, and ideas for how to improve the platform.

We expect to release v1.0 in the coming months, which will include additional features, performance optimizations, benchmarks, and bug fixes. We are also planning to release a series of blog posts and tutorials that will help users get started with ScalarLM and build their own LLMs with advanced reasoning capabilities.

Lessons Learned

Over the last ten years, we have built seven generations of training and inference stacks at Baidu, Apple, AI Fund, MLCommons, several startups including Lamini, and multiple enterprise partners. We have learned a lot about what works and what doesn’t. ScalarLM is a clean slate rewrite of a unified inference and training stack incorporating the lessons learned. Here are some of the key lessons:

• Training and Inference are both essential, but most frameworks focus on one or the other.
• SLURM and Kubernetes have different worldviews, but pieces of both are useful.
• A model exchange format has not yet emerged, and the best we have is the huggingface model hub.
• The main dependency is the GPU, and it it possible to achieve a simpler and more secure design by cutting out the rest.

These infrastructure issues are so difficult that many organizations focus on either training or inference, leaving significant accuracy on the table. Our vision for ScalarLM is to enable organizations to build LLM systems that leverage backpropagation in addition to in-context learning, just like humans, enabling them to solve more challenging problems.

Unified Training and Inference

Training is essential because it enables models to learn. Training creates flywheels. A machine learning model without training could never be updated and could not learn.

Inference is essential for models to have practical value. Inference is used when a model is deployed in production to serve predictions to users.

ScalarLM unifies training and inference into a single platform. This enables users to simultaneously train and deploy models using the same GPU cluster, and to easily switch between the two modes without changing their code.

It does this by building on pytorch models stored in the huggingface model hub. This allows training models that are supported by huggingface. The inference engine is based on vLLM, which loads the saved huggingface models. It is possible to train new models from scratch, or to continue training existing models.

A ScalarLM cluster will spin up inference and training workers as needed. Training jobs are submitted to the SLURM queue and pulled out by the training workers as batch jobs. Inference requests are submitted to a persistent queue and pulled out by the inference workers as needed. As checkpoints are produced by the training workers, they are pickup up by the inference workers and new inference requests can immediately use the new model.

ScalarLM follows the API design of the Llama Stack, which is a set of APIs that are designed to be easy to build Agent applications on top of open source LLMs. The Llama Stack includes guards including AILuminate benchmark suite, which is a set of benchmarks that are designed to measure the performance of LLMs. ScalarLM is designed to be compatible with the Llama Stack, and to work well with the AILuminate benchmarks.

SLURM and Kubernetes are Better Together

Many software developers are familiar with Kubernetes for cloud and enterprise apps. So the natural temptation is to use Kubernetes as the job queue, e.g. with the MPI Operator to get the network to work. This leaves network performance on the table and forces the user to set up the RDMA network stack manually.

Kubernetes without SLURM can work, but it forces the infra team to reinvent the wheel.

Some HPC clusters run SLURM on bare metal. This also works and simplifies the network configuration, but deployments without containers create version and dependency management hell.

In ScalarLM, we run a virtual SLURM cluster inside of a kubernetes application, providing the best of both worlds. We are closely monitoring development on the SLINKY project and expect to fold in better approaches as they become available.

Let’s look at SLURM and Kubernetes in more detail.

SLURM is a batch job scheduler that is widely used in high performance computing environments. It is designed to run large numbers of jobs on a cluster of machines with a high performance network. The high performance network is essential for training large models, because it allows the GPUs to communicate with each other quickly using RDMA primitive operations that are composed into collective operations like allreduce and allgather. Superficially it seems like SLURM only handles job scheduling, but it also sets up the MPI environment that brings up the high performance network that distributed training frameworks like Megatron-LM tap into.

Kubernetes is a container orchestration system that is widely used in cloud computing environments. It handles the problem of deploying containers onto a cluster of machines. Containers are essential to manage the complex software dependencies. If you want to upgrade the software on the cluster, you can just upgrade the container image and restart the container. Kubernetes allows you to do this without affecting the other containers running on the cluster.

Putting them together may seem difficult, and indeed, there are many subtle details that need to be taken care of including how SLURM daemons manage hostnames and handle cgroups which make conflicting assumptions with the lower layers of kubernetes networking, containerd, and accelerator management. However, SLURM is a essentially set of lightweight C programs that coordinate over TCP sockets. They tap into rdma interfaces exposed by linux network drivers. It is possible to entirely encapsulate these SLURM daemons in containers, and run those containers on Kubernetes. ScalarLM does this.

When the ScalarLM worker containers start, they have access to a virtual SLURM network and control plane. We build the orchestration, distributed training framework, and collectives on top of these interfaces. So a user can helm install cray, and get a fully functional SLURM cluster.

Huggingface is the model exchange format

New models come up almost daily, and it is important to be able to train and deploy them with minimal porting effort.

The huggingface model hub is the closest thing we have to a model exchange format. The transformers library has pytorch source code implementations of hundreds of classes of models. Even though most LLMs are based on transformers, there is no standard implementation of a transformer. Different models might choose different position embeddings or layer shapes. One of the biggest points of friction between training and inference teams in large companies, is resolving these differences. ScalarLM modifies Megraton-LM to use the huggingface model hub as the source of truth for model implementations, which makes it possible to immediately load most trained models directly into the inference engine.

Secure Enclave

Multiple previous generations of ScalarLM used a variety of cloud services including parallel distributed object storage such as Azure Blob or GCS, database services like BigQuery, cloud load balancers, etc. These lead to significant complexity and reduced portability between supercomputers, which are often optimized for efficiency and have different cloud stacks. The current generation of ScalarLM requires GPU compute nodes, a high performance interconnect, and nothing else.

This allows ScalarLM to be deployed in a secure enclave, with no external dependencies.

We had to address several design issues to make this possible.

Training State Versioning without a Database or Object Store

Training jobs need to produce a model that is passed to the inference engine and state needs to be tracked between the training and inference workers. We uniquely identify each model by a hash of the training data, the model code, and the hyperparameters. This state is saved to a shared filesystem that is accessible by both the training and inference workers. The training system picks up this training state and launches a SLURM job to train the model. As it trains, it produces checkpoints that are saved to the shared filesystem. The inference system picks up these checkpoints as they are made available. This enables fault tolerance, because if a training worker fails, training system can restart from a checkpoint. It also avoids the need for a database to track the state of the training jobs.

Massive Inference without Cloud Queues or Load Balancers

LLMs are big and expensive, even on GPUs. Generating thousands of tokens can take minutes. A user who submits a few requests from their laptop in a loop can quickly overwhelm a GPU cluster that costs millions of dollars. Submitting these requests as HTTP REST requests leads to load imbalance and timeouts. Submitting too few requests leads to underutilization. vLLM uses continuous batching to handle load imbalance, but that can still lead to overwhelming the workers. ScalarLM pushes inference requests into a persistent queue, which can be very large in size. The requests are pulled out by vLLM workers as they have capacity. The user client polls to get the results from the queue. This allows efficient and lossless processing of very high numbers of inference requests, common to inference pipelines.

Datasets without an Object Store

Training requires training data. Typically this requires uploading a training dataset. The natural place to put it is in a cloud object store. It initially seems hard to get around this. But how hard it is really?

How big is each data item that an LLM is trained on? A few hundred tokens.

How many bytes does that take up? Let’s say one byte per token compressed, so a few hundred bytes per item.

How many training items are there for a dataset like Alpaca? 52K.

So how big is that? 52K (Q&A pairs) * ~300 (bytes per pair) = 15.6 MB

How long does it take to upload 15.6MB over a cloud internet connection, e.g. 1 Gbs? = 0.124 seconds

How long does it take to train a 70B parameter model on that data on an H100 at 40% MFU? = 6.5 hours

So the upload perf isn’t nearly as big of a practical problem as it seems.

We built a fast file uploader in ScalarLM that can handle 10s to 100s of GBs of data over common cloud connections. It hashes the data to avoid uploading duplicates. No need for an object store.

Our main insight is as follows. If you are running ScalarLM, you built a supercomputer capable of training an LLM. That computer can handle moving around a few GB, so could most laptops today. An object store managing thousands of parallel disks is overkill.

Future

Our target for v1.0 release is to include performance optimizations and kernel level benchmarks on modern GPUs including MI300X and H200.

We plan to keep ScalarLM up to date with the latest advances in LLM research, including new models, new training stack optimizations, and new inference stack optimizations.

We are excited to build a new generation of LLM supercomputers that democratize access to training language models. We are in particular interested in applications that build on top of ScalarLM to support LLM agents that train and deploy themselves.

Contributing

ScalarLM is an open source project developed by a small team at TensorWave.

• Greg Diamos, PhD, is a computer architect and deep learning researcher. He has worked at NVIDIA.
• Sudnya Diamos is a 3x startup founding engineer, and former NVIDIA GPU architdect.
• Naila Farooqui, PhD, is a deep learning researcher. She is the co-author of NVIDIA Cutlass.
• Suhabe Bugrara, PhD, is a 3x CTO.

Contact Us to explore opportunities to work with us.

• Check out the examples - Example - Coming Soon
• Read deeper into Tokenformer Adaptors, a new method for efficiently training LLMs: Tokenformer Blog - Coming Soon

We welcome community contributions. Some areas that could use help:

1. Help completing the Llama Stack client
2. Running AILuminate evals
3. Performance optimizations in the training and inference engines
4. New deployment targets, benchmarks, and regression tests
5. Hosted scalarlm clusters on your cloud
6. A refactor of the unified vLLM build that makes it easier to accept upstream changes
7. Integration with SLINKY

Transformers, which form the backbone of most foundation models, face significant challenges when it comes to scaling. This is primarily due to their reliance on a set number of parameters within linear projections. When changes are made to the architecture, such as adjusting channel dimensions, it typically necessitates retraining the entire model from the beginning. This becomes increasingly problematic as models grow larger, as the computational costs associated with retraining become prohibitively expensive and unsustainable.

To address these scaling issues, a new architecture called Tokenformer has been developed. Tokenformer is a fully attention-based model that expands upon the traditional use of attention mechanisms. In addition to using attention for computations between input tokens, Tokenformer also applies attention to the interactions between tokens and the model’s parameters.

The key innovation of Tokenformer lies in its treatment of model parameters as if they were tokens themselves. This approach allows for the replacement of all linear projections found in traditional Transformers with a novel token-parameter attention layer. In this layer, the input tokens function as queries, while the model parameters serve as keys and values.

This reformulation of the Transformer architecture offers a significant advantage: it enables the model to be scaled up progressively and efficiently without the need for complete retraining from scratch. This ability to scale without full retraining represents a major step forward in addressing the computational challenges associated with developing and expanding large language models.

To apply tokenformer to a model in , we follow an approach similar to LoRA (Low-Rank Adaptation) where token-parameter attention layers are added in parallel to the existing attention layers. This allows for the model to be incrementally scaled up without the need for full retraining. The number of parameters in the key-value pairs can be adjusted after training to control the model’s capacity and performance.

Tokenformer is an innovative, fully attention-based architecture that addresses scaling challenges in traditional Transformers:

1. Replaces linear projections with token-parameter attention layers
2. Treats model parameters as tokens, with input tokens as queries and parameters as keys/values
3. Enables incremental scaling without full retraining
4. Functions as an extreme Mixture of Experts (MoE) model
5. Facilitates efficient parameter tuning for new tasks

Key advantages:

1. Progressive scaling from 124M to 1.4B parameters and beyond
2. Performance comparable to full retraining
3. Over 50% reduction in training costs

Tokenformer preserves inter-token computations while extending the Transformer architecture, offering a more flexible and efficient approach to large language model development.

LoRA (Low-Rank Adaptation) versus Tokenformer

Tokenformer and LoRA are two different approaches for adapting large language models (LLMs).

LoRA (Low-Rank Adaptation) is a parameter-efficient fine-tuning method that adds small trainable rank decomposition matrices to existing weights. It freezes the pretrained model weights and injects trainable rank decomposition matrices into each layer of the transformer architecture. This approach significantly reduces the number of trainable parameters while maintaining performance comparable to full training.

In contrast, Tokenformer is a more recent development that only leverages the attention mechanism for both input token computations and token-parameter interactions. This reformulation results in a natively scalable architecture, allowing for progressive and efficient scaling.

In summary, while both methods aim to improve LLM adaptation, LoRA focuses on parameter efficiency and versatility, while Tokenformer prioritizes flexibility and scalability. The key advantage of Tokenformer is its ability to scale up models without full retraining, or the need to keep adding new LoRA adaptors to increase capacity.

Our Implementation

is a unified framework for training and inference of language models. Its primary objective is to seamlessly support novel modeling techniques across both training and inference phases. The framework’s inference engine is based on vLLM, which imposes certain constraints on achieving this unified approach. This discussion outlines the implementation strategy for integrating Tokenformer into ScalarLM and provides specific details for both training and inference processes.

Tokenformer’s architecture requires the integration of cross-attention based adapters to either the feed-forward (MLP) layers, the existing attention layers, or both. To accomplish this, a visitor pattern is implemented. This pattern, implemented by the TokenformerSurgeon class, traverses the underlying PyTorch graph, identifying MLP and attention layers within any given model network and wrapping them with Tokenformer adapters.

The current implementation of Tokenformer in specifically targets Llama-based models.

Training with Tokenformer

The implementation of Llama models in the HuggingFace transformers library presents a challenge for the visitor strategy. This is because attention layers are not always encapsulated as separate modules. Instead, the library often utilizes torch.nn.functional attention mechanisms, such as scaled_dot_product_attention, within encompassing modules.

To address this issue, the implementation extends the HuggingFace transformers Llama library. This extension encapsulates attention mechanisms into their own module classes, thereby enabling the seamless application of the visitor pattern.

vLLM Inference with Tokenformer

vLLM, the basis for ‘s inference engine, does not rely on HuggingFace’s transformers implementations for many of its models. Instead, it provides its own implementations in the vllm/model_executor/models directory.

This architectural difference creates an incompatibility between vLLM and HuggingFace transformers. As a result, separate Tokenformer adapters were implemented for vLLM and HuggingFace transformers. These distinct implementations are necessary due to the different APIs for forward calls that need to be supported.

The implementation details are visually represented in a UML diagram, which illustrates the relationships and structures of the components involved in supporting Tokenformer within the framework.

Conclusion

Tokenformer introduces a naturally scalable architecture that leverages the attention mechanism to facilitate not only inter-token computations but also interactions between tokens and model parameters, thereby enhancing architectural flexibility. By representing model parameters as tokens, we can replace all linear projection layers in the Transformer with attention layers, allowing for seamless and efficient incremental scaling without the need for retraining from scratch. This architecture offers greater flexibility than traditional Transformers and will further contribute to the development of foundation models.