nvalchemi.distributed.reshard_by_destination

nvalchemi.distributed.reshard_by_destination(tensor, destinations, mesh)

Redistribute tensor elements to new ranks based on per-element destinations.

Unlike ShardTensor.redistribute() which changes placement strategy, this physically moves elements between ranks based on a destination map. Returns a plain torch.Tensor with the received elements.

Parameters:

• tensor (torch.Tensor) – Local tensor, shape (N_local, ...).

• destinations (torch.Tensor) – (N_local,) int tensor where destinations[i] is the rank that should own element i after resharding.

• mesh (DeviceMesh) – 1D device mesh for communication.

Returns:

Received elements, shape (N_new_local, ...).

Return type:

torch.Tensor