DeepSpeed

84fd92aa - Fix device mismatch in test_gate_up_partition_covers_the_whole_weight (#8308)

Fix device mismatch in test_gate_up_partition_covers_the_whole_weight (#8308) `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 ```python 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. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>