lit-gptでLlama2を学習してみたい

2024年1月5日 17:31

注）結局学習が上手く行ったわけではないです。悪しからず。

(参考記事)
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

４日後…

学習結果

あまりlossが下がらなかった. シンプルにデータ量や学習の不足な気がする。
納得のできる回答はほとんどできない。

ファインチューニング

（公式解説）https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/finetune_lora.md

この記事が気に入ったらサポートをしてみませんか？