CodeRabbitのすべてのレビューは、最初にルーティング判断から始まります。私たちのシステムはコード変更を評価し、モデルの出力をレビュー設定に変換します。ルーティングはCodeRabbitで最も実行回数の多いモデルタスクのひとつなので、より小さなモデルがレビュー工程の実際の一部を担えるかを試すのに適しています。
NVIDIA Nemotron 3.5 Lightningはコード関連の応答が優れているうえ、より小規模なインフラでも効率よく動かせるサイズだったため、候補として際立っていました。私たちは、事後学習によってCodeRabbitのルーティングルールを学習させられるかを検証したいと考えました。
そのために、CodeRabbitはNVIDIAおよびBasetenと協力し、2段階の事後学習を実施しました。
- 正しいルーティング判断の例を使った教師ありファインチューニング（SFT）
- CodeRabbitのルーティングポリシーで各回答を採点する、検証可能な報酬による強化学習（RLVR）
このプロジェクトでは、NVIDIA Nemotron 3.5 Lightningモデルや、NVIDIA NeMo AutoModelおよびNVIDIA NeMo RLを含む事後学習フレームワーク、BasetenのマネージドH100 Training JobsとA100 Dedicated Inference、そしてCodeRabbitのコードレビュー用データと評価基盤を組み合わせました。
結果
- Nemotron 3.5 Lightningの事後学習は非常に容易で、実験にかかった時間は3時間未満、コストは100ドル未満
- 事後学習後のモデルは、従来のGPTクラスのモデルよりも約4%高い精度を達成
- 推定推論コストも約50%削減
固定された1,000タスクの評価では、正確なルート一致率がGPTベースラインモデルの75.8%から、SFT後に80.4%、SFTとRLVR後に80.7%へ向上しました。Cohenのkappaで測定した出力一致度は、SFT後の0.461から0.544に達しました。
図1. 各候補モデルがコード変更に複雑度タグを付与し、CodeRabbitのスコアラーがそのタグをレビュー設定に変換
実験
まず、公開リポジトリから39,566件のサンプルを用意しました。学習データと評価データのリークを防ぐため、リポジトリ単位でデータを分割しています。フィルタリングと厳密なトークナイザー確認を行った後、主要なSFTセットは9,996件、固定評価データセットは1,000件になりました。
学習サンプルの作成には知識蒸留を用いています。より強力なモデルが、それぞれのコード変更に対する参照ルーティング判断を生成しました。その出力をフィルタリングし、残ったサンプルを使って、CodeRabbitがレビューをどうルーティングするかをNVIDIA Nemotron 3.5 Lightningに学習させました。
ステージ1: 蒸留による教師ありファインチューニング
最初に教師ありファインチューニングを行っています。BasetenのマネージドNVIDIA H100 Training Jobs上で、PyTorchネイティブのオープントレーニングライブラリであるNVIDIA NeMo AutoModelを使い、rank-8の軽量LoRAアダプターを1エポック学習しました。3つの学習率を試したところ、2e-4の試行が5つすべてのチェックポイントで最も低い検証損失を記録したため、そのstep-1,249アダプターを次の段階に選びました。
グラフ1. 教師ありファインチューニング中の学習損失と検証損失（低いほど良好）
固定された1,000タスクの評価では、教師ありファインチューニングによって正確なルート一致率が75.8%から80.4%に向上しています。また、Cohenのkappaで測定した出力一致度は0.429から0.461に向上しています。これが、報酬学習に進めるチェックポイントになりました。
教師ありチェックポイントを選択した後、最も難しい例での判断を改善するため、報酬学習へ進みました。
ステージ2: 検証可能な報酬による判断の改善
BasetenインフラのNVIDIA H100 GPU上で効率的な事後学習を行うため、NVIDIA NeMo RLオープンライブラリを使いました。ステージ1のLoRAをベースモデルの重みにマージしたうえで、ステージ1の調整済み出力にGRPOによるRLVRを適用しています。910個の選択済みプロンプトで学習し、25ステップごとに評価して、検証結果が頭打ちになったところで停止しました。
結果
Cohenのkappaは偶然を超えた一致度を測る指標で、高いほど良好です。ここでは、モデルの出力が参照出力にどれだけ近いかを示します。
SFTのみと比べると、SFTとRLVRの組み合わせはkappaを+0.0835、ルート一致率を+0.3パーセントポイント改善しました。ルート差分の対応信頼区間はゼロをまたいだため、このルート結果は統計的に決定的な改善ではなく、非劣化として扱っています。RLVRのより明確な成果は、出力一致度の向上です。
| 候補 | 出力一致度（Cohenのkappa） | 正確なルート一致率 |
| ベースラインモデル | 0.429 | 75.8% |
| Nemotron 3.5 Lightning + SFT | 0.461 | 80.4% |
| Nemotron 3.5 Lightning + SFT + RLVR | 0.544 | 80.7% |
最終モデルは1,000件すべてのリクエストを完了し、空の出力や異常に短い出力はありませんでした。これは明らかな崩壊がないことを示していますが、継続負荷や本番トラフィックでの確認は次の段階に残しています。
NVIDIA A100システムでの推論
Nemotron 3.5 Lightningの利点のひとつはコンパクトなサイズで、H100ではなくNVIDIA A100システムでもホストできることです。
サービングでは、Baseten Dedicated InferenceがSFTとRLVRを組み合わせた最終出力を、元のNVIDIA Nemotron 3.5 Lightningモデル上のrank-16 LoRAとして読み込みました。最終ポリシーは、1基の80 GiB A100上でvLLMを通じて動作しています。
測定されたスループットは、8つの同時リクエストで合計314.82出力トークン/秒でした。
ピーク時のサービングコスト
同じ評価ワークロードについて、OpenAIモデルとチューニング済みモデルの推論コストを比較したのが以下の内容です。
- OpenAIの公開価格では、総コストは2.34ドル
- Basetenの公開リソース表では、NVIDIA A100は1基あたり1分0.06667ドル。測定されたピーク合計レートでは、総コストは1.16ドル
- 推定削減額は1.18ドル、つまり50.4%削減
チューニング済みモデルは、このタスクにおいてベースラインモデルよりも出力トークン数が63.4%少なくなりました。
同じ1,000タスクのピークサービングコスト（飽和スループット時に50.4%の潜在的削減）
まとめ
この実験により、NVIDIA Nemotron 3.5 Lightningはファインチューニングにうまく応答することがわかりました。SFTは明確な品質向上をもたらし、RLVRはルート精度を維持しながら出力一致度を改善しています。NVIDIAのNemotron 3.5 Lightningモデルとトレーニングスタック、Basetenのマネージドトレーニングおよびサービングインフラを組み合わせることで、実験全体を現実的に進めることができています。この結果により、LightningはCodeRabbitのパイプライン内にある、範囲が狭く高頻度な他の処理にも有力な候補となりました。私たちはNVIDIAおよびBasetenとともに、こうした機会をさらに探っていくことを楽しみにしています。
