Add Phi 3.5 MoE

[DePasqualeOrg]

This is my attempt to port the Phi 3.5 MoE model from the Python implementation. Unfortunately I can't test it myself, since my MacBook doesn't have enough RAM. I marked two places in PhiMoE.swift with comments starting with !! which need to be checked. You can test this with ModelConfiguration.phi3_5MoE. Go ahead and make any necessary changes if you'd like, since I won't be able to run this myself.

[davidkoski]

I will give it a try on Tuesday!

[davidkoski]

It looks like there is a problem loading the weights:

Error: Mismatched parameter weight shape. Actual [16, 6400, 512], expected [16, 6400, 4096]

this is on SwitchLinear (maybe this error needs some more context).

SwitchLinear(bias=nil, inputDims=4096, numExperts=16, outputDims=6400)

And actually the parameters has a biases which is not expected here:

(lldb) po parameters.mapValues { $0.shape }
▿ [
  biases: [16, 6400, 64],
  scales: [16, 6400, 64],
  weight: [16, 6400, 512]
]

I need to look into this further

[davidkoski]

OK, I think SwitchLinear isn't quite right -- it is missing the bias (at the very least):

• https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/switch_layers.py#L87

and I think the dimension mismatch comes down to quantization -- the SwitchLinear isn't being replaced by SwitchLinearQuantized because it doesn't implement the protocol:

• https://github.com/ml-explore/mlx-swift/blob/main/Source/MLXNN/Quantized.swift#L14

and I think QuantizedSwitchLinear will need to be a subtype of SwitchLinear for the replacement to happen -- this @ModuleInfo(key: "gate_proj") var gateProj: SwitchLinear requires that the type be SwitchLinear or a subtype. Linear/QuantizedLinear are modeled the same way.

[DePasqualeOrg]

Thanks for that feedback. I've tried to make those changes, although I don't know how helpful this will be, since I unfortunately can't test it myself. If it's too much trouble and you'd prefer to focus on other priorities, don't worry about it.

[davidkoski]

OK, I will see if I can test/finish this -- it may be a few days before I get a chance.