[Feature] Distributed layerwise offload with DP multi-concurrency + mmap weight… (#5397)
Signed-off-by: evanchueng <evanchueng@outlook.com> Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: david6666666 <530634352@qq.com>
This commit is contained in:
1 parent
b3f4fbf9e2
commit
744c65b730
23 files changed
+3586
-78
No files matched your search
@@ -102,7 +102,10 @@ def _start_worker(req_q, res_q, count=2):
|
||||
|
||||
def _run():
|
||||
for _ in range(count):
|
||||
req = req_q.get(timeout=10)
|
||||
try:
|
||||
req = req_q.get(timeout=10)
|
||||
except queue.Empty:
|
||||
break
|
||||
method = req.get("method", "")
|
||||
args = req.get("args", ())
|
||||
if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput):
|
||||
@@ -146,7 +149,7 @@ def _inject_interleave(executor):
|
||||
orig_enqueue(item)
|
||||
if threading.current_thread().name == "thread_a":
|
||||
a_enqueued.set() # tell B: "A has enqueued"
|
||||
b_complete.wait(5) # block A until B finishes
|
||||
b_complete.wait(1) # block A until B finishes
|
||||
|
||||
executor._broadcast_mq.enqueue = _controlled
|
||||
return a_enqueued, b_complete
|
||||
@@ -173,8 +176,8 @@ class TestConcurrentRequestExecution:
|
||||
results["B"] = engine.add_req_and_wait_for_response(_mock_request("B"))
|
||||
b_complete.set() # release A
|
||||
|
||||
ta = threading.Thread(target=_a, name="thread_a")
|
||||
tb = threading.Thread(target=_b, name="thread_b")
|
||||
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
|
||||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||||
ta.start()
|
||||
tb.start()
|
||||
ta.join(10)
|
||||
@@ -274,8 +277,8 @@ class TestConcurrentCollectiveRpc:
|
||||
)
|
||||
b_complete.set()
|
||||
|
||||
ta = threading.Thread(target=_a, name="thread_a")
|
||||
tb = threading.Thread(target=_b, name="thread_b")
|
||||
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
|
||||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||||
ta.start()
|
||||
tb.start()
|
||||
ta.join(10)
|
||||
@@ -311,8 +314,8 @@ class TestConcurrentRequestExecutionAndCollectiveRpc:
|
||||
)
|
||||
b_complete.set()
|
||||
|
||||
ta = threading.Thread(target=_a, name="thread_a")
|
||||
tb = threading.Thread(target=_b, name="thread_b")
|
||||
ta = threading.Thread(target=_a, name="thread_a", daemon=True)
|
||||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||||
ta.start()
|
||||
tb.start()
|
||||
ta.join(10)
|
||||
@@ -683,7 +686,10 @@ class TestMultiprocExecutorStepStreamingOutput:
|
||||
|
||||
def _worker():
|
||||
for runner_output in runner_outputs:
|
||||
req_q.get(timeout=10)
|
||||
try:
|
||||
req_q.get(timeout=10)
|
||||
except queue.Empty:
|
||||
break
|
||||
res_q.put(runner_output)
|
||||
|
||||
thread = threading.Thread(target=_worker, daemon=True)
|
||||
|
||||
Reference in new issue
Block a user