Agent Lightning v1.0.3: 低精度RLの補正、MoEルーティングのメモリ削減

1. 低精度推論で報酬学習の安定化を試みる

今回のリリースはQwen3.5-9Bの数学学習例に、スコアセンタリングを追加した。回答生成モデルと学習モデルの確率差が生む、平均的な学習方向の偏りを差し引く補正だ。低精度で回答を作るモデルと学習中のモデルでは、次のトークンを選ぶ確率が異なり得る。スコアセンタリングは、この差によって学習の更新方向が偏る問題を狙った補正方法である。

この例では生成位置ごとに、トークンを抽出した時点の上位128候補の確率を保存し、現在の学習モデルの確率と照合して補正値を計算する。記録しなかった残りの候補の確率は、現在のモデルに比例すると近似する。

同社の発表によると、スコアセンタリングは学習後半の報酬を約0.6に維持した。応答グループの報酬を比較して学習する既存のGRPOを補正なしで使ったベースライン、つまり比較用の基準結果では、約45ステップ目に報酬が0に落ち、そのまま推移した。

この結果は、B200を1基使い、Qwen/Qwen3.5-9Bで数学を学習した例の学習曲線、つまり学習に伴う性能変化である。回答生成にはINT8の重み・活性値とINT4 KVキャッシュを使い、学習計算はBF16で行った。報酬はmath_verifyが最終回答を判定し、正解なら1、不正解なら0とし、途中で打ち切られた生成も0点とする。

2. 実行例では推論精度と学習精度を分ける

リリースにはモデル準備と精度検証のツール、学習重みをINT8推論モデルに同期する経路も含まれる。数値表現のビット数を減らす量子化の対象は、入力を重み付き和に変換するニューラルネットワークの線形層だ。トークンを数値ベクトルに変換した表現である埋め込みと、言語モデルの出力層はBF16のままとなる。設定は次のように推論部分に置く。

actor_rollout_ref:
  rollout:
    dtype: bfloat16
    quantization: int8
    quantization_config_file: ${oc.env:SC_INT8_CONFIG,outputs/precision_validation/Qwen3.5-9B-W8A8-INT8/config.json}
    engine_kwargs:
      vllm:
        linear_backend: triton
        attention_backend: TRITON_ATTN
        kv_cache_dtype: int4_per_token_head

dtypeは周辺の計算精度を定め、明示された量子化・キャッシュ設定がINT8とINT4の経路を選択する。この例の学習全体の設定は、プロンプト32件に対し各4件の応答、200回の更新である。簡易チェックではプロンプト2件に対し各4件の応答で、3ステップ実行する。

3. MoEルーティング再生は一時配列を順次処理

MoEは複数のエキスパートから一部を選び、トークンを処理するモデル構造だ。R3ルーティング再生は、回答生成時に選んだエキスパート情報を学習過程に渡し、同じ経路を再生する。以前はバッチ内のルーティング配列をすべてデコードしてから最終テンソルにコピーしていたが、今回の修正では一度に1つの配列だけをデコードし、コピー後に解放する。

改善提案者が公開したローカルCPUでの測定は、合成生成データ64件を処理した結果である。下表のように、ルーティング情報を展開する一時領域と、プロセスの最大常駐メモリ（RSS）がともに減った。

表：R3ルーティング再生のホストメモリ

条件：ローカルCPUベンチマーク、合成rollout、モデルが生成した応答グループ64件

単位：MB

削減幅は一時デコードバッファで最も大きく、最終的な学習テンソルは変更前後とも128 MBである。結果テンソルの形式とパディング・切り詰めの動作を保ちながら、中間処理の領域を減らした改善だ。

4. 削減の代償と適用範囲

低精度推論はメモリ・計算効率を狙うが、スコアセンタリングは候補確率の記録・受け渡しと、学習モデルの全語彙の確率計算を追加する。未記録の分布を近似するため、補正値には誤差の可能性が残る。この例のGRPOとスコアセンタリングは、グループ標準偏差による報酬正規化、つまり報酬をグループ内の変動幅で調整する方式を使う。論文の平均だけを使う設定や、勾配に沿って損失を下げる学習法である勾配降下法と同一の実験再現ではない。

リリースには報酬の小数値の保持と正規化精度の改善も含まれる。報酬を先にBF16に変換し、1.0と1.001が同じ値として処理されていた問題を修正した。正規化の分母のオーバーフロー、つまり表現範囲を超える数値エラーと、トークン数の丸めも防いだ。強化学習の更新に使う入力値が、低精度変換や大きな数の計算で失われる状況を減らす修正である。

今回のバージョンは低精度推論とMoEルーティング再生のボトルネックをそれぞれ狙い、メモリ削減は最終テンソルではなく中間バッファの処理から生まれる。