Unverified Commit 9069884b authored by Cheng Wan's avatar Cheng Wan Committed by GitHub
Browse files

[ci] disable memory imbalance check for draft worker (#8108)

parent 8a7a7770
......@@ -561,7 +561,7 @@ class ModelRunner:
# Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id)
if self.tp_size > 1:
if self.tp_size > 1 and not self.is_draft_worker:
if min_per_gpu_memory < local_gpu_memory * 0.9:
if get_bool_env_var("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK"):
logger.warning(
......
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