見出し画像

Codex MobileからGemma4 E2B 2048に挑んだら、Colab CLIのkeep-alive問題まで見えた

2048を解くLLMを育てるつもりで、Unsloth公式のGemma4 E2B 2048 GRPO実験をColabで回しました。

操作はいつものように、Codex MobileからMac上のCodexへ指示を出し、Google Colab CLIやColabブラウザUIを動かす形です。スマホから実験を指示し、ColabのT4やL4でUnsloth公式Notebookを走らせ、学習前後で2048戦略関数がどう変わるのかをログで追う。やりたかったことは、そこでした。

前回までは、Codex MobileからGoogle Colab CLIでT4を確保し、画像生成やUnslothのファインチューニングを試していました。

今回はその続きです。

ただし、結論から言うと、今回の成果は「2048が上達した」ではありません。

本当に価値があったのは、その手前にある問題を分解できたことでした。Unsloth公式Notebookの学習経路、実行環境に合わせた暫定パッチ、公式Notebook内の60 step / 600 stepの食い違い、Colab CLIのセッション維持、Google Drive保存の扱い。このあたりを、ログを見ながら一つずつ切り分けました。

きっかけは2048の公式GRPO実験

対象にしたのは、UnslothのGemma4 E2B Reinforcement Learning 2048 Game Notebookです。

https://raw.githubusercontent.com/unslothai/notebooks/main/nb/Gemma4_(E2B)_Reinforcement_Learning_2048_Game.ipynb

Notebookの目的は、Gemma4 E2Bに2048の戦略関数を出力させ、GRPOで報酬を与えながら改善していくことです。

モデルは2048の盤面を受け取り、W, A, S, D のどれかを返すPython関数を書くよう求められます。報酬関数では、関数として動くか、禁止されたimportを使っていないか、実際にゲームを進められるか、2048に到達できるかが見られます。

報酬はざっくり次の構造でした。

success: 20.0
failed but function works: 2.0
timeout: -1.0
exception: -3.0

最初に想定していた流れはシンプルです。

  • 公式Notebookの経路を確認する

  • Colab T4でモデルをロードする

  • GRPOを回す

  • LoRA adapterを保存する

  • before / after評価を取る

  • 2048戦略関数のvalid率やtimeout率を可視化する

ところが、実際に動かしてみると、学習以前にいくつかの互換性問題と運用上の落とし穴が見えてきました。

学習経路は通った。ただし上達とは別問題

最初の山は、Unsloth側のGemma4 loaderと、TRL / PEFT周りの互換性でした。

公式Notebookの経路をそのまま辿ろうとすると、最初はGemma4 loader側で止まりました。

原因は、Unsloth Zoo側のGemma4 temporary patchでした。

_Gemma4KVSharedSafeProxy.__getattr__num_kv_shared_layers に対して AttributeError を返す。その状態でTransformers 5.9.0側のconfig生成や __repr__ / to_diff_dict が走ると、config構築の途中で落ちる。

ここで当てた暫定パッチは、これです。

import unsloth_zoo.temporary_patches.gemma4 as gemma4_patch

proxy_cls = gemma4_patch._Gemma4KVSharedSafeProxy
original_getattr = proxy_cls.__getattr__

def patched_getattr(self, name):
    if name == "num_kv_shared_layers":
        return 0
    return original_getattr(self, name)

proxy_cls.__getattr__ = patched_getattr
proxy_cls._codex_num_kv_patch = True

変更点はかなり限定的です。

num_kv_shared_layers が存在しない扱いで例外になる箇所にだけ、0 を返すようにしました。

これで、Gemma4のconfig生成が先へ進み、公式Notebook由来のモデルロード経路を検証できるようになりました。

ただし、これは恒久修正ではありません。公式Notebookの2048 GRPO経路を、現在のColab環境でどこまで通せるかを確認するための検証用パッチです。

公式Notebookでは beta が明示されていません。一方、今回の環境では現在のTRL側の挙動により、参照モデル側のadapter経路に入り、Unsloth側の Gemma4ClippableLinear にPEFT adapterを挿そうとして落ちました。

エラーはこの形です。

ValueError: Target module Gemma4ClippableLinear is not supported.

ここで入れた変更は、GRPOConfigbeta=0.0 を明示することです。

grpo_kwargs = dict(
    max_steps=MAX_STEPS,
    max_completion_length=4094,
    logging_steps=1,
    output_dir=str(OUT_DIR),
)

grpo_kwargs["beta"] = 0.0
training_args = GRPOConfig(**grpo_kwargs)

これで、TRLがref adapter経路に入り、PEFTが Gemma4ClippableLinear にadapterを挿そうとして落ちる経路を避けました。

ここまでで入れた重要な変更は2つです。

  • num_kv_shared_layers だけ 0 を返すGemma4 loader暫定パッチ

  • TRL / PEFTのref adapter経路に入らないようにする beta=0.0 の明示

この2つで、ようやくGRPOのtrain開始まで進めました。

この時点で、T4でもモデルロード、GRPOのtrain開始、checkpoint保存、LoRA adapter保存までは到達しています。実測上もT4のVRAM内に収まっており、少なくとも「T4だからモデルが載らない」という見方は主因ではありませんでした。

最初の確認では、T4で train 1 step、LoRA adapter保存、成果物パッケージングまで到達しました。before evalでは、モデルが関数コードではなく説明文を出しており、rewardは missing function で0.0でした。

その後、公式パラメータ寄りの設定でも確認し、T4では checkpoint-1、L4では checkpoint-3 まで進みました。VRAM/RAMのログを見る限り、停止原因を単純なメモリ不足とは言えませんでした。

その後、ColabブラウザUI側のT4実行では、60 stepの学習完走まで確認しました。

global_step=60
train_runtime=20642.6487 sec
adapter_saved

約5時間44分です。

ここだけを見ると、「学習できた」と言いたくなります。けれど、実験としてはここで止めるべきではありませんでした。

止まっているのか、計算中なのかを見えるようにした

途中で大きな問題になったのが、「止まっているのか、計算中なのか分からない」ことでした。

そこで、30秒ごとにheartbeatを出すようにし、処理段階ごとにログを残しました。

  • model load

  • before eval

  • generation

  • trainer.train

  • adapter save

  • artifact package

各段階で、GPU名、VRAM使用量、RAM使用量、nvidia-smi、経過秒を記録しました。あわせて、次のJSONLを保存対象にしました。

  • progress.jsonl

  • reward_events.jsonl

  • eval_samples.jsonl

これで、単に「落ちた」「固まった」と言うのではなく、どの段階で、どのリソース状態で、どのログまで進んでいたかを追えるようになりました。

60 step完走後の報酬ログは、かなり厳しい結果でした。

reward rows: 120
-1.0 timeout: 82
-3.0 exception: 30
2.0 failed but function works: 8
20.0 success: 0

関数らしきものは出ています。しかし、多くはtimeoutか例外です。関数として動いたものも、ゲーム状態はfailedで、2048成功には届いていません。

この時点で言えるのは、ここまでです。

  • 学習処理としての trainer.train() は完走した

  • LoRA adapterは保存できた

  • reward logは取れた

  • ただし、2048戦略が上達したとは言えない

trainer.train() が終わることと、目的タスクが改善することは別です。特に今回のようなGRPO実験では、報酬が上がっているか、validな関数の比率が増えているか、timeoutやexceptionが減っているかを見ないと、成功とは言えません。

公式Notebookのステップ数を読み違えていた

今回もっとも大きな修正点は、学習ステップ数の読み方でした。

公式Notebookの現在の設定セルには、max_steps = 60 と書かれていました。そのため最初は、この60 stepを公式設定として扱っていました。

しかし、Notebook内に保存されている実行出力を見ると、別の数字が出てきます。

Total steps = 600

進捗バーも 434/600 付近まで進んでいました。

さらにNotebook本文には、次の注意書きがあります。

You might have to wait 150 to 200 steps for any action.
You'll probably get 0 reward for the first 100 steps.
Please be patient!

現在セルの max_steps=60 だけを見て「公式通り」と判断するのは不十分でした。

公式Notebookの保存済み出力は、少なくとも600 step系の実行を示しています。しかもNotebook自身が、最初の100 stepは報酬0かもしれない、150から200 step待つ必要があるかもしれない、と書いています。

60 step完走は、経路確認としては意味があります。

でも、公式Notebookの成功例と同じ土俵に立ったとは言えません。

現在セルだけを見て、保存済み出力、進捗バー、本文の警告まで合わせて読めていなかった。ここは今回の検証で明確に直すべき点でした。

以後の再現では、Notebookの現在セルだけでなく、保存済み出力、Total steps、進捗バー、Markdownの注意書き、成功例のstep数を照合してから学習長を決めます。

セッション消失の正体はkeep-aliveだった

もう一つ大きかったのが、Colab CLIのセッション消失です。

途中で何度も、CLI側では次の状態になりました。

No active sessions found on server.

これだけを見ると、Colab側のランタイムが落ちたのか、GPUが足りないのか、RAMが足りないのか、自分で止めたのかが分かりません。

そこでログを追うと、Colab CLIのkeep-aliveが失敗していることが分かりました。

Colab CLIは、作成したランタイムを維持するために、裏でkeep-alive RPCを叩いています。エンドポイントはこれです。

POST https://colab.pa.googleapis.com/$rpc/google.internal.colab.v1.RuntimeService/KeepAliveAssignment

このkeep-aliveが、繰り返し403で失敗していました。

403 Forbidden
USER_PROJECT_DENIED
Caller does not have required permission to use project 1014160490159.
Grant the caller roles/serviceusage.serviceUsageConsumer or serviceusage.services.use.

この時点で、少なくとも「VRAM不足で落ちた」とは言えません。

学習コードが落ちたのでも、GPUにモデルが載らなかったのでもありません。CLIがランタイムを維持するkeep-alive経路で弾かれていました。

ヘッダを外しても直らない

google-colab-cli 0.5.9 の実装も確認しました。

keep_alive_assignment の中で、次のヘッダが固定送信されていました。

x-goog-user-project: 1014160490159

では、これを外せば直るのか。

ここでは、原因を切り分けるための診断パッチを当てました。

ただし、これは「修正パッチ」ではありません。原因を切り分けるために、x-goog-user-project ヘッダを送らない場合にエラーがどう変わるかを見るためのA/Bパッチです。

結果はこうでした。

400 CONSUMER_INVALID
The API Key and the authentication credential are from different projects.

単純に x-goog-user-project を消せば直る、という話ではありません。

固定ヘッダありでは、Colab内部project 1014160490159 を使う権限がないと言われる。固定ヘッダなしでは、API keyと認証credentialのprojectが違うと言われる。

このA/Bで見えたのは、Colab CLIのpublic API key側のprojectと、OAuth credential側のconsumer project / 権限の噛み合わせが問題になっている、ということです。

同じGoogleアカウントでも、ColabブラウザUIではDrive mountが成功しています。

Mounted at /content/drive

しかし、CLIのkeep-aliveだけは失敗する。

ここが今回の本質でした。

上流issueにも追加情報を残した

この問題は、こちらの環境だけではありませんでした。

googlecolab/google-colab-cli のissue #14に、同じような報告がありました。

そこでも、KeepAliveAssignment403project 1014160490159serviceusage.serviceUsageConsumer が出ています。同じGoogleアカウントでColab Web UIは使えるのに、CLIだけkeep-aliveで失敗する、という点も似ています。

そこで、こちらで確認したA/B結果もコメントしました。

コメントした内容は、主に次の点です。

  • google-colab-cli 0.5.9

  • OAuth scopes

  • Colab Web UIでは同じアカウントでDrive mount成功

  • CLIでは KeepAliveAssignment403 USER_PROJECT_DENIED

  • x-goog-user-project: 1014160490159 固定送信を確認

  • 固定ヘッダを外すと 400 CONSUMER_INVALID に変化

  • 単純な診断パッチでは解決しない

この時点では、手元で直ったとは言えません。上流側で確認する価値のある再現情報として残しました。

Driveを本線にする

もう一つ、運用としてはっきりさせた点があります。

Colabで長時間学習を回すなら、成果物はColab内でGoogle Driveに直接保存するべきです。

手元へdownloadして、またuploadする経路は時間がかかります。何より、Colab実験の自然な運用から外れます。

今回、一部の未マウント /content 成果物は緊急回収しました。ただし、それを通常経路にはしません。

今後のルールは明確です。

  • checkpoint

  • LoRA adapter

  • progress.jsonl

  • reward_events.jsonl

  • eval_samples.jsonl

  • metadata

こうした長時間実験の成果物は、Colab内で /content/drive/MyDrive/... に直接保存する。

Drive mountが通っていないなら、長時間学習に進まない。

これは単なる好みではなく、再現性と作業時間の問題です。

今回わかったこと

今回成功したことはあります。

  • Unsloth公式Notebookのタスク構造を確認した

  • Unsloth側のGemma4 loader問題を暫定パッチで突破した

  • Gemma4 loader / TRL beta / PEFT adapter問題を切り分けた

  • beta=0.0 で学習開始に進めることを確認した

  • 30秒heartbeatと段階別ログで、処理中か停止かを見分けられるようにした

  • T4で1 step、L4でcheckpoint-3までの途中経路を確認した

  • T4でもVRAM/RAM不足が主因ではないことを確認した

  • T4ブラウザUIで60 step完走した

  • adapter保存とreward log取得まで到達した

  • 公式Notebookの60 step / 600 stepの食い違いを確認した

  • Colab CLIの固定 x-goog-user-project ヘッダを診断パッチで外し、403から400へ変わることを確認した

  • CLI keep-alive 403の原因をログとソースから追った

  • 上流issueに追加証拠をコメントした

一方で、未達もはっきりしています。

  • 2048戦略が上達したとはまだ言えない

  • 60 stepは公式成功例と同等ではない

  • 600 step級の公式保存出力再現はまだできていない

  • Colab CLI keep-alive問題は診断パッチだけでは解決していない

  • CLI経由の長時間実行はまだ信頼できない

  • Unsloth loader側の暫定パッチも、上流互換性まで解決したものではない

ここを混ぜると、記事としても実験としても危ない。

今回の成果は、「2048に勝つAIができた」ではありません。

今回の成果は、そこへ行くために必要な実行経路の問題をかなり潰したことです。

次に試すこと

次は、Colab CLIのkeep-aliveに頼らず、ColabブラウザUIで実行します。

そのうえで、最初にDrive mountを確認します。

from google.colab import drive
drive.mount('/content/drive')

出力先は、最初からMyDrive配下にします。

/content/drive/MyDrive/...

学習ステップ数は、現在セルの max_steps=60 だけでは決めません。公式Notebookの保存済み出力や注意書きを踏まえて決めます。

現実的には、まず200 step前後で報酬が動き始めるかを見る。その後、公式保存出力に合わせるなら600 step級で再現する。

評価では、単にadapterが保存されたかではなく、次を見ます。

  • validな def strategy(board): の比率

  • timeout率

  • exception率

  • failed but function works の比率

  • success 報酬の有無

  • before / afterで同じ評価条件を使った比較

ここまで見て、初めて「2048が少しうまくなったか」を話せます。

再現用コードとNotebook

今回の元になったNotebookは、Unsloth公式のこちらです。

https://raw.githubusercontent.com/unslothai/notebooks/main/nb/Gemma4_(E2B)_Reinforcement_Learning_2048_Game.ipynb

ただし、今回そのまま公開済みのfork Notebookを作って実行したわけではありません。

実際の実行は、公式Notebookを取得し、Colab CLIやColabブラウザUI上でPythonスクリプトとして実行・診断する形でした。

再現用コードとパッチの詳細は、公開gistにまとめています。

中核になったのは、gist内の gemma4_e2b_2048_official_diagnose.py です。

このコードは、次の役割を持たせました。

  • Unsloth公式NotebookをURLから取得する

  • 公式NotebookのセルをPython側で順番に実行する

  • Gemma4 loader問題を回避する暫定パッチを当てる

  • beta=0.0 を明示して、TRL / PEFTのref adapter経路を避ける

  • 30秒前後のheartbeatで、model load / eval / generation / train / saveの進捗を出す

  • GPU名、VRAM、RAM、nvidia-smi、経過秒を progress.jsonl に残す

  • reward計算のイベントを reward_events.jsonl に残す

  • 評価サンプルを eval_samples.jsonl に残す

  • adapter、checkpoint、ログを成果物として保存する

  • Drive mount済みなら /content/drive/MyDrive/... を優先し、未マウント時の /content 出力は緊急回収扱いにする

重要な変更は2つです。

  • Gemma4 loader側では、num_kv_shared_layers だけ 0 を返すmonkey patchを入れた

  • GRPO側では、beta=0.0 を明示してref adapter経路を避けた

実コード全体はgistに置いています。抜粋だけでなく、runner全体を確認できます。

再現する場合は、ColabでGPUランタイムを選び、Driveをmountしてから、gistの gemma4_e2b_2048_official_diagnose.py を置いて実行します。

from google.colab import drive
drive.mount("/content/drive")
import os, subprocess, sys

os.environ["RUN_NAME"] = "gemma4-e2b-2048-repro"
os.environ["MAX_STEPS"] = "60"
os.environ["EVAL_SAMPLES"] = "0"
os.environ["HEARTBEAT_SECONDS"] = "30"
os.environ["STOP_AFTER_TRAIN_ARTIFACT"] = "1"
os.environ["OMIT_BETA"] = "0"

subprocess.check_call([sys.executable, "gemma4_e2b_2048_official_diagnose.py"])

実行部分の要点は、次のような構成です。

from google.colab import drive
drive.mount("/content/drive")

# 公式Notebook由来のGRPO設定を使う
# ただし現行TRL/PEFT経路ではref adapter問題を踏むため beta=0.0 を明示
training_args = GRPOConfig(
    max_steps=60,
    max_completion_length=4094,
    beta=0.0,
    logging_steps=1,
    output_dir="/content/drive/MyDrive/...",
)

# 進捗とリソースをJSONLへ残す
log_event("trainer_train_start", gpu=gpu_snapshot())
trainer.train()
log_event("trainer_train_done", gpu=gpu_snapshot())

model.save_pretrained("/content/drive/MyDrive/.../adapter")
tokenizer.save_pretrained("/content/drive/MyDrive/.../adapter")

これは要点の抜粋です。実際のコードでは、heartbeat、例外時のtraceback保存、nvidia-smi、成果物tar化、reward/evalログ保存まで入れています。

以後は、Drive mountが確認できていない状態では長時間学習に進まない方針にします。未マウント /content からの手元回収は、すでに生成してしまった成果物の緊急回収に限定します。

また、Colab CLIのkeep-alive原因調査では、google-colab-cli 0.5.9 を読み、keep_alive_assignmentx-goog-user-project: 1014160490159 を固定送信していることを確認しました。そのうえで、診断パッチとしてこのヘッダを外し、エラーが 403 USER_PROJECT_DENIED から 400 CONSUMER_INVALID に変わることを確認しました。

今回使ったコードは、大きく3系統です。

  • 公式Unsloth Notebook

  • 公式Notebookを実験用に動かす再現用Python runner

  • Colab CLI keep-alive原因を切り分ける診断パッチ

この検証をそのまま開けるColabノートブックは、まだ公開していません。公開しているのは、元のUnsloth公式Notebookと、再現用runnerを置いたgistです。

まとめ

Codex MobileからGoogle Colab CLIを使い、Unsloth公式Gemma4 E2B 2048 GRPOに挑みました。

学習経路としては、Unsloth側のGemma4 loader問題を暫定パッチで突破し、beta=0.0 でTRL / PEFTのref adapter互換問題を回避したうえで、T4でもモデルロード、GRPO実行、60 step完走、adapter保存、reward log取得まで到達しました。

ただし、それは2048戦略が上達したことを意味しません。

公式Notebookの保存済み出力は600 step系で、本文にも150から200 step待つ必要があるかもしれないと書かれていました。60 step完走は、公式成功例の再現ではなく、経路確認に近い結果です。

さらに、Colab CLIのセッション消失は、VRAM不足やRAM不足ではなく、KeepAliveAssignment403 USER_PROJECT_DENIED が主な原因として見えてきました。

x-goog-user-project: 1014160490159 を外す診断パッチを当てても 400 CONSUMER_INVALID になるため、単純なパッチでは解決しません。この点は上流issueにもコメント済みです。

今回いちばん価値があったのは、「できた」「落ちた」を雑に言わず、どこまで通って、どこから先が未達なのかを分解できたことでした。

次は、ColabブラウザUI + Drive mountを前提に、200 step以上、できれば600 step級の再現へ進めます。

評価はadapter保存ではなく、2048戦略関数のvalid率と報酬分布で見る。

ここからが本番です。

いいなと思ったら応援しよう!