Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai
This commit is contained in:
Dobromir Popov
@@ -134,8 +134,9 @@ class TorchModelShard:
|
||||
self.model.to(self.device)
|
||||
except Exception as exc:
|
||||
if _looks_like_oom(exc):
|
||||
memory_kind = "VRAM" if self.device.type == "cuda" else "RAM"
|
||||
raise InsufficientVRAMError(
|
||||
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
|
||||
f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} "
|
||||
f"with {quantization} quantization; choose a smaller shard or lower quantization"
|
||||
) from exc
|
||||
raise
|
||||
@@ -411,7 +412,7 @@ def _should_partial_materialize_shard(
|
||||
return False
|
||||
if total_layers_hint is None:
|
||||
return False
|
||||
return not (shard_start == 0 and shard_end >= total_layers_hint - 1)
|
||||
return True
|
||||
|
||||
|
||||
def _load_partial_model_from_snapshot(
|
||||
@@ -476,7 +477,7 @@ def _load_partial_model_from_snapshot(
|
||||
)
|
||||
|
||||
with init_empty_weights_fn():
|
||||
model = auto_model_for_causal_lm.from_config(cfg, torch_dtype=dtype)
|
||||
model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype)
|
||||
tie_weights = getattr(model, "tie_weights", None)
|
||||
if callable(tie_weights):
|
||||
tie_weights()
|
||||
@@ -498,7 +499,7 @@ def _load_partial_model_from_snapshot(
|
||||
for tensor_name in names:
|
||||
set_tensor_fn(
|
||||
model,
|
||||
tensor_name,
|
||||
_checkpoint_tensor_name_for_model(model, tensor_name),
|
||||
device,
|
||||
value=handle.get_tensor(tensor_name),
|
||||
dtype=dtype,
|
||||
@@ -569,38 +570,74 @@ def _native_torch_dtype(cfg: Any, torch: Any) -> Any:
|
||||
return torch.bfloat16
|
||||
|
||||
|
||||
def _causal_lm_config(cfg: Any) -> Any:
|
||||
"""Use the text decoder config for composite VLM/MoE presets."""
|
||||
get_text_config = getattr(cfg, "get_text_config", None)
|
||||
if callable(get_text_config):
|
||||
try:
|
||||
return get_text_config()
|
||||
except Exception:
|
||||
pass
|
||||
text_config = getattr(cfg, "text_config", None)
|
||||
if text_config is not None:
|
||||
return text_config
|
||||
return cfg
|
||||
|
||||
|
||||
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
|
||||
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
|
||||
inner = getattr(model, "model", None)
|
||||
if inner is not None and hasattr(inner, "language_model"):
|
||||
return tensor_name
|
||||
if ".language_model." in tensor_name:
|
||||
return tensor_name.replace(".language_model.", ".")
|
||||
return tensor_name
|
||||
|
||||
|
||||
def _transformer_backbone(model: Any) -> Any:
|
||||
if hasattr(model, "model"):
|
||||
inner = model.model
|
||||
language_model = getattr(inner, "language_model", None)
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
return inner
|
||||
if hasattr(model, "transformer"):
|
||||
return model.transformer
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no transformer backbone found"
|
||||
)
|
||||
|
||||
|
||||
def _model_layers(model: Any) -> Any:
|
||||
if hasattr(model, "model") and hasattr(model.model, "layers"):
|
||||
return model.model.layers
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
|
||||
return model.transformer.h
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("layers", "h", "blocks"):
|
||||
layers = getattr(backbone, attr, None)
|
||||
if layers is not None:
|
||||
return layers
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no transformer layers found"
|
||||
)
|
||||
|
||||
|
||||
def _embed_tokens(model: Any) -> Any:
|
||||
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
|
||||
return model.model.embed_tokens
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
|
||||
return model.transformer.wte
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("embed_tokens", "wte"):
|
||||
embed = getattr(backbone, attr, None)
|
||||
if embed is not None:
|
||||
return embed
|
||||
raise ModelBackendError(
|
||||
"unsupported HuggingFace model architecture: no token embeddings found"
|
||||
)
|
||||
|
||||
|
||||
def _position_embeddings(model: Any) -> Any | None:
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
|
||||
return model.transformer.wpe
|
||||
return None
|
||||
backbone = _transformer_backbone(model)
|
||||
return getattr(backbone, "wpe", None)
|
||||
|
||||
|
||||
def _rotary_embedding_module(model: Any) -> Any | None:
|
||||
if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
|
||||
return model.model.rotary_emb
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
|
||||
return model.transformer.rotary_emb
|
||||
return None
|
||||
backbone = _transformer_backbone(model)
|
||||
return getattr(backbone, "rotary_emb", None)
|
||||
|
||||
|
||||
def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]:
|
||||
@@ -627,10 +664,11 @@ def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> l
|
||||
|
||||
|
||||
def _final_norm(model: Any) -> Any | None:
|
||||
if hasattr(model, "model") and hasattr(model.model, "norm"):
|
||||
return model.model.norm
|
||||
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
|
||||
return model.transformer.ln_f
|
||||
backbone = _transformer_backbone(model)
|
||||
for attr in ("norm", "ln_f", "final_layer_norm"):
|
||||
norm = getattr(backbone, attr, None)
|
||||
if norm is not None:
|
||||
return norm
|
||||
return None
|
||||
|
||||
|
||||
@@ -743,7 +781,12 @@ def _looks_like_oom(exc: BaseException) -> bool:
|
||||
current: BaseException | None = exc
|
||||
while current is not None:
|
||||
text = str(current).lower()
|
||||
if "out of memory" in text or "cuda error: out of memory" in text:
|
||||
if (
|
||||
"out of memory" in text
|
||||
or "cuda error: out of memory" in text
|
||||
or "paging file is too small" in text
|
||||
or "os error 1455" in text
|
||||
):
|
||||
return True
|
||||
current = current.__cause__ or current.__context__
|
||||
return False
|
||||
|
||||
@@ -995,6 +995,20 @@ def run_startup(
|
||||
)
|
||||
if user_pinned_shard:
|
||||
shard_label = f"{shard_label} (pinned)"
|
||||
if user_pinned_shard and assigned_total_layers and assignment_bytes_per_layer:
|
||||
pinned_layers = shard_end - shard_start + 1
|
||||
max_layers = _max_assignable_layers(
|
||||
memory_budget_mb,
|
||||
assigned_total_layers,
|
||||
assignment_bytes_per_layer,
|
||||
)
|
||||
if pinned_layers > max_layers:
|
||||
raise ValueError(
|
||||
f"Pinned shard layers {shard_start}–{shard_end} ({pinned_layers} layers) exceed "
|
||||
f"the {memory_budget_mb / 1024:.1f} GB {memory_budget_source} budget "
|
||||
f"(fits up to {max_layers}/{assigned_total_layers} layers at bfloat16). "
|
||||
"Drop --shard-start/--shard-end to let the tracker auto-assign, or pin a smaller range."
|
||||
)
|
||||
print(f" Shard: {shard_label}", flush=True)
|
||||
|
||||
# 4. Download shard
|
||||
@@ -1020,7 +1034,77 @@ def run_startup(
|
||||
)
|
||||
print(f" Cached at: {shard_path}", flush=True)
|
||||
|
||||
# 5. Start HTTP server
|
||||
# 5. Start HTTP server — real HF weights use TorchNodeServer; stub-model stays stub.
|
||||
_node_start_time = time.monotonic()
|
||||
if hf_repo and assigned_model != "stub-model":
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=hf_repo,
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
cache_dir=shard_path,
|
||||
debug=debug,
|
||||
max_loaded_shards=max_loaded_shards,
|
||||
)
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers
|
||||
shard_label = _format_shard_label(shard_start, shard_end, total_layers, model_name=assigned_model)
|
||||
if user_pinned_shard:
|
||||
shard_label = f"{shard_label} (pinned)"
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
address,
|
||||
local_base_url,
|
||||
endpoint,
|
||||
relay_url=relay_url,
|
||||
)
|
||||
_attach_relay_bridge(node, relay_bridge)
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_model,
|
||||
"hf_repo": hf_repo,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"downloaded_models": downloaded_models,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
"model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
|
||||
**registration_capabilities,
|
||||
**relay_fields,
|
||||
}
|
||||
tracker_node_id = _register_with_tracker(
|
||||
tracker_url, reg_payload, node, _node_start_time,
|
||||
)
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Model ID: {hf_repo}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||||
f" Hardware: {_hardware_label(device, gpu_name)}\n"
|
||||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
return node
|
||||
|
||||
is_last = shard_end >= assignment.get("model_layers_end", shard_end)
|
||||
node = StubNodeServer(
|
||||
host=host,
|
||||
@@ -1031,7 +1115,6 @@ def run_startup(
|
||||
model=assigned_model,
|
||||
shard_path=shard_path,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
|
||||
@@ -8,14 +8,20 @@
|
||||
:root { --bg:#0d1117; --panel:#161b22; --border:#30363d; --fg:#c9d1d9;
|
||||
--dim:#8b949e; --accent:#58a6ff; --ok:#3fb950; --bad:#f85149; --warn:#d29922; }
|
||||
* { box-sizing:border-box; }
|
||||
html, body { height:100%; }
|
||||
body { margin:0; background:var(--bg); color:var(--fg);
|
||||
font:13px/1.5 ui-monospace,SFMono-Regular,Menlo,Consolas,monospace; }
|
||||
body.chat-tab-active { overflow:hidden; height:100dvh; display:flex; flex-direction:column; }
|
||||
header { display:flex; align-items:baseline; gap:14px; padding:14px 20px;
|
||||
border-bottom:1px solid var(--border); }
|
||||
border-bottom:1px solid var(--border); flex-shrink:0; }
|
||||
header h1 { font-size:16px; margin:0; color:var(--accent); }
|
||||
header .meta { color:var(--dim); font-size:12px; }
|
||||
main { display:grid; grid-template-columns:repeat(auto-fit,minmax(340px,1fr));
|
||||
gap:14px; padding:14px 20px; }
|
||||
body.chat-tab-active main {
|
||||
flex:1; min-height:0; display:flex; flex-direction:column;
|
||||
padding:0; gap:0; overflow:hidden;
|
||||
}
|
||||
section { background:var(--panel); border:1px solid var(--border);
|
||||
border-radius:8px; padding:12px 14px; min-height:80px; }
|
||||
section h2 { margin:0 0 8px; font-size:12px; text-transform:uppercase;
|
||||
@@ -53,27 +59,116 @@
|
||||
.tabs { display:flex; gap:10px; margin-bottom:8px; }
|
||||
.tabs a { color:var(--dim); cursor:pointer; }
|
||||
.tabs a.active { color:var(--accent); border-bottom:1px solid var(--accent); }
|
||||
.dashboard-tabs { display:flex; gap:10px; padding:10px 20px 0; border-bottom:1px solid var(--border); }
|
||||
.dashboard-tabs { display:flex; gap:10px; padding:10px 20px 0; border-bottom:1px solid var(--border); flex-shrink:0; }
|
||||
.dashboard-tabs button { border:0; border-bottom:1px solid transparent; border-radius:0;
|
||||
background:transparent; color:var(--dim); padding:5px 0 8px; }
|
||||
.dashboard-tabs button.active { color:var(--accent); border-bottom-color:var(--accent); }
|
||||
.wide { grid-column:1 / -1; }
|
||||
section[hidden] { display:none !important; }
|
||||
.chat-shell { display:grid; grid-template-columns:minmax(0, 1.35fr) minmax(320px, 0.65fr); gap:12px; }
|
||||
.chat-pane { display:flex; flex-direction:column; gap:10px; min-width:0; }
|
||||
.chat-panel { background:var(--bg); border:1px solid var(--border); border-radius:6px; padding:10px; }
|
||||
.chat-controls { display:flex; gap:10px; align-items:end; flex-wrap:wrap; }
|
||||
.chat-controls label { display:flex; flex-direction:column; gap:4px; color:var(--dim); }
|
||||
.chat-controls select { min-width:220px; }
|
||||
.chat-history { display:flex; flex-direction:column; gap:8px; min-height:220px; max-height:420px; overflow:auto; }
|
||||
.chat-message { border:1px solid #21262d; border-radius:6px; padding:8px 10px; background:#10151d; }
|
||||
.chat-role { color:var(--dim); font-size:11px; text-transform:uppercase; letter-spacing:.06em; margin-bottom:4px; }
|
||||
.chat-role-user { color:var(--accent); }
|
||||
.chat-role-assistant { color:var(--ok); }
|
||||
.chat-role-error { color:var(--bad); }
|
||||
.chat-compose { display:flex; flex-direction:column; gap:8px; }
|
||||
.chat-compose textarea { min-height:112px; resize:vertical; width:100%; }
|
||||
.chat-status { color:var(--dim); font-size:12px; }
|
||||
section.chat-section {
|
||||
padding:0; border:0; border-radius:0; background:var(--bg); min-height:0;
|
||||
}
|
||||
body.chat-tab-active section.chat-section {
|
||||
flex:1; display:flex !important; flex-direction:column; min-height:0;
|
||||
}
|
||||
.chat-app {
|
||||
display:grid; grid-template-columns:260px minmax(0, 1fr); gap:0;
|
||||
flex:1; min-height:0; overflow:hidden; background:var(--bg);
|
||||
}
|
||||
.chat-sidebar {
|
||||
display:flex; flex-direction:column; min-height:0;
|
||||
border-right:1px solid var(--border); background:var(--panel);
|
||||
}
|
||||
.chat-new-btn {
|
||||
margin:12px; width:calc(100% - 24px); text-align:left;
|
||||
border:1px solid var(--border); border-radius:8px; padding:10px 12px;
|
||||
background:transparent; color:var(--fg);
|
||||
}
|
||||
.chat-new-btn:hover { background:#10151d; border-color:var(--accent); }
|
||||
.chat-session-list {
|
||||
flex:1; overflow:auto; padding:0 8px 12px; display:flex; flex-direction:column; gap:2px;
|
||||
}
|
||||
.chat-session-list.empty-state {
|
||||
justify-content:center; align-items:center; color:var(--dim); font-style:italic;
|
||||
padding:24px 12px;
|
||||
}
|
||||
.chat-session-item {
|
||||
position:relative; display:block; width:100%; text-align:left;
|
||||
padding:10px 32px 10px 12px; border:1px solid transparent; border-radius:8px;
|
||||
background:transparent; color:var(--fg); cursor:pointer;
|
||||
}
|
||||
.chat-session-item:hover { background:#10151d; }
|
||||
.chat-session-item.active { background:#10151d; border-color:#30363d; }
|
||||
.chat-session-title {
|
||||
font-size:13px; white-space:nowrap; overflow:hidden; text-overflow:ellipsis;
|
||||
}
|
||||
.chat-session-meta {
|
||||
margin-top:3px; font-size:11px; color:var(--dim);
|
||||
white-space:nowrap; overflow:hidden; text-overflow:ellipsis;
|
||||
}
|
||||
.chat-session-delete {
|
||||
position:absolute; top:50%; right:6px; transform:translateY(-50%);
|
||||
padding:2px 6px; min-width:0; border:0; border-radius:4px;
|
||||
background:transparent; color:var(--dim); line-height:1.2; opacity:0;
|
||||
}
|
||||
.chat-session-item:hover .chat-session-delete,
|
||||
.chat-session-item.active .chat-session-delete { opacity:1; }
|
||||
.chat-session-delete:hover { color:var(--bad); background:#1a1012; }
|
||||
.chat-main { display:flex; flex-direction:column; min-height:0; min-width:0; }
|
||||
.chat-toolbar {
|
||||
display:flex; gap:12px; align-items:center; flex-shrink:0;
|
||||
padding:10px 16px; border-bottom:1px solid var(--border); background:var(--panel);
|
||||
}
|
||||
.chat-toolbar label {
|
||||
display:flex; align-items:center; gap:8px; color:var(--dim); margin:0;
|
||||
}
|
||||
.chat-toolbar select { min-width:220px; max-width:min(420px, 50vw); }
|
||||
.chat-status { color:var(--dim); font-size:12px; margin-left:auto; }
|
||||
.chat-messages {
|
||||
flex:1; overflow:auto; padding:24px 16px; min-height:0;
|
||||
}
|
||||
.chat-messages-inner {
|
||||
max-width:768px; margin:0 auto; display:flex; flex-direction:column; gap:20px;
|
||||
}
|
||||
.chat-messages.empty .chat-messages-inner {
|
||||
min-height:100%; justify-content:center; align-items:center;
|
||||
color:var(--dim); font-size:14px;
|
||||
}
|
||||
.chat-row { display:flex; width:100%; }
|
||||
.chat-row.user { justify-content:flex-end; }
|
||||
.chat-row.assistant, .chat-row.error { justify-content:flex-start; }
|
||||
.chat-bubble {
|
||||
max-width:85%; padding:12px 14px; border-radius:16px; line-height:1.55;
|
||||
white-space:pre-wrap; word-break:break-word; font-size:13px;
|
||||
}
|
||||
.chat-bubble.user {
|
||||
background:#1f3a5f; border:1px solid #264a73; border-bottom-right-radius:4px;
|
||||
}
|
||||
.chat-bubble.assistant {
|
||||
background:transparent; border:0; padding-left:0; padding-right:0; max-width:100%;
|
||||
}
|
||||
.chat-bubble.error {
|
||||
background:#1a1012; border:1px solid #5c2020; color:#ffb4b4; border-bottom-left-radius:4px;
|
||||
}
|
||||
.chat-compose-wrap {
|
||||
flex-shrink:0; padding:12px 16px 16px; border-top:1px solid var(--border);
|
||||
background:var(--panel);
|
||||
}
|
||||
.chat-compose {
|
||||
display:flex; gap:8px; align-items:flex-end; max-width:768px; margin:0 auto;
|
||||
padding:10px 12px; border:1px solid var(--border); border-radius:16px;
|
||||
background:var(--bg);
|
||||
}
|
||||
.chat-compose:focus-within { border-color:var(--accent); }
|
||||
.chat-compose textarea {
|
||||
flex:1; min-height:24px; max-height:200px; resize:none; width:auto;
|
||||
border:0; background:transparent; padding:2px 0; outline:none;
|
||||
}
|
||||
.chat-compose button {
|
||||
flex-shrink:0; min-width:36px; height:36px; padding:0;
|
||||
border-radius:8px; border:1px solid var(--border);
|
||||
}
|
||||
.chat-compose button:disabled { opacity:.45; cursor:not-allowed; }
|
||||
.console {
|
||||
background:var(--bg); border:1px solid var(--border); border-radius:6px;
|
||||
min-height:160px; max-height:280px; overflow:auto; padding:7px 9px;
|
||||
@@ -111,27 +206,27 @@
|
||||
<section data-tab="overview"><h2>Nodes & coverage</h2><div id="nodes" class="empty">loading…</div></section>
|
||||
<section data-tab="overview"><h2>Model usage (RPM)</h2><div id="stats" class="empty">loading…</div></section>
|
||||
<section data-tab="overview" class="wide"><h2>Call wall</h2><div id="call-wall" class="empty">loading...</div></section>
|
||||
<section data-tab="chat" class="wide">
|
||||
<h2>Chat / inference</h2>
|
||||
<div class="chat-shell">
|
||||
<div class="chat-pane">
|
||||
<div class="chat-panel chat-controls">
|
||||
<section data-tab="chat" class="wide chat-section">
|
||||
<div class="chat-app">
|
||||
<aside class="chat-sidebar">
|
||||
<button type="button" class="chat-new-btn" onclick="createNewChatSession()">+ New chat</button>
|
||||
<div id="chat-session-list" class="chat-session-list empty-state">No chats yet</div>
|
||||
</aside>
|
||||
<div class="chat-main">
|
||||
<div class="chat-toolbar">
|
||||
<label>Model
|
||||
<select id="chat-model" onchange="selectChatModel(this.value)"></select>
|
||||
</label>
|
||||
<button class="small" onclick="clearChatHistory()">clear history</button>
|
||||
</div>
|
||||
<div class="chat-panel chat-compose">
|
||||
<textarea id="chat-prompt" placeholder="Ask a question or describe the task"></textarea>
|
||||
<div class="form-row">
|
||||
<button onclick="sendChat()" id="chat-send">Send</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="chat-pane">
|
||||
<div class="chat-panel">
|
||||
<div id="chat-status" class="chat-status">select a model to start</div>
|
||||
<div id="chat-history" class="chat-history empty">no messages yet</div>
|
||||
</div>
|
||||
<div id="chat-history" class="chat-messages empty">
|
||||
<div class="chat-messages-inner">Send a message to start this conversation.</div>
|
||||
</div>
|
||||
<div class="chat-compose-wrap">
|
||||
<div class="chat-compose">
|
||||
<textarea id="chat-prompt" placeholder="Message…" rows="1" aria-label="Message"></textarea>
|
||||
<button type="button" onclick="sendChat()" id="chat-send" title="Send (Enter)">↑</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -604,17 +699,214 @@ let lastStats = null;
|
||||
let availableModels = [];
|
||||
let chatHistory = [];
|
||||
let chatBusy = false;
|
||||
let chatSessions = [];
|
||||
let activeChatSessionId = "";
|
||||
let selectedChatModel = localStorage.getItem("meshnet_chat_model") || "";
|
||||
const CHAT_SESSIONS_KEY = "meshnet_chat_sessions_v1";
|
||||
const CHAT_ACTIVE_SESSION_KEY = "meshnet_chat_active_session_v1";
|
||||
const CHAT_SESSIONS_LIMIT = 50;
|
||||
|
||||
function newChatSessionId() {
|
||||
if (window.crypto && crypto.randomUUID) return crypto.randomUUID();
|
||||
return "chat-" + Date.now().toString(36) + "-" + Math.random().toString(36).slice(2, 8);
|
||||
}
|
||||
|
||||
function loadChatSessionsStore() {
|
||||
try {
|
||||
const raw = localStorage.getItem(CHAT_SESSIONS_KEY);
|
||||
const parsed = raw ? JSON.parse(raw) : [];
|
||||
return Array.isArray(parsed) ? parsed : [];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
function saveChatSessionsStore() {
|
||||
localStorage.setItem(CHAT_SESSIONS_KEY, JSON.stringify(chatSessions));
|
||||
if (activeChatSessionId) {
|
||||
localStorage.setItem(CHAT_ACTIVE_SESSION_KEY, activeChatSessionId);
|
||||
}
|
||||
}
|
||||
|
||||
function chatSessionTitle(session) {
|
||||
const firstUser = (session.messages || []).find(msg => msg.role === "user");
|
||||
if (!firstUser || !firstUser.content) return "New chat";
|
||||
const text = String(firstUser.content).trim().replace(/\s+/g, " ");
|
||||
return text.length > 42 ? text.slice(0, 42) + "…" : text;
|
||||
}
|
||||
|
||||
function formatSessionTime(iso) {
|
||||
if (!iso) return "";
|
||||
const date = new Date(iso);
|
||||
if (Number.isNaN(date.getTime())) return "";
|
||||
const now = new Date();
|
||||
const sameDay = date.toDateString() === now.toDateString();
|
||||
if (sameDay) return date.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" });
|
||||
return date.toLocaleDateString([], { month: "short", day: "numeric" });
|
||||
}
|
||||
|
||||
function getActiveChatSession() {
|
||||
return chatSessions.find(session => session.id === activeChatSessionId) || null;
|
||||
}
|
||||
|
||||
function persistActiveChatSession() {
|
||||
const session = getActiveChatSession();
|
||||
if (!session) return;
|
||||
session.messages = chatHistory.slice();
|
||||
session.model = selectedChatModel || session.model || "";
|
||||
session.title = chatSessionTitle(session);
|
||||
session.updatedAt = new Date().toISOString();
|
||||
chatSessions.sort((a, b) => String(b.updatedAt).localeCompare(String(a.updatedAt)));
|
||||
if (chatSessions.length > CHAT_SESSIONS_LIMIT) {
|
||||
chatSessions = chatSessions.slice(0, CHAT_SESSIONS_LIMIT);
|
||||
if (!chatSessions.some(item => item.id === activeChatSessionId)) {
|
||||
activeChatSessionId = chatSessions[0].id;
|
||||
chatHistory = chatSessions[0].messages.slice();
|
||||
}
|
||||
}
|
||||
saveChatSessionsStore();
|
||||
renderChatSessionList();
|
||||
}
|
||||
|
||||
function clearChatPrompt() {
|
||||
const promptEl = $("chat-prompt");
|
||||
if (!promptEl) return;
|
||||
promptEl.value = "";
|
||||
promptEl.style.height = "auto";
|
||||
}
|
||||
|
||||
function createNewChatSession() {
|
||||
if (chatBusy) return;
|
||||
const session = {
|
||||
id: newChatSessionId(),
|
||||
title: "New chat",
|
||||
model: selectedChatModel || "",
|
||||
messages: [],
|
||||
createdAt: new Date().toISOString(),
|
||||
updatedAt: new Date().toISOString(),
|
||||
};
|
||||
chatSessions.unshift(session);
|
||||
activeChatSessionId = session.id;
|
||||
chatHistory = [];
|
||||
clearChatPrompt();
|
||||
saveChatSessionsStore();
|
||||
renderChatSessionList();
|
||||
renderChatHistory();
|
||||
renderChatAuthHint();
|
||||
const promptEl = $("chat-prompt");
|
||||
if (promptEl) promptEl.focus();
|
||||
}
|
||||
|
||||
function selectChatSession(sessionId) {
|
||||
if (chatBusy) return;
|
||||
const session = chatSessions.find(item => item.id === sessionId);
|
||||
if (!session) return;
|
||||
if (sessionId === activeChatSessionId) return;
|
||||
activeChatSessionId = session.id;
|
||||
chatHistory = (session.messages || []).slice();
|
||||
clearChatPrompt();
|
||||
if (session.model) {
|
||||
selectedChatModel = session.model;
|
||||
localStorage.setItem("meshnet_chat_model", selectedChatModel);
|
||||
const select = $("chat-model");
|
||||
if (select) select.value = selectedChatModel;
|
||||
}
|
||||
localStorage.setItem(CHAT_ACTIVE_SESSION_KEY, activeChatSessionId);
|
||||
renderChatSessionList();
|
||||
renderChatHistory();
|
||||
renderChatAuthHint();
|
||||
}
|
||||
|
||||
function deleteChatSession(sessionId, event) {
|
||||
if (event) {
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
}
|
||||
if (chatBusy) return;
|
||||
const index = chatSessions.findIndex(item => item.id === sessionId);
|
||||
if (index < 0) return;
|
||||
chatSessions.splice(index, 1);
|
||||
if (activeChatSessionId === sessionId) {
|
||||
if (chatSessions.length) {
|
||||
activeChatSessionId = chatSessions[0].id;
|
||||
chatHistory = (chatSessions[0].messages || []).slice();
|
||||
clearChatPrompt();
|
||||
if (chatSessions[0].model) {
|
||||
selectedChatModel = chatSessions[0].model;
|
||||
localStorage.setItem("meshnet_chat_model", selectedChatModel);
|
||||
}
|
||||
} else {
|
||||
saveChatSessionsStore();
|
||||
createNewChatSession();
|
||||
return;
|
||||
}
|
||||
}
|
||||
saveChatSessionsStore();
|
||||
renderChatSessionList();
|
||||
renderChatHistory();
|
||||
renderChatModels();
|
||||
}
|
||||
|
||||
function initChatSessions() {
|
||||
chatSessions = loadChatSessionsStore();
|
||||
activeChatSessionId = localStorage.getItem(CHAT_ACTIVE_SESSION_KEY) || "";
|
||||
const active = chatSessions.find(session => session.id === activeChatSessionId);
|
||||
if (!active) {
|
||||
if (chatSessions.length) {
|
||||
activeChatSessionId = chatSessions[0].id;
|
||||
chatHistory = (chatSessions[0].messages || []).slice();
|
||||
if (chatSessions[0].model) selectedChatModel = chatSessions[0].model;
|
||||
} else {
|
||||
createNewChatSession();
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
chatHistory = (active.messages || []).slice();
|
||||
if (active.model) selectedChatModel = active.model;
|
||||
}
|
||||
renderChatSessionList();
|
||||
renderChatHistory();
|
||||
}
|
||||
|
||||
function renderChatSessionList() {
|
||||
const list = $("chat-session-list");
|
||||
if (!list) return;
|
||||
if (!chatSessions.length) {
|
||||
list.className = "chat-session-list empty-state";
|
||||
list.innerHTML = "No chats yet";
|
||||
return;
|
||||
}
|
||||
list.className = "chat-session-list";
|
||||
list.innerHTML = chatSessions.map(session => {
|
||||
const active = session.id === activeChatSessionId ? " active" : "";
|
||||
const title = esc(chatSessionTitle(session));
|
||||
const when = esc(formatSessionTime(session.updatedAt || session.createdAt));
|
||||
const id = JSON.stringify(session.id);
|
||||
return `<div class="chat-session-item${active}" role="button" tabindex="0"` +
|
||||
` onclick="selectChatSession(${id})"` +
|
||||
` onkeydown="if(event.key==='Enter'||event.key===' '){event.preventDefault();selectChatSession(${id});}">` +
|
||||
`<div class="chat-session-title">${title}</div>` +
|
||||
(when ? `<div class="chat-session-meta">${when}</div>` : "") +
|
||||
`<button type="button" class="chat-session-delete" title="Delete chat"` +
|
||||
` onclick="deleteChatSession(${id}, event)">×</button>` +
|
||||
`</div>`;
|
||||
}).join("");
|
||||
}
|
||||
|
||||
function switchDashboardTab(name) {
|
||||
if (name === "admin" && !isAdmin) name = "overview";
|
||||
if (name === "billing" && !isLoggedIn) name = "overview";
|
||||
dashboardTab = name;
|
||||
document.body.classList.toggle("chat-tab-active", name === "chat");
|
||||
updateSectionVisibility();
|
||||
for (const tabName of ["overview", "chat", "billing", "admin"]) {
|
||||
const button = $("tab-" + tabName);
|
||||
if (button) button.classList.toggle("active", tabName === dashboardTab);
|
||||
}
|
||||
if (name === "chat") {
|
||||
const promptEl = $("chat-prompt");
|
||||
if (promptEl) promptEl.focus();
|
||||
}
|
||||
}
|
||||
|
||||
function updateSectionVisibility() {
|
||||
@@ -632,18 +924,18 @@ function renderChatStatus(text) {
|
||||
|
||||
function renderChatHistory() {
|
||||
const history = $("chat-history");
|
||||
if (!history) return;
|
||||
if (!chatHistory.length) {
|
||||
history.classList.add("empty");
|
||||
history.innerHTML = "no messages yet";
|
||||
history.className = "chat-messages empty";
|
||||
history.innerHTML = '<div class="chat-messages-inner">Send a message to start this conversation.</div>';
|
||||
return;
|
||||
}
|
||||
history.classList.remove("empty");
|
||||
history.innerHTML = chatHistory.map(msg => {
|
||||
const roleClass = msg.role === "user" ? "chat-role-user" : msg.role === "assistant" ? "chat-role-assistant" : "chat-role-error";
|
||||
const label = msg.role === "user" ? "user" : msg.role === "assistant" ? "assistant" : "error";
|
||||
const meta = msg.model ? ` <span class="dim">· ${esc(short(msg.model, 24))}</span>` : "";
|
||||
return `<div class="chat-message"><div class="chat-role ${roleClass}">${label}${meta}</div><div>${esc(msg.content)}</div></div>`;
|
||||
history.className = "chat-messages";
|
||||
const rows = chatHistory.map(msg => {
|
||||
const role = msg.role === "user" ? "user" : msg.role === "assistant" ? "assistant" : "error";
|
||||
return `<div class="chat-row ${role}"><div class="chat-bubble ${role}">${esc(msg.content)}</div></div>`;
|
||||
}).join("");
|
||||
history.innerHTML = `<div class="chat-messages-inner">${rows}</div>`;
|
||||
history.scrollTop = history.scrollHeight;
|
||||
}
|
||||
|
||||
@@ -674,12 +966,13 @@ function renderChatModels() {
|
||||
function selectChatModel(value) {
|
||||
selectedChatModel = value || "";
|
||||
localStorage.setItem("meshnet_chat_model", selectedChatModel);
|
||||
}
|
||||
|
||||
function clearChatHistory() {
|
||||
chatHistory = [];
|
||||
renderChatHistory();
|
||||
renderChatStatus("history cleared");
|
||||
const session = getActiveChatSession();
|
||||
if (session) {
|
||||
session.model = selectedChatModel;
|
||||
session.updatedAt = new Date().toISOString();
|
||||
saveChatSessionsStore();
|
||||
renderChatSessionList();
|
||||
}
|
||||
}
|
||||
|
||||
function chatAuthToken() {
|
||||
@@ -935,8 +1228,10 @@ async function sendChat() {
|
||||
chatBusy = true;
|
||||
$("chat-send").disabled = true;
|
||||
promptEl.value = "";
|
||||
promptEl.style.height = "auto";
|
||||
chatHistory.push({ role: "user", content: prompt, model: selectedChatModel });
|
||||
renderChatHistory();
|
||||
persistActiveChatSession();
|
||||
renderChatStatus("sending request…");
|
||||
const r = await apiCall("/v1/chat/completions", "POST", body, bearerToken);
|
||||
chatBusy = false;
|
||||
@@ -947,6 +1242,7 @@ async function sendChat() {
|
||||
: "request failed";
|
||||
chatHistory.push({ role: "error", content: error, model: selectedChatModel });
|
||||
renderChatHistory();
|
||||
persistActiveChatSession();
|
||||
renderChatStatus(error);
|
||||
promptEl.focus();
|
||||
return;
|
||||
@@ -959,12 +1255,29 @@ async function sendChat() {
|
||||
model: selectedChatModel,
|
||||
});
|
||||
renderChatHistory();
|
||||
persistActiveChatSession();
|
||||
renderChatStatus(usage
|
||||
? `done: ${usage.total_tokens ?? "?"} tokens`
|
||||
: "done");
|
||||
promptEl.focus();
|
||||
}
|
||||
|
||||
function bindChatPromptShortcuts() {
|
||||
const promptEl = $("chat-prompt");
|
||||
if (!promptEl || promptEl.dataset.bound === "1") return;
|
||||
promptEl.dataset.bound = "1";
|
||||
promptEl.addEventListener("keydown", event => {
|
||||
if (event.key === "Enter" && !event.shiftKey) {
|
||||
event.preventDefault();
|
||||
sendChat();
|
||||
}
|
||||
});
|
||||
promptEl.addEventListener("input", () => {
|
||||
promptEl.style.height = "auto";
|
||||
promptEl.style.height = Math.min(promptEl.scrollHeight, 200) + "px";
|
||||
});
|
||||
}
|
||||
|
||||
async function renderAdminPanel() {
|
||||
const r = await apiCall("/v1/admin/accounts");
|
||||
if (!r.ok) { setAdminMode(false); return; }
|
||||
@@ -1034,6 +1347,8 @@ async function refresh() {
|
||||
$("refreshed").textContent = "refreshed " + new Date().toLocaleTimeString();
|
||||
}
|
||||
refresh();
|
||||
initChatSessions();
|
||||
bindChatPromptShortcuts();
|
||||
renderAccountPanel();
|
||||
renderChatModels();
|
||||
renderChatHistory();
|
||||
|
||||
@@ -4864,6 +4864,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
"model": resolved_name,
|
||||
"model_layers_end": required_end,
|
||||
"peers": peers,
|
||||
"bytes_per_layer": _preset_bytes_per_layer(preset),
|
||||
"model_sources": self._model_sources(
|
||||
resolved_name,
|
||||
preset,
|
||||
|
||||
@@ -1646,6 +1646,106 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch):
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_preset_startup_rejects_pinned_shard_above_memory_budget(tmp_path, monkeypatch):
|
||||
"""Pinned layer ranges that exceed the node memory budget fail before model load."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 8 * 1024},
|
||||
)
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"big-model": {
|
||||
"layers_start": 0,
|
||||
"layers_end": 39,
|
||||
"hf_repo": "org/big-model",
|
||||
"bytes_per_layer": {"bfloat16": 2 * 1024 * 1024 * 1024},
|
||||
},
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
with pytest.raises(ValueError, match="Pinned shard layers 0–39"):
|
||||
run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model="big-model",
|
||||
shard_start=0,
|
||||
shard_end=39,
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
cache_dir=tmp_path / "shards",
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_preset_model_with_hf_repo_loads_torch_backend(tmp_path, monkeypatch, capsys):
|
||||
"""Named presets that advertise hf_repo must load TorchNodeServer, not the stub server."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 16
|
||||
|
||||
torch_calls: list[dict] = []
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
torch_calls.append(kwargs)
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
self.chat_completion_count = 0
|
||||
self.tracker_node_id = None
|
||||
|
||||
def start(self):
|
||||
self.port = 7002
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "StubNodeServer", lambda **_kw: (_ for _ in ()).throw(AssertionError("preset with hf_repo must not use StubNodeServer")))
|
||||
|
||||
model_dir = tmp_path / "node-shards" / "tiny-llama"
|
||||
model_dir.mkdir(parents=True)
|
||||
(model_dir / "config.json").write_text('{"num_hidden_layers": 16}')
|
||||
monkeypatch.setattr(startup_mod, "download_shard", lambda *_a, **_kw: model_dir)
|
||||
|
||||
tracker = TrackerServer(model_presets={
|
||||
"tiny-llama": {"layers_start": 0, "layers_end": 15, "hf_repo": "org/tiny-llama-shards"}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model="tiny-llama",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
cache_dir=tmp_path / "node-shards",
|
||||
)
|
||||
try:
|
||||
assert len(torch_calls) == 1
|
||||
assert torch_calls[0]["model_id"] == "org/tiny-llama-shards"
|
||||
assert torch_calls[0]["cache_dir"] == model_dir
|
||||
output = capsys.readouterr().out
|
||||
assert "Loading real PyTorch model shard..." in output
|
||||
assert "Model ID: org/tiny-llama-shards" in output
|
||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
||||
registered = network_map["nodes"][0]
|
||||
assert registered["hf_repo"] == "org/tiny-llama-shards"
|
||||
assert registered["num_layers"] == 16
|
||||
finally:
|
||||
node.stop()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_torch_startup_retries_registration_when_tracker_unreachable(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
|
||||
@@ -17,6 +17,7 @@ from meshnet_node.model_backend import (
|
||||
TensorPayload,
|
||||
TorchModelShard,
|
||||
_call_layer,
|
||||
_checkpoint_tensor_name_for_model,
|
||||
_load_partial_model_from_snapshot,
|
||||
_should_partial_materialize_shard,
|
||||
_decoder_attention_mask,
|
||||
@@ -429,7 +430,7 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho
|
||||
39,
|
||||
total_layers_hint=40,
|
||||
uses_quantized_weights=False,
|
||||
) is False
|
||||
) is True
|
||||
assert _should_partial_materialize_shard(
|
||||
str(snapshot_dir),
|
||||
4,
|
||||
@@ -446,6 +447,118 @@ def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapsho
|
||||
) 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_materializes_only_assigned_tensors(tmp_path):
|
||||
snapshot_dir = tmp_path / "snapshot"
|
||||
snapshot_dir.mkdir()
|
||||
|
||||
Reference in New Issue
Block a user