Fix update weight error for blackwell DeepGEMM (#13910)

This commit is contained in:
fzyzcjy
2025-11-25 13:28:12 -08:00
committed by GitHub
parent 1ab6ce0e62
commit 64a11303ce
2 changed files with 7 additions and 5 deletions
+4 -2
View File
@@ -736,7 +736,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter):
weight_block_size = self.quant_method.quant_config.weight_block_size
block_n, _ = weight_block_size[0], weight_block_size[1]
raw_block_n, _ = weight_block_size[0], weight_block_size[1]
block_n = 1 if getattr(param, "format_ue8m0", False) else raw_block_n
shard_offset = (
(sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n
) // self.tp_size
@@ -966,7 +967,8 @@ class QKVParallelLinear(ColumnParallelLinear):
if isinstance(param, BlockQuantScaleParameter):
weight_block_size = self.quant_method.quant_config.weight_block_size
block_n, _ = weight_block_size[0], weight_block_size[1]
raw_block_n, _ = weight_block_size[0], weight_block_size[1]
block_n = 1 if getattr(param, "format_ue8m0", False) else raw_block_n
shard_offset = (shard_offset + block_n - 1) // block_n
shard_size = (shard_size + block_n - 1) // block_n