Repository navigation
[Bugfix][EPLB] Share one NIXL agent per worker across elastic EP reconfigurations, enable MI355 scaling test - #63134
stefankoncarevic wants to merge 16 commits into
Conversation
Each NixlEplbCommunicator built its own NIXL agent and registered the expert weights with it again. A scale up builds the new communicator while the old one is still alive, so a worker ended up holding two and then three agents over the identical pages. Whichever agent was collected first deregistered memory the others could still be reading, and on gfx950 the second registration of the same pages is refused outright with HSA_STATUS_ERROR_INVALID_ARGUMENT. Key the agent by (device, arena start, arena end) and reuse it, so each page is registered once and deregistered only when expert parallelism itself goes away. Also guard __del__, which can run after the module globals its body needs are already None. Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
Mirrors the H100 nightly step onto mi355_4 upstream, and adds an MI355 step next to the existing MI300 one internally. Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ci run |
|
❌ @stefankoncarevic, A reviewer with write access must run |
stefankoncarevic
added
the
ready
|
/ci run |
|
✅ @stefankoncarevic, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #96664 for commit |
| Two agents over one arena is a use after free waiting to happen: both hold | ||
| the identical pages, and whichever is collected first deregisters memory | ||
| the other may still be reading, which surfaces as a memory fault partway | ||
| through the expert reads. |
There was a problem hiding this comment.
I don't think NIXL behaves like this. Each agent has its own UCX context, and deregistering only drops that agent's own memh; another agent's registration over the same pages is unaffected. On ROCm, rocm_ipc dereg is just ucs_free(key). The more likely issue is that every registration calls hsa_amd_ipc_memory_create on the base allocation again, which gfx950 rejects. Could you reword this, and share the UCX/HSA logs from the failing run if you have them?
| # reconfiguration gathers metadata this agent already knows. Adding | ||
| # it twice would mean two entries over one peer, and removing | ||
| # either would pull the rkeys out from under the live one. |
There was a problem hiding this comment.
nit: loadRemoteMD is idempotent for an already known agent. It merges into the same entry, and conn info is a no-op if unchanged, so there's never two entries. remote_names is fine to keep, but this comment is inaccurate.
Replace the per-arena agent registry with passing the previous communicator's NixlWrapper to the new one. Expert weights and expert_buffer keep their storage across reconfigurations, so the shared agent's registrations stay valid; Python refcounting decides when the agent dies and NIXL deregisters/disconnects in its destructor. Re-adding a known peer's metadata is idempotent in NIXL, so no remote-agent cache is needed either. Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
NickLucche
left a comment
NickLucche
left a comment
There was a problem hiding this comment.
hey, I think Arena abstraction might be overkill, can you check https://github.com/stefankoncarevic/vllm/pull/1/changes ?
|
You are right on both counts, I will reword those comments. My own early debug run shows the same thing: 7 On the memory fault, it is likely that my own diagnostics and the workarounds I layered on top of the refusal Thanks for writing the simplification. I am merging your branch now and running it on MI355 and H100, and will report back. |
|
/ci run |
|
✅ Triggered Buildkite CI #96848 for commit |
|
Both steps pass with your branch merged in. MI355 is 4 passed, 1 skipped, H100 is 5 passed, and Initialized NIXL agent now counts 16 across 16 processes on MI355 and 20 across 20 on H100, none holding more than one, with zero The only thing left in the H100 log is two AsyncLLM output_handler failed tracebacks at shutdown. That race is on main and not related to this change, I will cover that in a separate PR. |
stefankoncarevic
changed the title
|
The |
|
This pull request has merge conflicts that must be resolved before it can be |
|
Two things I haven't seen raised yet: 1. The CI part conflicts with #63609, and half of it merges silently. #63609 (merged 10-25) commented out both
The install isn't needed on ROCm anyway. The ci_base image already installs the pinned Minor: the rationale in the
2. The MI355 results predate #62827, which changes the transport (relevant to @itayalroy's scale-down point). Build 96848 ran on an image with the old UCX pin ( This also affects the scale-down disconnect. At the new pin ( |
Conflicts: - .buildkite/test_areas/expert_parallelism.yaml: the H100 Elastic EP step this job mirrored from is commented out on main, so the MI355 job becomes a standalone step. The kv_connectors install and the pytest timeout flags are dropped: the image already carries tblib 3.1.0 and lm-eval 0.4.12 pinned from requirements/test/rocm.txt, and run-amd-test.sh already exports PYTEST_TIMEOUT=2400 and --timeout-method=thread. - vllm/distributed/eplb/eplb_communicator.py: keep both sides, expert_buffer is now optional and prev_communicator is still threaded through. Co-authored-by: Cursor <cursoragent@cursor.com>
Sharing one agent across reconfigurations means a replacement communicator calls add_remote_agent on an agent the outgoing one may still be running transfers on. That happens during EEP preparation, well before EPLB is drained at commit, and without a listener thread NIXL defaults to NIXL_THREAD_SYNC_NONE, which leaves agent state unlocked. No measurable cost on MI355. Under sustained load with step_interval=10, giving 164 and 166 EPLB rearranges in two minutes, NONE reached 8.39 req/s at p50 2.207 s and STRICT reached 8.63 req/s at p50 2.144 s. The full elastic EP suite is 4 passed, 1 skipped in 762 s. Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
You're right, and the window is specific to this PR. The agent runs with NIXL_THREAD_SYNC_NONE (no listener thread, so that's the default), and add_remote_agent is called from _prepare_eplb_communicator during preparation, while the server is still serving and EPLB is only drained later at commit. Before this change each communicator built its own agent, so peers were added to an agent with nothing in flight; now a replacement communicator adds them to an agent the outgoing one may still be using. On the cost of the stricter modes, I couldn't measure one. Under sustained load with step_interval=10, so 164 and 166 EPLB rearranges in two minutes, NONE gave 8.39 req/s at p50 2.207 s and STRICT gave 8.63 req/s at p50 2.144 s, with STRICT coming out marginally ahead, which puts the difference inside run-to-run noise. So I've switched the EPLB agent to NIXL_THREAD_SYNC_STRICT. It only applies to the agent I create, so nothing else using NIXL is affected. Full suite on MI355 with the change is 4 passed, 1 skipped in 762 s, with scale and downtime numbers unchanged. |
All of this is in. expert_parallelism.yaml now has a standalone device: mi355_4 step instead of the mirror block, and the uv pip install line and the --timeout flags are gone from it and from test-amd.yaml. The allowlist entry didn't need touching: #64010 already dropped it when it re-enabled the MI300 step. Reran on rocm/vllm-ci:test-deps-rocm-f7d949c58b19bd0a, which is on the new pin: 4 passed, 1 skipped in 762 s. Your reading of rocm_ipc_cache.c is correct, remove_remote_agent included. Numbers are in @itayalroy's thread, and the residue turns out to accumulate per scale cycle. |
|
/ci run |
|
✅ Triggered Buildkite CI #99545 for commit |
|
/ci retry |
|
@itayalroy PTAL |
|
✅ Triggered Buildkite CI #99742 for commit |
| with contextlib.suppress(Exception): | ||
| for agent_name in self._remote_agents.values(): | ||
| with contextlib.suppress(Exception): |
There was a problem hiding this comment.
@stefankoncarevic The UCX IPC cache memleak is a separate issue which will be resolved in NIXL 1.6.0 (see openucx/ucx#12108), but as I said, even with this change, releasing IPC cache entries from the IPC cache requires disconnecting from the departing peers (and this is also the correct thing to do lifecycle-wise, our NIXL agent shouldn't keep connections to removed peers). We're not planning to disable the IPC cache in vLLM as it affects performance.
Sharing the NIXL agent means the outgoing communicator cannot drop its own connections, since the agent outlives it. Let the communicator that takes over release the peers that are gone, once the reshuffle is done, matching them by agent name because scaling down renumbers the survivors. Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #100259 for commit |
|
This pull request has merge conflicts that must be resolved before it can be |
Overview
NixlEplbCommunicatorcreated its own NIXL agent and registered the expert weights with it. Areconfiguration builds the replacement communicator while the old one is still alive, so a worker ended up
holding several agents with the same expert pages registered. On MI355 the extra IPC export is refused and
the test fails.
This hands the replacement communicator the agent it is replacing, so the pages are registered once. Two
things follow from sharing an agent and are part of this change: the agent runs in strict thread sync mode,
and the replacement communicator disconnects from the ranks that scaling down removed. The Elastic EP
Scaling step on MI355 is added.
Claims
tests/distributed/test_elastic_ep.pypasses on MI355, which it did not before.every backend; it is only fatal on some.
Validation
Agent count per worker process, counted from
Initialized NIXL agentlines in CI logs oftests/distributed/test_elastic_ep.py, which runs four scale cycles:Build 92231 ran this step on #56679, which adds Buildkite yaml and nothing else, so that is
unmodified code. There MI355 reports
hipIpcGetMemHandle failed : invalid argument6 and 8 times and thetest fails, while H100 passes in the same builds. The counts are lower than 16 because those runs failed
part way through, so fewer workers reached the point of creating an agent.
The MI300X row is an unmodified run on hardware that tolerates the extra registrations. It shows the extra
registration is not specific to the platform where it breaks.
The test. The MI355 step failed on every attempt before this change. With it, the two most recent runs
were 4 passed, 1 skipped, in 857 s and 853 s, with zero
hipIpc* failed, and the H100 step green in the samebuilds. The fifth case skips on ROCm because NIXL EP is not installed in that image. Accuracy stays in the
H100 band: the largest deviation from the pre-scaling GSM8K score was 0.036 on MI355 and 0.035 on H100,
against a tolerance of 0.080.
Disconnecting the departed. Four scale downs across two MI355X runs, each printing
EPLB disconnected from 2 departed peers.from every survivor, which is exactly the number of ranks the4 to 2 step removes, and no failed disconnect in any of them. The cost is not measurable: for
enforce_eager_nonethe scale down took 0.81 s with 0.51 s of downtime, against 0.85 s and 0.46 s on thissame branch before the disconnect was added.
tests/distributed/test_eplb_communicator.pycovers theselection and the guards on CPU, 6 passed in 2.19 s.
Details
Root cause. A NIXL registration belongs to the agent that made it, and
num_local_expertsis fixed, sothe expert weights and the receive buffers keep their storage across a reconfiguration. Building a fresh
agent for the replacement communicator therefore exports IPC handles for memory that is already exported,
and leaves the worker holding two and then three agents over the same pages.
The fix.
EplbState.create_communicatoralready holds the communicator being replaced, so it passes itdown and the NIXL backend reuses its agent, registering memory only when it creates a fresh one. The agent
is dropped when the last communicator referencing it is collected, and NIXL deregisters the memory from its
destructor.
add_remote_agentfor a peer the shared agent already knows is a no op, sinceloadRemoteMDmerges into the existing entry and the connection info is unchanged.
This is @NickLucche's implementation, from a branch he wrote against this one. My first version kept a module
level registry keyed by the expert address range, which was more machinery for the same result.
Thread sync. A replacement communicator adds its peers to an agent the outgoing one may still be running
transfers on, and the NIXL default of
NIXL_THREAD_SYNC_NONEleaves that unguarded. The shared agent iscreated with
NIXL_THREAD_SYNC_STRICT, falling back to no config when the installed NIXL does not exposethe enum.
Departed peers. Sharing the agent also means the outgoing communicator can no longer drop its own
connections, since the agent outlives it.
EplbCommunicator.release_departed_peersis a no op by defaultand the NIXL backend implements it, called from
update_communicatoronce the replacement is live. Thatpoint is after the reshuffle, so nothing is still reading the memory, which is the constraint raised in
review. Peers are compared by agent name rather than by rank, because scaling down renumbers the survivors
while the shared agent keeps each worker's name.
Memory release. Dropping the connection is what frees the imported IPC mapping on
cuda_ipc, throughopenucx/ucx#12108, which releases the unused mapping at rkey destroy. The ROCm transport reaches the same
end differently in openucx/ucx#12077, where the owner's free is what detaches the importers. Neither is in
the image this was tested on, so this PR does not claim a memory reduction; it makes the disconnect happen,
which is a precondition on one of the two and correct lifecycle either way.
The
__del__guard. Separate and small, in the same teardown path.__del__can run during interpretershutdown, when module globals are already being set to
None, including thecontextlibthat the teardownbody itself uses. The cleanup then raises and every worker prints a traceback on exit. The body is now
wrapped in a plain
trythat depends on no module.On the ROCm side. A probe in the worker showed that at the moment a 6 MiB export is refused a 2 MiB one
from the same process still succeeds, and that shortly before the refusal the same process had completed 64
consecutive 6 MiB exports. So the refusal is narrower than the duplicate export alone would suggest, and
that part is being looked at separately in ROCm/rocm-systems#11892. This PR does not depend on that issue and
does not work around it: it removes the duplicate registration, which was worth removing on the vLLM side
regardless of how any runtime responds to it.
CI steps. The H100 elastic EP step is commented out until the KV connector installation is locked, so
there is nothing to mirror. The MI355 step is therefore a standalone step in the same file, nightly only on
mi355_4, and it needs no connector install because the ROCm image already ships NIXL. A matching step isadded next to the existing MI300 one in the AMD pipeline. The new CPU test joins the existing EPLB Algorithm
group, which already depends on
vllm/distributed/eplb.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.