fix: add @ModuleInfo to all Module-typed properties for weight update support

Module-typed properties declared as plain `let` without @ModuleInfo
cannot be updated via Module.update(modules:verify:), which is called
during quantization. This causes a fatal crash when loading quantized
models:

  MLXNN.UpdateError.needModuleInfo("Unable to get @ModuleInfo for
  BertModel.pooler -- must be wrapped to receive updates")

The quantize() function replaces Linear modules with QuantizedLinear
via update(modules:), which requires @ModuleInfo setters. Without
them, the non-throwing update(modules:) wrapper hits try! on the
thrown error → SIGABRT.

Fixed 36 properties across MLXEmbedders, MLXLLM, and MLXVLM:
- MLXEmbedders: BertModel.pooler, NomicBertModel.pooler,
  Qwen3 TransformerBlock.mlp, Qwen3ModelInner.norm
- MLXLLM: norm properties in 29 model classes, plus
  FalconH1Mixer.conv1d
- MLXVLM: Idefics3 TransformerBlock.mlp, Idefics3 LanguageModel.norm,
  Ministral3ModelInner.norm, LanguageModelInner.norm

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

Author: jowharshamshiri

Libraries/MLXEmbedders/Models/Bert.swift

@@ -259,7 +259,7 @@ public class BertModel: Module, EmbeddingModel {
     @ModuleInfo(key: "embeddings") fileprivate var embedder: BertEmbedding
 
     /// A linear layer used to "pool" the [CLS] token into a single sentence vector.
-    let pooler: Linear?
+    @ModuleInfo var pooler: Linear?
 
     /// The stack of Transformer layers.
     fileprivate let encoder: Encoder
@@ -280,10 +280,10 @@ public class BertModel: Module, EmbeddingModel {
 
         if lmHead {
             _lmHead.wrappedValue = LMHead(config)
-            self.pooler = nil
+            _pooler.wrappedValue = nil
         } else {
             // Pooler projects the [CLS] token to a hidden state of the same size
-            pooler = Linear(config.embedDim, config.embedDim)
+            _pooler.wrappedValue = Linear(config.embedDim, config.embedDim)
             _lmHead.wrappedValue = nil
         }
     }

Libraries/MLXEmbedders/Models/NomicBert.swift

@@ -671,7 +671,7 @@ public class NomicBertModel: Module, EmbeddingModel {
 
     /// The optional pooler layer.
     /// Used to extract a single vector representation for the whole sequence (usually from the [CLS] token).
-    let pooler: Linear?
+    @ModuleInfo var pooler: Linear?
 
     /// The stack of Transformer blocks.
     fileprivate let encoder: Encoder
@@ -696,9 +696,9 @@ public class NomicBertModel: Module, EmbeddingModel {
 
         // Initialize Pooler (for sentence embeddings)
         if pooler {
-            self.pooler = Linear(config.embedDim, config.embedDim, bias: false)
+            _pooler.wrappedValue = Linear(config.embedDim, config.embedDim, bias: false)
         } else {
-            self.pooler = nil
+            _pooler.wrappedValue = nil
         }
 
         // Initialize LM Head (for training/masked prediction)

Libraries/MLXEmbedders/Models/Qwen3.swift

@@ -169,7 +169,7 @@ private class TransformerBlock: Module {
     @ModuleInfo(key: "self_attn") var attention: Attention
 
     /// The feed-forward network (SwiGLU).
-    let mlp: MLP
+    @ModuleInfo var mlp: MLP
 
     /// Normalization applied before the attention layer.
     @ModuleInfo(key: "input_layernorm") var inputLayerNorm: RMSNorm
@@ -182,7 +182,7 @@ private class TransformerBlock: Module {
     public init(_ args: Qwen3Configuration) {
         // Initialize the two main processing sub-layers
         _attention.wrappedValue = Attention(args)
-        self.mlp = MLP(dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize)
+        _mlp.wrappedValue = MLP(dimensions: args.hiddenSize, hiddenDimensions: args.intermediateSize)
 
         // Initialize RMSNorm layers with the specified epsilon for numerical stability
         _inputLayerNorm.wrappedValue = RMSNorm(
@@ -236,7 +236,7 @@ private class Qwen3ModelInner: Module {
     fileprivate let layers: [TransformerBlock]
 
     /// The final normalization layer applied after all transformer blocks.
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     /// Initializes the model backbone.
     /// - Parameter args: Configuration containing `vocabularySize`, `hiddenLayers`, and `hiddenSize`.
@@ -255,7 +255,7 @@ private class Qwen3ModelInner: Module {
         }
 
         // 3. Initialize final RMSNorm
-        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     /// Forward pass through the model backbone.

Libraries/MLXLLM/Models/Apertus.swift

@@ -303,7 +303,7 @@ private class ApertusModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
     let layers: [ApertusBlock]
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     public init(_ args: ApertusConfiguration) {
         precondition(args.vocabSize > 0)
@@ -313,7 +313,7 @@ private class ApertusModelInner: Module {
             dimensions: args.hiddenSize
         )
         self.layers = (0 ..< args.numHiddenLayers).map { _ in ApertusBlock(args) }
-        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     public func callAsFunction(

Libraries/MLXLLM/Models/BaichuanM1.swift

@@ -203,7 +203,7 @@ public class BaichuanM1ModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
     fileprivate let layers: [BaichuanM1DecoderLayer]
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     init(_ config: BaichuanM1Configuration) {
         self.args = config
@@ -212,7 +212,7 @@ public class BaichuanM1ModelInner: Module {
         self.layers = (0 ..< config.hiddenLayers).map {
             BaichuanM1DecoderLayer(config, layerIdx: $0)
         }
-        norm = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: config.hiddenSize, eps: config.rmsNormEps)
     }
 
     func callAsFunction(

Libraries/MLXLLM/Models/BailingMoe.swift

@@ -313,7 +313,7 @@ class BailingMoeTransformerBlock: Module {
 public class BailingMoeModelInner: Module {
     @ModuleInfo(key: "word_embeddings") var embedTokens: Embedding
     let layers: [BailingMoeTransformerBlock]
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     init(_ args: BailingMoeConfiguration) {
         precondition(args.vocabularySize > 0)
@@ -322,7 +322,7 @@ public class BailingMoeModelInner: Module {
         self.layers = (0 ..< args.hiddenLayers).map {
             BailingMoeTransformerBlock(args, layerIdx: $0)
         }
-        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {

Libraries/MLXLLM/Models/Bitnet.swift

@@ -405,7 +405,7 @@ public class BitnetModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
     fileprivate let layers: [BitnetTransformerBlock]
-    var norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     init(_ args: BitnetConfiguration) {
         precondition(args.vocabularySize > 0)
@@ -417,7 +417,7 @@ public class BitnetModelInner: Module {
         layers = (0 ..< args.hiddenLayers).map { _ in
             BitnetTransformerBlock(args)
         }
-        norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {

Libraries/MLXLLM/Models/Cohere.swift

@@ -121,7 +121,7 @@ public class CohereModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
     fileprivate let layers: [CohereTransformerBlock]
-    let norm: LayerNorm
+    @ModuleInfo var norm: LayerNorm
 
     public init(_ args: CohereConfiguration) {
         precondition(args.vocabularySize > 0)
@@ -133,7 +133,7 @@ public class CohereModelInner: Module {
             .map { _ in
                 CohereTransformerBlock(args)
             }
-        self.norm = LayerNorm(dimensions: args.hiddenSize, eps: args.layerNormEps)
+        _norm.wrappedValue = LayerNorm(dimensions: args.hiddenSize, eps: args.layerNormEps)
     }
 
     public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {

Libraries/MLXLLM/Models/Ernie4_5.swift

@@ -173,7 +173,7 @@ class Ernie45DecoderLayer: Module {
 public class Ernie45ModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
     let layers: [Ernie45DecoderLayer]
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     public init(_ args: Ernie45Configuration) {
         self._embedTokens.wrappedValue = Embedding(
@@ -182,7 +182,7 @@ public class Ernie45ModelInner: Module {
         self.layers = (0 ..< args.numHiddenLayers).map { _ in
             Ernie45DecoderLayer(args)
         }
-        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {

Libraries/MLXLLM/Models/Exaone4.swift

@@ -141,7 +141,7 @@ public class Exaone4ModelInner: Module {
     @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
 
     fileprivate let layers: [Exaone4TransformerBlock]
-    let norm: RMSNorm
+    @ModuleInfo var norm: RMSNorm
 
     public init(_ args: Exaone4Configuration) {
         precondition(args.vocabularySize > 0)
@@ -162,7 +162,7 @@ public class Exaone4ModelInner: Module {
                 }
                 return Exaone4TransformerBlock(args, isLocal: isLocal)
             }
-        self.norm = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
+        _norm.wrappedValue = RMSNorm(dimensions: args.hiddenSize, eps: args.rmsNormEps)
     }
 
     public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]? = nil) -> MLXArray {