Fix device mismatch in test_gate_up_partition_covers_the_whole_weight

[alanhuangyoo]

tests/unit/module_inject/test_fused_repartition.py::test_gate_up_partition_covers_the_whole_weight fails on any accelerator-backed machine:

assert torch.equal(torch.cat([gate, up], dim=0), full_weight)
E  RuntimeError: Expected all tensors to be on the same device, but got other
   is on cpu, different from other tensors on cuda:0

GateUpPack_LinearLayer._tp_partition() finishes with self.move(_partition), and move() resolves its target as

device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name()

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):

tests/unit/module_inject/  ->  36 passed, 1 failed

After:

tests/unit/module_inject/test_fused_repartition.py  ->  5 passed
tests/unit/module_inject/                           ->  37 passed

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.

[tohtana]

Thank you for the catch, @alanhuangyoo! Looks good to me.