![見出し画像](https://assets.st-note.com/production/uploads/images/126705011/rectangle_large_type_2_fb7acba404b795639650365a9df3f656.png?width=1200)
lit-gptでLlama2を学習してみたい
注)結局学習が上手く行ったわけではないです。悪しからず。
(参考記事)
https://zenn.dev/if001/articles/6c507e15cd958b
https://zenn.dev/syoyo/articles/52f1d0d62fcad5 ←似た取り組み
https://github.com/OpenLLM-France/Lit-Claire/blob/main/pretrain.py#L245
(環境)
NVIDIA A100 (80GB) 4台
CUDA Version 12.2
torch version = 2.3.0.dev20231221+cu121
(ツール) lit-gpt GitHubはこちら
挑戦①
1. データを準備する
lit-gptのscripts/prepare-*.pyを参考に必要なデータを用意するprepare_hoge.pyを作成する。
python scripts/prepare_hoge.py
--tokenizer_path <トークナイザのディレクトリ>
--destination_path <データの保存先のディレクトリ>
2.pretrainコードを作成する
pretrain/redpajama.pyを参考にhoge.pyを作成する。
- model_name = "Llama-2-7b-chat-hf"に変更する。
ハイパーパラメータ
- micro_batch_size = 1にする。
- その他はデフォルトのまま。
3.pretrainを実行する
python pretrain/hoge.py \
--devices 4 \
--train_data_dir <データの保存先のディレクトリ> \
--precision 16-mixed
メモリ対策(公式ページ)
自分の場合・・・
モデルサイズ7B, micro_batch_size = 1, precision=16-mixedにすることで
47GB VRAM消費 x 4台.
学習のスピード
Total parameters 6,738,415,616
Estimated TFLOPs: 838.34
Measured TFLOPs: 755.06
1iterにつき2秒。125iterで1step。(batch_size=125)
結果
lossが10前後から下がらない。。。
でも7Bがこれくらいの計算リソースでちゃんと学習できるのは, Mistral8x7BなどのMoEが今後台頭してきそうなことを鑑みると嬉しい気がする。
推論する場合
lit-gptのchat/base.pyを利用しLlama2-7Bで推論
A100 1枚, 15GB RAMくらい消費。
めっちゃ高速
(例)Time for inference: 6.26 sec total, 42.02 tokens/sec, 263 tokens
日本語は出力できない
NGワードも制限されている
気を取り直してTinyLlamaでの再挑戦を試みる
挑戦②(リベンジ)
1.データセットの準備 .
"izumi-lab/wikinews-ja-20230728",
# "izumi-lab/wikinews-en-20230728",
"izumi-lab/wikipedia-ja-20230720",
# "izumi-lab/wikipedia-en-20230720",
# 'izumi-lab/open-text-books',
'if001/oscar_2023_filtered',
'if001/aozorabunko-clean-sin'
和泉研の日本語コーパスを拝借。めちゃくちゃbinファイルがdata/以下に生成される。三番目の日本語wikipediaで730個のファイルに。
トークナイザは、ELYZAを利用。
トークン数は、tokens 2.93M , tokens 1.53B, tokens 7.54B, tokens 210.88M
total tokens 9.29B
(ちなみに元論文は700B+250B=950B x 3epochs なので, 100分の1の規模です…)
vaildデータは、aozorabunkoとwikinewsに設定。
2.モデルのロード
python scripts/download.py --repo_id TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T
checkpoints/TinyLlama/ にbinファイルを保存。3Tトークン学習した1.1Bモデルだが、Chat用にファインチューニングしたモデルとどちらがベースとして良いのかはよくわからない。
3.モデルの変換
python scripts/convert_hf_checkpoint.py --checkpoint_dir checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T/
すぐ終わった。
4.学習
continual pretrainingがいいのか, fine-tuningがいいのか… 前者は公式でサポートされてないので自信がない。
今回はTinyLlamaがSlimpajamaとStarcoderで事前学習されているので、できれば日本語Wikipediaで追加で事前学習したい。
resumeオプションについて
--resumeオプションで2.のモデルを指定すればcontinual pretrainingをやっていることになるのかな🤔
usage: tinyllama.py [-h] [--config CONFIG] [--print_config[=flags]]
[--resume RESUME]
使うデータセットの書き換え
論文にもtrain:validの分割については言及がない. テキトーでいいのか. 実装を見るとSlimpajamaのほうだけvalを切り出して使っているっぽい. 結構val多いな.
Streaming Datasetについて
TinyLlamaの学習コード pretrain/tinyllama.py においては、データローダの作成にlightning.data.StreamingDataSetが用いられている.
これがあんまりうまくいかなかった…index.jsonがフォルダに無いよというエラーが発生。原因と対処がわからなかったので、素直にpretrain/redpajama.pyを参考にして使わないように修正。
コードは結局こうなった。
def create_dataloaders(batch_size: int, block_size: int) -> Tuple[DataLoader, DataLoader]:
from lightning.data import StreamingDataset
from lightning.data.streaming.item_loader import TokensLoader
# Increase by one because we need the next word as well
effective_block_size = block_size + 1
# train_datasets = [
# StreamingDataset(
# input_dir="data/slimpajama/train",
# item_loader=TokensLoader(block_size=effective_block_size),
# shuffle=True,
# drop_last=True,
# ),
# StreamingDataset(
# input_dir="data/starcoder",
# item_loader=TokensLoader(block_size=effective_block_size),
# shuffle=True,
# drop_last=True,
# ),
# ]
train_datasets = [
StreamingDataset(
input_dir="data/ja_data/train",
item_loader=TokensLoader(block_size=effective_block_size),
shuffle=True,
drop_last=True,
),
]
# Mix SlimPajama data and Starcoder data with these proportions:
#weights = (0.693584, 0.306416)
weights = [1.0]
combined_dataset = CombinedDataset(datasets=train_datasets, seed=42, weights=weights)
train_dataloader = DataLoader(
combined_dataset, batch_size=batch_size, pin_memory=True, num_workers=8, drop_last=True
)
# val_dataset = StreamingDataset(
# input_dir="data/slimpajama/val",
# item_loader=TokensLoader(block_size=effective_block_size),
# shuffle=True,
# # Consider setting to False, but we would lose some samples due to truncation when world size > 1
# drop_last=True,
# )
val_dataset = StreamingDataset(
input_dir="data/ja_data/val",
item_loader=TokensLoader(block_size=effective_block_size),
shuffle=True,
# Consider setting to False, but we would lose some samples due to truncation when world size > 1
drop_last=True,
)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, pin_memory=True, num_workers=8, drop_last=True)
return train_dataloader, val_dataloader
あと、resumeした時にデータを見送らないように以下をコメントアウト。
# resume data loader state by fast-forwarding through all seen batches
# drop this once streaming dataset supports proper resuming
# if resume:
# resume_t0 = time.perf_counter()
# for resume_iter in range(initial_iter):
# next(train_iterator)
# if resume_iter % 1000 == 0:
# fabric.print(f"Resuming dataset: {resume_iter} / {initial_iter}")
# fabric.barrier()
# fabric.print(
# f"Resuming data loader finished. Took {time.perf_counter() - resume_t0:.1f} seconds to reach iteration"
# f" {initial_iter}, epoch {train_iterator.epoch}."
# )
モデルの読み込み
ここでpthファイルを探すみたい。
if resume is True:
resume = max(out_dir.glob("*.pth"), key=(lambda p: int(p.name.split("-")[1])))
if resume:
fabric.print(f"Resuming training from {resume}")
fabric.load(resume, state)
なので、学習スクリプト実行時に、--resumeでそのpthファイル(lit_model.pth)を指定してみる。
python pretrain/tinyllama.py \
--resume "checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T/lit_model.pth"
エラー
Resuming training from checkpoints/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T/lit_model.pth
...
KeyError: 'model'
pretrainしてないlit_model.pthを指定したので, fabric.loadで形式が合ってない。
同様のissue
上記を参考にfabric.load_raw(path, model)に変更してみると・・・
今度は["_orig_mod."]の差によりstate_dictが合わない
これに関しては
model = torch.compile(model)
が原因のよう。(参考)
紆余曲折の結果、こうなった。
if resume is True:
resume = max(out_dir.glob("*.pth"), key=(lambda p: int(p.name.split("-")[1])))
if resume:
fabric.print(f"Resuming training from {resume}")
#from lit_gpt.utils import load_checkpoint
#load_checkpoint(fabric, model, resume, strict=True) #modelには_orig_modが全て含まれる. resumeのpthにはそれがない.
fabric.load_raw(resume,model)
#fabric.load(resume, state) ###ここでエラー
model = torch.compile(model)
学習中
Measured TFLOPs: 239.66
iter 32 | step 1: loss 4.0068, iter time: 1545.82 ms (optimizer.step), remaining time: 3.48 days
・・・
Validating ...
iter 288000: val loss 1.9565, val time: 36661.23 ms
Saving checkpoint to ‘out/lit-tiny-llama-1.1b/step-00009000.pth’
・・・
iter 305152 | step 9536: loss 1.4478
4日後…
学習結果
あまりlossが下がらなかった. シンプルにデータ量や学習の不足な気がする。
納得のできる回答はほとんどできない。
ファインチューニング
(公式解説)https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/finetune_lora.md
この記事が気に入ったらサポートをしてみませんか?