diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 959d71c0e..0087545df 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -344,7 +344,10 @@ def _select_tensor_parallel_shard(self, key, tensor): if tensor.shape[0] % self.tp_size: raise ValueError(f"Cannot column-shard {key} shape {tuple(tensor.shape)} across TP size {self.tp_size}") - return torch.chunk(tensor, self.tp_size, dim=0)[self.tp_rank].contiguous() + # A dim-0 chunk is already contiguous, so ``.contiguous()`` would keep + # it as a view backed by the full checkpoint tensor. Materialize the + # local shard so CPU loading does not retain peer ranks' storage. + return torch.chunk(tensor, self.tp_size, dim=0)[self.tp_rank].clone(memory_format=torch.contiguous_format) def _should_load_weights(self): # Each TP rank reads its own slices. This is also correct for TP+SP, diff --git a/test_cases/test_minimax_h3_tp_shards.py b/test_cases/test_minimax_h3_tp_shards.py new file mode 100644 index 000000000..84015a1ca --- /dev/null +++ b/test_cases/test_minimax_h3_tp_shards.py @@ -0,0 +1,21 @@ +import torch + +from lightx2v.models.networks.minimax_h3.model import MiniMaxH3Model + + +def test_column_tp_shard_owns_only_local_storage(): + model = MiniMaxH3Model.__new__(MiniMaxH3Model) + model.config = {"tensor_parallel": True} + model.tp_rank = 2 + model.tp_size = 4 + + source = torch.arange(32, dtype=torch.float32).reshape(8, 4) + shard = model._select_tensor_parallel_shard( + "transformer_blocks.0.attn.to_q.weight", + source, + ) + + torch.testing.assert_close(shard, source[4:6]) + assert shard.is_contiguous() + assert shard.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size()