OpenELM・SmolLM等小規模LLM（SLM）の各種チューニング手法について

　昨年末にAppleがOpenELMをオープンソースに近いラセンス（apple-sample-code-license）でHuggingFace（以下HF）にリリースして以降、今年に入ってから立て続けにMetaがMobileLLMを、HF自身がSmolLMを、オープンソースの小規模のLLM（SLM）としてリリースしました。各小規模モデルともスマートフォン上でのローカルな実行を念頭に置き、パラメーター数は1B（10億）未満のモデルもラインナップされています。特に日本語対応を前提にした参考記事が少なかったので、多少の不正確さを勘弁して頂く前提で、これらのSLMモデルのマージやファインチューニングの手法を試した範囲で紹介したいと思います。
　既存の日本語対応LLMで1B（10億）レベルのモデルである、llm-jp,karasuなどと比較して下記の表にまとめてみました。

| Model Name | Organization|     model size    | Architecture　|License|
|:-----------:|:------------:|:-----------------:|:-----------------:|:------------:|
| [OpenELM](https://huggingface.co/collections/apple/openelm-instruct-models-6619ad295d7ae9f868b759ca)    |    Apple    | 270M/450M/1.1B/3B | OpenELMForCausalLM |apple-sample-code-license|
| [MobileLLM](https://github.com/facebookresearch/MobileLLM/tree/main)  |    Meta     | 125M/350M/600M/1B/1.5B | LlamaForCausalLM |CC-BY-NC 4.0|
| [SmolLM](https://huggingface.co/collections/HuggingFaceTB/smollm-6695016cad7167254ce15966)  |HuggingFace  | 135M/360M/1.7B | LlamaForCausalLM |Apache 2.0|
|  [llm-jp](https://huggingface.co/llm-jp)   |   llm-jp    | 1.3B/13B         | GPT2LMHeadModel |Apache 2.0|
| [karasu](https://huggingface.co/collections/lightblue/karasu-658e78e9188f3934c6fb701e) |  lightblue  | 1.1B/7B/14B(qarasu) | LlamaForCausalLM |Apache 2.0|

　上記モデルのうちMetaのMobileLLMのみHFではなくGitHubに開発環境が公開されており、事前学習を自前で行う前提なので、pre-trainedのモデルは提供されていません。事前学習にはモデルのサイズが小さくてもそれなりのコンピュータリソースが必要になるので、今後の先達の方々の日本語学習モデル公開に期待して、今回はpre-trainedのモデルが提供されているOpenELMとSmolLMを中心に紹介させてもらいます。

１. OpenELMのファインチューニング
　Appleが公開しているLLM関連の開発環境の主なものはmlxとcorenetの2つで、元々OpenELMはcorenet側でサポートされていました。ただmlx側もサポートを始めたので今回は手順の簡単なmlxを使ったファインチューニング手法を紹介します。但し結果をHF等で公開する場合は、ライセンスの"apple-sample-code-license"を確認してからの方が良いと思います。また、以下はMac前提のmlxの内容になりますが、OpenELMはtransformersのpeftやtrlでも扱えますので、Mac以外の方は後半のpeft/trl部分を参照して下さい。
[▶corenetのOpenELMサイト](https://github.com/apple/corenet/tree/main/mlx_examples/open_elm)
[▶ml-explore/mlx-examplesのOpenelm.py](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/models/openelm.py)
　mlxの基本的な使い方は別の方の記事が参考になりますので下記のリンクを参照して下さい。
 [① A simple guide to local LLM fine-tuning on a Mac with MLX](https://apeatling.com/articles/simple-guide-to-local-llm-fine-tuning-on-a-mac-with-mlx/)
 [② MLX で Mac で LLM fine-tuning](https://zenn.dev/seya/articles/c06d4d5459f938)
 [③ MLXを使ってMacのローカル環境で”ござる”loraを作成してみる](https://zenn.dev/michy/articles/041d0db2ae9895)
　mlx側は他のモデルと同様の扱いでLORA(QLORA)によるファインチューニング(FT)の対象モデルの一つに含まれているだけなので、LORA用の設定ファイルを下記のように用意すればFT出来るようになります。
  ```openelm_config.yaml
# The path to the local model directory or Hugging Face repo.
model: "mlx-community/OpenELM-450M-8bit"
#model: "mlx-community/OpenELM-1_1B-Instruct-8bit"
#model: "mlx_model" #自分でconvertした量子化モデル
# Whether or not to train (boolean)
train: true

# Directory with {train, valid, test}.jsonl files
data: "dataset"

# The PRNG seed
seed: 0

# Number of layers to fine-tune
lora_layers: 20 #1B:28,450M:20, 270M:16, 3B:36

# Minibatch size.
batch_size: 1 #1, 4, 8

# Iterations to train for.
iters: 400 # 1000

# Number of validation batches, -1 uses the entire validation set.
val_batches: -1 #default 25

# Adam learning rate.
learning_rate: 2e-5

# Number of training steps between loss reporting.
steps_per_report: 100

# Number of training steps between validations.
steps_per_eval: 100

# Load path to resume training with the given adapter weights.
resume_adapter_file: null

# Save/load path for the trained adapter weights.
adapter_path: "adapters"

# Save the model every N iterations.
save_every: 200

# Evaluate on the test set after training
test: false

# Number of test set batches, -1 uses the entire test set.
test_batches: -1 #defaulr 100

# Maximum sequence length.
max_seq_length: 4096 #450M:1536, 270M:1280

# Use gradient checkpointing to reduce memory use.
grad_checkpoint: false

# Use DoRA instead of LoRA.
use_dora: false #true false

# LoRA parameters can only be specified in a config file
lora_parameters:
  # The layer keys to apply LoRA to.
  # These will be applied for the last lora_layers
  keys: ["attn.qkv_proj", "attn.out_proj", "ffn.proj_1", "ffn.proj_2"]
  rank: 32
  alpha: 32.0
  scale: 10.0
  dropout: 0.05

# Schedule can only be specified in a config file, uncomment to use.
lr_schedule:
  name: cosine_decay
  warmup: 100 # 0 for no warmup
  warmup_init: 1e-5 # 0 if not specified
  arguments: [1e-5, 400, 1e-7] # passed to scheduler
```
　モデルの指定は、サイズに応じてmlx用に量子化されたモデルがHF上の"[mlx-community](https://huggingface.co/mlx-community)"で公開されているので、これを利用するのが一番簡単ですが、自分でmlxのconvertで変換(量子化)したモデルをローカルで使うことも出来ます。あとはデータセットをmlx用に変換したものを用意して読み込みますが、変換の仕方はここを参考して下さい。
 [mlx-examples: Fine-Tuning with LoRA or QLoRA](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#Data)
 [mlx-examples: sample dataset](https://github.com/ml-explore/mlx-examples/tree/main/lora/data)
基本は下記のようなJSONL形式に変換したものを、test.jsonl,train.jsonl,valid.jsonlの３種類のファイルでディレクトリ配下に用意して、ディレクトリを指定して読み込む感じです。
```dataset/train.jsonlの例
{"text": "USER:SFファンタジーとは？ ASSISTANT:SFファンタジーは、SFとファンタジーの両方のトロフィーや要素を組み合わせたハイブリッドなジャンルである。"}
...
```
　またFTの対象とするtarget_moduleは、mlxではlora_parametersの中のkeysで設定しますが、上記の通りで大丈夫だと思います。その他のパラメータは、自分のコンピュータリソースと相談しながら所要時間を推定して修正して下さい。なお、"Maximum sequence length"は、データセットの文章の長さによって警告が出る場合もあるので、今回は大きめの設定としてます。
　FTの開始は下記コマンドになります。
```
mlx_lm.lora -c openelm_config.yaml　
```
　FTが終了すると、adapterフォルダにadapterファイルが生成されます。mlxで直接出力を確認したい場合は、generateを下記のように使えば出来ます。
```
 mlx_lm.generate --model mlx-community/OpenELM-450M-8bit --adapter-path adapters --prompt "日本の首都は？"
 ```
　次にこの生成されたadapterファイルを使って、fuseコマンドで通常のmodel.safetensorsファイルを生成すれば完成です。
```
mlx_lm.fuse --model mlx-community/OpenELM-450M-8bit　--de-quantize
```
　adapterの指定は省略できるようなので、元のモデルを指定して逆量子化(de-quantize、元のモデルが量子化モデルで無ければ不要)を指定すれば、lora_fused_modelフォルダにmodel.safetensors等のモデルファイルが保存されます。後はHFにアップロードしたり、llama.cpp（Branch3324以降のバージョンが必要）を使ってGGUF化・量子化してOllama等で利用して下さい。
 
　以上はMac前提のmlxによるOpenELMのファインチューニング方法の紹介でしたが、前段でも触れた通り、OpenELM自体はpeftやtrlでもFT可能です。ポイントはpeftのLoraConfigのtarget_modulesを上記のmlxと同じものを設定することです。データセットはmlxのtrain.jsolをそのまま流用して使えるように下記の変換処理入れるか、peft用に準備したものを使って下さい。
```peft.の場合にmlx用のデータセットを流用するための変換処理部分
import datasets
# データセットのJSONLを開く
import json
outputs = []
with open("train.jsonl", "r", encoding="utf-8") as f:
    lines = f.readlines()
    for line in lines:
        pair = json.loads(line)
        outputs.append(pair["text"])
# outputsをdatasets.DatasetDictに変換
data = datasets.Dataset.from_dict({"text": outputs})
data = datasets.DatasetDict({"train": data})
data = data.map(lambda samples: tokenizer(samples["text"]), batched=True)
# dataを９0%の訓練データと10%の検証データに分割
data = data['train']
data = data.train_test_split(test_size=0.1)
train_dataset = data["train"]
eval_dataset = data["test"]
```
 残念ながら現時点でMergeKitはOpenELMには対応していないようです。"OpenELMForCausalLM"のアーキテクチャモデルは他にあまり無いようなのでニーズがないのかもしれません。
 
２. SmolLMのファインチューニング
　OpenELMと異なる部分を中心に記載します。まずはmlxの場合のyamlファイルですが、変更箇所は下記の通り3箇所です。
```smollm_config.yaml
①model: "mlx-community/SmolLM-135M-Instruct-8bit"
②lora_layers: 30 #135M:30,360M:32, 1.7B:24
lora_parameters:
③keys: ['mlp.gate_proj', 'mlp.down_proj', 'self_attn.q_proj', 'mlp.up_proj', 'self_attn.o_proj','self_attn.v_proj', 'self_attn.k_proj']
```
　これで基本的には動きます。その他のパラメータは自分の環境に合わせて調整して下さい。
 ```
mlx_lm.lora -c smollm_config.yaml　
```
　generateによる出力確認やfuseによるモデル生成も同様です。またpeftについても、LoraConfigのtarget_modulesを上記の③keysと同じ設定すればOKです。
　SmolLMはModel Architectureが"LlamaForCausalLM"なので、マージに関しては敷居が低いのですが、レイヤ数が135M:30,360M:32,1.7B:24とパラメータ規模での標準よりかなり大きく、vocab_sizeも49152と大きくなっているため、簡単には他の同じArchitectureモデルとマージは出来ません。お試しするなら、まずはMergeKitを使ってSlomLMのサイズの同じモデルのbase modelとinstruct model同士で、下記のsample_yaml_configを使ってマージしてみて下さい。この際更にもう一工夫必要で、"mergekit/mergekit/_data/architectures/llama.json"の最後のllm_headの記述を削除して下さい。これで動くはずです。
```sample_yaml_config.
slices:
- sources:
  - layer_range: [0, 24]
    model: HuggingFaceTB/SmolLM-1.7B-Instruct
    parameters:
      density: [1, 0.7, 0.1]
      weight: 1.0
  - layer_range: [0, 24]
    model: HuggingFaceTB/SmolLM-1.7B
    parameters:
      density: 0.33
      weight:
        - filter: mlp
          value: 0.5
        - value: 0
merge_method: dare_ties
base_model: HuggingFaceTB/SmolLM-1.7B-Instruct
parameters:
  normalize: true
  int8_mask: true
dtype: bfloat16
```
```llama.json
...

"post_weights": [
        {
            "name": "model.norm.weight",
            "input_space": "running_residual"
        },　# ←最後のカンマも削除
        #以下のlm_head部分を削除
        {
            "name": "lm_head.weight",
            "input_space": "running_residual",
            "is_embed":true,
            "aliases": [
                "model.lm_head.weight"
            ]
        }
        #削除はここまで
    ]
}
```
　削除しないと"RuntimeError: Tensor lm_head.weight required but not present"が出ます。
 最後にpeft, trlについてですが、特に注意点は無くそのままで動くと思います。datasetをmlxと共通化する場合は、OpenELMで記載した上記変換処理がそのまま使えますので、target_modulesの設定をmlxのkeysと同じ設定にしてトライしてみて下さい。

【参考】　llm-jpのファインチューニング
　llm-jpさんが下記のリンクでllm-jpモデルのSFT（Supervised Fine-Tuning）用のコードを公開してくれていますので、peft, trlでFTする場合はこれを使わしてもらうのがベストですね。
[LLM-jp SFTのGitHubサイト](https://github.com/llm-jp/llm-jp-sft)
　MacユーザがmlxでFTしたい場合、llm-jpのモデルがpytorch_model.bin形式なのでそのままは動いてくれません。方法は2つで、①safetensors_utilというツールでpytorch_model.binをsafetensorsに変換する、②MergeKitで一旦マージモデルを作成し、HFにアップロードしたものを使う（MergeKitが自動的にsafetensorsに変換してくれます）。②は周りくどいという方は、下記のsafetensors_utilで変換されることをお勧めします。
 [safetensors_utilのGitHubサイト](https://github.com/by321/safetensors_util)
 　マージに関しては、下記の記事がとても参考になります（特に注釈部分が肝です）。こちらの方法で私も何とか成功しました。MergeKitでそのままマージするとエラーが出て失敗するので、SmolLMと似た対応ですが、"mergekit/mergekit/_data/architectures/gpt2.json"の中で定義されている全てのパラメータ名の前に"transformer"を付けることで解決できます。瓦さんに感謝です。
[mergekit でモデルマージを試してみる](https://zenn.dev/fusic/articles/try-llmjp-merge)
　出来上がったマージ・FTのモデルはllama.cppでそのままGGUF化出来ないので、今度は[mmngaさんのサイト](https://huggingface.co/mmnga/llm-jp-1.3b-v1.0-gguf)から変換スクリプトを拝借して、これを使ってGGUF化すれば後は普通に量子化できます。

３. まとめ
　今回はSLM（小規模LLM）のマージやFTの手法について、実際に試行錯誤した結果うまく行ったケースをまとめて紹介させてもらいました。日本でもっとSLMに関する話題が盛り上がる一助になれば幸いです。また、先達方々には是非OpenELMやSmolLM、そしてMobileLLMの日本語対応を一考して頂けると嬉しいです。ご一読頂きありがとうございました。コメントや修正点等のご指摘お待ちしておりますので、お気軽に投稿して下さい。