feat: harden node placement and partial model loading

2026-07-13 21:58:08 +02:00
parent a6bcc69288
commit 5d87e81bc9
21 changed files with 497 additions and 55 deletions
--- a/tests/test_real_model_backend.py
+++ b/tests/test_real_model_backend.py
@@ -1075,6 +1075,190 @@ def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path):
    assert model.model.rotary_emb.to_calls == ["cpu:0"]


+def _build_partial_snapshot_fixture(snapshot_dir: Path, *, num_layers: int = 4) -> dict[str, dict[str, int]]:
+    """Minimal HF snapshot with per-tensor numel metadata for memory-scaling tests."""
+    weight_map: dict[str, str] = {
+        "model.embed_tokens.weight": "shard-head.safetensors",
+        "model.norm.weight": "shard-tail.safetensors",
+        "lm_head.weight": "shard-tail.safetensors",
+    }
+    tensor_numel: dict[str, dict[str, int]] = {
+        "shard-head.safetensors": {"model.embed_tokens.weight": 10_000},
+        "shard-tail.safetensors": {
+            "model.norm.weight": 512,
+            "lm_head.weight": 10_000,
+        },
+    }
+    for layer in range(num_layers):
+        rel = f"shard-layer-{layer}.safetensors"
+        weight_map[f"model.layers.{layer}.self_attn.q_proj.weight"] = rel
+        weight_map[f"model.layers.{layer}.mlp.down_proj.weight"] = rel
+        per_layer = 1_000 * (layer + 1)
+        tensor_numel[rel] = {
+            f"model.layers.{layer}.self_attn.q_proj.weight": per_layer,
+            f"model.layers.{layer}.mlp.down_proj.weight": per_layer,
+        }
+        (snapshot_dir / rel).write_bytes(b"stub")
+
+    (snapshot_dir / "config.json").write_text(json.dumps({"num_hidden_layers": num_layers}))
+    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))
+    (snapshot_dir / "shard-head.safetensors").write_bytes(b"stub")
+    (snapshot_dir / "shard-tail.safetensors").write_bytes(b"stub")
+    return tensor_numel
+
+
+def _partial_load_materialized_numel(
+    snapshot_dir: Path,
+    shard_start: int,
+    shard_end: int,
+    tensor_numel: dict[str, dict[str, int]],
+) -> int:
+    """Return the summed numel of checkpoint tensors assigned to one shard load."""
+
+    class FakeModule:
+        def to(self, device):
+            return self
+
+    class FakeModel:
+        def __init__(self, num_layers: int):
+            self.model = types.SimpleNamespace(
+                embed_tokens=FakeModule(),
+                layers=[FakeModule() for _ in range(num_layers)],
+                rotary_emb=FakeModule(),
+                norm=FakeModule(),
+            )
+            self.lm_head = FakeModule()
+
+        def tie_weights(self):
+            pass
+
+    class AutoConfigStub:
+        @staticmethod
+        def from_pretrained(model_id):
+            assert model_id == str(snapshot_dir)
+            return types.SimpleNamespace(num_hidden_layers=num_layers)
+
+    class AutoModelStub:
+        @staticmethod
+        def from_config(cfg, torch_dtype=None):
+            return FakeModel(cfg.num_hidden_layers)
+
+    class UnusedContext:
+        def __enter__(self):
+            return None
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+    num_layers = json.loads((snapshot_dir / "config.json").read_text())["num_hidden_layers"]
+    loaded_numel = 0
+
+    def fake_set_tensor(module, tensor_name, device, value=None, dtype=None):
+        nonlocal loaded_numel
+        loaded_numel += int(getattr(value, "numel", lambda: 0)())
+
+    class FakeTensor:
+        def __init__(self, numel: int):
+            self._numel = numel
+
+        def numel(self) -> int:
+            return self._numel
+
+    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 FakeTensor(tensor_numel[self.filename][tensor_name])
+
+    _load_partial_model_from_snapshot(
+        AutoConfigStub,
+        AutoModelStub,
+        types.SimpleNamespace(),
+        str(snapshot_dir),
+        shard_start,
+        shard_end,
+        "bf16",
+        "cpu:0",
+        init_empty_weights_fn=UnusedContext,
+        set_tensor_fn=fake_set_tensor,
+        safe_open_fn=FakeSafeOpen,
+    )
+    return loaded_numel
+
+
+def test_partial_snapshot_resident_weight_numel_scales_with_shard(tmp_path):
+    "Partial load materializes only assigned-layer weights, not the full checkpoint.\n\nTags: model, node, real-inference"
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    tensor_numel = _build_partial_snapshot_fixture(snapshot_dir, num_layers=4)
+
+    middle_numel = _partial_load_materialized_numel(snapshot_dir, 1, 1, tensor_numel)
+    full_numel = _partial_load_materialized_numel(snapshot_dir, 0, 3, tensor_numel)
+
+    # Layer 1 only: two tensors at 2000 numel each.
+    assert middle_numel == 4_000
+    # Head embed + four layers (two tensors each, increasing sizes) + tail norm/lm_head.
+    assert full_numel == 10_000 + 512 + 10_000 + 2_000 + 4_000 + 6_000 + 8_000
+    assert middle_numel < full_numel // 4
+
+
+@pytest.mark.integration
+def test_partial_snapshot_materialized_param_count_with_real_torch(tmp_path):
+    "When torch is installed, partial load leaves unassigned params on meta device.\n\nTags: model, node, real-inference"
+    torch = _require_functional_torch()
+    pytest.importorskip("transformers")
+    pytest.importorskip("safetensors")
+    pytest.importorskip("accelerate")
+    from safetensors.torch import save_file
+    from transformers import AutoConfig, AutoModelForCausalLM, GPT2Config
+
+    n_layer = 4
+    n_embd = 16
+    snapshot_dir = tmp_path / "snapshot"
+    snapshot_dir.mkdir()
+    config = GPT2Config(n_layer=n_layer, n_embd=n_embd, n_head=2, n_positions=32)
+    (snapshot_dir / "config.json").write_text(config.to_json_string())
+
+    weight_map: dict[str, str] = {}
+    for layer in range(n_layer):
+        key = f"transformer.h.{layer}.attn.c_attn.weight"
+        rel = f"layer-{layer}.safetensors"
+        weight_map[key] = rel
+        save_file({key: torch.ones(n_embd, 3 * n_embd)}, snapshot_dir / rel)
+    (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map}))
+
+    def count_materialized(model) -> tuple[int, int]:
+        materialized = 0
+        meta = 0
+        for param in model.parameters():
+            if param.device.type == "meta":
+                meta += param.numel()
+            else:
+                materialized += param.numel()
+        return materialized, meta
+
+    middle = _load_partial_model_from_snapshot(
+        AutoConfig,
+        AutoModelForCausalLM,
+        torch,
+        str(snapshot_dir),
+        1,
+        1,
+        torch.float32,
+        torch.device("cpu"),
+    )
+    mid_mat, mid_meta = count_materialized(middle)
+    assert mid_mat == n_embd * 3 * n_embd
+    assert mid_meta > mid_mat * 10
+
+
 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"