[AMD][Fix] Restore aiter-padded MoE weight dims for serialized checkpoints (#33090)
AMD only change
This commit is contained in:
@@ -196,9 +196,18 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
|||||||
extra_weight_attrs["weight_loader"] = weight_loader
|
extra_weight_attrs["weight_loader"] = weight_loader
|
||||||
|
|
||||||
# WEIGHTS
|
# WEIGHTS
|
||||||
|
# Serialized MXFP4 must keep the aiter-aligned dims from
|
||||||
|
# get_moe_weight_sizes(): w13_weight_scale / w2_weight_scale below are
|
||||||
|
# still sized off w13_up_dim / w2_down_dim, and `weight_padded` is
|
||||||
|
# advertised to the loader, so allocating the raw (unpadded) dims here
|
||||||
|
# desyncs the weights from their scales.
|
||||||
w13_shape = (
|
w13_shape = (
|
||||||
num_experts,
|
num_experts,
|
||||||
2 * intermediate_size_per_partition,
|
(
|
||||||
|
w13_up_dim
|
||||||
|
if self.is_checkpoint_mxfp4_serialized
|
||||||
|
else 2 * intermediate_size_per_partition
|
||||||
|
),
|
||||||
hidden_size // 2 if self.is_checkpoint_mxfp4_serialized else hidden_size,
|
hidden_size // 2 if self.is_checkpoint_mxfp4_serialized else hidden_size,
|
||||||
)
|
)
|
||||||
w13_weight = torch.nn.Parameter(
|
w13_weight = torch.nn.Parameter(
|
||||||
@@ -217,7 +226,7 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
|||||||
num_experts,
|
num_experts,
|
||||||
hidden_size,
|
hidden_size,
|
||||||
(
|
(
|
||||||
intermediate_size_per_partition // 2
|
w2_down_dim
|
||||||
if self.is_checkpoint_mxfp4_serialized
|
if self.is_checkpoint_mxfp4_serialized
|
||||||
else intermediate_size_per_partition
|
else intermediate_size_per_partition
|
||||||
),
|
),
|
||||||
|
|||||||
Reference in New Issue
Block a user