fix(reshard): validate published elsize against its dtype

modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py

@@ -165,9 +165,21 @@ def decode_shard_table(blob: bytes) -> list:
 
 
 def _torch_dtype(label: str):
+    """Resolve a shard-table dtype label to a ``torch.dtype``.
+
+    Publishers emit either ``"torch.bfloat16"`` or ``"bfloat16"``, so the prefix
+    is optional. The label must name a dtype rather than merely some torch
+    attribute. A fixed allowlist of names is deliberately avoided: the publish
+    format is an external contract, and a hardcoded list would reject dtypes a
+    newer torch supports.
+    """
     import torch
 
-    return getattr(torch, label.split(".")[-1])
+    name = label.split(".")[-1]
+    dtype = getattr(torch, name, None)
+    if not isinstance(dtype, torch.dtype):
+        raise ValueError(f"unsupported dtype label {label!r} in shard table")
+    return dtype
 
 
 def build_sources(tensors: list) -> tuple:
@@ -182,6 +194,12 @@ def build_sources(tensors: list) -> tuple:
     session_to_device = {}
     for t in tensors:
         dtype = _torch_dtype(t.dtype)
+        if t.elsize != dtype.itemsize:
+            raise ValueError(
+                f"tensor {t.name!r} published elsize {t.elsize} disagrees with dtype "
+                f"{t.dtype} (itemsize {dtype.itemsize}); elsize drives raw address "
+                f"arithmetic, so a mismatch would read the wrong bytes"
+            )
         shards = []
         for s in t.shards:
             session = s.agent_name
@@ -219,10 +237,14 @@ def merge_shard_tables(tables: list) -> list:
                     t.name, t.dtype, t.elsize, t.full_shape, list(t.shards)
                 )
                 continue
-            if cur.full_shape != t.full_shape or cur.dtype != t.dtype:
+            if (
+                cur.full_shape != t.full_shape
+                or cur.dtype != t.dtype
+                or cur.elsize != t.elsize
+            ):
                 raise ValueError(
-                    f"tensor {t.name!r} published with inconsistent shape/dtype across ranks: "
-                    f"{cur.full_shape}/{cur.dtype} vs {t.full_shape}/{t.dtype}"
+                    f"tensor {t.name!r} published with inconsistent shape/dtype/elsize across ranks: "
+                    f"{cur.full_shape}/{cur.dtype}/{cur.elsize} vs {t.full_shape}/{t.dtype}/{t.elsize}"
                 )
             cur.shards.extend(t.shards)
     return list(merged.values())

modelexpress_client/python/tests/test_reshard_refit_rendezvous.py

@@ -13,6 +13,8 @@
     PublishedShard,
     PublishedTensor,
     _mx_version,
+    build_sources,
+    merge_shard_tables,
     wrap_rendezvous_blob,
 )
 
@@ -291,3 +293,62 @@ def publish_metadata(self, *_args, **_kwargs):
 
     with pytest.raises(ValueError, match="must be positive"):
         rendezvous.publish(b"registered")
+
+
+def _tensor(dtype="torch.bfloat16", elsize=2, name="weight"):
+    return PublishedTensor(
+        name=name,
+        dtype=dtype,
+        elsize=elsize,
+        full_shape=(4, 4),
+        shards=[
+            PublishedShard(
+                agent_name="trainer-agent",
+                device_id=0,
+                addr=4096,
+                shard_offset=(0, 0),
+                shape=(4, 4),
+            )
+        ],
+    )
+
+
+def test_a_published_elsize_that_disagrees_with_its_dtype_is_rejected():
+    # elsize drives raw address arithmetic in the slice plan, so a wrong value
+    # reads the wrong bytes rather than failing.
+    with pytest.raises(ValueError, match="disagrees with dtype"):
+        build_sources([_tensor(dtype="torch.bfloat16", elsize=4)])
+
+
+def test_a_published_elsize_matching_its_dtype_is_accepted():
+    sources, _, _ = build_sources([_tensor(dtype="torch.bfloat16", elsize=2)])
+    assert sources["weight"].elsize == 2
+
+
+def test_a_stripped_dtype_label_resolves_the_same_as_a_prefixed_one():
+    stripped, _, _ = build_sources([_tensor(dtype="bfloat16", elsize=2)])
+    prefixed, _, _ = build_sources([_tensor(dtype="torch.bfloat16", elsize=2)])
+    assert stripped["weight"].dtype == prefixed["weight"].dtype
+
+
+def test_a_dtype_label_naming_a_non_dtype_torch_attribute_is_rejected():
+    # getattr(torch, "load") resolves to a function; without an allowlist it
+    # would be accepted as a dtype.
+    with pytest.raises(ValueError, match="unsupported dtype label"):
+        build_sources([_tensor(dtype="torch.load", elsize=2)])
+
+
+def test_ranks_publishing_the_same_tensor_with_different_elsize_are_rejected():
+    with pytest.raises(ValueError, match="inconsistent shape/dtype/elsize"):
+        merge_shard_tables([[_tensor(elsize=2)], [_tensor(elsize=4)]])
+
+
+def test_ranks_publishing_a_consistent_tensor_merge_their_shards():
+    merged = merge_shard_tables([[_tensor()], [_tensor()]])
+    assert len(merged) == 1
+    assert len(merged[0].shards) == 2
+
+
+def test_a_dtype_label_that_names_nothing_in_torch_is_rejected():
+    with pytest.raises(ValueError, match="unsupported dtype label"):
+        build_sources([_tensor(dtype="torch.not_a_real_dtype", elsize=2)])