PEFT documentation
学習可能なトークン
学習可能なトークン
Trainable Tokensメソッドは、埋め込み行列全体を学習したり、埋め込み行列にアダプターを適用したりすることなく、特定のトークン埋め込みをファインチューニングの対象とする方法を提供します。この手法は、こちらの初期実装に基づいています。
このメソッドは特定のトークンにのみ対象を絞り、指定したトークンインデックスを選択的に学習します。その結果、必要なRAMは低くなり、完全なファインチューニング済み埋め込み行列を保存する場合と比べて、ディスクメモリも大幅に低くなります。
gemma-2-2b（比較的大きな埋め込み行列を持つ）では、このスクリプトで取得した予備的なベンチマークの結果から、埋め込み行列を完全にファインチューニングするよりも、Trainable Tokens を使用することで約 4 GiB の VRAM を節約できることが示唆されています。LoRA も同程度の VRAM を使用しますが、変更したくないトークンも対象となる可能性があります。なお、これらはあくまで目安であり、埋め込み行列のサイズによってこれらの数値が多少変動する可能性があります。
このメソッドはトークンを追加しません。トークナイザーにトークンを追加し、モデルの埋め込み行列をそれに応じてリサイズする必要があります。このメソッドは、指定したトークンの埋め込みのみを再学習します。このメソッドはLoRAレイヤーと併用することもできます。LoRAのドキュメントを参照してください。
新しいトークンを追加する場合、save_pretrained() でモデルを保存したり、get_peft_model_state_dict() で状態辞列を取得したりすると、埋め込み行列がリサイズされたため、差分のみではなく埋め込み行列全体を保存する可能性があります。これは予防策です。容量を節約するには、
save_pretrainedを呼び出す際にsave_embedding_layers=Falseを設定して、この動作を無効にできます。埋め込み行列を他の方法で変更しない限り、この変更は学習可能なトークンによって追跡されないため、安全に実行できます。
ベンチマーク概要
API
TrainableTokensConfig
クラス peft.TrainableTokensConfig
< source >( task_type: Optional[Union[str, TaskType]] = Nonepeft_type: Optional[Union[str, PeftType]] = Noneauto_mapping: Optional[dict] = Nonepeft_version: Optional[str] = Nonebase_model_name_or_path: Optional[str] = Nonerevision: Optional[str] = Noneinference_mode: bool = Falsetoken_indices: list[int] = <factory>target_modules: Optional[Union[list[str], str]] = Noneinit_weights: bool = True )
パラメータ
- token_indices (
list[int]) — 学習可能なトークンのインデックスを示す整数のリストです。トークナイザーでトークンのインデックスを調べるには、文字列をトークン化して返されるinput_idsを確認できます。インデックスの数がトークンの総数に近づくほど、この手法の効率は低下します。 - target_modules (
Optional[Union[list[str], str]]) —TrainableTokensLayerに置き換えるモジュール名のリスト、またはモジュール名の正規表現です。 定義されていない場合、モデルにget_input_embeddingsメソッドがある場合（transformer モデルには通常あります）、そのモデルの入力埋め込み層を取得しようとします。それが失敗した場合、デフォルトは ‘embed_tokens’ です。 その他のターゲットの例としては、embedding、encoder.embeddings、またはdecoder.embeddingsがあります。 - init_weights (
bool) — デフォルトでは、新しいトークンの重みは、対応するトークン埋め込みと同じ値で初期化されます。これにより、学習を行わない場合、TrainableTokensは何もしない（no-op）状態になります。Falseに設定すると、重みはランダムな値になります。正確に意図していることが明確でない限り、この設定を変更しないでください。
TrainableTokensメソッドの設定。
新しいトークンの学習（および既存トークンの再学習）を、埋め込み行列全体を学習せずに実行できます。一部の選択したトークン（インデックスで識別）を学習可能にし、残りは変更しないことで、この手法はメモリを節約しながら新しいトークンを追加したり、既存トークンの埋め込みを変更したりするために使用できます。埋め込み行列を完全に学習する場合と比べて、ストレージおよびワーキングメモリの使用量が削減されます。
FSDP/DeepSpeedでのトレーニングがまだ完全にサポートされていない可能性があることに注意してください。
TrainableTokensModel
クラス peft.TrainableTokensModel
< source >( modelpeft_config: Union[PeftConfig, dict[str, PeftConfig]]adapter_name: strlow_cpu_mem_usage: bool = Falsestate_dict: Optional[dict[str, torch.Tensor]] = None )