Support sanity checking weight consistency especially for RL

[fzyzcjy]

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @fzyzcjy, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a robust mechanism for verifying the consistency of model weights within the system. It provides a new API endpoint that allows users to perform sanity checks on the model's parameters, enabling operations like snapshotting the current weights, resetting them to random values, and comparing the current state against a previously saved snapshot. This functionality is particularly valuable for debugging and ensuring stability in dynamic model environments, such as those found in Reinforcement Learning.

Highlights

• New API Endpoint for Weight Checks: Introduced a "/check_weights" HTTP endpoint to trigger weight consistency checks, enabling external systems to initiate these operations.
• WeightChecker Utility Class: Added a new WeightChecker class that encapsulates the logic for snapshotting, resetting, and comparing model weights, centralizing this functionality.
• Integrated Weight Management: The ModelRunner now utilizes the WeightChecker to perform these operations, and the Scheduler and TokenizerManager are updated to expose this functionality throughout the system.
• New Data Structures for Requests: Defined CheckWeightsReqInput and CheckWeightsReqOutput for structured communication regarding weight checks, ensuring clear data exchange.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[fzyzcjy]

/tag-and-rerun-ci

[gemini-code-assist]

Code Review

This pull request introduces a valuable feature for sanity checking model weight consistency, which is particularly useful in Reinforcement Learning workflows. The implementation is well-structured, adding a new WeightChecker utility, exposing it through a new /check_weights HTTP endpoint, and wiring it through the TokenizerManager and Scheduler. The code is clear and follows existing patterns in the codebase. I have one suggestion to fix a minor bug in the random tensor generation logic.

[slin1237]

do we have to add an endpoint for this?
can this be part of model info endpoint?
such that we don't add a new resource/endpoint, likewise the verb in front of endpoint
Existing SGLang frontend has too many of those endpoints

[fzyzcjy]

it needs to mutate state for the checker, i.e. the json payload can be snapshot,reset_params,check. while I guess model info endpoint is purely read w/o any mutations.

[slin1237]

ref to comments

[fzyzcjy]

/tag-and-rerun-ci