Apriel 2

[tscholak]

This PR enables declarative architecture surgery on Apriel2 models. You can now convert a standard attention-based model into exotic architectures (stochastic supernets, hybrid Mamba/GDN models) with a single command and a YAML config.

To convert Apriel-1.5-15b-Thinker into a stochastic supernet with 3 mixer types per layer:

python convert.py ServiceNow-AI/Apriel-1.5-15b-Thinker /tmp/apriel2-supernet \
    --surgery examples/stochastic_supernet.yaml --verbose

The converter shows you exactly what it's doing with a visual plan tree:

============================================================
CONVERSION PLAN
============================================================
├── model/
│   ├── [embed_tokens, norm]/
│   │   └── weight ← language_model.model.[embed_tokens, norm].weight
│   ├── vision_encoder/
│   │   └── encoder/blocks/[0..23]/...
│   └── decoder/
│       └── blocks/
│           └── [0..47]/
│               ├── mixer/
│               │   └── mixers/
│               │       ├── attention/
│               │       │   └── self_attn/[q,k,v,o]_proj/weight ← ...self_attn.[q,k,v,o]_proj.weight
│               │       ├── sliding_window/
│               │       │   └── self_attn/[q,k,v,o]_proj/weight ← ...self_attn.[q,k,v,o]_proj.weight
│               │       └── gated_delta_net/
│               │           └── gdn/
│               │               ├── in_proj_qkvz/weight = [Q[:160,:]; K[:160,:]; V slices...; 𝟎(640×5120)]
│               │               ├── in_proj_ba/weight = 𝟎(64×5120)
│               │               ├── out_proj/weight ← ...self_attn.o_proj.weight
│               │               ├── conv1d/weight = scaled_identity_conv(7680×1×4)
│               │               ├── A_log = A_log(32)
│               │               └── norm/weight = 𝟏(160)
│               ├── mlp/[gate,up,down]_proj/weight ← ...
│               └── [input_layernorm, post_attention_layernorm]/weight ← ...
└── lm_head/weight ← language_model.lm_head.weight
============================================================
Converting: 100%|████████████████████████████████| 1185/1185 [00:40<00:00]
Conversion complete! Output saved to /tmp/apriel2-supernet

The result is a sharded Apriel2 checkpoint with attention + sliding window + gated delta net mixers in every layer.

The surgery config for this looks like this:

decoder:
  type: fixed
  block:
    mixer:
      type: stochastic
      main_mixer_name: attention
      sampling_strategy: uniform
      mixers:
        # Main attention mixer - inherits config and weights from source
        attention:
          type: attention
          init: transfer

        # Sliding window - same architecture with window size override
        sliding_window:
          type: attention
          init: transfer
          sliding_window: 4096

        # Gated delta net - DIL initialization maps Q/K/V/O -> GDN projections
        gated_delta_net:
          type: gated_delta_net
          init: transfer
          conv_kernel_size: 4  # Only required param - rest derived from source

    # MLP and normalization transfer from source
    mlp:
      init: transfer

    normalization:
      init: transfer

The converter supports the following better-than-random initializations for linear mixers:

• MIL: Q→C, K→B, V→x, O→out_proj
• DIL: Q/K/V→in_proj_qkvz with GQA-aware tiling