← All courses ← Course home

Lesson 03 of 05

Register a custom architecture

The loader keys on config.json model_type, not the Hugging Face id. Register a missing family such as Gemma 4 before you load weights.

Agent brief (llms.md)

You finished environment setup. Now you can write the first load. Open the model's config.json. Find model_type. That string is the key. The Hugging Face repo id is not.

flowchart LR
  C[config.json] --> T[model_type]
  T --> R[Type registry]
      

mlx-swift-lm keeps a type registry. On current releases that registry lives on LLMTypeRegistry.shared, which is a ModelTypeRegistry. The factory (LLMModelFactory.shared) reads model_type, asks the registry to build the Swift class, then loads safetensors. If the type is missing, load throws an unsupported-model-type error. A stock call with only a repo id does not invent the architecture for you.

Gemma 4 is the usual example. Current mlx-swift-lm registers strings such as gemma4, gemma4_text, and gemma4_unified. An older pin does not. A Gemma 4 assistant or MTP drafter may use a different model_type that is still missing. Do not assume the type you need is in your pin. Check.

flowchart LR
  T[model_type] --> R[Type registry]
  R -->|missing| E[unsupported model type]
  R -->|present| L[Load weights]
      

Verify against the docs you pinned

Public names have moved. Older docs show a class. Current mlx-swift-lm treats ModelTypeRegistry as an actor with registerModelType and contains. The helper that decodes a configuration and calls YourModel.init is often private inside the package. In your app, write a small creator that decodes Data and returns the model. If a call will not compile, open the DocC for MLXLMCommon and MLXLLM on the tag you resolved. Do not invent a factory method.

Useful docs to keep open:

Check, register, then load

The sketch below follows the current public pattern. Confirm await if the registry is an actor in your pin. Confirm the configuration and model type names in MLXLLM. If Gemma 4 already exists in your pin, contains is true and you skip register.

import MLXLLM
import MLXLMCommon

func ensureArchitectureRegistered(modelType: String) async {
    let registry = LLMTypeRegistry.shared
    if await registry.contains(modelType) {
        return
    }

    // Sketch: decode the config type that matches this model_type,
    // then return the matching Module + LLMModel class from MLXLLM
    // or from your port. Confirm types in the package you pinned.
    await registry.registerModelType(modelType) { data in
        let configuration = try JSONDecoder().decode(
            Gemma4Configuration.self,
            from: data
        )
        return Gemma4Model(configuration)
    }
}

func loadRegisteredModel(id: String) async throws -> ModelContainer {
    let configuration = ModelConfiguration(id: id)
    return try await LLMModelFactory.shared.loadContainer(
        configuration: configuration
    )
}

If the family is not in MLXLLM at all, you port it. The public shape is a configuration struct that matches config.json, plus a top-level class that conforms to Module and LLMModel (and usually KVCacheDimensionProvider). Follow the porting guide. Registration is still the same: map the model_type string to that class. You can also add a ModelConfiguration id to the model registry so callers can look the repo up by name. That id step is convenience. The type step is required.

flowchart LR
  T[model_type] --> C[contains?]
  C -->|no| REG[registerModelType]
  C -->|yes| L[loadContainer]
  REG --> L
      

Load weights only after the type is registered. Do not start generate in this lesson. Route work to the right Apple model is the other pick: Foundation Models versus a custom Core AI file. Here the pick already happened. You are making the MLX family you chose actually construct.

Rule. Read model_type. Check the registry. Register if missing. Then load. Confirm every name in the docs for your pin.

Next, gate generate on scenePhase so Metal does not run in the background.

Key concepts

  • The factory keys on config.json model_type, not the Hugging Face id.
  • A missing type throws unsupported model type. Gemma 4 and newer strings need a registry entry.
  • contains then registerModelType then loadContainer is the production order.
  • API names vary by release. Verify against the mlx-swift-lm docs you pinned.

Takeaways

  • Do not debug Metal until the architecture constructs.
  • Port a class only when MLXLLM does not already ship that type.
  • Keep the register call next to app startup, before the first load.