Add update_weights_from_disk endpoint to Engine (#4102)

Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
This commit is contained in:
Jhin
2025-03-05 14:25:18 -06:00
committed by GitHub
parent ef9d3b3c2c
commit 70b3c6eeb1
3 changed files with 239 additions and 31 deletions

View File

@@ -44,6 +44,7 @@ from sglang.srt.managers.io_struct import (
InitWeightsUpdateGroupReqInput,
ReleaseMemoryOccupationReqInput,
ResumeMemoryOccupationReqInput,
UpdateWeightFromDiskReqInput,
UpdateWeightsFromDistributedReqInput,
UpdateWeightsFromTensorReqInput,
)
@@ -302,6 +303,27 @@ class Engine:
self.tokenizer_manager.update_weights_from_tensor(obj, None)
)
def update_weights_from_disk(
self,
model_path: str,
load_format: Optional[str] = None,
):
"""Update the weights from disk inplace without re-launching the engine.
This method allows updating the model weights from disk without restarting
the engine. It can be used to load a different model or update weights with
new training.
"""
obj = UpdateWeightFromDiskReqInput(
model_path=model_path,
load_format=load_format,
)
loop = asyncio.get_event_loop()
return loop.run_until_complete(
self.tokenizer_manager.update_weights_from_disk(obj, None)
)
def get_weights_by_name(self, name: str, truncate_size: int = 100):
"""Get weights by parameter name."""
obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size)