Shard weights along vocab dim (same as vLLM's VocabParallelEmbedding.weight_loader).
(weights: torch.Tensor, tp: TPInfo)
| 92 | hidden_states = torch.nn.functional.pad(hidden_states, (0, pad)) |
| 93 | return weights, hidden_states |
| 94 | |
| 95 | |
| 96 | def shard_weights(weights: torch.Tensor, tp: TPInfo) -> torch.Tensor: |
| 97 | """Shard weights along vocab dim (same as vLLM's VocabParallelEmbedding.weight_loader).""" |
| 98 | if tp.size == 1: |
| 99 | return weights # early return for single-GPU case |
| 100 | shard_size = weights.shape[0] // tp.size |
| 101 | start_idx = tp.rank * shard_size |
| 102 | shard = weights.narrow(0, start_idx, shard_size) |
| 103 | assert shard.is_contiguous() |
| 104 | return shard |
| 105 |
no outgoing calls
no test coverage detected