Models fail to load with keyNotFound error

[atdrendel]

Related to #214, I believe.

Some models fail with keyNotFound errors. The Qwen1.5 and 2.5 errors were solved in #210. Some errors still remain, though.

• Phi-3.5-MoE-instruct-4bit: keyNotFound(base: "SuScaledRotaryEmbedding", key: "_freqs")
• OpenELM-270M-Instruct: keyNotFound(base: "Linear", key: "weight")

[davidkoski]

I will take a look -- probably a variation on the same issue.

[davidkoski]

Phi-3.5-MoE-instruct-4bit is interesting -- it has a rope layer which is SuScaledRotaryEmbedding, but it is not meant to be parameters. Indeed it is not marked with @ParameterInfo, it is just a constant value.

po self.rope.parameters()
▿ [
  
]
  - contents : 0 elements

and that is because:

    public func parameters() -> ModuleParameters {
        filterMap(filter: Self.filterValidParameters, map: Self.mapParameters())
    }

filters out parameters with leading _ names. I think we should probably not consider those for validation.

That will be fixed with ml-explore/mlx-swift#200

[davidkoski]

OpenELM-270M-Instruct has a couple of cases where layers should be optional but are not:

        var out = transformer(inputs, cache: cache)
        if shareInputOutputLayers {
            out = matmul(out, transformer.embedTokens.weight.T)
        } else {
            out = lmHead(out)
        }

lmHead is always present but ignored depending on this config value. Should be:

        var out = transformer(inputs, cache: cache)
        if let lmHead {
            out = lmHead(out)
        } else {
            out = matmul(out, transformer.embedTokens.weight.T)
        }

there are some other layers in the model set up the same way.

[davidkoski]

I inspected the rest of the models for similar patterns but didn't see any.