📖

特化型llm(Doujinshi-1.8b)の開発報告書④:DeepSpeedを用いた継続事前学習

2025/03/19に公開

LLM

 はじめに沼津高専のpuwaerです。この度、R18に特化した大規模言語モデル（LLM）、Doujinshi-1.8bを開発しました。

この記事では、DeepSpeedを用いた継続事前学習の作成方法をプログラムを交えて解説します。
本記事では、DeepSpeedを用いた継続事前学習の具体的な手法について詳しく解説します。

また、プログラムは以下の記事を参考にして作成しています。

HuggingFaceとDeepSpeedで実践継続事前学習
参考にした記事のプログラムがhuggingfaceに公開されているデータセットを使用して学習されていました。

継続事前学習を行うにあたり、huggingfaceの公開データセットを使用するのではなく、ローカル環境にある parquet や json 形式のデータを学習に利用できるように使いやすくしたプログラムを紹介します。
下の流れでする合成データの具体的な手法について解説します。

1.使用機器

2.環境構築

3.各プログラムの役割

4.プログラムの詳細

5.設定ファイルの解説
本記事で利用するプログラムは以下のGitHubリポジトリで管理しています。

https://github.com/puwaer/continual_pretrain_deepspeed
このプログラムで作成したモデル
huggingface: Doujinshi-1.8b

 1.使用機器以下のようなpcを使い開発しました。

高専4年生でまだ研究室配属されていないため教授に頼みpcを貸してもらいました。

また、個人のpcはdeepspeedを用いたllmの継続事前学習の検証用にramが増設されています。


機器
CPU
RAM
GPU
VRAM


研究室PC
Intel i7-12700
64GB
NVIDIA RTX A6000
48GB

個人PC
Ryzen 7 5700X
96GB
NVIDIA GeForce RTX 3060
12GB


 2. 環境構築
 2.1 Docker環境の構築まず、Dockerコンテナをビルドし、起動します。

 2.1.0. リポジトリをクローンするgit clone https://github.com/puwaer/continual_pretrain_deepspeed.git

 2.1.1. Dockerイメージのビルドcd continual-pretrain/docker
chmod +x build.sh
./build.sh

 2.1.2. Dockerコンテナの起動cd continual-pretrain/docker
docker compose up -d

 2.1.3. コンテナへアクセスdocker exec -it LLM bash

 2.2 依存ライブラリのインストール以下のコマンドで必要なライブラリをインストール、またはアップグレードしてください。
pip install --upgrade huggingface_hub datasets transformers
Qwenを学習に使用する場合は、追加で以下のライブラリをインストールしてください。
pip install --upgrade 'transformers>=4.36.0'
pip install --upgrade 'accelerate>=0.26.0'
pip install --upgrade wandb

 3. 各プログラムの役割本プロジェクトのディレクトリ構成と、各フォルダ・ファイルの役割を説明します。

continual-pretrain/configs/deepspeed : DeepSpeedの設定ファイル

continual-pretrain/configs/train_configs : 学習の設定ファイル

continual-pretrain/src : Pythonで書かれた学習コードが格納されている
.continual-pretrain
├── abci
│   └── requirements.txt
├── configs
│   ├── deepspeed
│   │   ├── ds_config_zero2.json
│   │   ├── ds_config_zero3.json
│   │   ├── ds_config_zero_llmjp.json
│   │   └── ds_config_zero_qwen.json
│   └── train_configs
│       ├── train_abci.yaml
│       ├── train_base.yaml
│       ├── train_test.yaml
│       ├── train_test_llmjp_json.yaml
│       └── train_test_qwen.yaml
├── data
│   └── test_data.json
├── docker
│   ├── Dockerfile
│   ├── build.sh
│   └── docker-compose.yml
├── script
│   └── continual_pretrain_abci.sh
└── src
    ├── train_deepspeed.py
    ├── train_deepspeed_1.py
    ├── train_deepspeed_2.py
    ├── train_deepspeed_llmjp.py
    ├── train_deepspeed_llmjp_json.py
    ├── train_deepspeed_llmjp_tokenizer.py
    ├── train_deepspeed_qwen.py
    └── utils.py


 4.プログラムの詳細
 学習の実行学習を実行する際には、目的に応じたコマンドを実行してください。

 一般的なLLMの事前学習deepspeed src/train_deepspeed.py --train_config ./configs/train_configs/train_base.yaml
このスクリプトは基本的なLLMの学習を行います。train_base.yamlではデータセットのパス、バッチサイズ、学習率などを指定します。

データセットはhuggingfaceで公開されているものを使用、モデルもhuggingfaceで公開されているものを使用
deepspeed src/train_deepspeed_2.py --train_config ./configs/train_configs/train_test.yaml
このスクリプトは基本的なLLMの学習を行います。train_test.yamlではデータセットのパス、バッチサイズ、学習率などを指定します。

データセットはローカル環境にあるparquetファイルを使用、モデルもhuggingfaceで公開されているものを使用

 Qwenを使用する場合deepspeed src/train_deepspeed_qwen.py --train_config ./configs/train_configs/train_test_qwen.yaml
このスクリプトは基本的なLLMの学習を行います。train_test_qwen.yamlではデータセットのパス、バッチサイズ、学習率などを指定します。

データセットはローカル環境にあるparquetファイルを使用、モデルもhuggingfaceで公開されているものを使用

 JSONファイルの学習データを使用する場合deepspeed src/train_deepspeed_llmjp_json.py --train_config ./configs/train_configs/train_test_llmjp_json.yaml
このスクリプトは基本的なLLMの学習を行います。train_test_llmjp_json.yamlではデータセットのパス、バッチサイズ、学習率などを指定します。

データセットはローカル環境にあるjsonファイルを使用、モデルもhuggingfaceで公開されているものを使用

 学習に使用するjsonファイルの構成学習に使用するjsonファイルの構成を以下に示す
[
    {
      "id": "1",
      "text": "任意のテキストデータ"
    },
    {
      "id": "2",
      "text": "任意のテキストデータ"
    }
]

 5.設定ファイルの解説学習を動かすための設定ファイルを詳しく解説します。

基本的には使用しているgpuのvramの容量とramの容量に応じてパラメータを指定してください

 5.1 DeepSpeed設定ファイル (continual-pretrain/configs/deepspeed/ds_config_zero3.json)以下の設定ファイルは、DeepSpeedのゼロ冗長性最適化（ZeRO）Stage 3を使用して、メモリ負荷を軽減しつつ学習を進める構成になっています。
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",        
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": false,
        "contiguous_gradients": true,
        "sub_group_size": 5e8,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 5e8,
        "stage3_max_reuse_distance": 5e8,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

 5.2 DeepSpeed設定ファイルのパラメータの解説

設定項目
説明


fp16
自動で16bit浮動小数点を有効化し、メモリ削減と高速化を図る設定。

optimizer

AdamW で重み減衰を行いながら学習を安定させる最適化アルゴリズム。

scheduler
ウォームアップ学習率 (WarmupLR) を導入し、初期学習率を滑らかに増加させて過学習や収束の不安定化を防ぐ。

zero_optimization
Stage 3 によりモデルパラメータ・オプティマイザ状態をCPUオフロードし、VRAM消費を大幅に抑える大規模学習用の設定。

stage: 3
Stage 3では、モデルのパラメータ・勾配・Optimizer状態を細かく分割して管理し、必要なときだけGPUにロードする仕組みを持つ。

offload_optimizer
Optimizer状態をCPUメモリに退避させることで、VRAMを削減。 pin_memory: true で高速化も狙う。

device

cpuまたはnoneを指定可能。noneを選択するとGPUを使用します。


 メモリ消費の内訳（例）パラメータ w ：2バイト（FP16）
勾配  ∆W𝐿：2バイト（FP16）
Optimizer状態：3 × 4バイト（FP32）
1つのパラメータにつき m（移動平均）, v（二次モーメント）, v_hat（修正後モーメント）の3つの状態を保持するため合計12バイト


 補足説明Zero Stage 3では、これらを個別にgpuにオフロードすることができるため、大規模モデルでもGPUのメモリが少なくても学習が可能です。しかし、データ通信における損失が生じ、GPUの性能の約9割しか活用できません。この問題を改善するためには、train_base.yamlのper_device_train_batch_size、per_device_eval_batch_size、およびgradient_accumulation_stepsの設定をgpuに応じて調整することが有効です。

 5.3 学習設定ファイル (continual-pretrain/configs/train_configs/train_base.yaml)次に、学習用の設定ファイルを見ていきます。
model:
  model: llm-jp/llm-jp-3-1.8b
  tokenizer: llm-jp/llm-jp-3-1.8b
  use_cache: False
  max_length: 512

train:
  output_dir: ../outputs
  evaluation_strategy: steps
  logging_strategy: steps
  save_strategy: steps
  learning_rate: 1e-4
  num_train_epochs: 1
  per_device_train_batch_size: 3
  per_device_eval_batch_size: 3
  gradient_accumulation_steps: 2
  gradient_checkpointing: True
  weight_decay: 0.01
  warmup_ratio: 0.1
  optim: adamw_torch
  fp16: True
  bf16: False
  dataloader_num_workers: 4
  eval_steps: 50
  save_steps: 50
  logging_steps: 50
  run_name: test
  save_total_limit: 2
  save_on_each_node: False
  neftune_noise_alpha: 5
  deepspeed: ./configs/deepspeed/ds_config_zero3.json
  report_to: wandb

seed: 42

dataset:
  path: hotchpotch/wikipedia-ja-20231030
  subset: chunked
  split: train

 5.4 学習設定ファイルのパラメータの解説
 model セクション

設定項目
説明


model
使用するモデルを指定。ここでは llm-jp/llm-jp-3-1.8b（1.8Bパラメータの日本語LLM）を採用。

tokenizer
トークナイザを指定。モデルと同じ llm-jp/llm-jp-3-1.8b を使用し、入力データのトークン化を行う。

max_length
入力シーケンスの最大長を512トークンに設定。長いシーケンスは切り捨てられ、短いものはパディングされる。


 train セクション

設定項目
説明


output_dir
学習結果（モデルチェックポイントやログ）の保存先ディレクトリ。../outputs に保存される。

evaluation_strategy
評価のタイミングを指定。steps は一定ステップごと（eval_steps で定義）に評価を行う。

logging_strategy
ログ記録のタイミング。steps で一定ステップごと（logging_steps で定義）にログを記録。

save_strategy
モデル保存のタイミング。steps で一定ステップごと（save_steps で定義）にチェックポイントを保存。

learning_rate
学習率の設定。一般的に事前学習では 1e-4 を使用する

num_train_epochs
学習エポック数を1に設定。言語モデルの事前学習では全データを1回だけ学習することが多い。

per_device_train_batch_size
1GPUごとの訓練バッチサイズを3に設定。GPUメモリに応じた調整が必要。

per_device_eval_batch_size
1GPUごとの評価バッチサイズを3に設定。訓練と同様にメモリ制約を考慮。

gradient_accumulation_steps
勾配蓄積ステップ数を2に設定。バッチサイズを効果的に増やし（3×2=6）、メモリ負荷を軽減。

gradient_checkpointing
勾配チェックポイントを有効化（True）。中間計算を保存せず再計算することでメモリ使用量を削減。

weight_decay
重み減衰を 0.01 に設定。過学習を防ぎつつモデルの汎化性能を向上させる正則化手法。

warmup_ratio
学習率ウォームアップの割合を 0.1 に設定。全体ステップの10%で学習率を徐々に増加させる。

optim
最適化アルゴリズムとして adamw_torch（PyTorch版AdamW）を指定。重み減衰付きのAdamを使用。

fp16
16bit浮動小数点演算を有効化（True）。メモリ使用量を削減し、計算を高速化。

bf16
BF16演算を無効化（False）。FP16を優先して使用。

dataloader_num_workers
データローダーのワーカー数を4に設定。データの前処理を並列化し、学習のボトルネックを軽減。

run_name
実行名を test に設定。ログや可視化ツール（W&Bなど）での識別用。

deepspeed
DeepSpeed設定ファイルのパスを指定。ds_config_zero3.json を使用してメモリ効率を最適化。

report_to
学習ログを wandb（Weights & Biases）に送信し、リアルタイムで可視化・追跡。


 dataset セクション

設定項目
説明


path
データセットのパスを指定。hotchpotch/wikipedia-ja-20231030 は日本語Wikipediaデータを使用。


 補足説明gpuのスペックやデータセットのサイズに応じて、train_base.yamlのper_device_train_batch_size、per_device_eval_batch_size、およびgradient_accumulation_stepsを調整し、最適化してください。

 おわりに本記事では、DeepSpeedを用いた継続事前学習の手法を解説しました。

書くことが思いのほかあり、後半の方は疲れて雑になり、説明の不足しているところが多くあると思います。

パラメータなど分からないことがありましたら、気軽にTwitterのDMに質問してください。

 開発支援のお願い現在、開発を続けていますが、クラウドGPUの価格が高く、十分な計算リソースを確保できずにいます。そのため、思い通りに開発が出来ていません。

また、オープンソースの理念を大切にしており、プログラム・データセット・モデルを有料で公開するつもりはありません。そのため、金銭的に余裕のある方に支援していただけると大変助かります。

TwitterのDMやご支援いただける方は、以下のプラットフォームよりお願いいたします。
Twitter
Zenn
github sponsors

設定項目	説明
`model`	使用するモデルを指定。ここでは `llm-jp/llm-jp-3-1.8b`（1.8Bパラメータの日本語LLM）を採用。
`tokenizer`	トークナイザを指定。モデルと同じ `llm-jp/llm-jp-3-1.8b` を使用し、入力データのトークン化を行う。
`max_length`	入力シーケンスの最大長を512トークンに設定。長いシーケンスは切り捨てられ、短いものはパディングされる。

設定項目	説明
`output_dir`	学習結果（モデルチェックポイントやログ）の保存先ディレクトリ。`../outputs` に保存される。
`evaluation_strategy`	評価のタイミングを指定。`steps` は一定ステップごと（`eval_steps` で定義）に評価を行う。
`logging_strategy`	ログ記録のタイミング。`steps` で一定ステップごと（`logging_steps` で定義）にログを記録。
`save_strategy`	モデル保存のタイミング。`steps` で一定ステップごと（`save_steps` で定義）にチェックポイントを保存。
`learning_rate`	学習率の設定。一般的に事前学習では `1e-4` を使用する
`num_train_epochs`	学習エポック数を1に設定。言語モデルの事前学習では全データを1回だけ学習することが多い。
`per_device_train_batch_size`	1GPUごとの訓練バッチサイズを3に設定。GPUメモリに応じた調整が必要。
`per_device_eval_batch_size`	1GPUごとの評価バッチサイズを3に設定。訓練と同様にメモリ制約を考慮。
`gradient_accumulation_steps`	勾配蓄積ステップ数を2に設定。バッチサイズを効果的に増やし（3×2=6）、メモリ負荷を軽減。
`gradient_checkpointing`	勾配チェックポイントを有効化（`True`）。中間計算を保存せず再計算することでメモリ使用量を削減。
`weight_decay`	重み減衰を `0.01` に設定。過学習を防ぎつつモデルの汎化性能を向上させる正則化手法。
`warmup_ratio`	学習率ウォームアップの割合を `0.1` に設定。全体ステップの10%で学習率を徐々に増加させる。
`optim`	最適化アルゴリズムとして `adamw_torch`（PyTorch版AdamW）を指定。重み減衰付きのAdamを使用。
`fp16`	16bit浮動小数点演算を有効化（`True`）。メモリ使用量を削減し、計算を高速化。
`bf16`	BF16演算を無効化（`False`）。FP16を優先して使用。
`dataloader_num_workers`	データローダーのワーカー数を4に設定。データの前処理を並列化し、学習のボトルネックを軽減。
`run_name`	実行名を `test` に設定。ログや可視化ツール（W&Bなど）での識別用。
`deepspeed`	DeepSpeed設定ファイルのパスを指定。`ds_config_zero3.json` を使用してメモリ効率を最適化。
`report_to`	学習ログを `wandb`（Weights & Biases）に送信し、リアルタイムで可視化・追跡。

特化型llm(Doujinshi-1.8b)の開発報告書④:DeepSpeedを用いた継続事前学習

はじめに

1.使用機器

2. 環境構築

2.1 Docker環境の構築

2.1.0. リポジトリをクローンする

2.1.1. Dockerイメージのビルド

2.1.2. Dockerコンテナの起動

2.1.3. コンテナへアクセス

2.2 依存ライブラリのインストール

3. 各プログラムの役割

4.プログラムの詳細

学習の実行

一般的なLLMの事前学習

Qwenを使用する場合

JSONファイルの学習データを使用する場合

学習に使用するjsonファイルの構成

5.設定ファイルの解説

5.1 DeepSpeed設定ファイル (`continual-pretrain/configs/deepspeed/ds_config_zero3.json`)

5.2 DeepSpeed設定ファイルのパラメータの解説

メモリ消費の内訳（例）

補足説明

5.3 学習設定ファイル (`continual-pretrain/configs/train_configs/train_base.yaml`)

5.4 学習設定ファイルのパラメータの解説

`model` セクション

`train` セクション

`dataset` セクション

補足説明

おわりに

開発支援のお願い

Discussion

機器	CPU	RAM	GPU	VRAM
研究室PC	Intel i7-12700	64GB	NVIDIA RTX A6000	48GB
個人PC	Ryzen 7 5700X	96GB	NVIDIA GeForce RTX 3060	12GB

設定項目	説明
fp16	自動で16bit浮動小数点を有効化し、メモリ削減と高速化を図る設定。
optimizer	`AdamW` で重み減衰を行いながら学習を安定させる最適化アルゴリズム。
scheduler	ウォームアップ学習率 (`WarmupLR`) を導入し、初期学習率を滑らかに増加させて過学習や収束の不安定化を防ぐ。
zero_optimization	Stage 3 によりモデルパラメータ・オプティマイザ状態をCPUオフロードし、VRAM消費を大幅に抑える大規模学習用の設定。
stage: 3	Stage 3では、モデルのパラメータ・勾配・Optimizer状態を細かく分割して管理し、必要なときだけGPUにロードする仕組みを持つ。
offload_optimizer	Optimizer状態をCPUメモリに退避させることで、VRAMを削減。 `pin_memory: true` で高速化も狙う。
device	`cpu`または`none`を指定可能。`none`を選択するとGPUを使用します。