fix(refit): stream a parameter larger than the IPC staging buffer - #3356
fix(refit): stream a parameter larger than the IPC staging buffer#3356terrykong wants to merge 1 commit into
Conversation
stream_weights_via_ipc_zmq_impl sizes its two ping-pong staging buffers
from *free memory* (NRL_REFIT_BUFFER_MEMORY_RATIO, default 0.3, halved
again for ping-pong) with no floor at the largest parameter. When one
parameter does not fit a single buffer it cannot be packed at all, and
the refit aborted on an assertion.
DeepSeek-V3 hits this: model.embed_tokens.weight is 129280 x 7168 bf16 =
1.73 GiB against a 1.65 GiB buffer, only 4.6% over, so
AssertionError: Parameter model.embed_tokens.weight too large for
buffer: 1853358080 > 1771261132
kills every colocated DSv3 run during the first refit, long after the
32-node allocation has been paid for.
Ship such a parameter on its own in a buffer sized to fit instead:
flush whatever group is pending, send the oversized parameter alone,
and wait for its ACK before freeing the one-off buffer (unlike the
ping-pong pair it is not kept alive across the next send). This mirrors
the HTTP streaming path, which already gives an oversized parameter a
bucket of its own rather than raising.
The new branch is only reachable where the old code raised, so refits
that work today are byte-for-byte unaffected; a regression test pins
that parameters which fit are still batched many per group.
Tested on CPU with a stub socket -- no GPU needed, since the packing and
hand-off logic is plain Python around a byte buffer. The three oversized
tests fail against the unpatched code with the same assertion seen in
CI, so they are not vacuous.
Signed-off-by: Terry Kong <terryk@nvidia.com>
|
/ok to test 831fa73 |
|
Closing as superseded — folded into #3280 (commit Rationale: validating DeepSeek-V3 requires the vLLM 0.25 node-spanning startup fixes (#3350, #3355) and this refit fix together. Keeping them on separate branches would have meant running the DSv3 perf suite twice, on two SKUs, for no added signal — so both live on one branch and get one validation cycle. The change itself is unmodified from what was green here (90 pass / 0 fail): Worth restating since it now rides a much larger PR: the new branch is reachable only where the old code raised |
1 participant
The bug
stream_weights_via_ipc_zmq_implsizes its two ping-pong staging buffers from free memory (NRL_REFIT_BUFFER_MEMORY_RATIO, default0.3, halved again for ping-pong) with no floor at the largest parameter. A parameter that does not fit a single buffer cannot be packed at all, and the refit aborted on an assertion.DeepSeek-V3 hits it.
model.embed_tokens.weightis 129280 x 7168 bf16 = 1.73 GiB against a 1.65 GiB buffer — only 4.6% over:raised from
MegatronPolicyWorker.stream_weights_via_ipc_zmq. It kills every colocated DSv3 run during the first refit, i.e. long after a 32-node allocation has been paid for. This is pre-existing onmainand unrelated to any vLLM version.The fix
Ship such a parameter on its own in a buffer sized to fit: flush whatever group is pending, send the oversized parameter alone, and consume its ACK before freeing the one-off buffer — unlike the ping-pong pair, it is not kept alive across the next send.
This mirrors the HTTP streaming path, which already gives an oversized parameter a bucket of its own (
if bucket and bucket_size + param_size >= buffer_size_bytes) instead of raising. The two sibling implementations now agree.Deliberately not fixed by raising the buffer size — neither by flooring the auto-computed value nor by setting
refit_buffer_size_gbper recipe. The 70% of free memory that the buffer does not take is reserved for parameter all-gather across the expert-parallel dimension, so growing the buffer trades a clean assertion for an OOM, and a per-recipe constant has to be re-tuned per model and cluster.Blast radius: none for working refits
The new branch is reachable only where the old code raised
AssertionError— a guaranteed failure. Any refit that works today takes byte-for-byte the same path.test_parameters_that_fit_are_still_batchedpins that.Validation
tests/unit/models/policy/test_refit_oversized_param.py— 5 tests, ~1.5s, no GPU and no Ray: the packing and hand-off logic is plain Python around a byte buffer, so CPU tensors and a stub socket exercise it faithfully.Against the unpatched code the suite splits exactly as intended, so it is not vacuous:
test_oversized_parameter_is_streamed_instead_of_abortingAssertionError: Parameter model.embed_tokens.weight too large for buffer: 3072 > 2048test_oversized_parameter_alonetest_consecutive_oversized_parameterstest_parameters_that_fit_are_still_batchedtest_alignment_is_what_decides_oversizedThe failure signature matches production exactly, at reduced scale.
Writing the harness surfaced one protocol detail worth recording: the receiver ACKs the end-of-stream
COMPLETEmarker as well as each data group, so a fake socket that only ACKs data groups deadlocks the streamer.Not covered here
Not yet validated on a real refit. That needs a cluster allocation. DeepSeek-V3 is the natural end-to-end check, but it currently also needs the vLLM 0.25 node-spanning startup fixes from #3350 and #3355, so a green DSv3 run is not the success criterion for this PR on its own.
🤖 Generated with Claude Code