[None][feat] share MNNVL MoE workspace checkpoint lifecycle - #1
Conversation
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
e064a65 to
a8bc80a
Compare
| class _MnnvlAlltoAllWorkspaceLifecycle: | ||
| """Own checkpoint and watchdog transitions for one shared MoE workspace. | ||
|
|
||
| The checkpoint coordinator must prevent new dispatches from starting while |
There was a problem hiding this comment.
I agree with keeping admission control outside this workspace owner, but could we make the precondition more explicit?
The caller must already have atomically stopped new engine/executor work and drained or aborted in-flight requests before calling this lifecycle. The local-client and rank-wide idle checks below are useful validation/defense-in-depth, but they do not themselves prevent a dispatch from starting after the idle vote.
I expect the authoritative gate to live in a companion top-level TRT-LLM checkpoint lifecycle, following the layering in Dynamo #12226: engine pause/admission first, resource hooks second, serving resume last. vLLM #46877 is also a useful reference for bubbling resource hooks to the engine, although it still relies on external orchestration for quiescence.
For the initial pre-serving Snapshot scenario discussed in #16635, the admission/drain implementation can initially be an assertion or no-op, but the top-level lifecycle seam should still exist. Could we document that dependency here and describe this idle vote as preflight rather than admission control?
| self._coordinator = self._create_coordinator() | ||
| torch.cuda.synchronize() | ||
| comm.barrier() | ||
| self._start_watchdog() |
There was a problem hiding this comment.
Could we add a final all-rank readiness agreement after the fallible local restore work and before publishing MAPPED?
The current barrier happens before _start_watchdog() and each client's _mnnvl_checkpoint_reset(). If either operation fails on one rank after the barrier, that rank enters BROKEN while its peers can call _checkpoint_restore_complete() and return MAPPED. The next collective can then run on only the successful ranks.
I suggest keeping every allocation in RESTORING, catching errors from the post-barrier local initialization, and exchanging a per-rank readiness result. If any rank reports failure, every rank should fail closed; only an all-rank success result should transition to MAPPED. Ideally, no fallible operation remains after that agreement.
The explicit RESTORING state added by this PR is a good improvement—we should use it to make publication rank-consistent.
| self._workspace_lifecycle = None | ||
| self._dispatch_state = {"phase": "destroyed"} | ||
|
|
||
| def __del__(self) -> None: |
There was a problem hiding this comment.
Should __del__() call destroy() here?
Destruction and cache eviction are process-local, but construction of a missing workspace key performs rank collectives. Python GC is not synchronized across ranks, so the last wrapper could be finalized and evict the cache on rank 0 while rank 1 still retains its cached workspace. A later synchronized construction would then take the cache-miss/allocation path only on rank 0, while rank 1 takes the cache-hit path and skips the matching collective.
I think explicit, all-rank destroy() should be the only path that evicts a collectively constructed workspace. If a finalizer is necessary, it should be limited to non-collective local bookkeeping and leave the symmetric workspace cached until coordinated teardown or process exit.
| raise RuntimeError(f"Cannot restore MNNVL allocation in {record.state.value} state") | ||
| comm_size = comm.Get_size() | ||
| comm_rank = comm.Get_rank() | ||
| if comm_size != record.comm_size or comm_rank != record.comm_rank: |
There was a problem hiding this comment.
Are rank and size sufficient to prove that the replacement communicator matches the graph-visible layout?
A communicator can have the same local rank and size but different participants or a different rank ordering. Handle exchange would then map the new group's rank slots into virtual addresses whose CUDA graphs and MoE routing still interpret them according to the original ordered group.
Could we either:
- persist and validate an ordered membership signature, such as stable world-rank IDs; or
- make “same ordered process group” an explicit precondition validated by the future engine-level checkpoint coordinator?
I would also avoid publishing the replacement communicator globally until the complete frontend restore has succeeded.
| if all(workspace is None or not workspace.mapped for workspace in workspaces): | ||
| MnnvlMoe.checkpoint_prepare() | ||
| return | ||
| local_clients_idle = not any(instance._dispatch_state for instance in self._INSTANCES) |
There was a problem hiding this comment.
Could the two-sided path adapt to the same workspace lifecycle owner rather than independently implementing instance registration, idle voting, idempotency, and client reset?
The resource shape differs—it has a primary and prepare allocation plus a different protocol initializer—but the lifecycle policy is the same. It seems like the common owner could accept a memory bundle and protocol-initialization callback:
- one-sided: one allocation, metadata validation, watchdog hooks;
- two-sided: two allocations,
moe_initialize_workspace, no watchdog.
That would keep rank readiness, failure publication, client registration, and restore ordering consistent between the two implementations. I view this as a design/follow-up suggestion rather than a request to put engine admission control in this owner.
There was a problem hiding this comment.
I think this can be deferred -- they share the low-level memory mapping, but own different allocation bundles and protocol initialization.
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
4f63b09 to
8fb01da
Compare
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
|
Thanks for keeping the shared-lifecycle work separate while it is under review. Could you rebase this PR onto the current NVIDIA#16632 head or merge it to NVIDIA#16632? |
|
@chienchunhung just did! |
83d5492
into
hannahz/dep-1082-port-flashinfer-stable-va-lifecycle-for-native-moe-all-to
3 participants
@coderabbitai summary
Description
Introduce a workspace-scoped lifecycle owner for native one-sided MNNVL MoE All-to-All allocations. It coordinates all local frontends sharing an allocation and performs a rank-wide readiness check before checkpoint teardown. During restore, it suspends and recreates the shared watchdog, refreshes metadata and coordinator state, resets each frontend’s protocol state, and publishes the allocation only after frontend initialization and rank synchronization.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.