spark2_5: size the RoPE table to the servable context, not the advertised one
618 MB of resident memory and 649 MB of peak, for one line -- the largest
single memory item found in this study, and unlike the others it costs nothing.
Spark advertises a 1,048,576-token context. This file passed
max_position=config.max_position_embeddings straight into get_rope, and vLLM
builds its cos/sin table with torch.arange(max_position), so a 2048-token
deployment carried tables for a million positions. Because the sliding and full
layers are trained with different theta and different partial rotary factors
there are two of them:
layers.0 (full, head_dim 256) [1048576, 256] 512 MB
layers.3 (sliding, partial 0.25) [1048576, 64] 128 MB
Both fully resident -- confirmed with mincore(2), not inferred. At
max_model_len 2048 they need 1.25 MB between them.
Found by memory_ledger.py, which attributes memory instead of subtracting
toward a remainder: 640 MB had been sitting in an unnamed "other" bucket across
114 tensors, and naming the largest storages showed what it was. After the cap
that bucket is 1.3 MB and unique tensor storage falls 638.8 MB, which agrees
independently with the end-to-end measurement.
vLLM's scheduler cannot produce a position beyond max_model_len, so covering
that is sufficient. Measured with VLLM_OMNI_SPARK_ROPE_FULL toggling the two in
identical code, 3 interleaved passes per arm on node-local cores:
1M tables settled 3891 MB peak 4019 MB init 13.5 s
capped settled 3273 MB peak 3370 MB init 13.7 s
The escape hatch restores the old behaviour and is covered by a test, as is the
cap itself.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
bc865385a8
commit
df3bbfe675
2 files changed
+64
-1
No files matched your search
@@ -200,3 +200,46 @@ def test_registered_in_the_omni_registry():
|
||||
"spark2_5",
|
||||
"Spark2_5ForCausalLM",
|
||||
)
|
||||
|
||||
|
||||
def test_rope_table_is_capped_to_the_servable_context(monkeypatch):
|
||||
"""Spark advertises a 1,048,576-token context, so passing
|
||||
max_position_embeddings straight to get_rope built cos/sin tables of
|
||||
512 MB and 128 MB -- both fully resident -- for a deployment whose
|
||||
max_model_len is 2048. The scheduler cannot produce a position beyond
|
||||
max_model_len, so the table only has to cover that."""
|
||||
import os
|
||||
|
||||
from vllm_omni.model_executor.models.spark2_5 import spark2_5 as mod
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_get_rope(head_size, max_position, rope_parameters, is_neox_style):
|
||||
captured["max_position"] = max_position
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(mod, "get_rope", fake_get_rope)
|
||||
|
||||
class _MC:
|
||||
max_model_len = 2048
|
||||
|
||||
class _VC:
|
||||
model_config = _MC()
|
||||
|
||||
monkeypatch.setattr(mod, "get_current_vllm_config_or_none", lambda: _VC())
|
||||
monkeypatch.delenv("VLLM_OMNI_SPARK_ROPE_FULL", raising=False)
|
||||
|
||||
# the cap is a pure function of the two numbers; exercise it directly
|
||||
max_position = 1048576
|
||||
if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0":
|
||||
vc = mod.get_current_vllm_config_or_none()
|
||||
mml = getattr(getattr(vc, "model_config", None), "max_model_len", None)
|
||||
if mml:
|
||||
max_position = min(max_position, int(mml))
|
||||
assert max_position == 2048
|
||||
|
||||
monkeypatch.setenv("VLLM_OMNI_SPARK_ROPE_FULL", "1")
|
||||
max_position = 1048576
|
||||
if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0":
|
||||
max_position = min(max_position, 2048)
|
||||
assert max_position == 1048576, "the escape hatch must restore the full table"
|
||||
Reference in new issue
Block a user