Pre-Training Hybrid LLMs from Scratch in Primus on AMD Instinct™ GPUs

9 min read. | 2201 total words.

In this blog you will learn how to pre-train hybrid LLMs from scratch with Primus on AMD Instinct™ MI355X GPUs. In particular, we will cover how to declare a hybrid layer stack with a single configuration string, how three linear-recurrent mixers (Mamba2, Gated DeltaNet, and Kimi Delta Attention) trade off throughput with training quality, which recipe settings determine whether your run fits in memory in the first place, and which environment variables can quietly change your results. Along the way you will receive the released configurations behind every number in this post, so that you can reproduce the runs yourself. Ready to train your own hybrid? Pull the Primus image and follow along — everything here runs on a single eight-GPU node.

Why Train Hybrid Models on Primus?#

Language model architecture is moving beyond the pure Transformer. Standard self-attention remains the primary bottleneck for long-context windows: its compute cost scales quadratically \(O(N^2)\) with sequence length, and its Key-Value (KV) cache grows linearly with sequence length. Conversely, linear-recurrent blocks offer the opposite trade-off—processing sequences through a fixed-size state at linear cost \(O(N)\). However, on their own, linear-recurrent layers struggle with precise, non-local retrieval. Hybrid models bridge the gap between these trade-offs by combining both types of layers. Recent designs such as Jamba, Samba, Qwen3-Next, and Kimi Linear[1] interleave a small number of attention layers within a stack dominated by linear-recurrent blocks. The attention layers give direct access to earlier tokens, which a fixed-size recurrent state cannot guarantee. Since many of those layers can be linear-recurrent, hybrid architectures reduce memory and compute overhead while retaining retrieval where it matters.

Supporting these architectures in a training framework presents a distinct engineering problem. A hybrid model is a heterogenous stack assembled from distinct sublayer types, each with its own custom kernels, memory characteristics, and numerical stability demands. A production-grade framework must address three critical requirements:

Primus AMD’s unified, backend-agnostic training framework for foundation models on AMD Instinct GPUs, is a platform that addresses the aforementioned requirements, supporting backends such as Megatron-LM[2] and TorchTitan[3]. This article describes the design and implementation choices made to enable hybrid model pre-training from scratch on AMD GPUs. Furthermore, we will expand on the architectures Primus supports, the cost associated with each design choice, and the results we obtained by training five models to completion on a single node of eight MI355X GPUs.

We trained five models to completion on one node of eight MI355X GPUs, all at sequence length 2048 in bfloat16 on FineWeb-Edu[4]. The label “75% hybrid” refers to the *-M-M-M-*-M-M-M-*-M-M-M- pattern, in which three MLA attention blocks[5] are distributed among twelve mixers.

Table 1: The five pre-training runs that produced the results in this article. Each completed with no NaN values and no skipped iterations, and the configurations used are included in the repository.

A hybrid model in Primus is a single layer stack constructed from three kinds of sublayer, which may be interleaved in any order.

Table 2: The three sublayer types from which a hybrid stack is assembled.

This 24-character pattern describes twelve mixer blocks alternating with twelve MLP blocks, with attention placed at mixer positions 0, 4 and 8 and a linear mixer at every other position. A purely linear model is simply a pattern containing no * characters, which may also be expressed as hybrid_attention_ratio: 0.0.

For all mixers, the surrounding SwiGLU MLP and RMSNorm sublayers are identical, as is the training recipe. Because only the token mixer changes, swapping one for another takes just a single-line configuration change. The three mixers differ fundamentally in how they determine what to retain in their fixed-size state.

Mamba2[6] is a selective state-space model (SSM) that decays its state using one scalar gate per head.

Mamba2[6] is a selective state-space model (SSM) that decays its state using one scalar gate per head.

Gated DeltaNet (GDN)[7] builds on Mamba2 by adding the delta rule, a form of associative memory that updates the stored association for the current key rather than only accumulating into the state, while keeping one scalar forget gate per head.

Gated DeltaNet (GDN)[7] builds on Mamba2 by adding the delta rule, a form of associative memory that updates the stored association for the current key rather than only accumulating into the state, while keeping one scalar forget gate per head.

Kimi Delta Attention (KDA)[1] extends GDN by refining that gate to one value per key dimension, computed through a low-rank bottleneck projection. This optimization provides finer control over what the state retains, at a small cost in throughput.

Kimi Delta Attention (KDA)[1] extends GDN by refining that gate to one value per key dimension, computed through a low-rank bottleneck projection. This optimization provides finer control over what the state retains, at a small cost in throughput.

mamba_370M-pretrain.yaml

mamba_370M-pretrain.yaml

gdn_{300M,1B}_BF16-pretrain.yaml

gdn_{300M,1B}_BF16-pretrain.yaml

zebra_llama_gdn_1B_BF16-pretrain.yaml

zebra_llama_gdn_1B_BF16-pretrain.yaml

kda_{300M,1B}_BF16-pretrain.yaml

kda_{300M,1B}_BF16-pretrain.yaml

zebra_llama_kda_1B_BF16-pretrain.yaml

zebra_llama_kda_1B_BF16-pretrain.yaml

Table 3: Released configurations, located in examples/megatron/configs/MI355X/. Hybrid configurations are also provided for MI300X, and the Mamba2 configurations for MI325X.

In these filenames the zebra_llama_ prefix records the lineage of the hybrid model, and the suffix names the linear mixer and size, as in zebra_llama_gdn_1B. The particular combination of MLA attention with linear blocks used here follows the layer plans established by Zebra-Llama[8] and extended by HyLo[9]. Note that although both papers train hybrid models by upcycling existing Transformer checkpoints through post-training, the work described in this article pre-trains the Zebra-Llama architectures from scratch through random initialization.

The default released configurations include all of the settings below with the exception of the batch size and iteration count, which are deliberately set small so that an initial run completes quickly.

1024 (2.10M tokens/iter)

1024 (2.10M tokens/iter)

512 (1.05M tokens/iter)

512 (1.05M tokens/iter)

beta = (0.9, 0.95), weight decay 0.01, gradient clip 1.0

beta = (0.9, 0.95), weight decay 0.01, gradient clip 1.0

beta = (0.9, 0.95), weight decay 0.01, gradient clip 1.0

beta = (0.9, 0.95), weight decay 0.01, gradient clip 1.0

Table 4: The complete pre-training recipe. Each column is a self-contained setup.

Two entries in this table are operationally critical, because they determine whether a run executes at all rather than merely how it converges.

Chunked fused cross-entropy prevents the logits tensor from dominating activation memory at a micro-batch of this size. At these shapes it is the difference between a run that fits and one that raises an out-of-memory error.

Chunked fused cross-entropy prevents the logits tensor from dominating activation memory at a micro-batch of this size. At these shapes it is the difference between a run that fits and one that raises an out-of-memory error.

The four use_fla_* flags replace the elementwise SwiGLU, normalization and short-convolution paths with Triton kernels from FLA. Disabling them reduces both throughput and available headroom.

The four use_fla_* flags replace the elementwise SwiGLU, normalization and short-convolution paths with Triton kernels from FLA. Disabling them reduces both throughput and available headroom.

The remaining settings constitute a conventional GPT-style pre-training recipe and may be treated as such.

Access to the tokenizer repository is a strict requirement. The released configurations tokenize with meta-llama/Llama-3.2-1B, a gated Hugging Face repository, and without valid credentials the launcher terminates with an HTTP 401 error during setup. This occurs even when training on mock data, which does not otherwise require a tokenizer.

The allocator setting has a greater effect than its brevity suggests. These models allocate a small number of very large tensors, and without expandable segments the caching allocator becomes fragmented around them. Reserved memory then grows substantially beyond what is actually in use, and a run with genuine headroom remaining can still fail because no single contiguous block is available. With the setting enabled, reserved memory tracks allocated memory closely. It costs approximately 4% in throughput and is a prerequisite for the larger micro-batch sizes.

PRIMUS_FLA_MLA_ATTN applies only to the hybrid models. Setting it to 0 keeps the MLA blocks on the native attention path, which delivers the highest throughput at the shapes used here. The setting must be supplied as an environment variable, because the ambient value takes precedence over the YAML key of the same name.

Finally, if a repository clone is mounted from the host rather than created inside the image, git will decline to inspect a tree owned by another user, which causes the editable installations to fail during setup. The safe.directory entry addresses this.

Primus separates configuration into two layers, so that the training recipe can be modified without altering the architecture. The model definition, under primus/configs/models/megatron/, specifies layer count, hidden size, the hybrid pattern and the mixer hyperparameters. The experiment configuration references a model definition and supplies everything else:

The released configurations sets mock_data: true, allowing a run to be launched immediately in order to validate the software stack without any dataset. For training on real data, prepare a Megatron-format dataset once, then set mock_data: false and point train_data_path at the result:

Primus launches every run through a single entry point, primus-cli. From the Primus repo inside the container, start a single-node run in direct mode and pass the experiment configuration. GPUS_PER_NODE defaults to 8, and logs and checkpoints are written under ./output:

Changing the above run from GDN to KDA requires a single edit: point EXP at kda_300M_BF16-pretrain.yaml instead. On a Slurm cluster, launch from a login node with the Primus checkout on a shared filesystem. primus-cli starts a container on each node, mounts the checkout into it and forwards HF_TOKEN; pass the other two settings from the environment section with --env. Without Slurm, run the direct-mode command inside the container on every node with the rendezvous variables set:

The 300M configurations ship with a deliberately small recipe, consisting of micro-batch 2, global batch 16, and 50 iterations, so that training itself completes promptly and confirms that the software stack is functioning. Be aware that the first launch is expected to take longer because the launcher’s setup hooks install the training dependencies before training begins as these dependencies not cached.

To reproduce the results reported below, increase micro_batch_size to 128 and global_batch_size to 1024, which corresponds to eight GPUs at a micro-batch of 128 with no gradient accumulation. Each configuration documents the largest micro-batch that has been verified to fit on a single MI355X.

All results were obtained on a single node of eight AMD Instinct MI355X GPUs at sequence length 2048 in bfloat16, using the released configurations together with the environment settings described above.

Reproducing them depends on two runtime settings. The hybrid configurations ship with the chunked fused_ce_mode: 1, which prioritizes memory efficiency; set fused_ce_mode: 2 for the higher throughput used in these results, and PRIMUS_FLA_MLA_ATTN=0.

Figure 1: Training loss for the four 300M models over 10B tokens. The final 2B tokens are shown separately on the right, where the curves begin to separate. The 75% hybrid GDN model reaches the lowest final loss.#

Figure 1: Training loss for the four 300M models over 10B tokens. The final 2B tokens are shown separately on the right, where the curves begin to separate. The 75% hybrid GDN model reaches the lowest final loss.#

Figure 2: Pure GDN 1B over 100B tokens, covering 95,368 iterations from a loss of 12.17 to 2.487 and ending at a gradient norm of 0.102.#

Figure 2: Pure GDN 1B over 100B tokens, covering 95,368 iterations from a loss of 12.17 to 2.487 and ending at a gradient norm of 0.102.#

Figure 3: Pure GDN 300M over 10B tokens.#

Figure 3: Pure GDN 300M over 10B tokens.#

Figure 4: Pure KDA 300M over 10B tokens.#

Figure 4: Pure KDA 300M over 10B tokens.#

Figure 5: 75% Hybrid GDN 300M over 10B tokens, which achieved the lowest final loss of the four 300M runs.#

Figure 5: 75% Hybrid GDN 300M over 10B tokens, which achieved the lowest final loss of the four 300M runs.#

Figure 6: 75% Hybrid Mamba2 300M over 10B tokens.#

Figure 6: 75% Hybrid Mamba2 300M over 10B tokens.#

All five curves descend smoothly, with no loss spikes, no divergence and no recovery events. The largest single-step increase anywhere in the 1B run is 0.03.

The table below reports zero-shot accuracy from lm-evaluation-harness[10]. The purely linear models are evaluated after conversion to Hugging Face format using the converters in tools/hybrid/, and the hybrid models are evaluated directly from the Megatron checkpoint.

Table 5: Zero-shot accuracy for all five models.

These results summarize the practical conclusion of the study. At a fixed budget of 10B tokens, the choice of mixer has little effect on downstream accuracy while driving as much as a 1.75x spread in throughput, which makes it primarily a performance decision rather than a quality one. What does improve quality is the token budget: the 1B run at 100B tokens exceeds the 300M models by roughly ten to twenty points on most tasks. The strategy this suggests is to select the mixer that trains most efficiently on the available hardware, then reinvest the savings in additional tokens.

The first lesson is that, in Primus, a hybrid layer stack is simply a configuration string. Attention, linear-recurrent and MLP sublayers are interleaved through a single hybrid_override_pattern field, so trying a new architecture means editing one line rather than writing new model code.

That flexibility matters because, at a fixed token budget, the choice of mixer turned out to be a performance decision rather than a quality one. The four 300M models finished within 0.03 loss of one another and within two points on nearly every downstream task, yet differed by as much as 1.75x in throughput. What did improve quality was the token budget. The 1B run at 100B tokens exceeds the 300M models by roughly ten to twenty points on most tasks, considerably more than any architectural choice made at 10B tokens.

Adding a small amount of attention also proved inexpensive and worthwhile. Placing three MLA blocks among twelve mixers reduced throughput by less than 3%, and that model produced the lowest loss of the 300M runs.

The recipe held up at scale as well. The 1B run descended from a loss of 12.17 to 2.487 across 95,368 iterations, with a largest single-step increase of 0.03, no NaN values and no skipped iterations.

Finally, the environment is part of the recipe. The allocator setting and the MLA wrapper variable each affected throughput or capacity more than most hyperparameters, and neither appears in a configuration file. When you reproduce these runs, treat those settings with the same care as the YAML itself.

In this blog you explored how Primus pre-trains hybrid language models, understood as arbitrary interleavings of MLA attention, linear-recurrent mixers, and MLP blocks, from a single configuration string while retaining the fused kernels that make each sublayer type efficient. You walked through the recipe, the environment settings that silently shape results, and five completed runs on one node of eight AMD Instinct MI355X GPUs. Across those runs, the choice of mixer changed throughput by as much as 1.75x while having little effect on loss or downstream accuracy. A 75% hybrid introduced three MLA blocks for less than 3% in throughput and achieved the lowest loss of the group, and a 1B model completed 100B tokens with no NaN values and no skipped iterations. The configurations behind every reported number are included in the repository.

Now train a hybrid of your own. Start from one of the released configurations, swap the mixer with a single-line edit, and compare your throughput and loss against the numbers in this post. For the upcycling counterpart to this work, which converts existing Transformer checkpoints into the same hybrid architectures, see the companion blog Zebra-HyLo: Upcycling Transformers into Long-Context Hybrid LLMs on AMD Instinct™ GPUs. Our team is continuing this line of work: researching efficient model architectures, advancing hybrid-model and upcycling research, and making Primus even faster on AMD Instinct GPUs. Watch for follow-up posts on ROCm Blogs.

Base image: rocm/primus:v26.7 on Docker Hub

Base image: rocm/primus:v26.7 on Docker Hub

The information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes. THIS INFORMATION IS PROVIDED ‘AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. AMD, the AMD Arrow logo, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. © 2026 Advanced Micro Devices, Inc. All rights reserved

Third-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT.