Files
neuron-tai/tests/test_real_model_backend.py

1262 lines
41 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 (
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
def _require_functional_torch():
"""Skip tensor-behaviour tests when the installed torch namespace is incomplete."""
torch = pytest.importorskip("torch")
required = ("tensor", "zeros", "ones", "arange", "bfloat16", "long")
if not all(hasattr(torch, name) for name in required):
pytest.skip("requires a functional PyTorch tensor runtime")
return torch
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 session=" in out
assert "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 = _require_functional_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 = _require_functional_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 = _require_functional_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 = _call_layer(
NeedsPositionEmbeddings(),
"hidden",
attention_mask=None,
position_ids="positions",
position_embeddings="rotary",
)
assert hidden == "hidden"
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 = _require_functional_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 = _require_functional_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 = _require_functional_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: types.SimpleNamespace(type=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 = _require_functional_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()