[Feature] Integration SGLang into OpenRLHF

[zhaochenyang20]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

We've implemented the weight update API for RLHF pipeline:

api:

sglang/python/sglang/srt/server.py

Lines 214 to 239 in 21e9e63

	@app.post("/init_weights_update_group")
	async def init_weights_update_group(
	obj: InitWeightsUpdateGroupReqInput, request: Request
	):
	"""Initialize the parameter update group."""
	success, message = await tokenizer_manager.init_weights_update_group(obj, request)
	content = {"success": success, "message": message}
	if success:
	return ORJSONResponse(content, status_code=200)
	else:
	return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
	@app.post("/update_weights_from_distributed")
	async def update_weights_from_distributed(
	obj: UpdateWeightsFromDistributedReqInput, request: Request
	):
	"""Update model parameter from distributed online."""
	success, message = await tokenizer_manager.update_weights_from_distributed(
	obj, request
	)
	content = {"success": success, "message": message}
	if success:
	return ORJSONResponse(content, status_code=200)
	else:
	return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)

test case / usage:

https://github.com/sgl-project/sglang/blob/main/test/srt/test_update_weights_from_distributed.py

We will integrated SGLang into OpenRLHF this week. Here is the data for our accuracy and speed test.

https://huggingface.co/datasets/OpenRLHF/prompt-collection-v0.1-dev-rand5k

https://huggingface.co/datasets/OpenRLHF/prompt-collection-v0.1-dev-100k

Typically, 50K data requires several hours.

Related resources

See above.