AMD Primus Pre-Trains Hybrid LLMs From Scratch on MI355X GPUs

AMD published details on pre-training hybrid large language models from scratch using Primus, its unified training framework for foundation models on AMD Instinct GPUs. The work ran on a single node of eight MI355X GPUs, at sequence length 2048 in bfloat16, on the FineWeb-Edu dataset. Five models were trained to completion: a 1B pure Gated DeltaNet model on 100B tokens, a 300M pure GDN model on 10B tokens, a 300M pure Kimi Delta Attention model on 10B tokens, a 300M 75% hybrid Mamba2 model on 10B tokens, and a 300M 75% hybrid GDN model on 10B tokens. All runs completed with no NaN values and no skipped iterations, and the configurations are released in the Primus repository.

Hybrid Stack Definition

A hybrid model in Primus is a single layer stack assembled from three sublayer types: attention blocks using multi-head latent attention, linear-recurrent mixers (Mamba2, Gated DeltaNet, or Kimi Delta Attention), and standard SwiGLU feed-forward blocks. The arrangement is declared with a single configuration string. The pattern "*-M-M-M-*-M-M-M-*-M-M-M-" describes twelve mixer blocks alternating with twelve MLP blocks, with attention placed at mixer positions 0, 4, and 8. A purely linear model is a pattern with no attention characters, which can also be expressed with a hybrid attention ratio of 0.0.

The three linear mixers differ in how they retain information in their fixed-size state. Mamba2 is a selective state-space model that decays its state using one scalar gate per head. Gated DeltaNet adds 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 extends GDN by refining that gate to one value per key dimension, computed through a low-rank bottleneck projection, providing finer control over what the state retains at a small cost in throughput.

The hybrid configurations follow the layer plans established by Zebra-Llama and extended by HyLo. Both papers train hybrid models by upcycling existing Transformer checkpoints through post-training, but the work described here pre-trains the Zebra-Llama architectures from scratch through random initialization.

Recipe and Memory Considerations

The released configurations include a complete pre-training recipe. For the 300M models, the global batch is 1024 with a micro-batch of 128 per GPU, running 4,768 iterations for 10B tokens. The 1B model uses a global batch of 512 with a micro-batch of 64 per GPU, running 95,368 iterations for 100B tokens. Both use cosine learning rate schedules, AdamW with beta values of 0.9 and 0.95, weight decay of 0.01, and gradient clip of 1.0. Parallelism is pure data parallel with tensor, pipeline, and expert parallelism all set to 1. The 300M model uses replicated optimizer state, while the 1B model uses ZeRO-1 sharded optimizer state.

Two settings are operationally critical for whether a run executes at all. Chunked fused cross-entropy prevents the logits tensor from dominating activation memory at the micro-batch size used, and 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.

Environment variables can quietly change results. The PyTorch caching allocator setting for expandable segments is a prerequisite for the larger micro-batch sizes, because these models allocate a small number of very large tensors and the caching allocator becomes fragmented around them without it. Reserved memory then grows substantially beyond what is actually in use, and a run with genuine headroom remaining can still fail. The setting costs approximately 4% in throughput. For hybrid models, the PRIMUS_FLA_MLA_ATTN environment variable must be set to 0 to keep MLA blocks on the native attention path, which delivers the highest throughput at the shapes used. Access to the gated tokenizer repository is a strict requirement, and the launcher terminates with an HTTP 401 error during setup without valid credentials, even when training on mock data.