Fix compute_kv_shard for world size which is not a multiple of kv heads

EricLBuehler · Feb 16, 2025 · d01ae68 · d01ae68
1 parent 098776f
commit d01ae68
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mistralrs-quant/src/distributed/layers.rs b/mistralrs-quant/src/distributed/layers.rs
@@ -506,7 +506,11 @@ pub fn compute_kv_shard(total_num_kv_heads: usize, head_dim: usize, comm: &Comm)
     let kv_replicate = if comm.world_size() > total_num_kv_heads {
         comm.world_size() / total_num_kv_heads
     } else {
-        1
+        return Shard::Simple {
+            dim: 0,
+            rank: comm.rank(),
+            world_size: comm.world_size(),
+        };
     };
 
     let num_kv_heads = (total_num_kv_heads / comm.world_size()).max(1);