Add Transformers support with AutoModelForCausalLM and preset selection

- Add AutoModelForCausalLM to auto_map in all config.json files (main + 9 presets)
- Add config_class = Apriel2Config to Apriel2ForCausalLM for compatibility
- Add Apriel2PatternMixerAdapter to match supernet checkpoint weight paths
(blocks.{i}.mixer.mixers.{type}.{param}) when loading pattern preset configs
- Update README: move Transformers section after vLLM, add install deps,
document preset selection workflow with config.json copy

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

README.md

@@ -107,51 +107,7 @@ SuperApriel-15b-Instruct is trained in two stages:
 
 ## How to Use
 
-Install dependencies:
-
-```bash
-pip install transformers
-```
-
-> **🔴 TODO: There is currently no mechanism to select a preset when using Transformers directly. The model defaults to the all-attention preset during inference. Preset selection requires vLLM with the Fast-LLM plugin (see below). We need to add a Transformers API for preset switching.**
-
-Basic usage with Transformers (all-attention preset):
-
-```python
-import re
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-model_name = "ServiceNow-AI/SuperApriel-15b-Instruct"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto",
-    trust_remote_code=True,
-)
-
-prompt = "Positive real numbers $x$ and $y$ satisfy $y^3=x^2$ and $(y-x)^2=4y^2$. What is $x+y$?\nMark your solution with \\boxed"
-messages = [
-    {"role": "user", "content": prompt}
-]
-
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True,
-    tools=[]
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-generated_ids = model.generate(**model_inputs, max_new_tokens=1024)
-output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-
-response = re.findall(r"\[BEGIN FINAL RESPONSE\](.*?)\[END FINAL RESPONSE\]", output, re.DOTALL)[0].strip()
-print("response:", response)
-```
-
-See [Apriel-1.6-15b-Thinker](https://huggingface.co/ServiceNow-AI/Apriel-1.6-15b-Thinker) for recommended inference settings.
+The recommended serving backend is **vLLM** with the Fast-LLM plugin, which supports preset selection and runtime switching. For simpler use cases, Transformers is also supported (see [Use with Transformers](#use-with-transformers) below).
 
 ## Use with vLLM
 
@@ -265,6 +221,70 @@ The supernet mode above switches the active placement globally for all requests.
 
 > **🔴 TODO: Add per-request placement selection via the vLLM serving API (e.g. `placement_id` field in the request body). This is managed separately from the global `collective_rpc` switching shown above.**
 
+## Use with Transformers
+
+The model works with `AutoModelForCausalLM` for standard Transformers inference. Each preset config selects one mixer per layer, and the model automatically loads only the relevant weights from the supernet checkpoint.
+
+### Installation
+
+```bash
+pip install transformers torch einops accelerate
+```
+
+Presets that include GDN or KDA mixers additionally require:
+
+```bash
+pip install causal-conv1d mamba-ssm
+```
+
+Attention-only presets (`all-attention`, and others using only FA/SWA) work without these extra packages.
+
+### Usage
+
+The repo ships with the supernet config as default (for vLLM). Before loading with Transformers, you must select a preset by copying its `config.json`:
+
+```python
+import re
+import shutil
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_dir = "ServiceNow-AI/SuperApriel-15b-Instruct"  # or local path
+
+# Select a preset (required — the default supernet config is for vLLM only)
+preset = "all-attention"  # or any preset from preset_configs/
+shutil.copy(f"{model_dir}/preset_configs/{preset}/config.json",
+            f"{model_dir}/config.json")
+
+tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_dir,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+prompt = "Positive real numbers $x$ and $y$ satisfy $y^3=x^2$ and $(y-x)^2=4y^2$. What is $x+y$?\nMark your solution with \\boxed"
+messages = [{"role": "user", "content": prompt}]
+
+text = tokenizer.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True, tools=[]
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+generated_ids = model.generate(**model_inputs, max_new_tokens=1024)
+output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+response = re.findall(
+    r"\[BEGIN FINAL RESPONSE\](.*?)\[END FINAL RESPONSE\]", output, re.DOTALL
+)[0].strip()
+print("response:", response)
+```
+
+Available presets: `all-attention`, `Reg_Lklhd-26`, `Reg_Lklhd-18`, `Reg_Lklhd-13`, `Reg_Lklhd-10`, `Idealized_All-18`, `Idealized_All-6`, `Idealized_Lklhd-6`, `extra_bayesian-mix-7`.
+
+See [Apriel-1.6-15b-Thinker](https://huggingface.co/ServiceNow-AI/Apriel-1.6-15b-Thinker) for recommended inference settings.
+
 ## Chat Template
 
 ```

config.json

@@ -1,151 +1,152 @@
 {
- "architectures": [
-  "Apriel2ForConditionalGeneration"
- ],
- "auto_map": {
-  "AutoConfig": "configuration_apriel2.Apriel2Config",
-  "AutoModel": "modeling_apriel2.Apriel2Model",
-  "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
- },
- "bos_token_id": 1,
- "decoder": {
-  "block": {
-   "mixer": {
-    "main_mixer_name": "attention",
-    "mixers": {
-     "attention": {
-      "add_linear_biases": false,
-      "head_groups": 8,
-      "head_size": 128,
-      "heads": 32,
-      "rotary": {
-       "theta": 1000000000.0,
-       "type": "mistral_1d"
-      },
-      "type": "attention",
-      "window_size": null
-     },
-     "gdn": {
-      "convolution_layer": {
-       "kernel_size": 4
+  "architectures": [
+    "Apriel2ForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_apriel2.Apriel2Config",
+    "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
+    "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
+  },
+  "bos_token_id": 1,
+  "decoder": {
+    "block": {
+      "mixer": {
+        "main_mixer_name": "attention",
+        "mixers": {
+          "attention": {
+            "add_linear_biases": false,
+            "head_groups": 8,
+            "head_size": 128,
+            "heads": 32,
+            "rotary": {
+              "theta": 1000000000.0,
+              "type": "mistral_1d"
+            },
+            "type": "attention",
+            "window_size": null
+          },
+          "gdn": {
+            "convolution_layer": {
+              "kernel_size": 4
+            },
+            "key_head_dim": 128,
+            "key_heads": 8,
+            "type": "gdn",
+            "value_head_dim": 128,
+            "value_heads": 32
+          },
+          "kda": {
+            "convolution_layer": {
+              "kernel_size": 4
+            },
+            "head_dim": 128,
+            "heads": 32,
+            "normalization": {
+              "epsilon": 1e-05
+            },
+            "type": "kda"
+          },
+          "sliding_window": {
+            "add_linear_biases": false,
+            "head_groups": 8,
+            "head_size": 128,
+            "heads": 32,
+            "rotary": {
+              "theta": 1000000000.0,
+              "type": "mistral_1d"
+            },
+            "type": "attention",
+            "window_size": 4096
+          }
+        },
+        "sampling_strategy": "full_layout",
+        "type": "stochastic"
       },
-      "key_head_dim": 128,
-      "key_heads": 8,
-      "type": "gdn",
-      "value_head_dim": 128,
-      "value_heads": 32
-     },
-     "kda": {
-      "convolution_layer": {
-       "kernel_size": 4
+      "mlp": {
+        "activation": "silu",
+        "add_linear_biases": false,
+        "gated": true,
+        "intermediate_size": 14336,
+        "type": "mlp"
       },
-      "head_dim": 128,
-      "heads": 32,
       "normalization": {
-       "epsilon": 1e-05
-      },
-      "type": "kda"
-     },
-     "sliding_window": {
-      "add_linear_biases": false,
-      "head_groups": 8,
-      "head_size": 128,
-      "heads": 32,
-      "rotary": {
-       "theta": 1000000000.0,
-       "type": "mistral_1d"
-      },
-      "type": "attention",
-      "window_size": 4096
-     }
+        "epsilon": 1e-05,
+        "type": "rms_norm"
+      }
     },
-    "sampling_strategy": "full_layout",
-    "type": "stochastic"
-   },
-   "mlp": {
-    "activation": "silu",
-    "add_linear_biases": false,
-    "gated": true,
-    "intermediate_size": 14336,
-    "type": "mlp"
-   },
-   "normalization": {
-    "epsilon": 1e-05,
-    "type": "rms_norm"
-   }
-  },
-  "num_blocks": 48,
-  "type": "fixed"
- },
- "embeddings": {
-  "max_position_embeddings": 120000
- },
- "eos_token_id": 2,
- "head": {
-  "normalization": {
-   "epsilon": 1e-05,
-   "type": "rms_norm"
-  }
- },
- "hidden_size": 5120,
- "image_token_index": 10,
- "model_type": "apriel2",
- "tie_word_embeddings": false,
- "transformers_version": "4.57.3",
- "use_cache": true,
- "vision_encoder": {
-  "adapter": {
-   "activation": "gelu_pytorch_tanh",
-   "add_linear_biases": true,
-   "gated": false,
-   "intermediate_size": 5120,
-   "type": "mlp"
+    "num_blocks": 48,
+    "type": "fixed"
   },
   "embeddings": {
-   "input_channels": 3,
-   "normalization": {
-    "epsilon": 1e-05,
-    "type": "rms_norm"
-   },
-   "patch_height": 16,
-   "patch_width": 16
+    "max_position_embeddings": 120000
   },
-  "encoder": {
-   "block": {
-    "mixer": {
-     "add_linear_biases": false,
-     "causal": false,
-     "cross_document_attention": false,
-     "head_groups": 16,
-     "head_size": 64,
-     "heads": 16,
-     "rotary": {
-      "max_image_size": 1024,
-      "patch_size": 16,
-      "theta": 10000.0,
-      "type": "pixtral_2d"
-     },
-     "type": "attention"
-    },
-    "mlp": {
-     "activation": "silu",
-     "add_linear_biases": false,
-     "gated": true,
-     "intermediate_size": 4096,
-     "type": "mlp"
-    },
+  "eos_token_id": 2,
+  "head": {
     "normalization": {
-     "epsilon": 1e-05,
-     "type": "rms_norm"
+      "epsilon": 1e-05,
+      "type": "rms_norm"
     }
-   },
-   "num_blocks": 24,
-   "type": "fixed"
   },
-  "hidden_size": 1024,
-  "num_channels": 3,
-  "num_hidden_layers": 24,
-  "patch_size": 16
- },
- "vocab_size": 131072
-}
+  "hidden_size": 5120,
+  "image_token_index": 10,
+  "model_type": "apriel2",
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.3",
+  "use_cache": true,
+  "vision_encoder": {
+    "adapter": {
+      "activation": "gelu_pytorch_tanh",
+      "add_linear_biases": true,
+      "gated": false,
+      "intermediate_size": 5120,
+      "type": "mlp"
+    },
+    "embeddings": {
+      "input_channels": 3,
+      "normalization": {
+        "epsilon": 1e-05,
+        "type": "rms_norm"
+      },
+      "patch_height": 16,
+      "patch_width": 16
+    },
+    "encoder": {
+      "block": {
+        "mixer": {
+          "add_linear_biases": false,
+          "causal": false,
+          "cross_document_attention": false,
+          "head_groups": 16,
+          "head_size": 64,
+          "heads": 16,
+          "rotary": {
+            "max_image_size": 1024,
+            "patch_size": 16,
+            "theta": 10000.0,
+            "type": "pixtral_2d"
+          },
+          "type": "attention"
+        },
+        "mlp": {
+          "activation": "silu",
+          "add_linear_biases": false,
+          "gated": true,
+          "intermediate_size": 4096,
+          "type": "mlp"
+        },
+        "normalization": {
+          "epsilon": 1e-05,
+          "type": "rms_norm"
+        }
+      },
+      "num_blocks": 24,
+      "type": "fixed"
+    },
+    "hidden_size": 1024,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "patch_size": 16
+  },
+  "vocab_size": 131072
+}

modeling_apriel2.py

@@ -972,6 +972,32 @@ def create_mixer(mixer_config: dict, hidden_size: int, layer_idx: int, config, a
         return mixer_class(hidden_size, mixer_config, layer_idx=layer_idx)
 
 
+class Apriel2PatternMixerAdapter(nn.Module):
+    """Adapter that wraps a single mixer under mixers.{name} to match supernet weight paths.
+
+    The supernet checkpoint stores weights as blocks.{i}.mixer.mixers.{type}.{param},
+    but a bare mixer creates blocks.{i}.mixer.{param}. This adapter adds the intermediate
+    mixers.{name} level so pattern configs can load from supernet checkpoints.
+    """
+
+    def __init__(self, mixer_name: str, mixer: nn.Module):
+        super().__init__()
+        self.mixers = nn.ModuleDict({mixer_name: mixer})
+        self._mixer_name = mixer_name
+
+    def forward(self, *args, **kwargs):
+        return self.mixers[self._mixer_name](*args, **kwargs)
+
+    def preprocess(self, *args, **kwargs):
+        return self.mixers[self._mixer_name].preprocess(*args, **kwargs)
+
+    @classmethod
+    def setup(cls, mixer_name: str, mixer_config: dict, hidden_size: int, max_position_embeddings: int) -> nn.ModuleDict:
+        mixer_type = mixer_config.get("type", "attention")
+        mixer_class = get_mixer_class(mixer_type)
+        return mixer_class.setup(mixer_config, hidden_size, max_position_embeddings)
+
+
 class Apriel2Mamba(nn.Module):
     """Mamba mixer."""
 
@@ -1906,14 +1932,16 @@ class Apriel2BlockSequence(nn.Module):
 
         blocks = []
         for layer_idx in range(num_blocks):
-            # Get block_config for this layer
+            # Get block_config and block_name for this layer
             if seq_type == "fixed":
                 block_config = self.sequence_config.get("block", {})
+                block_name_for_layer = None  # No adapter needed for fixed type
             elif seq_type == "pattern":
                 pattern = self.sequence_config.get("pattern", [])
                 blocks_config = self.sequence_config.get("blocks", {})
                 block_name = pattern[layer_idx % len(pattern)]
                 block_config = blocks_config[block_name]
+                block_name_for_layer = block_name  # Pass to Apriel2Block for weight path matching
             else:
                 raise ValueError(f"Unknown sequence type: {seq_type}")
 
@@ -1925,6 +1953,7 @@ class Apriel2BlockSequence(nn.Module):
                     layer_idx=layer_idx,
                     rms_norm_eps=rms_norm_eps,
                     config=self.config,
+                    block_name=block_name_for_layer,
                 )
             )
 
@@ -2031,6 +2060,7 @@ class Apriel2Block(nn.Module):
         layer_idx: int,
         rms_norm_eps: float,
         config: Apriel2TextConfig,
+        block_name: Optional[str] = None,
     ):
         """
         Args:
@@ -2039,6 +2069,7 @@ class Apriel2Block(nn.Module):
             layer_idx: Layer index in the sequence
             rms_norm_eps: Epsilon for RMS normalization
             config: Model config (passed to mixers that need it)
+            block_name: For pattern configs, the mixer name (e.g. "attention") to match supernet weight paths
         """
         super().__init__()
         self.hidden_size = hidden_size
@@ -2046,7 +2077,13 @@ class Apriel2Block(nn.Module):
 
         # Create mixer based on type
         mixer_config = block_config.get("mixer", {"type": "attention"})
-        self.mixer = create_mixer(mixer_config, hidden_size, layer_idx, config, allow_stochastic=True)
+        raw_mixer = create_mixer(mixer_config, hidden_size, layer_idx, config, allow_stochastic=True)
+
+        # For pattern configs, wrap in adapter to match supernet checkpoint weight paths
+        if block_name is not None:
+            self.mixer = Apriel2PatternMixerAdapter(block_name, raw_mixer)
+        else:
+            self.mixer = raw_mixer
 
         # Create MLP
         mlp_config = block_config.get("mlp", {"type": "mlp"})
@@ -2435,6 +2472,7 @@ class Apriel2TextModel(Apriel2PreTrainedModel):
 class Apriel2ForCausalLM(Apriel2PreTrainedModel, GenerationMixin):
     """Apriel2 model with a language modeling head (text-only)."""
 
+    config_class = Apriel2Config
     _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config: Apriel2TextConfig):

preset_configs/Idealized_All-18/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Idealized_All-6/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Idealized_Lklhd-6/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Reg_Lklhd-10/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Reg_Lklhd-13/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Reg_Lklhd-18/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/Reg_Lklhd-26/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/all-attention/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,

preset_configs/extra_bayesian-mix-7/config.json

@@ -5,6 +5,7 @@
   "auto_map": {
     "AutoConfig": "configuration_apriel2.Apriel2Config",
     "AutoModel": "modeling_apriel2.Apriel2Model",
+    "AutoModelForCausalLM": "modeling_apriel2.Apriel2ForCausalLM",
     "AutoModelForImageTextToText": "modeling_apriel2.Apriel2ForConditionalGeneration"
   },
   "bos_token_id": 1,