Add plamo-2-1b model

[mitmul]

This PR adds the latest SLM from Preferred Networks, PLaMo-2-1B.

Currently it can successfully generate correct responses when I use the following code:

import mlx.core as mx
from mlx_lm.utils import load
from mlx_lm.models import cache

model, tokenizer = load(
    "../plamo-2-1b-mlx",
    tokenizer_config={
        "trust_remote_code": True,
    },
    model_config={
        "trust_remote_code": True,
    },
)

user_input = "美味しいカレーの作り方のレシピを紹介します。"

input_ids = tokenizer.encode(user_input)
input_ids = mx.array(input_ids, dtype=mx.int64)

y = input_ids
prompt_cache = cache.make_prompt_cache(model)

for _ in range(100):
    logits = model(inputs=y[None], cache=prompt_cache)
    logits = logits[:, -1, :]
    logprobs = logits - mx.logsumexp(logits, keepdims=True)
    next_y = mx.argmax(logprobs, axis=-1)

    response = tokenizer.decode(next_y.item())
    print(response, end="")
    y = mx.concat([y, mx.array(next_y)])

However, when I run the following command from my terminal:

python -m mlx_lm.generate \
--model plamo-2-1b-mlx \
--prompt '美味しいカレーの作り方のレシピを紹介します。' \
--verbose true \
--max-tokens 128 \
--ignore-chat-template

the outputs are totally corrupted.
Now I'm investigating the cause of why mlx_lm.generate doesn't work for this model.
I guess it relates to the cache implementation in this code.

[mitmul]

I found a bug in my code that performs causal-conv1d update, which was related to the difference between the channel-first conv1d in torch and channel-last conv1d in MLX. I fixed it and confirmed that it runs properly as expected, so I changed this PR's state to review-ready 🙏

[awni]

I tried running:

mlx_lm.generate --model pfnet/plamo-2-1b --prompt "The talles mountain in"

But it crashes with missing layers. Did you encounter that?

ValueError: Received parameters not in model: model.layers.layers.13.post_mixer_norm.weight model.layers.layers.10.mixer.A_log model.layers.layers.6.mixer.dt_bias model.layers.layers.14.mlp.gate_up_proj.weight model.layers.layers.4.mixer.dt_norm_weight model.layers.layers.3.post_mlp_norm.weight model.layers.layers.8.mixer.A_log model.layers.layers.15.mixer.qkv_proj.weight model.layers.layers.3.pre_mlp_norm.weight ...

[mitmul]

Thank you for checking. That's because the config.json in pfnet/plamo-2-1b has "model_type": "plamo" and the old architecture for PLaMo-13B has existed in mlx_lm as plamo, so the older one is selected in your case instead of the right one, plamo2 (Samba architecture).

Now I'm discussing with the members of the model development in Preferred Networks if we can change the value of model_type in config.json on the huggingface repo: pfnet/plamo-2-1b.