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"
|
||||
@@ -44,6 +44,7 @@ from vllm.model_executor.layers.linear import (
|
||||
)
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.config import get_current_vllm_config_or_none
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
@@ -194,9 +195,28 @@ class Spark2_5Attention(nn.Module):
|
||||
# different partial-rotary factors; both live under rope_parameters.
|
||||
rope_parameters = dict(config.rope_parameters[layer_type])
|
||||
rope_parameters.setdefault("rope_type", "default")
|
||||
|
||||
# Spark advertises a 1 048 576-token context, and vLLM builds the RoPE
|
||||
# cos/sin table over the whole of it -- `torch.arange(max_position)` --
|
||||
# regardless of how long a sequence the engine will actually accept.
|
||||
# Measured with mincore on a 2048-token deployment: 512 MB resident for
|
||||
# the full-attention layers (head_dim 256) and 128 MB for the sliding
|
||||
# ones (partial rotary, 64 wide), 640 MB for two tables that need
|
||||
# 1.25 MB between them. The table only has to cover positions the
|
||||
# scheduler can produce, which vLLM caps at max_model_len.
|
||||
max_position = config.max_position_embeddings
|
||||
if os.environ.get("VLLM_OMNI_SPARK_ROPE_FULL", "0") == "0":
|
||||
vllm_config = get_current_vllm_config_or_none()
|
||||
model_config = (
|
||||
getattr(vllm_config, "model_config", None) if vllm_config else None
|
||||
)
|
||||
max_model_len = getattr(model_config, "max_model_len", None)
|
||||
if max_model_len:
|
||||
max_position = min(max_position, int(max_model_len))
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
self.head_dim,
|
||||
max_position=config.max_position_embeddings,
|
||||
max_position=max_position,
|
||||
rope_parameters=rope_parameters,
|
||||
is_neox_style=True,
|
||||
)
|
||||
|
||||
Reference in new issue
Block a user