Optimum Neuronでresume_from_checkpointをするとTypeError None // int

先日の記事のように Optimun Neuronを使うとHuggingFace Trainerを簡単にAWS Neuron上で扱えます。

ですがまだあまり安定していないというのが正直なところです。今回の問題は0.0.23から0.0.25の範囲で確認済みです。記事に書いてしまっていますが、もし良ければちゃんと調査してIssueにしてください（すいません、時間が取れればやりたいと思ってます）。

Optimun Neuronでチェックポイントからの再開時にTypeError

Optimun Neuronでチェックポイントを残しつつ学習し、実際に resume_from_checkpoint=True をつけて再開すると以下の行で TypeError が発生します。

        train_bs_state = trainer_state.train_batch_size // max(1, training_args.n_gpu)

https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L1874

これは trainer_state.train_batch_size にNoneが設定されており、intとの比較でエラーとなっています。呼び出し元はこちら。

            self.compare_trainer_and_checkpoint_args(self.args, self.state)

https://github.com/huggingface/optimum-neuron/blob/v0.0.25/optimum/neuron/trainers.py#L900

Optimun Neuronの0.0.25ではNeuronTrainerの実装を見ても設定する処理がないようです。
https://github.com/huggingface/optimum-neuron/blob/v0.0.25/optimum/neuron/trainers.py

HuggingFace Trainerでは設定しています。
https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L2249

問題のとりあえずの回避方法

かなり雑な回避方法ですが、Optimun NeuronのNeuron Trainerでチェックポイントを今すぐ使いたい人は以下を参考にしてください（ちゃんと本家を修正しましょう。すいません）。この方法では読み込み時に、 training_state の値ではなく指定されたバッチサイズを使っています。

from optimum.neuron import NeuronTrainer

# Avoid a bug that NeuronTrainer does not store batch size to checkpoint
# And it cause None // int TypeError in compare_trainer_and_checkpoint_args.
class Trainer(NeuronTrainer):
    def compare_trainer_and_checkpoint_args(self, training_args, training_state):
        if training_state.train_batch_size is None:
            training_state.train_batch_size = self._train_batch_size
        super().compare_trainer_and_checkpoint_args(training_args, training_state)