Zero3 メモリオプティマイザによる GPT-NeOx の微調整
このスクリプトは、ゼロDPメモリ最適化機能を備えた複数のデバイス上でGPT-Neoxモデルのバイアスパラメータをトレーニングします。
14import datetime
15
16import torch
17import torch.distributed
18
19from labml import experiment, monit, tracker
20from labml.configs import option
21from labml.logger import inspect
22from labml_nn.neox.samples.finetune import PipelineParallelTrainerConfパイプラインパラレルトレーナーの設定を使用して、Zero3 メモリオプティマイザーに適応させてください。
27class Configs(PipelineParallelTrainerConf):
28 rank: int
29 world_size: int32@option(Configs.optimizer, 'Zero3Adam')
33def _optimizer(c: Configs):39 from labml_nn.optimizers.adam_fp16 import AdamFP16
40 return AdamFP16(c.model.get_trainable_chunk(), lr=c.learning_rate)Zero3 メモリオプティマイザーでモデルを作成
43@option(Configs.model, 'Zero3')
44def _model(c: Configs):48 from labml_nn.scaling.zero3 import Zero3Layer, Zero3Sequential微調整器がトレーニング可能なパラメーターを設定していることを確認するには
51 _ = c.fine_tunerレイヤーを以下のようにラッピングします。Zero3Layer
54 modules = []
55 for m in monit.iterate('Zero3', c.layers):
56 modules.append(Zero3Layer(m.to(c.device),
57 c.rank, c.world_size, c.device, c.dtype))シーケンシャルモデルの作成
60 model = Zero3Sequential(modules)63 return modelランクのあるノードでトレーニングを実行しますrank
。
66def main(rank: int, world_size: int, init_method: str = 'tcp://localhost:23456'):PyTorch 分散プロセスグループを初期化
71 with monit.section('Distributed'):
72 torch.distributed.init_process_group('nccl',
73 timeout=datetime.timedelta(seconds=30),
74 init_method=init_method,
75 rank=rank,
76 world_size=world_size)現在のデバイスを設定
79 device = torch.device(f'cuda:{rank}')
80 torch.cuda.set_device(device)実験を作成
83 experiment.create(name='zero3_neox', writers={'screen', 'labml'})
84 experiment.distributed(rank, world_size)構成の作成
87 conf = Configs()構成をロード
90 experiment.configs(conf, {
91 'model': 'Zero3',
92 'optimizer': 'Zero3Adam',
93
94 'device': device,
95 'rank': rank,
96 'world_size': world_size,
97
98 'learning_rate': 3e-4,
99 'max_seq_len': 128,
100 'batch_size': 16,
101 })実験を始める
104 with experiment.start():モデルを初期化します。これをループの前に行って、ログをクリーンアップしてください。
106 _ = conf.modelモデルのトレーニング
109 for epoch in monit.loop(conf.epochs):
110 conf.train_epoch()
111 tracker.new_line()115if __name__ == '__main__':マシン構成をログに記録する
117 inspect([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])
118 inspect(
119 n_gpus=torch.cuda.device_count(),
120 mpi=torch.distributed.is_mpi_available(),
121 nccl=torch.distributed.is_nccl_available(),
122 )
123
124 n_gpu = torch.cuda.device_count()GPU ごとにプロセスを開始します。複数のコンピューターを使用している場合は、別のランチャーが必要になります
。127 torch.multiprocessing.spawn(main, args=(n_gpu,), nprocs=n_gpu, join=True)