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