Fix reward/classification models broken by load_weights v2 dispatch (#28671)

[Jialin]

Motivation

#28671 refactored LlamaForCausalLM/Qwen2ForCausalLM.load_weights to dispatch to self._legacy_load_weights / self._load_weights_v2. The reward and classification wrappers (llama_reward, llama_classification, qwen2_rm, qwen2_classification) borrow the base loader via an unbound call Base.load_weights(self, ...) but inherit only nn.Module, so self has no _legacy_load_weights and weight load fails with AttributeError (default path, since v2 is opt-in). This breaks loading of every reward/classification model built on Llama or Qwen2 — e.g. LxzGordon/URM-LLaMa-3.1-8B crashes at startup with 'LlamaForSequenceClassificationWithNormal_Weights' object has no attribute '_legacy_load_weights'.
Call _legacy_load_weights directly, restoring the pre-#28671 behavior. These wrapped models are not AutoWeightLoader v2 citizens, so legacy is the correct path.

Modifications

• Changed the unbound base-loader borrow from Base.load_weights(self, ...) to Base._legacy_load_weights(self, ...) in the four affected wrappers: llama_reward.py, llama_classification.py, qwen2_rm.py, qwen2_classification.py.
• Added a one-line comment at each site explaining why the legacy loader is called directly (so it isn't "fixed" back to load_weights).
• gemma2 / internlm2 / mixtral-based wrappers are intentionally untouched: their base models were not converted to the v2 dispatch, so their load_weights is still self-contained and works.

Accuracy Tests

No output change — this restores the exact weight-loading path that ran before #28671 (the refactor renamed that body to _legacy_load_weights verbatim). The path is exercised by the existing reward/classification loading tests (test/registered/prefill_only/test_reward_models.py, test_score_engine.py, test_pooled_hidden_states.py), which fail on current trunk and pass with this change. black and py_compile verified locally.

Speed Tests and Profiling

N/A — weight-loading path only; no inference-speed impact.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. (Covered by existing reward/classification model tests.)
• Update documentation according to Write documentations. (N/A — bugfix only.)
• Provide accuracy and speed benchmark results. (N/A — behavior-preserving loader fix.)
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #29884063179
Latest PR Test (Extra): ⏳ Run #29884063091

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[b8zhong]

/rerun-test test/registered/prefill_only/test_reward_models.py

[github-actions]

Results for /rerun-test test/registered/prefill_only/test_reward_models.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/prefill_only/test_reward_models.py