From 1183e1e06b4fe955482a528ea86d5331e727cce7 Mon Sep 17 00:00:00 2001 From: Nicolas 'Pixel' Noble Date: Wed, 12 Aug 2026 14:53:08 -0700 Subject: [PATCH 1/2] fix(reshard): validate published elsize against its dtype The shard table carries dtype and elsize as independent fields and nothing reconciled them. elsize drives raw address arithmetic in the slice plan (src_addr = addr + offset * elsize, nbytes = n * elsize), so a publisher emitting a byte width that disagrees with its dtype produced reads at the wrong offset and the wrong length instead of an error. - build_sources now rejects a tensor whose elsize disagrees with its resolved dtype's itemsize. - merge_shard_tables now includes elsize in the cross-rank agreement check, which its docstring already claimed it did. - _torch_dtype requires the label to resolve to a torch.dtype. It previously accepted any string naming a torch attribute, so "torch.load" resolved to a function and was carried as a dtype. Signed-off-by: Nicolas 'Pixel' Noble --- .../modelexpress/refit/reshard/rendezvous.py | 30 +++++++-- .../tests/test_reshard_refit_rendezvous.py | 61 +++++++++++++++++++ 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py b/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py index d2b61ecd..01c6bd82 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py +++ b/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()) diff --git a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py b/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py index 8b330f7a..de5ff922 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py +++ b/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)]) From 6cbf264e0fc5dd52f25756a3d9d0c06d05c70c15 Mon Sep 17 00:00:00 2001 From: Nicolas 'Pixel' Noble Date: Fri, 14 Aug 2026 10:34:29 -0700 Subject: [PATCH 2/2] fix(reshard): compare canonicalized dtype labels across ranks The publish contract accepts both the prefixed and unprefixed spelling of a dtype label, but the cross-rank merge compared the raw strings, so two ranks agreeing on a dtype could be reported as an inconsistency. Compare on a canonical key instead. - add _dtype_key, shared by the merge check and dtype resolution so the prefix handling cannot drift between them - reject a non-string dtype label as invalid metadata rather than failing later on a missing attribute - cover both with regression tests, plus a control that genuinely differing dtypes are still rejected Signed-off-by: Nicolas 'Pixel' Noble --- .../modelexpress/refit/reshard/rendezvous.py | 36 +++++++++++++------ .../tests/test_reshard_refit_rendezvous.py | 24 +++++++++++++ 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py b/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py index 01c6bd82..dc69ad6e 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py @@ -164,19 +164,30 @@ def decode_shard_table(blob: bytes) -> list: return tensors -def _torch_dtype(label: str): - """Resolve a shard-table dtype label to a ``torch.dtype``. +def _dtype_key(label) -> str: + """Canonical comparison key for a shard-table dtype label. 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. + is optional and the two spellings of one dtype must not read as a + disagreement. The label comes straight from decoded JSON, so its type is a + publisher's claim rather than a guarantee; a non-string is rejected as + invalid metadata instead of failing later on an attribute it does not have. + """ + if not isinstance(label, str): + raise ValueError(f"unsupported dtype label {label!r} in shard table") + return label.split(".")[-1] + + +def _torch_dtype(label): + """Resolve a shard-table dtype label to a ``torch.dtype``. + + 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 - name = label.split(".")[-1] - dtype = getattr(torch, name, None) + dtype = getattr(torch, _dtype_key(label), None) if not isinstance(dtype, torch.dtype): raise ValueError(f"unsupported dtype label {label!r} in shard table") return dtype @@ -227,7 +238,12 @@ def build_sources(tensors: list) -> tuple: def merge_shard_tables(tables: list) -> list: """Merge per-rank ``list[PublishedTensor]`` into one, concatenating shards for the same source across ranks (reshard fans in cross-rank). full_shape / - dtype / elsize must agree across ranks for a given tensor name.""" + dtype / elsize must agree across ranks for a given tensor name. + + Dtype agreement is decided on the canonicalized label, since the publish + contract accepts both the prefixed and unprefixed spellings: ranks that + spell one dtype differently agree, and reporting them as inconsistent would + name a cross-rank disagreement that does not exist.""" merged: dict = {} for table in tables: for t in table: @@ -239,7 +255,7 @@ def merge_shard_tables(tables: list) -> list: continue if ( cur.full_shape != t.full_shape - or cur.dtype != t.dtype + or _dtype_key(cur.dtype) != _dtype_key(t.dtype) or cur.elsize != t.elsize ): raise ValueError( diff --git a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py b/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py index de5ff922..f02e2e50 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py +++ b/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py @@ -352,3 +352,27 @@ def test_ranks_publishing_a_consistent_tensor_merge_their_shards(): 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)]) + + +def test_ranks_spelling_one_dtype_differently_are_not_a_disagreement(): + # The publish contract accepts both spellings, so ranks that disagree only + # on the prefix agree on the dtype; rejecting them would name a cross-rank + # inconsistency that does not exist. + merged = merge_shard_tables( + [[_tensor(dtype="torch.bfloat16")], [_tensor(dtype="bfloat16")]] + ) + assert len(merged) == 1 + assert len(merged[0].shards) == 2 + + +def test_ranks_publishing_the_same_tensor_with_different_dtype_are_rejected(): + with pytest.raises(ValueError, match="inconsistent shape/dtype/elsize"): + merge_shard_tables( + [[_tensor(dtype="torch.bfloat16")], [_tensor(dtype="float16")]] + ) + + +def test_a_non_string_dtype_label_is_rejected_as_invalid_metadata(): + # dtype rides through decode untouched, so its type is a publisher's claim. + with pytest.raises(ValueError, match="unsupported dtype label"): + build_sources([_tensor(dtype=123, elsize=2)])