fix(reshard): validate published elsize against its dtype - #620
Open
nicolasnoble wants to merge 2 commits into
Open
fix(reshard): validate published elsize against its dtype#620nicolasnoble wants to merge 2 commits into
nicolasnoble wants to merge 2 commits into
Conversation
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 <nicolas@nobis-crew.org>
Contributor
WalkthroughThe reshard rendezvous now validates Torch dtype labels, tensor element sizes, and cross-rank shard metadata. Tests cover valid normalization, invalid attributes, inconsistent metadata, and successful shard merging. ChangesReshard metadata validation
Estimated code review effort: 3 (Moderate) | ~20 minutes Mergeability Score: 🟡 Moderate · up to The PR adds useful elsize validation, but cross-rank merging can still reject valid shard tables when equivalent dtype labels use different prefixes, and malformed labels are not rejected consistently. Merge should wait for the label handling to be corrected. Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py`:
- Around line 240-247: Normalize or resolve dtype labels before the cross-rank
consistency check in the rendezvous merge path, so prefixed and unprefixed
producer forms compare equivalently while preserving shape and element-size
validation. Add a regression test in test_reshard_refit_rendezvous.py covering a
merge with mixed dtype label forms.
Apply the same fix in
`@modelexpress_client/python/tests/test_reshard_refit_rendezvous.py` around lines
328 - 331.
- Around line 178-181: Update decode_shard_table to validate that each dtype
label is a string before calling split("."); raise ValueError for non-string
labels, while preserving the existing unsupported-string validation and dtype
lookup behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 846dbdd7-04b0-46e8-8e5a-c0e4e8b97655
📒 Files selected for processing (2)
modelexpress_client/python/modelexpress/refit/reshard/rendezvous.pymodelexpress_client/python/tests/test_reshard_refit_rendezvous.py
Comment on lines
+178
to
+181
| 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") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Reject non-string dtype labels with ValueError.
decode_shard_table copies JSON dtype values without runtime type validation. A malformed table can therefore pass a non-string value to label.split("."), which raises AttributeError instead of rejecting the label as invalid metadata. Check the label type before splitting it.
Proposed fix
def _torch_dtype(label: str):
import torch
+ if not isinstance(label, str):
+ raise ValueError(f"unsupported dtype label {label!r} in shard table")
name = label.split(".")[-1]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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") | |
| if not isinstance(label, str): | |
| raise ValueError(f"unsupported dtype label {label!r} in shard table") | |
| 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") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py` around
lines 178 - 181, Update decode_shard_table to validate that each dtype label is
a string before calling split("."); raise ValueError for non-string labels,
while preserving the existing unsupported-string validation and dtype lookup
behavior.
Comment on lines
+240
to
+247
| 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}" |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Normalize dtype labels before cross-rank comparison.
The merge path compares raw labels even though the producer contract accepts prefixed and unprefixed forms.
modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py#L240-L247: compare resolved dtypes or canonicalized labels.modelexpress_client/python/tests/test_reshard_refit_rendezvous.py#L328-L331: add a mixed-label merge regression test.
📍 Affects 2 files
modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py#L240-L247(this comment)modelexpress_client/python/tests/test_reshard_refit_rendezvous.py#L328-L331
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py` around
lines 240 - 247, Normalize or resolve dtype labels before the cross-rank
consistency check in the rendezvous merge path, so prefixed and unprefixed
producer forms compare equivalently while preserving shape and element-size
validation. Add a regression test in test_reshard_refit_rendezvous.py covering a
merge with mixed dtype label forms.
Apply the same fix in
`@modelexpress_client/python/tests/test_reshard_refit_rendezvous.py` around lines
328 - 331.
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 <nicolas@nobis-crew.org>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The reshard shard table carries
dtypeandelsizeas independent fields and nothing reconciled them.elsizeis what the slice plan does address arithmetic with,src_addr = addr + offset * elsizeandnbytes = n * elsize, so a publisher emitting a byte width that disagrees with its dtype silently produced reads at the wrong offset and the wrong length.plan_pullalready compares source and destination dtypes; the byte width behind it was checked against nothing.build_sourcesnow rejects a tensor whoseelsizedisagrees with its resolved dtype'sitemsize, andmerge_shard_tablesincludeselsizein the cross-rank agreement check, which its docstring already described it as doing. Both checks sit on the consumer side, since encode and decode stay dependency-free and a publish-side check would cover only the in-tree publishers._torch_dtypealso now requires its label to resolve to an actualtorch.dtype; it was a baregetattr, sotorch.loadresolved to a function and was carried forward as a dtype.Nothing in the repo can trigger this today, since both in-tree publishers derive
elsizefromtensor.element_size(). It seems worth guarding anyway because the publish path is the documented trainer integration contract, andelsizeis a required field on it with no invariant tying it todtype. The six added tests were confirmed failing against the unpatched module.Summary by CodeRabbit
Bug Fixes
Tests