Fix device mismatch in test_gate_up_partition_covers_the_whole_weight - #8308
Merged
tohtana merged 1 commit intoAug 24, 2026
Merged
Conversation
Copy link
Copy Markdown
Contributor
Sorry, something went wrong.
GateUpPack_LinearLayer._tp_partition() ends with self.move(), which relocates the shard to get_accelerator().current_device_name(). The test then compares the reassembled shards against full_weight, which stays on CPU, so torch.equal() raises on any accelerator-backed machine: RuntimeError: Expected all tensors to be on the same device, but got other is on cpu, different from other tensors on cuda:0 Compare against full_weight on the shards' device. No-op on CPU-only runs. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
alanhuangyoo
requested review from
loadams,
tjruwase and
tohtana
as code owners
August 24, 2026 11:07
tohtana
approved these changes
Aug 24, 2026
tohtana
left a comment
Copy link
Copy Markdown
Collaborator
There was a problem hiding this comment.
Thank you for the catch, @alanhuangyoo! Looks good to me.
Sorry, something went wrong.
tohtana
enabled auto-merge
August 24, 2026 16:53
Merged
via the queue into
deepspeedai:master
with commit 84fd92a
Aug 24, 2026
13 of 15 checks passed
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
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.
tests/unit/module_inject/test_fused_repartition.py::test_gate_up_partition_covers_the_whole_weight fails on any accelerator-backed machine:
GateUpPack_LinearLayer._tp_partition() finishes with self.move(_partition), and move() resolves its target as
so each shard comes back on the accelerator. full_weight is never moved, and torch.equal() refuses the cross-device comparison. The sibling test test_gate_up_partition_ignores_later_grain_size_changes compares two post-partition tensors, so both operands share a device and it passes — this is the only case in the file that mixes a partitioned tensor with the original.
This compares against full_weight on the shards' device instead. It is a no-op on CPU-only runs, which is why the test passes there.
Before, on master (8x H20, CUDA):
After:
pre-commit run --files tests/unit/module_inject/test_fused_repartition.py is clean.
The test was added in #8185 four days ago. It is not covered by the live CI: modal-torch-latest runs tests/unit/v1/ only, and the self-hosted GPU workflows that would run unit/ have not produced a run in a long time (nv-a6000 last ran 2025-08-01, nv-nightly 2026-01-15, nv-torch-latest-v100 and nv-inference have no runs listed).
Env: torch 2.13.0+cu130, H20 (sm90), single node.