端の知識の備忘録

技術メモになりきれない、なにものか達の供養先

Transformers Trainerを使ってみて分かりにくかった仕様など

まえがき

言語モデルを自分でガッツリ使う経験が今まで無かったので、勉強がてら先週火曜日まで開催されていたKaggle - LLM Science Exam というコンペに参加してました。

www.kaggle.com

そこそこ頑張った結果、過去最高の成績(49th/2663, Top2%)を取ることができ、銀メダルを取ることができた。SolutionなどはKaggleの方に書いたのでそちらを。

www.kaggle.com

で、この記事で触れるのは、今まで機械学習モデルのトレーニングにはPyTorchの標準的なやり方(ループ内で推論してloss計算してloss.backward()で逆伝播させて……というやつ)をなるべく使ってきたのだが、最近流行り始めているっぽいHuggingFaceのTransformers Trainer系APIを全面的に利用してみたので、そのときに困ったこととかをまとめてみるというお話。

まとめ

特にHuggingFace Transformers使う上で困ったのは次の4点。

  • wandbなどのロギングサービスとの連携がいらないときは、TrainingArgumentで、report_to="none"("none"が文字列であることに注意)を指定する。
    • 最初ローカルだけでサクッと動かしたいだけなのに、このオプションがデフォルトで有効("all")になっていて少しイラッとした。wandbが使ったら便利なサービスなのはわかるけど、デフォルトは無効になっていてほしい……。
  • モデルの中身で何しているのかを見たいときには、GitHubここから対象のモデルのフォルダの、modeling_XXX.pyを確認する
    • AutoModelForSequentialClassification()で作られるモデルが実際どのようにForward passの処理をしているのか?とかがドキュメントに書いていないので、ソースを見て納得したいようなときの話
  • PEFTのLoRAでモデルのFineTuningを行ったのだが、学習させたモデルを読み込むとClassification層の重みが初期化されてしまう現象に遭遇。
    • 次のGitHub Issueみたいな感じの現象。当然、ランダムな重みのFC層を通った結果が出てくるので、推論結果はナンセンスなものになる。
    • 結局、LoRAの重みを別に読み込んで推論用モデルの重みを上書きするのが一番確実な方法だと思う。

github.com

  • remove_unused_columnsがTrueになっているTrainerにtrain_dataset, eval_datasetを渡すと、モデルのForwardに渡されないdatasetの要素が消し去られてしまい想定しないエラーにつながることがある。
    • remove_unused_columnsはデフォルトTrueである。要注意!
    • 私の場合、使わない要素を自前で省く処理を入れたDataCollatorをTrainerに渡していたのだが、これがコンフリクトしてエラーを起こしてしまった。外で普通のtorch.utils.data.DataLoaderで使う場合エラーが出なかったので、何が原因か分からず数時間浪費した。

HuggingFace系のライブラリは便利は便利なのだが、暗黙的に行われる処理の中身が分からず何かトラブったときの対応が難しいと言うのが難点。エラーメッセージも問題の根本を直接指し示すものとはならないことが多く、GitHubとかで類似のIssueを探したり、アドホックにprint仕込んでデバックしたりする局面が何回かあった。

とはいえ、PyTorchを直接書くよりはコードの量は格段に少なくなるし、PEFTによるファインチューニングやbitsandbytesによる量子化、DeepSpeedによるVRAM使用量削減のような自前でPyTorchと合わせて使うには少し手間がかかるライブラリを使いたいとき、基本的な使い方であればHuggingFaceのオプションに指定するだけで使えたりする利便性もある。

HuggingFaceは簡単そうに見えて色々わかっていないと問題発生時に困るライブラリなので、自前のデータセットでLLMトレーニングさせたい!みたいな場合まずPyTorchを少し勉強してから使うべきものという印象でした。

PEFT関連の注意点をもう少し

まとめで書いた「LoRAの重みを別に読み込んで推論用モデルの重みを上書きする」のやり方は次のような感じ。

peft_model_id = "/kaggle/input/mistral-7b-lora-seq/checkpoint-5000/checkpoint-5000"
base_model = "/kaggle/input/mistral-7b-v0-1-starter/Mistral-7B-v0.1/Mistral-7B-v0.1"
config = LoraConfig.from_pretrained(peft_model_id)
inference_model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=5, quantization_config=bnb_config)

# ここでトレーニングさせたLoRAの重みを直接読み込む
adapters_weights = torch.load(os.path.join(peft_model_id, 'adapter_model.bin'))

# PeftModelで学習済みのLoRAのPATHを指定しても、分類層の重みが初期化されてしまう。
model = PeftModel.from_pretrained(inference_model, peft_model_id)

# PEFTモデルの分類層の重みをLoRAの重みで上書きする
model.score.original_module.weight = Parameter(adapters_weights['base_model.model.score.weight'])

また、PEFTでLoRAとかを実際どう学習させるかなどの実装例はここを見ると良い。ドキュメントにはさらっとしか書いていないのであまり役に立たない。

github.com

remove_unused_columnsの罠の例

# こんな感じのpreprocessで、['label', "input_ids", "attention_mask", "prompt"]を要素に持つデータセットを準備する
def preprocess(row, max_seq_length=2048):
    pre_context = tokenizer.bos_token + "[INST] Answer the given multiple choice Question using the Context provided. Contexts are automatically generated and may not be reliable. In such cases, ignore the context and answer! Answers must be a single letter of the alphabet, either A, B, C, D, or E. \nContext:"
    tokenized_pre_context = tokenizer(pre_context)
    len_tokens_pre_context = len(tokenized_pre_context["input_ids"])

    options = '\n Options: A: ' + row['A'] + ' B: ' + row['B'] + ' C: ' + row['C'] + ' D: ' + row['D'] + ' E: ' + row['E'] + "\n[/INST] Answer:"
    
    tokenized_options = tokenizer(options)
    len_tokens_options = len(tokenized_options["input_ids"])
    
    tokenized_prompt = tokenizer('\n Question: ' + row["prompt"])
    len_tokens_prompt = len(tokenized_prompt["input_ids"])
    
    context = row["third_context"] + "\n" + row["second_context"] + "\n" + row["first_context"]
    max_tokens_context = max_seq_length - len_tokens_prompt - len_tokens_pre_context - len_tokens_options
    tokenized_context = tokenizer(context, max_length=max_tokens_context, truncation=True)
    
    
    example_encoded = {
        "input_ids": torch.tensor(tokenized_pre_context["input_ids"] + tokenized_context["input_ids"] + tokenized_prompt["input_ids"] + tokenized_options["input_ids"]),
        "attention_mask": torch.tensor(tokenized_pre_context["attention_mask"] + tokenized_context["attention_mask"] + tokenized_prompt["attention_mask"] + tokenized_options["attention_mask"]),
        "prompt": pre_context + context  + '\n Question: ' + row["prompt"] + options
    }


    example_encoded['label'] = option_to_index[row['answer']]
    
    return example_encoded

#  _ = [feature.pop("prompt") for feature in features] という処理でforward passで利用しない"prompt"を省く処理をCollate Functionに入れているとする
@dataclass
class DataCollatorForMultipleChoice:
    tokenizer: PreTrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = 2048
    pad_to_multiple_of: Optional[int] = None
    
    def __call__(self, features):
        labels = [feature.pop("label") for feature in features]
        _ = [feature.pop("prompt") for feature in features]
        batch_size = len(features)

        flattened_features = [[{k: v for k, v in feature.items()} for feature in features]]
        flattened_features = sum(flattened_features, [])
        batch = self.tokenizer.pad(flattened_features, padding=self.padding, max_length=self.max_length,pad_to_multiple_of=self.pad_to_multiple_of,return_tensors="pt",)
        batch = {k: v.view(batch_size, -1) for k, v in batch.items()} 
        batch["labels"] = torch.tensor(labels, dtype=torch.int64)
        return batch

# 通常のtorch.utils.data.DataLoaderでこのCollate Functionを使うのは問題ない。
train_dataloader = torch.utils.data.DataLoader(tokenized_val_dataset, batch_size=2, collate_fn=collate)

# しかし、TrainerでDataCollatorForMultipleChoiceを利用するとエラーになる!暗黙的にモデルが利用する['label', "input_ids", "attention_mask"]以外の要素(今回の場合"prompt")が削除されるため。
trainer = Trainer(
    model=lora_model,
    args=training_args,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_val_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForMultipleChoice(tokenizer=tokenizer),
    compute_metrics=compute_metrics,
)

#これを防ぐためには、TrainingArgumentsでremove_unused_columns=False を指定する。

training_args = TrainingArguments(
...
    remove_unused_columns=False,
...
)