[事前学習] - 高度化WG12 Attention Checkpoint Tokenの検証

[gojiteji]

Overview

WG12で研究している、LLMの解釈性と高速化に関する研究 Attentionの集中箇所(Attention Checkpoint Tokenを)を強制的に作り、事前学習を行う。

Details

Attentionを可視化した際、部分部分にAttentionをまとめれそうな箇所が存在する。本実験では、Attentionをまとめるトークン（Attention Checkpoint Token）を外部から、一定間隔に意図的に設置することで、1. 情報のまとめる箇所を作り、atentionの解釈性を向上させる。2. inference時にまとめた部分だけ残すことで効率化を目指す。

• LlamaのAttentionの例：画像点線のような部分を、意図的に作る

  

• Google Slides

• LLMC Notion

Resources

• 計算機

  • クラスタ: llm-jp-nvlink
  • ノード種別: gpu
  • ノード台数: 4 or 8

• コード

  • リポジトリ: https://github.com/gojiteji/Megatron-LM/blob/feat/attention-checkpoint-token/scripts/mdx/1.8B/1.8B-dense-cp.sh
  • コミット: FIXME xxxxxx

• 入力データ: v3.0.0系　/model/koki-tan/llm-jp-corpus/v3.0.0/tokenized/interval_${ATTENTION_CHECKPOINT_INTERVAL}/v3.0b1 v3.1.0系　/model/koki-tan/llm-jp-corpus/v3.1.0/tokenized/interval_${ATTENTION_CHECKPOINT_INTERVAL}/v3.0b1 　ATTENTION_CHECKPOINT_INTERVAL（挿入間隔）は0、16、32を想定

• 出力データ:

  • 保存先: `llm-jp-nvlink: /model/experiments/0078_wg12_acttoken
  • データ内訳:
  • {name}: xxx TB （バッファ容量を含む）

• W&B ログ:

  • https://wandb.ai/gojiteji/llm-jp-wg12

• 開始日: 2024-11-DD

• 終了予定日: YYYY-11-DD （バッファ期間を含む） 　1.8Bモデルを20k stepsはminimumで走らせたい。 2-3weeks程度想定。

本ISSUEは事前学習に関するものですが、1.8Bモデルを20k steps程度の学習で、最初の結果が確認できると考えています