Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai
This commit is contained in:
@@ -482,12 +482,36 @@ def _load_partial_model_from_snapshot(
|
|||||||
if callable(tie_weights):
|
if callable(tie_weights):
|
||||||
tie_weights()
|
tie_weights()
|
||||||
|
|
||||||
|
# Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and
|
||||||
|
# multi-token-prediction tensors the text-only CausalLM never builds;
|
||||||
|
# transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected,
|
||||||
|
# so the manual loader must skip them too.
|
||||||
|
expected_keys = _model_state_dict_keys(model)
|
||||||
tensors_by_file: dict[str, list[str]] = {}
|
tensors_by_file: dict[str, list[str]] = {}
|
||||||
|
skipped: list[str] = []
|
||||||
for tensor_name in sorted(tensor_names):
|
for tensor_name in sorted(tensor_names):
|
||||||
rel_file = weight_map.get(tensor_name)
|
rel_file = weight_map.get(tensor_name)
|
||||||
if not isinstance(rel_file, str):
|
if not isinstance(rel_file, str):
|
||||||
continue
|
continue
|
||||||
|
if (
|
||||||
|
expected_keys is not None
|
||||||
|
and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys
|
||||||
|
):
|
||||||
|
skipped.append(tensor_name)
|
||||||
|
continue
|
||||||
tensors_by_file.setdefault(rel_file, []).append(tensor_name)
|
tensors_by_file.setdefault(rel_file, []).append(tensor_name)
|
||||||
|
if skipped:
|
||||||
|
preview = ", ".join(skipped[:3])
|
||||||
|
print(
|
||||||
|
f" Skipping {len(skipped)} checkpoint tensors absent from the causal LM "
|
||||||
|
f"(e.g. {preview})",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
if not tensors_by_file:
|
||||||
|
raise PartialModelLoadUnsupported(
|
||||||
|
f"no checkpoint tensors for layers {shard_start}-{shard_end} match the "
|
||||||
|
f"causal LM built from {snapshot_dir}"
|
||||||
|
)
|
||||||
|
|
||||||
for rel_file, names in tensors_by_file.items():
|
for rel_file, names in tensors_by_file.items():
|
||||||
checkpoint_file = snapshot_dir / rel_file
|
checkpoint_file = snapshot_dir / rel_file
|
||||||
@@ -584,6 +608,17 @@ def _causal_lm_config(cfg: Any) -> Any:
|
|||||||
return cfg
|
return cfg
|
||||||
|
|
||||||
|
|
||||||
|
def _model_state_dict_keys(model: Any) -> set[str] | None:
|
||||||
|
"""Expected parameter/buffer names, or None when the model can't report them."""
|
||||||
|
state_dict = getattr(model, "state_dict", None)
|
||||||
|
if not callable(state_dict):
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return set(state_dict().keys())
|
||||||
|
except Exception:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
|
def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str:
|
||||||
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
|
"""Map multimodal checkpoint keys onto text-only CausalLM modules when needed."""
|
||||||
inner = getattr(model, "model", None)
|
inner = getattr(model, "model", None)
|
||||||
|
|||||||
@@ -5,8 +5,9 @@
|
|||||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||||
<title>meshnet tracker</title>
|
<title>meshnet tracker</title>
|
||||||
<style>
|
<style>
|
||||||
:root { --bg:#0d1117; --panel:#161b22; --border:#30363d; --fg:#c9d1d9;
|
:root { --bg:#0d1117; --panel:#161b22; --border:#30363d; --fg:#e6edf3;
|
||||||
--dim:#8b949e; --accent:#58a6ff; --ok:#3fb950; --bad:#f85149; --warn:#d29922; }
|
--dim:#8b949e; --accent:#58a6ff; --ok:#3fb950; --bad:#f85149; --warn:#d29922;
|
||||||
|
--chat-input-bg:#21262d; --chat-user-bg:#1f4788; --chat-user-border:#388bfd; }
|
||||||
* { box-sizing:border-box; }
|
* { box-sizing:border-box; }
|
||||||
html, body { height:100%; }
|
html, body { height:100%; }
|
||||||
body { margin:0; background:var(--bg); color:var(--fg);
|
body { margin:0; background:var(--bg); color:var(--fg);
|
||||||
@@ -114,7 +115,7 @@
|
|||||||
.chat-session-item:hover .chat-session-delete,
|
.chat-session-item:hover .chat-session-delete,
|
||||||
.chat-session-item.active .chat-session-delete { opacity:1; }
|
.chat-session-item.active .chat-session-delete { opacity:1; }
|
||||||
.chat-session-delete:hover { color:var(--bad); background:#1a1012; }
|
.chat-session-delete:hover { color:var(--bad); background:#1a1012; }
|
||||||
.chat-main { display:flex; flex-direction:column; min-height:0; min-width:0; }
|
.chat-main { display:flex; flex-direction:column; min-height:0; min-width:0; color-scheme:dark; }
|
||||||
.chat-toolbar {
|
.chat-toolbar {
|
||||||
display:flex; gap:12px; align-items:center; flex-shrink:0;
|
display:flex; gap:12px; align-items:center; flex-shrink:0;
|
||||||
padding:10px 16px; border-bottom:1px solid var(--border); background:var(--panel);
|
padding:10px 16px; border-bottom:1px solid var(--border); background:var(--panel);
|
||||||
@@ -122,10 +123,15 @@
|
|||||||
.chat-toolbar label {
|
.chat-toolbar label {
|
||||||
display:flex; align-items:center; gap:8px; color:var(--dim); margin:0;
|
display:flex; align-items:center; gap:8px; color:var(--dim); margin:0;
|
||||||
}
|
}
|
||||||
.chat-toolbar select { min-width:220px; max-width:min(420px, 50vw); }
|
.chat-toolbar select {
|
||||||
|
min-width:220px; max-width:min(420px, 50vw);
|
||||||
|
color:var(--fg); background:var(--chat-input-bg); border:1px solid var(--border);
|
||||||
|
border-radius:6px; padding:6px 8px;
|
||||||
|
}
|
||||||
.chat-status { color:var(--dim); font-size:12px; margin-left:auto; }
|
.chat-status { color:var(--dim); font-size:12px; margin-left:auto; }
|
||||||
.chat-messages {
|
.chat-messages {
|
||||||
flex:1; overflow:auto; padding:24px 16px; min-height:0;
|
flex:1; overflow:auto; padding:24px 16px; min-height:0;
|
||||||
|
background:var(--bg); color:var(--fg);
|
||||||
}
|
}
|
||||||
.chat-messages-inner {
|
.chat-messages-inner {
|
||||||
max-width:768px; margin:0 auto; display:flex; flex-direction:column; gap:20px;
|
max-width:768px; margin:0 auto; display:flex; flex-direction:column; gap:20px;
|
||||||
@@ -139,13 +145,15 @@
|
|||||||
.chat-row.assistant, .chat-row.error { justify-content:flex-start; }
|
.chat-row.assistant, .chat-row.error { justify-content:flex-start; }
|
||||||
.chat-bubble {
|
.chat-bubble {
|
||||||
max-width:85%; padding:12px 14px; border-radius:16px; line-height:1.55;
|
max-width:85%; padding:12px 14px; border-radius:16px; line-height:1.55;
|
||||||
white-space:pre-wrap; word-break:break-word; font-size:13px;
|
white-space:pre-wrap; word-break:break-word; font-size:14px; color:var(--fg);
|
||||||
}
|
}
|
||||||
.chat-bubble.user {
|
.chat-bubble.user {
|
||||||
background:#1f3a5f; border:1px solid #264a73; border-bottom-right-radius:4px;
|
background:var(--chat-user-bg); border:1px solid var(--chat-user-border);
|
||||||
|
border-bottom-right-radius:4px; color:#f0f6fc;
|
||||||
}
|
}
|
||||||
.chat-bubble.assistant {
|
.chat-bubble.assistant {
|
||||||
background:transparent; border:0; padding-left:0; padding-right:0; max-width:100%;
|
background:var(--panel); border:1px solid var(--border);
|
||||||
|
border-bottom-left-radius:4px; max-width:100%; color:var(--fg);
|
||||||
}
|
}
|
||||||
.chat-bubble.error {
|
.chat-bubble.error {
|
||||||
background:#1a1012; border:1px solid #5c2020; color:#ffb4b4; border-bottom-left-radius:4px;
|
background:#1a1012; border:1px solid #5c2020; color:#ffb4b4; border-bottom-left-radius:4px;
|
||||||
@@ -157,16 +165,25 @@
|
|||||||
.chat-compose {
|
.chat-compose {
|
||||||
display:flex; gap:8px; align-items:flex-end; max-width:768px; margin:0 auto;
|
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;
|
padding:10px 12px; border:1px solid var(--border); border-radius:16px;
|
||||||
background:var(--bg);
|
background:var(--chat-input-bg);
|
||||||
|
}
|
||||||
|
.chat-compose:focus-within {
|
||||||
|
border-color:var(--accent);
|
||||||
|
box-shadow:0 0 0 1px var(--accent);
|
||||||
}
|
}
|
||||||
.chat-compose:focus-within { border-color:var(--accent); }
|
|
||||||
.chat-compose textarea {
|
.chat-compose textarea {
|
||||||
flex:1; min-height:24px; max-height:200px; resize:none; width:auto;
|
flex:1; min-height:24px; max-height:200px; resize:none; width:auto;
|
||||||
border:0; background:transparent; padding:2px 0; outline:none;
|
border:0; background:transparent; padding:4px 0; outline:none;
|
||||||
|
color:var(--fg); caret-color:var(--accent); font:inherit; font-size:14px; line-height:1.5;
|
||||||
}
|
}
|
||||||
|
.chat-compose textarea::placeholder { color:var(--dim); opacity:1; }
|
||||||
.chat-compose button {
|
.chat-compose button {
|
||||||
flex-shrink:0; min-width:36px; height:36px; padding:0;
|
flex-shrink:0; min-width:36px; height:36px; padding:0;
|
||||||
border-radius:8px; border:1px solid var(--border);
|
border-radius:8px; border:1px solid var(--chat-user-border);
|
||||||
|
background:var(--chat-user-bg); color:#f0f6fc;
|
||||||
|
}
|
||||||
|
.chat-compose button:hover:not(:disabled) {
|
||||||
|
border-color:var(--accent); background:#2563b8;
|
||||||
}
|
}
|
||||||
.chat-compose button:disabled { opacity:.45; cursor:not-allowed; }
|
.chat-compose button:disabled { opacity:.45; cursor:not-allowed; }
|
||||||
.console {
|
.console {
|
||||||
|
|||||||
@@ -559,6 +559,96 @@ def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path)
|
|||||||
assert set_calls == ["model.layers.1.self_attn.q_proj.weight"]
|
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):
|
def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
|
||||||
snapshot_dir = tmp_path / "snapshot"
|
snapshot_dir = tmp_path / "snapshot"
|
||||||
snapshot_dir.mkdir()
|
snapshot_dir.mkdir()
|
||||||
|
|||||||
Reference in New Issue
Block a user