AIが数学を解く能力向上

このチュートリアルでは、Tunix、JAX、LoRA、およびカスタム報酬関数を使用して、Gemma-3がGSM8Kの数学問題を推論するよう教えるエンドツーエンドのGRPOトレーニングワークフローを構築します。<br>まず、環境を準備し、Hugging Faceで認証し、Gemma-3モデルをロードし、GSM8Kの例を構造化された推論と最終的な数値回答の両方を必要とするプロンプト形式にラップします。<br>次に、形式の遵守と数学的正確さを評価する報酬関数を定義し、トレーニングを軽量に保つためにLoRAアダプターをアタッチし、ベースラインモデルを評価し、グループサンプリングされた生成を通じてポリシーを改善するためにGRPOを実行します。<br>これは、アダプターの重みのみをトレーニングし、ワークフローを単一アクセラレーターのセットアップに十分コンパクトに保つ強化学習チュートリアルを提供します。<br>TunixのインストールとGRPOトレーニングの設定 コードをコピーしました 別のブラウザを使用してください<br>import importlib.util, os, shutil as _sh<br>if importlib.util.find_spec("tunix") is None:<br> print("Installing Tunix + JAX ecosystem — this takes ~5-8 min…")<br> %pip install -q ipywidgets tensorboardX transformers grain nest_asyncio<br> %pip install -q datasets huggingface_hub "numpy>2"<br> %pip install -q tensorflow tensorflow_datasets<br> %pip install -q git+https://github.com/jax-ml/jax<br> %pip install -q git+https://github.com/google/tunix<br> %pip install -q git+https://github.com/google/qwix<br> %pip uninstall -q flax -y<br> %pip install -q git+https://github.com/google/flax<br> %pip uninstall -q wandb -y<br> print("<br><br> Install done. The runtime will RESTART now.")<br> print(" After it restarts, RUN THIS CELL AGAIN to start training.<br>")<br> os.kill(os.getpid(), 9)<br>import getpass, functools, json, re<br>import numpy as np<br>os.environ["WANDB_MODE"] = "disabled"<br>os.environ["TOKENIZERS_PARALLELISM"] = "false"<br>os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"<br>HF_TOKEN = os.environ.get("HF_TOKEN")<br>if not HF_TOKEN:<br> try:<br> from google.colab import userdata<br> HF_TOKEN = userdata.get("HF_TOKEN")<br> except Exception:<br> HF_TOKEN = getpass.getpass("Hugging Face token (needs Gemm")