Unverified Commit c726d44c authored by harrisonlimh's avatar harrisonlimh Committed by GitHub
Browse files

Recapture cuda graph after model weight update to resolve IMA error (#11780)

parent 283c8ba0
...@@ -981,6 +981,10 @@ class ModelRunner: ...@@ -981,6 +981,10 @@ class ModelRunner:
self.server_args.load_format = load_format self.server_args.load_format = load_format
self.load_config = load_config self.load_config = load_config
# Recapture device graph after model weight update.
if not self.server_args.disable_cuda_graph and self.device == "cuda":
self.init_device_graphs()
logger.info("Update weights end.") logger.info("Update weights end.")
return True, "Succeeded to update model weights." return True, "Succeeded to update model weights."
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment