staged training

ステージトレーニング

Transformer Language Modelsの段階的なトレーニングでは、小さなモデルから始まる段階的なトレーニングセットアップを提案し、モデルの深さと幅を増やすために「成長演算子」を適用することにより、トレーニングに使用される計算量を徐々に増やします。前のステージの出力で各ステージを初期化することにより、トレーニングプロセスは以前の段階から計算を効果的に再利用し、より効率的になります。

ここで、成長演算子と評価スクリプトの再現可能なコードをリリースします。

設定

このリポジトリのスクリプトには、Python 3.7以下が必要です。適切なPython環境を取得したら、公式の指示に従ってPytorch v1.9.0を最初にインストールします。その後、実行します

 pip install -r requirements.txt

成長オペレーター

成長オペレーター（幅/深さ）はそれぞれ、トレーニング状態全体（モデルパラメーター、オプティマイザー状態、学習レートスケジュールなどを含む）を入力し、トレーニングが続く新しいトレーニング状態を出力します。

対応するスクリプトの使用方法の詳細については、 scripts/cheatsheet.txt参照してください。

たとえば、幅オペレーターを以下で適用できます。

 CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/gpt_pretrain.py 
  --save_prefix final_gpt2_large_div2_width_check_bs512_lr0.0020_warmup3k_seqlen1024_debug 
  --gpu_count -1 
  --model gpt2  
  --tokenizer gpt2 
  --batch_size 4 
  --grad_accum 32  
  --lr 0.002006911598778545  
  --warmup_steps 3000   
  --train_steps 250000  
  --val_every 50  
  --val_batches 50 
  --fp16 
  --seqlen 1024 
  --log_rate 10 
  --num_workers 4 
  --size GPT2_large_div2_width 
  --random 
  --resume final_runs/final_gpt2_large_div2_width_check_bs512_lr0.0021_warmup3k_seqlen1024_debug/checkpoint-xxx.ckpt 
  --doubling weights

または次のような深度演算子

 CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/gpt_pretrain.py 
  --save_prefix final_gpt2_large_div2_depthx2_check_bs512_lr0.0020_warmup3k_seqlen1024_debug 
  --gpu_count -1 
  --model gpt2  
  --tokenizer gpt2 
  --batch_size 4 
  --grad_accum 32 
  --lr 0.002006911598778545 
  --warmup_steps 3000 
  --train_steps 250000 
  --val_every 50 
  --val_batches 50 
  --fp16 
  --seqlen 1024 
  --log_rate 10 
  --num_workers 4 
  --size GPT2_large_div2_depth 
  --random 
  --resume final_runs/final_gpt2_large_div2_depth_check_bs512_lr0.0020_warmup3k_seqlen1024_debug/checkpoint-epoch=0-step=6499.ckpt 
  --doubling layers

評価

evaluation/eval_wikitext.pyまたはevaluation/eval_lambada.pyを使用して、サポートされているデータセットの1つでGPT-2を評価します。例えば：

python evaluation/eval_wikitext.py

またはDockerを使用してください：

docker build -t evaluation:latest .
docker run --rm --gpus all evaluation:latest evaluation/eval_wikitext.py

参照

研究で段階的なトレーニングを使用している場合、またはここで公開されているベースラインの結果を参照する場合は、次のBibtexエントリを使用してください。

 @misc{shen2022staged,
    title={Staged Training for Transformer Language Models},
    author={Sheng Shen and Pete Walsh and Kurt Keutzer and Jesse Dodge and Matthew Peters and Iz Beltagy},
    year={2022},
    eprint={2203.06211},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}