[TRTLLM-12589][fix] Reset MoE A2A dispatch state on warmup OOM#14000
Conversation
c0e4eb5 to
2e0d90b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #47758 [ run ] triggered by Bot. Commit: |
|
PR_Github #47758 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #47806 [ run ] triggered by Bot. Commit: |
Barry-Delaney
changed the title
|
PR_Github #47806 [ run ] completed with state
|
|
Waiting on #14028 for fixing CI. |
2e0d90b to
41ab775
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48083 [ run ] triggered by Bot. Commit: |
|
PR_Github #48083 [ run ] completed with state
|
0a93d10 to
118e7a5
Compare
41ab775 to
5c45ae7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48504 [ run ] triggered by Bot. Commit: |
|
PR_Github #48504 [ run ] completed with state
|
The MoE all-to-all dispatch/combine state machine (MoeAlltoAll._state.phase and NVLinkOneSided._dispatch_state["phase"]) is left in "dispatched" when a forward raises between dispatch() and combine(). _general_warmup_impl's "try / except OutOfMemoryError" catches the OOM and clears CUDA cache but does not reset Python-side stateful objects, so the next warmup config's forward fails the dispatch() invariant with "dispatch called twice without an intervening combine" This converts the "skip this warmup shape and try a smaller one" behavior the framework intended into "die at server startup with a confusing assert that has nothing to do with memory". Fix: * Add reset_state() to MoeAlltoAll and NVLinkOneSided. * In _general_warmup_impl's OOM handler, walk self.model.modules() and call reset_state() on any moe_a2a / comm submodule that exposes it, before retrying the next warmup config. Non-OOM exceptions are deliberately not covered: they escape the `except OutOfMemoryError` block, propagate through PyExecutor.__init__ and kill the server before any further dispatch() runs, so the stuck state machine never gets exercised again -- guarding that path would violate "don't guard against scenarios that can't happen". This restores the documented "skip OOMing config" semantics. The underlying 4096-token warmup OOM at kv_cache_free_gpu_memory_fraction=0.9 is a separate KV pool sizing issue and is not addressed here. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
5c45ae7 to
08828db
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48576 [ run ] triggered by Bot. Commit: |
|
PR_Github #48576 [ run ] completed with state |
…A#14000) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…A#14000) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
…A#14000) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
2 participants
The MoE all-to-all dispatch/combine state machine (MoeAlltoAll._state.phase and NVLinkOneSided._dispatch_state["phase"]) is left in "dispatched" when a forward raises between dispatch() and combine(). _general_warmup_impl's "try / except OutOfMemoryError" catches the OOM and clears CUDA cache but does not reset Python-side stateful objects, so the next warmup config's forward fails the dispatch() invariant with
"dispatch called twice without an intervening combine"
This converts the "skip this warmup shape and try a smaller one" behavior the framework intended into "die at server startup with a confusing assert that has nothing to do with memory".
Fix:
Non-OOM exceptions are deliberately not covered: they escape the
except OutOfMemoryErrorblock, propagate through PyExecutor.init and kill the server before any further dispatch() runs, so the stuck state machine never gets exercised again -- guarding that path would violate "don't guard against scenarios that can't happen".This restores the documented "skip OOMing config" semantics. The underlying 4096-token warmup OOM at kv_cache_free_gpu_memory_fraction=0.9 is a separate KV pool sizing issue and is not addressed here.