1254 lines
40 KiB
Python
1254 lines
40 KiB
Python
"""US-012 tests for the real PyTorch node backend."""
|
|
|
|
from collections import OrderedDict
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import threading
|
|
import time
|
|
import types
|
|
import urllib.request
|
|
|
|
import pytest
|
|
|
|
from meshnet_node.model_backend import (
|
|
InsufficientVRAMError,
|
|
PartialModelLoadUnsupported,
|
|
KVCacheMiss,
|
|
TensorPayload,
|
|
TorchModelShard,
|
|
_call_layer,
|
|
_checkpoint_tensor_name_for_model,
|
|
_load_partial_model_from_snapshot,
|
|
_should_partial_materialize_shard,
|
|
_decoder_attention_mask,
|
|
_int_tensor_header,
|
|
_tensor_from_bfloat16_bytes,
|
|
_torch_cuda_is_executable,
|
|
build_quantization_config,
|
|
validate_quantization,
|
|
)
|
|
from meshnet_node.torch_server import TorchNodeServer
|
|
|
|
|
|
class _FakeBackend:
|
|
model_id = "fake-model"
|
|
total_layers = 12
|
|
is_head = True
|
|
is_tail = False
|
|
|
|
def encode_prompt(self, prompt: str) -> TensorPayload:
|
|
assert prompt == "The capital of France is"
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 6 * 8 * 2),
|
|
shape=[1, 6, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
)
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body,
|
|
shape,
|
|
attention_mask_header,
|
|
position_ids_header,
|
|
start_layer=None,
|
|
**kwargs, # noqa: ARG002
|
|
):
|
|
assert shape == [1, 6, 8]
|
|
return TensorPayload(
|
|
body=body,
|
|
shape=shape,
|
|
attention_mask_header=attention_mask_header,
|
|
position_ids_header=position_ids_header,
|
|
)
|
|
|
|
|
|
class _FakeTailBackend(_FakeBackend):
|
|
is_head = False
|
|
is_tail = True
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body,
|
|
shape,
|
|
attention_mask_header,
|
|
position_ids_header,
|
|
start_layer=None,
|
|
**kwargs, # noqa: ARG002
|
|
):
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
return " Paris"
|
|
|
|
|
|
class _FakeFullBackend(_FakeBackend):
|
|
is_head = True
|
|
is_tail = True
|
|
|
|
def generate_text(
|
|
self,
|
|
messages: list[dict],
|
|
max_new_tokens: int = 16,
|
|
temperature: float = 1.0,
|
|
top_p: float = 1.0,
|
|
) -> str:
|
|
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
|
assert max_new_tokens == 7
|
|
assert temperature == 1.0
|
|
assert top_p == 1.0
|
|
return "56"
|
|
|
|
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
|
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
|
return 8
|
|
|
|
def count_text_tokens(self, text: str) -> int:
|
|
assert text == "56"
|
|
return 1
|
|
|
|
|
|
class _FakeChatTokenizer:
|
|
eos_token = ""
|
|
|
|
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
|
assert add_generation_prompt is True
|
|
assert tokenize is False
|
|
return "debug prompt"
|
|
|
|
|
|
class _FakePipelineHeadBackend(_FakeBackend):
|
|
tokenizer = _FakeChatTokenizer()
|
|
|
|
def encode_prompt(self, prompt: str) -> TensorPayload:
|
|
assert prompt.startswith("debug prompt")
|
|
return TensorPayload(
|
|
body=b"\x00" * (1 * 6 * 8 * 2),
|
|
shape=[1, 6, 8],
|
|
attention_mask_header=None,
|
|
position_ids_header=None,
|
|
)
|
|
|
|
|
|
class _FakePipelineTailBackend(_FakeTailBackend):
|
|
def __init__(self) -> None:
|
|
self.start_layers: list[int | None] = []
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body,
|
|
shape,
|
|
attention_mask_header,
|
|
position_ids_header,
|
|
start_layer=None,
|
|
**kwargs, # noqa: ARG002
|
|
):
|
|
self.start_layers.append(start_layer)
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
return " token"
|
|
|
|
|
|
class _BlockingStreamingTailBackend(_FakeTailBackend):
|
|
def __init__(self, second_token_release: threading.Event) -> None:
|
|
self._release = second_token_release
|
|
self.calls = 0
|
|
|
|
def forward_bytes(
|
|
self,
|
|
body,
|
|
shape,
|
|
attention_mask_header,
|
|
position_ids_header,
|
|
start_layer=None,
|
|
**kwargs, # noqa: ARG002
|
|
):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return " first"
|
|
self._release.wait(timeout=3.0)
|
|
return " second"
|
|
|
|
|
|
def test_quantization_flag_validation():
|
|
"Quantization flag validation\n\nTags: model, node, real-inference"
|
|
assert validate_quantization("bfloat16") == "bfloat16"
|
|
assert validate_quantization("int8") == "int8"
|
|
assert validate_quantization("nf4") == "nf4"
|
|
with pytest.raises(ValueError, match="quantization"):
|
|
validate_quantization("float32")
|
|
|
|
|
|
def test_node_package_declares_torch_dependency():
|
|
"Node package declares torch dependency\n\nTags: model, node, real-inference"
|
|
pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")
|
|
|
|
assert '"torch>=' in pyproject
|
|
|
|
|
|
def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
|
|
"Bitsandbytes configs are created lazily\n\nTags: model, node, real-inference"
|
|
calls = []
|
|
|
|
class FakeBitsAndBytesConfig:
|
|
def __init__(self, **kwargs):
|
|
calls.append(kwargs)
|
|
|
|
monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16"))
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"transformers",
|
|
types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig),
|
|
)
|
|
|
|
assert build_quantization_config("bfloat16") is None
|
|
build_quantization_config("int8")
|
|
build_quantization_config("nf4")
|
|
|
|
assert calls == [
|
|
{"load_in_8bit": True},
|
|
{
|
|
"load_in_4bit": True,
|
|
"bnb_4bit_quant_type": "nf4",
|
|
"bnb_4bit_compute_dtype": "bf16",
|
|
},
|
|
]
|
|
|
|
|
|
def test_rocm_inventory_without_executable_kernels_is_not_used_as_cuda():
|
|
"Rocm inventory without executable kernels is not used as cuda\n\nTags: model, node, real-inference"
|
|
class FakeCuda:
|
|
@staticmethod
|
|
def is_available():
|
|
return True
|
|
|
|
@staticmethod
|
|
def synchronize():
|
|
raise AssertionError("synchronize should not run after empty() fails")
|
|
|
|
fake_torch = types.SimpleNamespace(
|
|
cuda=FakeCuda(),
|
|
empty=lambda *args, **kwargs: (_ for _ in ()).throw(
|
|
RuntimeError("HIP error: invalid device function")
|
|
),
|
|
)
|
|
|
|
assert _torch_cuda_is_executable(fake_torch) is False
|
|
|
|
|
|
def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
|
|
"Head forward accepts text prompt and returns bfloat16 activations\n\nTags: model, node, real-inference"
|
|
node = TorchNodeServer(backend=_FakeBackend())
|
|
port = node.start()
|
|
try:
|
|
payload = json.dumps({"prompt": "The capital of France is"}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/forward",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = resp.read()
|
|
headers = {key.lower(): value for key, value in resp.headers.items()}
|
|
|
|
assert len(body) == 1 * 6 * 8 * 2
|
|
assert headers["x-meshnet-shape"] == "1,6,8"
|
|
assert headers["x-meshnet-dtype"] == "bfloat16"
|
|
assert headers["x-meshnet-wire"] == "2"
|
|
finally:
|
|
node.stop()
|
|
|
|
|
|
def test_tail_forward_returns_text_completion_from_binary_activations():
|
|
"Tail forward returns text completion from binary activations\n\nTags: model, node, real-inference"
|
|
node = TorchNodeServer(backend=_FakeTailBackend())
|
|
port = node.start()
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/forward",
|
|
data=b"\x00" * (1 * 6 * 8 * 2),
|
|
headers={
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Shape": "1,6,8",
|
|
"X-Meshnet-Dtype": "bfloat16",
|
|
"X-Meshnet-Session": "session-1",
|
|
"X-Meshnet-Chunk-Index": "0",
|
|
"X-Meshnet-Chunk-Total": "1",
|
|
"X-Meshnet-Hop-Index": "1",
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body == {"text": " Paris"}
|
|
assert node.received_activations
|
|
assert node.forward_chunk_count == 1
|
|
finally:
|
|
node.stop()
|
|
|
|
|
|
def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys):
|
|
"Full model chat completion uses generation not single token decode\n\nTags: model, node, real-inference"
|
|
node = TorchNodeServer(backend=_FakeFullBackend())
|
|
port = node.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "What is 7 times 8?"}],
|
|
"max_tokens": 7,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Request-Id": "req-test-123",
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["choices"][0]["message"]["content"] == "56"
|
|
assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
|
|
finally:
|
|
node.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert " [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out
|
|
assert " [node] chat complete tokens=1 elapsed_s=" in out
|
|
|
|
|
|
def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
|
|
"Pipeline hop logs are suppressed without debug\n\nTags: model, node, real-inference"
|
|
tail_backend = _FakePipelineTailBackend()
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=tail_backend)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 1,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
body = json.loads(resp.read())
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert body["choices"][0]["message"]["content"] == " token"
|
|
assert tail_backend.start_layers == [22]
|
|
assert "pipeline hop 0:" not in out
|
|
assert "pipeline hop 0 returned text" not in out
|
|
|
|
|
|
def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
|
|
"Pipeline hop logs are enabled with debug\n\nTags: model, node, real-inference"
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
|
|
tail = TorchNodeServer(backend=_FakePipelineTailBackend())
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 1,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
json.loads(resp.read())
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert f" [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
|
|
assert " [node] pipeline hop 0 returned text=' token'" in out
|
|
|
|
|
|
def test_split_shard_chat_streams_each_generated_token_incrementally():
|
|
"Split shard chat streams each generated token incrementally\n\nTags: model, node, real-inference, streaming"
|
|
release_second = threading.Event()
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
response = None
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": True,
|
|
"max_tokens": 2,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
response = urllib.request.urlopen(req, timeout=5)
|
|
|
|
first_token_line = ""
|
|
deadline = time.time() + 2.0
|
|
while time.time() < deadline:
|
|
line = response.readline().decode()
|
|
if '"content": " first"' in line:
|
|
first_token_line = line
|
|
break
|
|
|
|
assert first_token_line
|
|
assert not release_second.is_set()
|
|
release_second.set()
|
|
rest = response.read().decode()
|
|
finally:
|
|
release_second.set()
|
|
if response is not None:
|
|
response.close()
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert '"content": " second"' in rest
|
|
assert "data: [DONE]" in rest
|
|
|
|
|
|
def test_current_requests_snapshot_while_generating():
|
|
"Current requests snapshot while generating\n\nTags: model, node, real-inference"
|
|
release_second = threading.Event()
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second))
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
response = None
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"stream": True,
|
|
"max_tokens": 2,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Request-Id": "req-live-1",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
response = urllib.request.urlopen(req, timeout=5)
|
|
deadline = time.time() + 2.0
|
|
while time.time() < deadline:
|
|
live = head.current_requests
|
|
if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1:
|
|
break
|
|
time.sleep(0.02)
|
|
assert head.current_requests
|
|
snap = head.current_requests[0]
|
|
assert snap["request_id"] == "req-live-1"
|
|
assert snap["tokens"] >= 1
|
|
assert snap["tokens_per_sec"] >= 0
|
|
assert snap["routing_complete"] is True
|
|
release_second.set()
|
|
response.read()
|
|
finally:
|
|
release_second.set()
|
|
if response is not None:
|
|
response.close()
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
assert head.current_requests == []
|
|
|
|
|
|
def test_distributed_generating_log_includes_tps(capsys):
|
|
"Distributed generating log includes tps\n\nTags: model, node, real-inference"
|
|
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
|
tail = TorchNodeServer(backend=_FakePipelineTailBackend())
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
payload = json.dumps({
|
|
"model": "fake-model",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 1,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-Meshnet-Route": json.dumps([
|
|
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
|
]),
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(req, timeout=5) as resp:
|
|
json.loads(resp.read())
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|
|
|
|
out = capsys.readouterr().out
|
|
assert "generating step=1/1" in out
|
|
assert " tps=" in out
|
|
assert "generation complete tokens=1" in out
|
|
assert out.count("generating step=1/1") == 1
|
|
|
|
|
|
def test_int_tensor_header_serializes_torch_tensors():
|
|
"Int tensor header serializes torch tensors\n\nTags: model, node, real-inference"
|
|
torch = pytest.importorskip("torch")
|
|
|
|
header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))
|
|
|
|
assert header.startswith("1,3:")
|
|
|
|
|
|
def test_bfloat16_wire_decode_views_owned_bytes_without_float32_round_trip():
|
|
"""Activation decode stays bf16 and does not clone bytes into bytearray.
|
|
|
|
Tags: model, performance, wire
|
|
"""
|
|
torch = pytest.importorskip("torch")
|
|
body = torch.tensor([[1, 2]], dtype=torch.bfloat16).view(torch.uint8).numpy().tobytes()
|
|
|
|
decoded = _tensor_from_bfloat16_bytes(body, [1, 2], torch)
|
|
|
|
assert decoded.dtype == torch.bfloat16
|
|
assert decoded.tolist() == [[1.0, 2.0]]
|
|
|
|
|
|
def test_decoder_attention_mask_is_causal_float_mask():
|
|
"Decoder attention mask is causal float mask\n\nTags: model, node, real-inference"
|
|
torch = pytest.importorskip("torch")
|
|
|
|
hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
|
|
mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)
|
|
|
|
assert mask.shape == (1, 1, 3, 3)
|
|
assert mask.dtype == torch.bfloat16
|
|
assert mask[0, 0, 0, 1] < 0
|
|
assert mask[0, 0, 2, 0] == 0
|
|
|
|
|
|
def test_call_layer_passes_rotary_position_embeddings():
|
|
"Call layer passes rotary position embeddings\n\nTags: model, node, real-inference"
|
|
class NeedsPositionEmbeddings:
|
|
def __call__(self, hidden_states, **kwargs):
|
|
assert kwargs["position_embeddings"] == "rotary"
|
|
return hidden_states
|
|
|
|
hidden, cache_state = _call_layer(
|
|
NeedsPositionEmbeddings(),
|
|
"hidden",
|
|
attention_mask=None,
|
|
position_ids="positions",
|
|
position_embeddings="rotary",
|
|
)
|
|
|
|
assert hidden == "hidden"
|
|
assert cache_state is None
|
|
|
|
|
|
def _fake_cache_shard(torch, *, max_sessions=16, ttl=600.0):
|
|
class RecordingLayer:
|
|
def __init__(self, index):
|
|
self.index = index
|
|
self.calls = []
|
|
|
|
def __call__(self, hidden_states, **kwargs):
|
|
self.calls.append({
|
|
"shape": tuple(hidden_states.shape),
|
|
"use_cache": kwargs.get("use_cache"),
|
|
"past_key_value": kwargs.get("past_key_value"),
|
|
})
|
|
present = {
|
|
"layer": self.index,
|
|
"shape": tuple(hidden_states.shape),
|
|
"opaque": object(),
|
|
}
|
|
return hidden_states + (self.index + 1), present
|
|
|
|
shard = object.__new__(TorchModelShard)
|
|
shard.shard_start = 0
|
|
shard.shard_end = 1
|
|
shard.torch = torch
|
|
shard.model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[]))
|
|
shard.layers = [RecordingLayer(0), RecordingLayer(1)]
|
|
shard._session_cache = OrderedDict()
|
|
shard._cache_max_sessions = max_sessions
|
|
shard._cache_ttl_seconds = ttl
|
|
return shard
|
|
|
|
|
|
def test_shard_cache_prefill_then_decode_reuses_opaque_layer_state():
|
|
"Shard cache prefill then decode reuses opaque layer state\n\nTags: cache, model, node, real-inference"
|
|
torch = pytest.importorskip("torch")
|
|
shard = _fake_cache_shard(torch)
|
|
|
|
prefill_hidden = torch.zeros((1, 4, 2), dtype=torch.bfloat16)
|
|
prefill_mask = torch.ones((1, 4), dtype=torch.long)
|
|
prefill_positions = torch.arange(4, dtype=torch.long).reshape(1, 4)
|
|
shard._run_layers(
|
|
prefill_hidden,
|
|
prefill_mask,
|
|
prefill_positions,
|
|
session_id="session-1",
|
|
cache_mode="prefill",
|
|
seq_len=4,
|
|
)
|
|
|
|
assert len(shard._session_cache) == 1
|
|
cached_states = next(iter(shard._session_cache.values())).layer_states
|
|
assert len(cached_states) == 2
|
|
assert cached_states[0]["shape"] == (1, 4, 2)
|
|
|
|
decode_hidden = torch.zeros((1, 1, 2), dtype=torch.bfloat16)
|
|
decode_mask = torch.ones((1, 5), dtype=torch.long)
|
|
decode_positions = torch.tensor([[4]], dtype=torch.long)
|
|
shard._run_layers(
|
|
decode_hidden,
|
|
decode_mask,
|
|
decode_positions,
|
|
session_id="session-1",
|
|
cache_mode="decode",
|
|
seq_len=5,
|
|
)
|
|
|
|
assert shard.layers[0].calls[-1]["shape"] == (1, 1, 2)
|
|
assert shard.layers[0].calls[-1]["past_key_value"] is cached_states[0]
|
|
assert shard.layers[1].calls[-1]["past_key_value"] is cached_states[1]
|
|
assert next(iter(shard._session_cache.values())).seq_len == 5
|
|
|
|
|
|
def test_shard_cache_decode_miss_is_explicit():
|
|
"Shard cache decode miss is explicit\n\nTags: cache, model, node, real-inference"
|
|
torch = pytest.importorskip("torch")
|
|
shard = _fake_cache_shard(torch)
|
|
|
|
with pytest.raises(KVCacheMiss):
|
|
shard._run_layers(
|
|
torch.zeros((1, 1, 2), dtype=torch.bfloat16),
|
|
torch.ones((1, 5), dtype=torch.long),
|
|
torch.tensor([[4]], dtype=torch.long),
|
|
session_id="missing",
|
|
cache_mode="decode",
|
|
seq_len=5,
|
|
)
|
|
|
|
|
|
def test_shard_cache_lru_bounds_sessions():
|
|
"Shard cache lru bounds sessions\n\nTags: cache, model, node, real-inference"
|
|
torch = pytest.importorskip("torch")
|
|
shard = _fake_cache_shard(torch, max_sessions=1)
|
|
|
|
for session in ("old", "new"):
|
|
shard._run_layers(
|
|
torch.zeros((1, 2, 2), dtype=torch.bfloat16),
|
|
torch.ones((1, 2), dtype=torch.long),
|
|
torch.arange(2, dtype=torch.long).reshape(1, 2),
|
|
session_id=session,
|
|
cache_mode="prefill",
|
|
seq_len=2,
|
|
)
|
|
|
|
assert list(shard._session_cache.keys()) == [("new", 0, 1)]
|
|
|
|
|
|
def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):
|
|
"Partial materialize guard requires local non full non quantized snapshot\n\nTags: model, node, real-inference"
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text("{}")
|
|
(snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')
|
|
|
|
assert _should_partial_materialize_shard(
|
|
str(snapshot_dir),
|
|
4,
|
|
7,
|
|
total_layers_hint=40,
|
|
uses_quantized_weights=False,
|
|
) is True
|
|
assert _should_partial_materialize_shard(
|
|
str(snapshot_dir),
|
|
0,
|
|
39,
|
|
total_layers_hint=40,
|
|
uses_quantized_weights=False,
|
|
) is True
|
|
assert _should_partial_materialize_shard(
|
|
str(snapshot_dir),
|
|
4,
|
|
7,
|
|
total_layers_hint=40,
|
|
uses_quantized_weights=True,
|
|
) is False
|
|
assert _should_partial_materialize_shard(
|
|
"repo/model",
|
|
4,
|
|
7,
|
|
total_layers_hint=40,
|
|
uses_quantized_weights=False,
|
|
) is False
|
|
|
|
|
|
def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm():
|
|
"Checkpoint tensor name remapped for text only causal lm\n\nTags: model, node, real-inference"
|
|
class TextOnlyModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(layers=[])
|
|
|
|
model = TextOnlyModel()
|
|
assert _checkpoint_tensor_name_for_model(
|
|
model,
|
|
"model.language_model.layers.0.mlp.gate.weight",
|
|
) == "model.layers.0.mlp.gate.weight"
|
|
assert _checkpoint_tensor_name_for_model(
|
|
model,
|
|
"model.language_model.embed_tokens.weight",
|
|
) == "model.embed_tokens.weight"
|
|
|
|
|
|
def test_checkpoint_tensor_name_kept_for_multimodal_backbone():
|
|
"Checkpoint tensor name kept for multimodal backbone\n\nTags: model, node, real-inference"
|
|
class MultimodalModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(language_model=types.SimpleNamespace())
|
|
|
|
model = MultimodalModel()
|
|
name = "model.language_model.layers.0.mlp.gate.weight"
|
|
assert _checkpoint_tensor_name_for_model(model, name) == name
|
|
|
|
|
|
def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path):
|
|
"Partial snapshot loader remaps language model checkpoint keys\n\nTags: model, node, real-inference"
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text(json.dumps({
|
|
"text_config": {"num_hidden_layers": 3},
|
|
}))
|
|
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
|
"weight_map": {
|
|
"model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
|
|
}
|
|
}))
|
|
(snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")
|
|
|
|
class FakeModule:
|
|
def __init__(self):
|
|
self.to_calls = []
|
|
|
|
def to(self, device):
|
|
self.to_calls.append(device)
|
|
return self
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(
|
|
layers=[FakeModule(), FakeModule(), FakeModule()],
|
|
rotary_emb=FakeModule(),
|
|
)
|
|
|
|
def tie_weights(self):
|
|
pass
|
|
|
|
class AutoConfigStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id):
|
|
return types.SimpleNamespace(
|
|
text_config=types.SimpleNamespace(num_hidden_layers=3),
|
|
get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
|
|
)
|
|
|
|
class AutoModelStub:
|
|
@staticmethod
|
|
def from_config(cfg, torch_dtype=None):
|
|
return FakeModel()
|
|
|
|
set_calls = []
|
|
|
|
def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
|
|
set_calls.append(tensor_name)
|
|
|
|
class FakeSafeOpen:
|
|
def __init__(self, filename, framework, device):
|
|
self.filename = Path(filename).name
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def get_tensor(self, tensor_name):
|
|
return tensor_name
|
|
|
|
class UnusedContext:
|
|
def __enter__(self):
|
|
return None
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
_load_partial_model_from_snapshot(
|
|
AutoConfigStub,
|
|
AutoModelStub,
|
|
types.SimpleNamespace(),
|
|
str(snapshot_dir),
|
|
1,
|
|
1,
|
|
"bf16",
|
|
"cpu:0",
|
|
init_empty_weights_fn=UnusedContext,
|
|
set_tensor_fn=fake_set_tensor,
|
|
safe_open_fn=FakeSafeOpen,
|
|
)
|
|
|
|
assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
|
|
|
|
|
|
def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path):
|
|
# Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.*
|
|
# tensors that the text-only CausalLM never builds — they must be skipped,
|
|
# not assigned (assignment raises AttributeError: 'mtp' / 'visual').
|
|
"Partial snapshot loader skips tensors absent from causal lm\n\nTags: model, node, real-inference"
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text(json.dumps({
|
|
"text_config": {"num_hidden_layers": 3},
|
|
}))
|
|
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
|
"weight_map": {
|
|
"model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
|
|
"mtp.layers.1.input_layernorm.weight": "shard-2.safetensors",
|
|
"model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors",
|
|
}
|
|
}))
|
|
(snapshot_dir / "shard-2.safetensors").write_bytes(b"stub")
|
|
|
|
class FakeModule:
|
|
def to(self, device):
|
|
return self
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(
|
|
layers=[FakeModule(), FakeModule(), FakeModule()],
|
|
rotary_emb=FakeModule(),
|
|
)
|
|
|
|
def tie_weights(self):
|
|
pass
|
|
|
|
def state_dict(self):
|
|
return {"model.layers.1.self_attn.q_proj.weight": None}
|
|
|
|
class AutoConfigStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id):
|
|
return types.SimpleNamespace(
|
|
text_config=types.SimpleNamespace(num_hidden_layers=3),
|
|
get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3),
|
|
)
|
|
|
|
class AutoModelStub:
|
|
@staticmethod
|
|
def from_config(cfg, torch_dtype=None):
|
|
return FakeModel()
|
|
|
|
set_calls = []
|
|
|
|
def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
|
|
set_calls.append(tensor_name)
|
|
|
|
class FakeSafeOpen:
|
|
def __init__(self, filename, framework, device):
|
|
pass
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def get_tensor(self, tensor_name):
|
|
return tensor_name
|
|
|
|
class UnusedContext:
|
|
def __enter__(self):
|
|
return None
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
_load_partial_model_from_snapshot(
|
|
AutoConfigStub,
|
|
AutoModelStub,
|
|
types.SimpleNamespace(),
|
|
str(snapshot_dir),
|
|
1,
|
|
1,
|
|
"bf16",
|
|
"cpu:0",
|
|
init_empty_weights_fn=UnusedContext,
|
|
set_tensor_fn=fake_set_tensor,
|
|
safe_open_fn=FakeSafeOpen,
|
|
)
|
|
|
|
assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
|
|
|
|
|
|
def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
|
|
"Partial snapshot loader materializes only assigned tensors\n\nTags: model, node, real-inference"
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text("{}")
|
|
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
|
"weight_map": {
|
|
"model.embed_tokens.weight": "shard-1.safetensors",
|
|
"model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors",
|
|
"model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors",
|
|
"model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors",
|
|
"model.norm.weight": "shard-3.safetensors",
|
|
"lm_head.weight": "shard-3.safetensors",
|
|
}
|
|
}))
|
|
for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"):
|
|
(snapshot_dir / rel).write_bytes(b"stub")
|
|
|
|
class FakeModule:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
self.to_calls = []
|
|
|
|
def to(self, device):
|
|
self.to_calls.append(device)
|
|
return self
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(
|
|
embed_tokens=FakeModule("embed"),
|
|
layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")],
|
|
rotary_emb=FakeModule("rotary"),
|
|
norm=FakeModule("norm"),
|
|
)
|
|
self.lm_head = FakeModule("lm_head")
|
|
self.tie_weights_called = 0
|
|
|
|
def tie_weights(self):
|
|
self.tie_weights_called += 1
|
|
|
|
class AutoConfigStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id):
|
|
assert model_id == str(snapshot_dir)
|
|
return types.SimpleNamespace(num_hidden_layers=3)
|
|
|
|
class AutoModelStub:
|
|
@staticmethod
|
|
def from_config(cfg, torch_dtype=None):
|
|
assert cfg.num_hidden_layers == 3
|
|
assert torch_dtype == "bf16"
|
|
return FakeModel()
|
|
|
|
class EmptyWeights:
|
|
def __init__(self):
|
|
self.entered = 0
|
|
self.exited = 0
|
|
|
|
def __call__(self):
|
|
return self
|
|
|
|
def __enter__(self):
|
|
self.entered += 1
|
|
return None
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
self.exited += 1
|
|
return False
|
|
|
|
init_empty_weights = EmptyWeights()
|
|
set_calls = []
|
|
|
|
def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
|
|
set_calls.append((tensor_name, device, value, dtype))
|
|
|
|
tensors = {
|
|
"shard-1.safetensors": {
|
|
"model.embed_tokens.weight": "embed",
|
|
"model.layers.0.self_attn.q_proj.weight": "layer0",
|
|
},
|
|
"shard-2.safetensors": {
|
|
"model.layers.1.self_attn.q_proj.weight": "layer1",
|
|
},
|
|
"shard-3.safetensors": {
|
|
"model.layers.2.self_attn.q_proj.weight": "layer2",
|
|
"model.norm.weight": "norm",
|
|
"lm_head.weight": "lm_head",
|
|
},
|
|
}
|
|
|
|
class FakeSafeOpen:
|
|
def __init__(self, filename, framework, device):
|
|
assert framework == "pt"
|
|
assert device == "cpu"
|
|
self.filename = Path(filename).name
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
def get_tensor(self, tensor_name):
|
|
return tensors[self.filename][tensor_name]
|
|
|
|
model = _load_partial_model_from_snapshot(
|
|
AutoConfigStub,
|
|
AutoModelStub,
|
|
types.SimpleNamespace(),
|
|
str(snapshot_dir),
|
|
1,
|
|
1,
|
|
"bf16",
|
|
"cpu:0",
|
|
init_empty_weights_fn=init_empty_weights,
|
|
set_tensor_fn=fake_set_tensor,
|
|
safe_open_fn=FakeSafeOpen,
|
|
)
|
|
|
|
assert init_empty_weights.entered == 1
|
|
assert init_empty_weights.exited == 1
|
|
assert model.tie_weights_called == 1
|
|
assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"]
|
|
assert model.model.layers[1].to_calls == ["cpu:0"]
|
|
assert model.model.layers[0].to_calls == []
|
|
assert model.model.layers[2].to_calls == []
|
|
assert model.model.embed_tokens.to_calls == []
|
|
assert model.model.norm.to_calls == []
|
|
assert model.lm_head.to_calls == []
|
|
assert model.model.rotary_emb.to_calls == ["cpu:0"]
|
|
|
|
|
|
def test_partial_snapshot_loader_requires_known_layer_count(tmp_path):
|
|
"Partial snapshot loader requires known layer count\n\nTags: model, node, real-inference"
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text("{}")
|
|
(snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({
|
|
"weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"}
|
|
}))
|
|
(snapshot_dir / "shard.safetensors").write_bytes(b"stub")
|
|
|
|
class AutoConfigStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id):
|
|
return types.SimpleNamespace()
|
|
|
|
class AutoModelStub:
|
|
@staticmethod
|
|
def from_config(cfg, torch_dtype=None):
|
|
raise AssertionError("from_config should not run without a known layer count")
|
|
|
|
class UnusedContext:
|
|
def __enter__(self):
|
|
return None
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"):
|
|
_load_partial_model_from_snapshot(
|
|
AutoConfigStub,
|
|
AutoModelStub,
|
|
types.SimpleNamespace(),
|
|
str(snapshot_dir),
|
|
0,
|
|
0,
|
|
"bf16",
|
|
"cpu:0",
|
|
init_empty_weights_fn=lambda: UnusedContext(),
|
|
set_tensor_fn=lambda *args, **kwargs: None,
|
|
safe_open_fn=lambda *args, **kwargs: None,
|
|
)
|
|
|
|
|
|
def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch):
|
|
"Torch model shard prefers partial loader for local snapshot\n\nTags: model, node, real-inference"
|
|
import meshnet_node.model_backend as backend
|
|
|
|
snapshot_dir = tmp_path / "snapshot"
|
|
snapshot_dir.mkdir()
|
|
(snapshot_dir / "config.json").write_text("{}")
|
|
(snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}')
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.model = types.SimpleNamespace(
|
|
layers=[object(), object(), object()],
|
|
embed_tokens=object(),
|
|
)
|
|
self.config = types.SimpleNamespace(hidden_size=8)
|
|
self.eval_called = 0
|
|
|
|
def eval(self):
|
|
self.eval_called += 1
|
|
|
|
fake_model = FakeModel()
|
|
partial_calls = []
|
|
|
|
class AutoConfigStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id, cache_dir=None):
|
|
return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16"))
|
|
|
|
class AutoModelStub:
|
|
@staticmethod
|
|
def from_pretrained(*args, **kwargs):
|
|
raise AssertionError("full model load should not run for partial local shards")
|
|
|
|
class AutoTokenizerStub:
|
|
@staticmethod
|
|
def from_pretrained(model_id, cache_dir=None):
|
|
assert model_id == str(snapshot_dir)
|
|
return types.SimpleNamespace()
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"torch",
|
|
types.SimpleNamespace(
|
|
cuda=types.SimpleNamespace(is_available=lambda: False),
|
|
device=lambda value: value,
|
|
bfloat16="bf16",
|
|
),
|
|
)
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"transformers",
|
|
types.SimpleNamespace(
|
|
AutoConfig=AutoConfigStub,
|
|
AutoModelForCausalLM=AutoModelStub,
|
|
AutoTokenizer=AutoTokenizerStub,
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
backend,
|
|
"_load_partial_model_from_snapshot",
|
|
lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model,
|
|
)
|
|
|
|
shard = TorchModelShard(
|
|
"repo/model",
|
|
1,
|
|
1,
|
|
quantization="auto",
|
|
cache_dir=snapshot_dir,
|
|
)
|
|
|
|
assert len(partial_calls) == 1
|
|
assert shard.model is fake_model
|
|
assert fake_model.eval_called == 1
|
|
assert shard.total_layers == 3
|
|
assert shard.is_head is False
|
|
assert shard.is_tail is False
|
|
|
|
|
|
@pytest.mark.integration
|
|
def test_two_node_gpt2_completion_is_deterministic():
|
|
"Two node gpt2 completion is deterministic\n\nTags: model, node, real-inference"
|
|
if os.environ.get("CI"):
|
|
pytest.skip("GPT-2 integration test is skipped in CI")
|
|
torch = pytest.importorskip("torch")
|
|
pytest.importorskip("transformers")
|
|
pytest.importorskip("safetensors")
|
|
pytest.importorskip("accelerate")
|
|
pytest.importorskip("bitsandbytes")
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("GPT-2 integration test requires a CUDA GPU")
|
|
|
|
head = TorchNodeServer(
|
|
model_id="openai-community/gpt2",
|
|
shard_start=0,
|
|
shard_end=6,
|
|
quantization="bfloat16",
|
|
)
|
|
tail = TorchNodeServer(
|
|
model_id="openai-community/gpt2",
|
|
shard_start=6,
|
|
shard_end=12,
|
|
quantization="bfloat16",
|
|
)
|
|
head_port = head.start()
|
|
tail_port = tail.start()
|
|
try:
|
|
prompt_req = urllib.request.Request(
|
|
f"http://127.0.0.1:{head_port}/forward",
|
|
data=json.dumps({"prompt": "The capital of France is"}).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(prompt_req, timeout=60) as resp:
|
|
activation = resp.read()
|
|
head_headers = resp.headers
|
|
|
|
tail_req = urllib.request.Request(
|
|
f"http://127.0.0.1:{tail_port}/forward",
|
|
data=activation,
|
|
headers={
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Shape": head_headers["X-Meshnet-Shape"],
|
|
"X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"],
|
|
"X-Meshnet-Session": "gpt2-session",
|
|
"X-Meshnet-Chunk-Index": "0",
|
|
"X-Meshnet-Chunk-Total": "1",
|
|
"X-Meshnet-Hop-Index": "1",
|
|
"X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"],
|
|
"X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"],
|
|
},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(tail_req, timeout=60) as resp:
|
|
body = json.loads(resp.read())
|
|
|
|
assert body["text"].strip()
|
|
assert body["text"] == " Paris"
|
|
finally:
|
|
head.stop()
|
|
tail.stop()
|