Add update_weights_from_disk endpoint to Engine (#4102)
Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user