Zebra: Hybrid Recurrent-Attention Models on AMD GPUs

Zebra: Hybrid Recurrent-Attention Models on AMD GPUs#

Zebra is the Primus family of hybrid models that combine recurrent layers (Mamba, KDA, or GDN) with Multi-Latent Attention (MLA) and SwiGLU MLP. Hybrid pretrain configs follow Megatron naming (zebra_llama_{mixer}_{size}_BF16-pretrain.yaml, where {mixer} is mamba, kda, or gdn). Pure recurrent models use architecture presets at kda_* / gdn_* / mamba_* with {model}_BF16-pretrain.yaml experiment configs. See the guides below.

This guide covers the complete workflow: environment setup, data preparation, pretraining, checkpoint conversion, and evaluation.

FLA-validated recipes: For runnable, FLA-parity-validated walkthroughs of the pure-recurrent variants, see Pure GDN guide and Pure KDA guide. For the exhaustive list of code/config/runtime changes required for exact parity with the Flash Linear Attention (FLA) reference implementation, see GDN ⇄ FLA parity and KDA ⇄ FLA parity.

Table of Contents#

Architecture Overview#

Zebra hybrid interleaves three types of layers in a repeating pattern:

[Attention] [MLP] [Recurrent] [MLP] [Recurrent] [MLP] ... [Attention] [MLP] ...

  • Recurrent layers — one of Mamba SSM, Kimi Delta Attention (KDA), or Gated Delta Net (GDN)

  • Attention layers — Multi-Latent Attention with YaRN rotary embeddings and LoRA-compressed KV

  • MLP layers — SwiGLU feed-forward with Transformer Engine fused norms

The hybrid_attention_ratio parameter controls what fraction of recurrent+attention layer pairs use attention (default 0.25 = 1 attention layer per 3 recurrent layers). Setting it to 0.0 yields a pure recurrent model (all KDA or GDN), while 1.0 yields a pure MLA attention model.

Available Configurations#

Hybrid models (Zebra)#

Recurrent mixer (Mamba2, KDA, or GDN) interleaved with MLA attention and SwiGLU MLP. Experiment configs use zebra_llama_{mixer}_{size}_BF16-pretrain.yaml; architecture presets use zebra_{mixer}_{size}_hybrid.yaml. In both, {mixer} is mamba, kda, or gdn.

Architecture preset

Pretrain config (MI300X)

Also on MI325X / MI355X

zebra_mamba_1B_hybrid.yaml

zebra_llama_mamba_1B_BF16-pretrain.yaml

MI325X, MI355X

zebra_kda_1B_hybrid.yaml

zebra_llama_kda_1B_BF16-pretrain.yaml

MI355X

zebra_gdn_1B_hybrid.yaml

zebra_llama_gdn_1B_BF16-pretrain.yaml

MI355X

zebra_mamba_300M_hybrid.yaml

zebra_llama_mamba_300M_BF16-pretrain.yaml

MI300X only

zebra_gdn_300M_hybrid.yaml

zebra_llama_gdn_300M_BF16-pretrain.yaml

MI300X only

zebra_mamba_3B_hybrid.yaml

zebra_llama_mamba_3B_BF16-pretrain.yaml

MI325X, MI355X

zebra_mamba_8B_hybrid.yaml

zebra_llama_mamba_8B_BF16-pretrain.yaml

MI325X, MI355X

Bridge SFT (Mamba hybrids): examples/megatron_bridge/configs/MI300X/zebra_mamba_{1B,3B,8B}_hybrid_sft_posttrain.yaml

Pretrain configs on MI300X (hybrid + pure)#

Config

Model

Recurrent Type

Seq Length

Params

Tokenizer

zebra_llama_mamba_1B_BF16-pretrain.yaml

1B (Mamba+MLA)

Mamba SSM

2048

~1B

meta-llama/Llama-3.2-1B

zebra_llama_kda_1B_BF16-pretrain.yaml

1B (KDA+MLA)

Kimi Delta Attention

8192

~1B

meta-llama/Llama-3.2-1B

kda_1B_BF16-pretrain.yaml

1B (pure KDA)

Kimi Delta Attention

2048

~1.2B

meta-llama/Llama-3.2-1B

kda_300M_BF16-pretrain.yaml

300M (pure KDA)

Kimi Delta Attention

2048

~300M

meta-llama/Llama-3.2-1B

zebra_llama_gdn_1B_BF16-pretrain.yaml

1B (GDN only)

Gated Delta Net

8192

~1B

fla-hub/gla-1.3B-100B

gdn_1B_BF16-pretrain.yaml

1B (pure GDN)

Gated Delta Net

2048

~1.2B

meta-llama/Llama-3.2-1B

gdn_300M_BF16-pretrain.yaml

300M (pure GDN)

Gated Delta Net

2048

~338M

meta-llama/Llama-3.2-1B

zebra_llama_mamba_3B_BF16-pretrain.yaml

3B (Mamba+MLA)

Mamba SSM

8192

~3B

meta-llama/Llama-3.2-3B

zebra_llama_mamba_8B_BF16-pretrain.yaml

8B (Mamba+MLA)

Mamba SSM

8192

~8B

meta-llama/Llama-3.1-8B

zebra_llama_mamba_300M_BF16-pretrain.yaml

300M (Mamba+MLA)

Mamba SSM

2048

~300M

meta-llama/Llama-3.2-1B

zebra_llama_gdn_300M_BF16-pretrain.yaml

300M (GDN+MLA)

Gated Delta Net

2048

~338M

meta-llama/Llama-3.2-1B

Model Configs (primus/configs/models/megatron/)#

Config

Layers

Hidden

FFN

Attention Ratio

Attention Type

zebra_mamba_1B_hybrid.yaml

32

2048

8192

0.25

MLA

kda_1B.yaml

32

2048

8192

0.0 (pure KDA)

None

zebra_gdn_1B_hybrid.yaml

32

2048

8192

0.0 (pure GDN)

None

gdn_1B.yaml

32 (16 GDN+16 MLP)

2048

8192

0.0 (pure GDN)

None

zebra_mamba_3B_hybrid.yaml

56

3072

8192

0.25

MLA

zebra_mamba_8B_hybrid.yaml

64

4096

14436

0.25

MLA

Note on pure KDA: The kda_1B config matches FLA’s kda_1B.json architecture (16 KDA layers, head_dim=32 for keys, head_dim=64 for values, tied embeddings, norm_eps=1e-6). It uses the FLA Triton kernel (use_fla_triton_kda: true) for fused forward+backward during training.

Note on pure GDN: The gdn_1B config matches FLA’s gated_deltanet_1B.json architecture (16 GDN + 16 MLP layers, num_heads=8, num_v_heads=16, short convolution with kernel size 4, tied embeddings). This config has been validated end-to-end against FLA on MI300X — the training loss curves match within ~1% across 76K steps on FineWeb-Edu 10BT. See Step 4 for conversion to FLA’s HuggingFace format.

Prerequisites#

  • Hardware: AMD Instinct MI300X (or compatible ROCm GPUs)

  • Software: ROCm drivers >= 7.0, Docker >= 24.0

  • HuggingFace Token: Required for gated tokenizers (HF_TOKEN)

  • Disk Space: ~50 GB for FineWeb-Edu 10BT tokenized data

Step 1: Environment Setup#

1.1 Pull the Docker Image#

docker pull docker.io/rocm/primus:v25.10

1.2 Clone the Repository#

git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git
cd Primus

1.3 Start a Development Container#

# Quick start (mounts Primus into /workspace/Primus)
bash tools/docker/start_container.sh

This creates a persistent container named dev_primus_<user>. You can customize it with environment variables:

DOCKER_IMAGE=docker.io/rocm/primus:v25.10 \
DATA_PATH=/path/to/data \
bash tools/docker/start_container.sh

Then exec into the container:

docker exec -it dev_primus_$(whoami) bash
cd /workspace/Primus

1.4 Install Python Dependencies (inside container)#

pip install -r requirements.txt

For GDN models (required for the Triton kernel and FLA model classes):

pip install flash-linear-attention

For evaluation, also install:

pip install lm-eval

Step 2: Dataset Preparation#

Zebra hybrid uses the FineWeb-Edu dataset, preprocessed into Megatron binary format.

2.1 Set Up Environment#

export HF_TOKEN="hf_your_token_here"
export PYTHONPATH="$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"

2.2 Run Data Preparation#

python examples/megatron/prepare_fineweb_edu.py \
    --primus-path . \
    --data-path ./data \
    --tokenizer-type HuggingFaceTokenizer \
    --tokenizer-model meta-llama/Llama-3.2-1B \
    --sample-size 10BT

This will:

  1. Download the FineWeb-Edu 10BT dataset from HuggingFace

  2. Tokenize it into Megatron binary format (.bin + .idx files)

  3. Output files to ./data/fineweb-edu-10BT/HuggingFaceTokenizer/

Available sample sizes: 10BT, 100BT, 350BT

The script uses all available CPU cores by default. To limit parallelism, add --workers N.

Note: The context length (sequence length) is not set during data prep. It is configured at training time via seq_length in your pretrain YAML.

2.3 Using a Different Tokenizer#

For the GDN config which uses fla-hub/gla-1.3B-100B:

python examples/megatron/prepare_fineweb_edu.py \
    --primus-path . \
    --data-path ./data \
    --tokenizer-type HuggingFaceTokenizer \
    --tokenizer-model fla-hub/gla-1.3B-100B \
    --sample-size 10BT

2.5 Update Data Paths in Config#

After preparation (standard or FLA-aligned), update the train_data_path in your pretrain config YAML to point to the generated files:

# Standard Megatron data prep (multiple shards)
train_data_path: >
  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_0_text_sentence
  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_1_text_sentence
  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_2_text_sentence
  /path/to/data/fineweb-edu-10BT/HuggingFaceTokenizer/fineweb_edu_10BT_3_text_sentence
mock_data: false

Step 3: Pretraining#

Single-Node (Local / Docker)#

Launch training inside a Docker container on a single node:

# Zebra hybrid 1B with KDA (Kimi Delta Attention)
export DATA_PATH=./data
GPUS_PER_NODE=8 HF_TOKEN=$HF_TOKEN \
./primus-cli container --volume "$DATA_PATH:$DATA_PATH" \
  --env DATA_PATH \
  -- train pretrain --config examples/megatron/configs/MI300X/zebra_llama_kda_1B_BF16-pretrain.yaml

Other model variants (same launcher, different --config):

# Zebra hybrid 1B with Mamba SSM
./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_mamba_1B_BF16-pretrain.yaml

# Zebra hybrid 1B with pure KDA (no attention layers)
./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/kda_1B_BF16-pretrain.yaml

# Zebra hybrid 1B with GDN (pure recurrent, no attention)
./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_gdn_1B_BF16-pretrain.yaml

# Zebra hybrid 1B pure GDN (FLA-validated, 4-GPU)
GPUS_PER_NODE=4 ./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/gdn_1B_BF16-pretrain.yaml

# Zebra hybrid 3B
./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_mamba_3B_BF16-pretrain.yaml

# Zebra hybrid 8B
./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_mamba_8B_BF16-pretrain.yaml

Multi-Node (Slurm)#

For multi-node training on a Slurm cluster:

export DATA_PATH=/shared/data
./primus-cli slurm srun -N 2 \
  -- container --volume "$DATA_PATH:$DATA_PATH" \
  --env DATA_PATH \
  -- train pretrain --config examples/megatron/configs/MI300X/zebra_llama_kda_1B_BF16-pretrain.yaml

Ensure the global_batch_size in your config is divisible by micro_batch_size * GPUS_PER_NODE * NNODES.

If Already Inside a Container#

If you are already inside a Docker container or on a bare-metal node with the environment set up, use direct mode instead of container — it skips the container launch and runs torchrun in place:

./primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_kda_1B_BF16-pretrain.yaml

ROCm FLA Triton note. FLA’s KDA intra-chunk kernels autotune over num_stages values that AMD Triton 3.6 cannot compile — the MLIR pass pipeline aborts inside tritonamdgpu-schedule-loops{num_stages=4}, so a KDA run dies during its first autotune sweep. Primus applies the workaround for you, inside the training process: the megatron.fla.kda_safe_autotune patch (primus/backends/megatron/patches/fla_kda_autotune_patches.py) drops the num_stages >= 3 candidates from the two kernels’ autotune space at startup, keeping the num_warps sweep. Nothing on disk is rewritten — the installed fla package is left pristine. It fires on ROCm whenever the run resolves to the FLA KDA backend (use_fla_triton_kda: true, or kda_backend: fla / use_kimi_k3_attention_backend: fla). GDN, mamba and plain-attention configs drive different FLA ops, compile fine, and are left untouched. On Triton 3.7 the full sweep compiles and the autotuner reaches the same configs on its own, so the narrowing costs at most ~1% of the KDA op on the largest shape measured; it will be removed once the supported toolchain floor reaches Triton 3.7.

Mock Data (Smoke Test)#

To quickly verify the model runs without real data, the 3B and 8B configs come with mock_data: true by default. For the 1B configs, you can override on the command line — every argument after --config is forwarded to the Primus Python CLI:

./primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/zebra_llama_kda_1B_BF16-pretrain.yaml \
  --mock_data true --train_iters 10

Key Training Parameters#

Parameter

Description

Typical Values

train_iters

Total training iterations

38147 (1B KDA pure), 400000 (1B Mamba)

micro_batch_size

Per-GPU batch size

4 (KDA/GDN), 16 (Mamba 1B)

global_batch_size

Total batch size across all GPUs

micro_batch_size * num_gpus

seq_length

Sequence length

2048, 4096, 8192

lr

Peak learning rate

2.0e-4

save_interval

Checkpoint save frequency

1000

auto_continue_train

Auto-resume from last checkpoint on crash

true / false

hybrid_attention_ratio

Fraction of attention layers (0.0 = pure recurrent)

0.0, 0.25

Step 4: Checkpoint Conversion to HuggingFace#

Convert a Megatron checkpoint to HuggingFace format for inference and evaluation.

4.1 Convert Checkpoint#

Pure GDN Models#

Pure GDN models use a dedicated converter that maps Primus’s fused projections to FLA’s native GatedDeltaNetForCausalLM format:

python tools/hybrid/convert_gdn_to_fla_hf.py \
    --checkpoint-path output/amd/root/gdn_1B_BF16-pretrain/checkpoints/iter_0076294 \
    --output-dir output/gdn_1B_fla_hf \
    --config /path/to/gated_deltanet_1B.json

This handles:

  • Splitting the fused in_proj (3104 → q/k/v/gate/beta/alpha projections)

  • Splitting the fused conv1d (q/k/v convolutions)

  • Splitting the fused SwiGLU fc1 (gate_proj + up_proj)

  • Mapping alternating GDN/MLP sublayers to combined FLA layers

  • Handling tied embeddings

After conversion, verify with the sanity check:

python tools/hybrid/verify_gdn_conversion.py --model-path output/gdn_1B_fla_hf

Expected output: Loss ~2-4, top prediction for “The capital of France is” should be “Paris”.

KDA / Hybrid Models#

The general converter auto-detects architecture from the checkpoint’s saved arguments:

# KDA+MLA hybrid model
python tools/hybrid/convert_zebra_llama_to_hf.py \
    --checkpoint-path output/zebra_llama_kda_1B_BF16-pretrain/iter_0028000 \
    --output-dir output/zebra_llama_kda_1B_hf_iter_0028000

# Pure KDA model
python tools/hybrid/convert_zebra_llama_to_hf.py \
    --checkpoint-path output/kda_1B_BF16-pretrain/iter_0038000 \
    --output-dir output/kda_1B_hf

The converter will:

  • Read the Megatron checkpoint and training arguments

  • Auto-detect architecture parameters (hybrid_attention_ratio, kda_num_heads, q_lora_rank, etc.)

  • Remap parameter names from Megatron conventions to HuggingFace conventions

  • Save pytorch_model.bin, config.json, and a model card README.md in the output directory

  • Copy modeling_zebra_llama.py into the output directory and record it under auto_map in config.json, so the directory loads with trust_remote_code=True

  • Not write tokenizer files: pass the base model to --tokenizer when evaluating or chatting

4.2 Verify Conversion#

The script prints a summary of missing, extra, and shape-mismatched keys. A successful conversion shows:

0 missing, 0 extra, 0 shape mismatches

4.3 Supported Architectures#

Architecture

hybrid_attention_ratio

Layer pattern

Pure KDA

0.0

All KDA + MLP

KDA + MLA hybrid

0.0 < r < 1.0

Mix of KDA and MLA + MLP

Pure MLA

1.0

All MLA + MLP

Pure GDN

0.0 (with GDN spec)

All GDN + MLP

Mamba + MLA hybrid

0.0 < r < 1.0 (with Mamba spec)

Mix of Mamba and MLA + MLP

Step 5: Evaluation with lm-eval-harness#

5.1 Pure GDN Models (FLA format)#

Pure GDN models use a dedicated eval wrapper (tools/hybrid/eval_gdn_lm_eval.py) that pre-registers FLA’s GatedDeltaNetForCausalLM with transformers’ AutoModel and patches compatibility issues with transformers >= 4.55:

python tools/hybrid/eval_gdn_lm_eval.py \
    --model hf \
    --model_args pretrained=output/gdn_1B_fla_hf,trust_remote_code=True,tokenizer=meta-llama/Llama-3.2-1B \
    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
    --batch_size auto \
    --output_path eval_results/gdn_1B

Note: Do not use lm_eval --model hf directly — it will fail because AutoConfig does not recognize gated_deltanet without FLA being imported first. The wrapper handles this. The tokenizer=meta-llama/Llama-3.2-1B argument is required since the converted model directory does not contain tokenizer files.

5.2 KDA / Hybrid Models (Zebra hybrid format)#

KDA and hybrid models use the custom ZebraLlamaForCausalLM architecture, which requires a dedicated lm-eval wrapper:

python3 tools/hybrid/lm_harness_eval.py \
    --model_path output/kda_1B_hf \
    --tokenizer meta-llama/Llama-3.2-1B \
    --dtype bfloat16 \
    --tasks arc_easy,arc_challenge,hellaswag,mmlu,openbookqa,piqa,race,winogrande \
    --batch_size auto

Pass --tokenizer explicitly: the converter writes weights, config and modeling code, but no tokenizer files.

5.3 Using the Eval Shell Script (KDA/Hybrid)#

This wrapper takes key=value assignments rather than GNU-style flags, and any it is not given fall back to the defaults shown by the script:

bash tools/hybrid/eval_zebra_llama_lm_eval.sh \
    model_path=output/kda_1B_hf \
    tokenizer=meta-llama/Llama-3.2-1B \
    batch_size=auto \
    output_path=eval_results/kda_1B

Important: both entry points import tools/hybrid/modeling_zebra_llama.py, and that import is what registers the custom architecture with the Auto* classes before lm_eval.simple_evaluate(model="hf", ...) runs. Do not re-register it yourself; a second registration raises a duplicate-registration ValueError. The task list in this wrapper is fixed in TASKS inside the script.

5.4 Available Benchmarks#

Task

Description

Metric

arc_easy

ARC Easy (science QA)

acc, acc_norm

arc_challenge

ARC Challenge (harder science QA)

acc, acc_norm

hellaswag

HellaSwag (commonsense NLI)

acc, acc_norm

mmlu

MMLU (57 subject knowledge benchmark)

acc

openbookqa

OpenBookQA

acc, acc_norm

piqa

PIQA (physical intuition QA)

acc, acc_norm

race

RACE (reading comprehension)

acc

winogrande

Winogrande (coreference resolution)

acc

5.5 Memory Considerations#

The pure-PyTorch KDA chunked attention is memory-intensive. If you encounter OOM errors:

  • Use --batch_size auto to let lm-eval find the largest fitting batch size

  • Reduce max_length (e.g., max_length=1024 in --model_args)

  • Reduce --batch_size to 1

Configuration Reference#

Hybrid Layer Specs#

The spec field in the pretrain config selects the layer arrangement:

Spec

Description

hybrid_stack_spec

Mamba SSM + MLA hybrid

kda_hybrid_stack_spec

KDA + MLA hybrid (or pure KDA with hybrid_attention_ratio: 0.0)

gdn_hybrid_stack_spec

GDN + MLA hybrid (or pure GDN with hybrid_attention_ratio: 0.0)

Environment Variables#

Variable

Description

Default

DOCKER_IMAGE

Docker image for training

docker.io/rocm/primus:v25.10

EXP

Path to experiment config YAML

examples/megatron/exp_pretrain.yaml

DATA_PATH

Path to dataset directory

./data

HF_TOKEN

HuggingFace API token

(required for gated models)

WANDB_API_KEY

Weights & Biases API key

(optional)

GPUS_PER_NODE

Number of GPUs per node

8

NNODES

Number of nodes

1

MASTER_ADDR

Master node address

localhost

MASTER_PORT

Master node port

1234

Troubleshooting#

OOM During Training#

  • Reduce micro_batch_size or seq_length

  • Enable activation checkpointing: add recompute_granularity: selective to the config

OOM During Evaluation#

  • Use --batch_size 1 or --batch_size auto

  • Add max_length=1024 to --model_args

ModuleNotFoundError: No module named 'megatron'#

Set the Python path before running data preparation:

export PYTHONPATH="$(pwd)/third_party/Megatron-LM:${PYTHONPATH}"

Checkpoint Conversion Shape Mismatches#

Ensure the modeling_zebra_llama.py model definition matches the architecture of your checkpoint (Mamba vs KDA vs GDN). The converter auto-detects architecture from checkpoint args, but the HF model code in tools/hybrid/modeling_zebra_llama.py must support the target architecture. Common causes of shape mismatches:

  • Mismatched hybrid_attention_ratio between config and checkpoint

  • Incorrect kda_num_heads or head dimension settings

  • Using a modeling_zebra_llama.py that doesn’t support the checkpoint’s attention type

ValueError: model type 'zebra_llama' not recognized#

This occurs when using lm_eval --model hf directly instead of the custom wrapper, so the zebra_llama architecture was never registered. Always use:

python3 tools/hybrid/lm_harness_eval.py --model_path <converted-checkpoint> --tokenizer <base-model> ...

Or the eval shell script, which handles this automatically.

Truncation Warnings During Eval#

Messages like Combined length of context and continuation exceeds model's maximum length mean some eval samples are being truncated. This has minimal impact on most benchmarks but can affect long-context tasks like RACE. To avoid truncation, increase max_length in --model_args.

NCCL / RCCL Timeout During Training#

On MI300X, intermittent RCCL hangs can occur (typically during checkpoint saves). Mitigations:

  • Set auto_continue_train: true in the pretrain config to auto-resume from the last checkpoint

  • Increase the heartbeat timeout: export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=7200

File Reference#

Primus/
├── examples/megatron/
│   ├── configs/MI300X/
│   │   ├── zebra_llama_mamba_1B_BF16-pretrain.yaml        # 1B Mamba+MLA
│   │   ├── zebra_llama_kda_1B_BF16-pretrain.yaml     # 1B KDA+MLA hybrid
│   │   ├── kda_1B_BF16-pretrain.yaml # 1B pure KDA
│   │   ├── zebra_llama_gdn_1B_BF16-pretrain.yaml     # 1B GDN
│   │   ├── gdn_1B_BF16-pretrain.yaml # 1B pure GDN (FLA-validated)
│   │   ├── zebra_llama_mamba_3B_BF16-pretrain.yaml         # 3B Mamba+MLA
│   │   └── zebra_llama_mamba_8B_BF16-pretrain.yaml         # 8B Mamba+MLA
│   ├── prepare_fineweb_edu.py                   # Data preparation script
│   ├── prepare_fineweb_edu.sh                   # Data prep shell wrapper
│   └── preprocess_data.py                       # Megatron tokenizer
├── primus/configs/models/megatron/
│   ├── zebra_mamba_1B_hybrid.yaml                      # 1B model architecture
│   ├── kda_1B.yaml             # 1B pure KDA architecture
│   ├── zebra_gdn_1B_hybrid.yaml                  # 1B GDN architecture
│   ├── gdn_1B.yaml             # 1B pure GDN (FLA-validated)
│   ├── zebra_mamba_3B_hybrid.yaml                      # 3B model architecture
│   └── zebra_mamba_8B_hybrid.yaml                      # 8B model architecture
├── tools/
│   ├── hybrid/
│   │   ├── convert_zebra_llama_to_hf.py         # Megatron → HF converter (KDA/hybrid)
│   │   ├── convert_gdn_to_fla_hf.py             # Megatron → FLA HF converter (pure GDN)
│   │   ├── verify_gdn_conversion.py             # Post-conversion sanity check (pure GDN)
│   │   ├── eval_gdn_lm_eval.py                  # lm-eval wrapper for GDN (registers FLA)
│   │   ├── convert_zebra_llama_to_hf.sh         # Converter shell wrapper
│   │   ├── modeling_zebra_llama.py              # HF model definition (KDA/hybrid)
│   │   ├── lm_harness_eval.py                   # lm-eval wrapper
│   │   ├── eval_zebra_llama_lm_eval.sh          # Eval shell wrapper
│   │   ├── run_zebra_eval.sh                    # Quick eval script
│   │   ├── chat_zebra_llama.py                  # Interactive chat
│   │   └── convert_fla_to_megatron.py           # FLA Arrow → Megatron binary converter
│   └── docker/start_container.sh                # Dev container launcher
├── runner/
│   ├── primus-cli                               # Unified launcher (direct/container/slurm)
│   └── helpers/envs/base_env.sh                 # NCCL/ROCm/cache env defaults
└── requirements.txt                             # Python dependencies
Contents