transformers
fa073e84 - tp loading: batched pre-slice via torch.chunk instead of per-rank loop

Commit
88 days ago
tp loading: batched pre-slice via torch.chunk instead of per-rank loop Replace the per-rank Python loop (N × tp_layer.rank mutation + shard_tensor dispatch) with a single _batch_shard_for_scatter() that uses torch.chunk / torch.split for GPU-native batched slicing. Falls back to the per-rank loop for unusual TP classes (EmbeddingParallel, MoeIdentityExpertParallel etc.) No measured perf change on NVLink — the bottleneck is the per-mapping sequential scatter calls, not the pre-slicing.
Author
Parents
Loading