Deep Speedについてまとめ
概要
Deep Speedについて学習した際のまとめです。
Getting Started
initialize
初期化処理
model_engine, optimizer, _, _ = deepspeed.initialize(args=cmd_args,
model=model,
model_parameters=params)
NCCL backend
Deep Speedでは基本的にNCCL backendを使い分散処理を行う。
訓練のコード
訓練コードは以下の様になる。
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()
学習率スケジューラーによって、並列学習、混合精度訓練を自動で行ってくれる。
DeepSpeed Configuration
設定ファイル。
{
"train_batch_size": 8,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.00015
}
},
"fp16": {
"enabled": true
},
"zero_optimization": true
}
hostfile
マルチノードで実行する際のノード一覧を記載。
worker-1 slots=4
worker-2 slots=4
ノードに存在するGPU数も記載する。
起動コマンド
deepspeed --hostfile=myhostfile <client_entry.py> <client args> \
--deepspeed --deepspeed_config ds_config.json
Getting Started with DeepSpeed for Inferencing Transformer based Models
互換性のあるモデル
こちらに定義されているモデルのみ互換性がある。（推論のみの話？）
推論コード
# create the model
if args.pre_load_checkpoint:
model = model_class.from_pretrained(args.model_name_or_path)
else:
model = model_class()
...
import deepspeed
# Initialize the DeepSpeed-Inference engine
ds_engine = deepspeed.init_inference(model,
mp_size=2,
dtype=torch.half,
checkpoint=None if args.pre_load_checkpoint else args.checkpoint_json,
replace_with_kernel_inject=True)
model = ds_engine.module
output = model('Input String')
互換性のないモデルで並列に推論したいケース
injection_policyを指定すれば互換性がなくてもできる？
# create the model
import transformers
from transformers.models.t5.modeling_t5 import T5Block
import deepspeed
pipe = pipeline("text2text-generation", model="google/t5-v1_1-small", device=local_rank)
# Initialize the DeepSpeed-Inference engine
pipe.model = deepspeed.init_inference(
pipe.model,
mp_size=world_size,
dtype=torch.float,
injection_policy={T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')}
)
output = pipe('Input String')
実行。
deepspeed --num_gpus 2 inference.py
GPT NEO 2.7Bで推論するサンプル
# Filename: gpt-neo-2.7b-generation.py
import os
import deepspeed
import torch
from transformers import pipeline
local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B',
device=local_rank)
generator.model = deepspeed.init_inference(generator.model,
mp_size=world_size,
dtype=torch.float,
replace_with_kernel_inject=True)
string = generator("DeepSpeed is", do_sample=True, min_length=50)
if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
print(string)
実行
deepspeed --num_gpus 2 gpt-neo-2.7b-generation.py
量子化
以下の様に指定する。
import deepspeed
model = deepspeed.init_inference(model,
checkpoint='./checkpoint.json',
dtype=torch.int8,
quantization_setting=(quantize_groups,
mlp_extra_grouping)
)
Training your large model with DeepSpeed
Overview
DeepSpeedは巨大なモデルでのトレーニングを可能にするための技術をまとめたもの。
巨大なモデルを訓練するための２つの方法
ZeRO based technologies
並列でGPUを動作させた時のメモリの効率化、ZeRO-Offload and ZeRO-Infinityという技術も使うことができる。
3D Parallelism based technologies
テンソルスライシング、パイプライン並列、データ並列（またはZeROによるデータ並列）という3つの異なる並列技術の組み合わせ。
使用するにはモデルのコードのリファクタリングが必要になってくる。
どの技術を使うか
3D Parallelism based technologiesを使うケース
GPT-2またはGPT-3系のモデルの場合は3D Parallelismへの対応がすでに終わっている可能性が高いのでこちらを使用すると良い。
詳細は以下を参照
ZeRO based technologiesを使うケース
3D Parallelismへの対応がされていないモデルの場合はこちらを使うのが良い。
リファクタリングを必要としないため。
この方法でうまくいったら3D parallelism対応のコード変更を行ってみるのも良い。
ZeRO-OffloadとZeRO-Infinityにもアクセスできるようになる。
ZeROについての詳細は以下を参照。
パフォーマンスとのトレードオフ
GPUあたりのバッチサイズが極端に小さくない場合、ZeROと3D並列の性能は概ね同等。
データ並列処理自体の通信コストとほとんど変わらない。
何千ものGPUでトレーニングする場合かつバッチサイズが小さい場合はZeROだけでは十分な性能を維持できない可能性がある。
そういう場合はバッチサイズを大きくするか3D parallelismにするのが良い。
Discussion