Support for loading LoRA layers from adapter config in Swift

[markrenaud]

I’m fine-tuning models with MLX QLoRA using the Python MLX APIs, and then integrating the resulting models into my apps via the Swift framework.

For smaller models (e.g. Gemma3 270M), I often need to train multiple projections to achieve suitable task performance.

In Python, this is straightforward to configure in the QLoRA YAML file, e.g.:

...

# LoRA parameters - rank 32 with all projections
lora_parameters:
  # The layer keys to apply LoRA to (all projections)
  keys: ["self_attn.q_proj", "self_attn.v_proj", "self_attn.k_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj"]
  
  # LoRA rank
  rank: 32
  
  # LoRA scale
  scale: 32.0
  
  # Dropout rate for LoRA layers
  dropout: 0.0

...

Loading the adapter for inference in Python works as expected with this configuration.

However, in Swift, the default model files (e.g. Gemma3Text) currently only have q_proj and v_proj hard-coded.

This means I need to create a local copy of the model definition and extend it manually, e.g.:

extension Gemma3TextModel: LoRAModel {
// ORIGINAL
//    public func loraLinearLayers() -> LoRALinearLayers {
//        model.layers.map { ($0.selfAttention, ["q_proj", "v_proj"]) }
//    }
    
    
    // CUSTOMIZED FOR EXTRA LORA LAYERS
    public func loraLinearLayers() -> LoRALinearLayers {
        let attnLayers = model.layers.map { ($0.selfAttention, ["q_proj", "v_proj", "k_proj", "o_proj"]) }
        let mlpLayers = model.layers.map { ($0.mlp, ["gate_proj", "up_proj", "down_proj"]) }
                                            
        return attnLayers + mlpLayers
    }

}

Since the generated adapter_config.json already specifies which projections the adapter targets, would it be possible for the Swift library to read this information at adapter load time and automatically apply the correct LoRA layers? (I'm not sure if the adapter_config.json has changed over time - if it has, the LoRALinearLayers could default back to q_proj and v_proj if the appropriate keys were not found in the json).

This would be especially useful when working with a shared base model while swapping in/out multiple adapters that target different sets of layers.

I am not sure if there is a technical or architectural reason that this hasn’t been implemented in the Swift package .

[davidkoski]

Yeah, this looks like a good thing to add. When LoRA was implemented on the swift side the python side hard coded the layer names, so it is somewhat equivalent.

[davidkoski]

I think this is resolved with #424, let me know if it is not and we can reopen. The LoRAConfiguration can indicate which layers to manipulate:

// Load LoRA adapter from directory or create a new one
let modelAdapter: ModelAdapter
do {
    modelAdapter = try LoRAContainer.from(directory: adapter)
} catch {
    modelAdapter = try await modelContainer.perform { context in
        return try LoRAContainer.from(
            model: context.model, configuration: LoRAConfiguration(numLayers: loraLayers))
    }
}