大規模モデルでdf.dropはまずい：Out of Memoryを解決した話

目次

• 初めに
• 学び
• OOMを解決した話
• OOMとは
• OOMが起きた場所
• OOMが起きた原因
• 対処法：ラベルベースのインデクシング
• 感想

初めに

2月に入り、手動正常性確認をパイプラインに組み込んでから初めてのモデル実行があった。

Slackに飛んできたのは「contact_rate Failed」の赤い通知。

原因を探るべく、私はアマゾン(AWS)の奥地へ向かった...

学び

• 大規模データをdropすると中間テーブルが生成されて瞬間的にメモリ使用量が増える
• ラベルベースのインデクシングで対処可能。カラムを先に指定しlist型で保持しておく
• int64は当然64bitなので無駄が大きい。ほぼint16以下で事足りる

OOMを解決した話

エラーの原因は137: OOM。
機械学習ではよくあるエラーだった。

重い処理が残りがちなモデルなのでちょくちょくリファクタリングをしているものの、まだ足りていなかったようです。

• OOMとは何か
• OOMの裏側で起きていること
• OOMが起きた箇所はどこか、なんで起きたのか
• 対処法

の順番で言語化していきます。

OOMとは

OOM = Out of Memory
その名の通り、使えるメモリ領域を超えたデータが降ってきてしまい、処理が継続できなくなった状態で起きるエラーです。

メモリはよく作業机に、データは書類などに例えられます。
通常書類はメモリ上に置かれてCPUがガンガン処理していくのですが、あまりに大きい書類がドカッと来てしまうと作業台に乗らなくなることがあります。

その時に出るエラーがOOMというわけです。

これに対処するには主に書類サイズを小さくすることを考えます。
無駄に書類をのっけていないか、のっけた書類を早期にどかせないか、載せる書類を小さくできないか、を考えて、それでも無理そうなら書類を分割して無理やり小さくするなどの対処法があります。

今回はそもそも無駄な書類をのっけていたので、そこに着手しています。

OOMが起きた場所

OOMが起きた場所は推論フェーズ。

本モデルは実績ありユーザーで学習して実績なしユーザーを推論するコールドスタートモデル。
実績なしユーザーの総数は1億以上。

推論テーブルは以前の記事で触れたように、
重い処理を回さないためにシングルトンを作って同じものを繰り返し使うようになっています。
https://zenn.dev/zenn_mita/articles/394dfe5e8703ef

ただ、シングルトンにするということはずっとメモリ上に保持しているということ。
サイズはおよそ30GB。無視できない大きさでした。

下記のような推論処理でエラーになっていそうでした。

error.py

def predict():
    transformed_input_df = input_df.drop(["acctid"], axis=1)
    transformed_input_df = transformed_input_df.astype("int")

OOMが起きた原因

上記の短い処理の中でもメモリを増やしてしまっている処理は3か所もありました。
それが「dropを使用していること」「別dfに格納していること」「astypeでint型に変更していること」の3点。

それぞれの問題点とそれがあることによって引き起こされる結果が下記のようになります。

■dropを使用していること

• 問題
  • dropをすると指定カラム以外がある中間テーブルが生まれてしまう
• 結果
  • 右辺のdfとほとんど同じ容量のdfができてしまい、容量が2倍になる

■別dfに格納している

• 問題
  • 使わないdfも残ってしまい、容量が増えたままになってしまう
• 結果
  • 後続処理で使えるメモリ量が減る

■astypeでint型に変更していること

• 問題
  • intを指定するとint64が選択されてしまう
• 結果
  • ワンホットベクトルのような1bitあれば十分なカラムに対しても8byte容量がかかるようになる

上記3つが重なって、100GB以上余裕のあったメモリがあっという間になくなったというわけでした。

対処法：ラベルベースのインデクシング

対処法は簡単で、dropを使わずラベルベースのインデクシングを使用すること、別dfを作らない or 使わないdfをすぐgc.collect()する、int16のような軽量な型を指定する、などなどでした。

今後頻繁に使いそうなのがラベルベースのインデクシング。
下記のような処理のことを指しています。

label_based_indexing.py

cols = [col for col in df.columns if col != "acctid"]
model.predict(input_df[cols])
...

簡単に言語化すると、欲しいカラムを先にリストで保持しておき、そのリストを使用してdfを絞り込むという処理をしています。
こうすると中間テーブルを作成せずに、今あるテーブルのみから欲しいテーブルが作れるのでメモリ削減に使われるということです。

注意点はdf = df[cols]のような代入をしてしまうと、たとえ右辺と左辺が同じdfだとしても一瞬中間テーブルが作成されてしまうので、OOMの可能性が上がってしまうことでしょうか。

メモリに余裕がある、かつ後続で頻繁に絞り込み後のdfを使用する場合など以外はインデクシングを使ったほうがいいかなと考えています。

感想

今回は処理の正当性を確保しつつ早く修正する必要があったためint32で型変換しています。
99%の確率でint16でもいいと思いますが、中間テーブル作り直しだったため、
以前取り組んだコスト削減が効かず万が一再々実行になったらまずいなと。
https://zenn.dev/zenn_mita/articles/bd54d81c4dab31

本当は型変換などせず取ってきたデータをそのまま使用できればいいと考えているので、
今後はテーブル定義を見直していく予定です。

ただ、源泉データから引っ張ってくるときすでにint64だったらどうしようもないな、と思っていたり。

自分にできることを最大限やりつつ、設計から関わるときは型を意識したものを作れるよう頑張ります。