1052 lines
34 KiB
Python
1052 lines
34 KiB
Python
"""US-012 tests for the real PyTorch node backend."""
|
|
|
|
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,
|
|
TensorPayload,
|
|
TorchModelShard,
|
|
_call_layer,
|
|
_checkpoint_tensor_name_for_model,
|
|
_load_partial_model_from_snapshot,
|
|
_should_partial_materialize_shard,
|
|
_decoder_attention_mask,
|
|
_int_tensor_header,
|
|
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):
|
|
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):
|
|
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):
|
|
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):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return " first"
|
|
self._release.wait(timeout=3.0)
|
|
return " second"
|
|
|
|
|
|
def test_quantization_flag_validation():
|
|
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():
|
|
pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8")
|
|
|
|
assert '"torch>=' in pyproject
|
|
|
|
|
|
def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
|
|
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_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
|
|
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():
|
|
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):
|
|
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):
|
|
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):
|
|
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():
|
|
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():
|
|
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):
|
|
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():
|
|
torch = pytest.importorskip("torch")
|
|
|
|
header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))
|
|
|
|
assert header.startswith("1,3:")
|
|
|
|
|
|
def test_decoder_attention_mask_is_causal_float_mask():
|
|
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():
|
|
class NeedsPositionEmbeddings:
|
|
def __call__(self, hidden_states, **kwargs):
|
|
assert kwargs["position_embeddings"] == "rotary"
|
|
return hidden_states
|
|
|
|
assert _call_layer(
|
|
NeedsPositionEmbeddings(),
|
|
"hidden",
|
|
attention_mask=None,
|
|
position_ids="positions",
|
|
position_embeddings="rotary",
|
|
) == "hidden"
|
|
|
|
|
|
def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):
|
|
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():
|
|
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():
|
|
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):
|
|
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').
|
|
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):
|
|
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):
|
|
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):
|
|
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():
|
|
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()
|