Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/supported-models.md
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ Implemented model families include:
- LFM2-VL (`lfm2_vl` / `lfm2-vl`): a SigLIP2-style packed-patch vision tower (native variable resolution, per-image bicubically-resampled position grid) + a pixel-unshuffle (space-to-depth) projector into the LFM2 hybrid text backbone. Each image is smart-resized so its post-downsample token count lands in `[min_image_tokens, max_image_tokens]`, packed at its native patch count (no padding), and its `ceil(h/f)*ceil(w/f)` projected tokens replace the `<image>` placeholder. The multi-tile `do_image_splitting` high-resolution path is a documented follow-up.
- PLaMo 2 (Preferred Networks: interleaves Mamba SSM and GQA attention layers by index; each block carries normformer-style pre/post offset RMSNorms, and the Mamba mixer derives B/C/dt from a post-conv projection). The architecture is validated against the mlx-lm reference at the token-id level (`tests/plamo2_parity.rs`). CLI text generation additionally needs support for PLaMo's custom `PlamoTokenizer` (the `tokenizer.jsonl` Unigram format), which the Rust tokenizer loader does not yet read.
- Kimi Linear, LongCat Flash, LongCat Flash N-gram
- GPT-2 (`gpt2`, OpenAI: the original decoder-only transformer, and the first family in this tree to use learned absolute position embeddings. `wpe(arange(seq_len) + cache_offset)` is added to the `wte` token embeddings at the input boundary and there is no RoPE anywhere, so the KV-cache offset enters the graph through the position lookup rather than through a rotation. The rest of the block is `LayerNorm` with bias (not RMSNorm), one fused `c_attn` QKV projection split three ways on the last axis, multi-head attention (`n_kv_heads == n_head`), a tanh-approximate GELU MLP at four times the model width, and a tied output head (`wte.as_linear`; the checkpoint has no `lm_head` tensor). Raw HuggingFace exports store `c_attn` / `c_proj` / `c_fc` in the `Conv1D` `[in, out]` layout, which the loader transposes to `[out, in]`; the transpose applies to weights only, never to the 1-D bias vectors, and the decision is taken from the `c_attn` weight shape so an already-sanitized MLX conversion loads unchanged. `h.N.attn.bias` is HuggingFace's `[1, 1, n_ctx, n_ctx]` registered causal-mask buffer rather than a projection bias; it is dropped at load and never reaches a linear layer. Validated checkpoints: `openai-community/gpt2` (f32, 124M, bare keys), `distilbert/distilgpt2` (`transformer.`-prefixed keys) and `mlx-community/gpt2-base-mlx` (`model.`-prefixed, already transposed); all three generate the same text for the same prompt. No GPT-2 checkpoint ships a chat template, so the CLI renders the prompt verbatim and `--no-chat-template` is optional. The `wpe` table has a hard 1024 rows: the default `-n -1` budget already resolves to `n_positions - prompt_len`, but an explicit `-n N` or a prompt longer than the context runs past the table, and every token past it is embedded at the last row (the loader clamps and warns once during generation rather than indexing out of bounds). Loading separately rejects a checkpoint whose `wpe` weight has fewer rows than `config.json`'s `n_positions` claims, and rejects a `c_attn` / `c_proj` / `c_fc` shape that disagrees with the layout `Gpt2Layout` detected from layer 0, since the position clamp and the Conv1D transpose both trust those config-derived and layer-zero-derived bounds for every later lookup.)
- GPT-OSS

Many of these families have checkpoint-specific config or weight-layout
Expand Down
4 changes: 4 additions & 0 deletions src/distributed/tensor_parallel/inference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,10 @@ fn fallback_architecture(model_type: ModelType) -> &'static str {
ModelType::Olmo => "olmo",
ModelType::Olmo2 => "olmo2",
ModelType::Olmo3 => "olmo3",
// GPT-2's fused `c_attn` and Conv1D key names have no shard rules, so
// the planner's supported-architecture validation rejects this string
// before any TP load is attempted. It keeps the dispatch table total.
ModelType::Gpt2 => "gpt2",
ModelType::StarCoder2 => "starcoder2",
ModelType::Mellum => "mellum",
ModelType::MiniCPM | ModelType::MiniCPMOVLM => "minicpm",
Expand Down
2 changes: 2 additions & 0 deletions src/loaded_model.rs
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ pub enum LoadedModel {
Olmo(models::OlmoModel),
Olmo2(models::OLMo2Model),
Olmo3(models::OLMo3Model),
Gpt2(models::Gpt2Model),
StarCoder2(models::StarCoder2Model),
// Mixed full/sliding caches use a wrapper
Mellum(models::MellumWrapper),
Expand Down Expand Up @@ -307,6 +308,7 @@ macro_rules! delegate_language_model {
LoadedModel::Olmo(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::Olmo2(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::Olmo3(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::Gpt2(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::StarCoder2(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::Mellum(inner) => LanguageModel::$method(inner, $($arg),*),
LoadedModel::MiniCPM(inner) => LanguageModel::$method(inner, $($arg),*),
Expand Down
1 change: 1 addition & 0 deletions src/model_metadata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,7 @@ macro_rules! for_each_model_registration {
Olmo => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::OlmoModel::load, args: models::olmo::ModelArgs, weight_builder: models::OlmoModel::from_weights, wrap: LoadedModel::Olmo } };
Olmo2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::OLMo2Model::load, args: models::olmo2::ModelArgs, weight_builder: models::OLMo2Model::from_weights, wrap: LoadedModel::Olmo2 } };
Olmo3 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::OLMo3Model::load, args: models::olmo3::OLMo3Config, weight_builder: models::OLMo3Model::from_weights, wrap: LoadedModel::Olmo3 } };
Gpt2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::Gpt2Model::load, args: models::gpt2::ModelArgs, weight_builder: models::Gpt2Model::from_weights, wrap: LoadedModel::Gpt2 } };
StarCoder2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::StarCoder2Model::load, args: models::starcoder2::StarCoder2Config, weight_builder: models::StarCoder2Model::from_weights, wrap: LoadedModel::StarCoder2 } };
Mellum => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MellumModel::load, args: models::mellum::ModelArgs, weight_builder: models::MellumModel::from_weights, wrap: |m| LoadedModel::Mellum(models::MellumWrapper::new(m)) } };
MiniCPM => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiniCPMModel::load, args: models::minicpm::ModelArgs, weight_builder: models::MiniCPMModel::from_weights, wrap: LoadedModel::MiniCPM } };
Expand Down
1 change: 1 addition & 0 deletions src/models/detection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,7 @@ pub fn get_model_type(model_path: &Path) -> Result<ModelType> {
"olmo" => Ok(ModelType::Olmo),
"olmo2" => Ok(ModelType::Olmo2),
"olmo3" => Ok(ModelType::Olmo3),
"gpt2" => Ok(ModelType::Gpt2),
"starcoder2" => Ok(ModelType::StarCoder2),
"mellum" => Ok(ModelType::Mellum),
"minicpm" => Ok(ModelType::MiniCPM),
Expand Down
29 changes: 29 additions & 0 deletions src/models/detection_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,35 @@ fn whisper_model_type_is_detected() {
fs::remove_dir_all(model_dir).unwrap();
}

#[test]
fn gpt2_model_type_is_detected() {
// GPT-2 configs use the original OpenAI field names (`n_embd` / `n_head` /
// `n_layer`), so detection must key off `model_type` alone.
let model_dir = temp_path("gpt2_text");
fs::create_dir_all(&model_dir).unwrap();
fs::write(
model_dir.join("config.json"),
r#"{
"model_type": "gpt2",
"architectures": ["GPT2LMHeadModel"],
"n_embd": 768,
"n_head": 12,
"n_layer": 12,
"n_positions": 1024,
"n_ctx": 1024,
"layer_norm_epsilon": 1e-05,
"vocab_size": 50257,
"activation_function": "gelu_new"
}"#,
)
.unwrap();

let detected = super::detection::get_model_type(&model_dir).unwrap();
assert_eq!(detected, ModelType::Gpt2);

fs::remove_dir_all(model_dir).unwrap();
}

#[test]
fn fastvlm_model_type_is_detected() {
for model_type in ["fastvlm", "llava_qwen2"] {
Expand Down
Loading