Gemma3-4b issue: Mismatched parameter scales shape. Actual [1024, 40], expected [256, 40]

[Remember2015]

Mismatched parameter scales shape. Actual [1024, 40], expected [256, 40]

The model is from mlx-community in hf: gemma-3-4b-it-4bit

[davidkoski]

I suspect it is a problem with quantization -- a MLXArray wasn't quantized when being loaded and it should.

[Remember2015]

I suspect it is a problem with quantization -- a MLXArray wasn't quantized when being loaded and it should.

ok, thanks for your reply, and is it the best way to quantize model by using mlx_lm.convert?

[davidkoski]

ok, thanks for your reply, and is it the best way to quantize model by using mlx_lm.convert?

I think I said it backwards -- the weights are quantized (you can see it in config.json). If the weights are quantized, then on load it will traverse the layers and quantize the initial weights in memory. Finally it will update the parameters with the weights loaded from the file:

public func loadWeights(
...
    // per-model cleanup
    weights = model.sanitize(weights: weights)

    // quantize if needed
    if quantization != nil || perLayerQuantization != nil {
        quantize(model: model) { path, module in
            if weights["\(path).scales"] != nil {
                if let perLayerQuantization {
                    return perLayerQuantization.quantization(layer: path)?.asTuple
                } else {
                    return quantization?.asTuple
                }
            } else {
                return nil
            }
        }
    }

    // apply the loaded weights
    let parameters = ModuleParameters.unflattened(weights)
    try model.update(parameters: parameters, verify: [.all])

The expected size is 1/4 the dimension of the actual, so that means scales was quantized in memory but it wasn't in the saved weights -- we will need to figure out why.

[davidkoski]

I wonder if this could possibly be #372

[Remember2015]

I wonder if this could possibly be #372

Let me try, thank you

[Remember2015]

I think the code works, now the model is running, but it did not stop generation when token <end_of_turn> generated

Who are you?<end_of_turn>
<end_of_turn>
<end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of_turn><end_of

I wonder if this could possibly be #372

Let me try, thank you

[davidkoski]

OK, if <end_of_turn> is supposed to stop but the tokenizer config doesn't configure it as such (e.g. via Tokenizer.eosTokenId) then you may have to use something like this:

    static public let phi3_5_4bit = ModelConfiguration(
        id: "mlx-community/Phi-3.5-mini-instruct-4bit",
        defaultPrompt: "What is the gravity on Mars and the moon?",
        extraEOSTokens: ["<|end|>"]
    )

This is done in the ModelConfiguration by id (the path on huggingface). Potentially llm-tool could let you override this but it does not do so today.

[Remember2015]

OK, if <end_of_turn> is supposed to stop but the tokenizer config doesn't configure it as such (e.g. via Tokenizer.eosTokenId) then you may have to use something like this:

static public let phi3_5_4bit = ModelConfiguration(
    id: "mlx-community/Phi-3.5-mini-instruct-4bit",
    defaultPrompt: "What is the gravity on Mars and the moon?",
    extraEOSTokens: ["<|end|>"]
)

This is done in the ModelConfiguration by id (the path on huggingface). Potentially llm-tool could let you override this but it does not do so today.

Bow to your greatness. Yeah after adding extra eos token, now gemma-3 is working well