[AMD] Fix DeepSeek block-FP8 loading on gfx94x (#38446)
Signed-off-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
@@ -628,7 +628,10 @@ class DeepseekV2WeightLoaderMixin:
|
||||
if hasattr(self_attn.kv_b_proj, "weight_scale")
|
||||
else self_attn.kv_b_proj.weight_scale_inv
|
||||
)
|
||||
if _is_fp8_fnuz:
|
||||
is_ue8m0_uint8 = (
|
||||
weight_scale.format_ue8m0 and weight_scale.dtype == torch.uint8
|
||||
)
|
||||
if _is_fp8_fnuz and not is_ue8m0_uint8:
|
||||
weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
|
||||
weight=w,
|
||||
weight_scale=weight_scale,
|
||||
@@ -638,7 +641,7 @@ class DeepseekV2WeightLoaderMixin:
|
||||
weight = w
|
||||
|
||||
# In multiple weight loading scenarios (e.g. RL), we need to inverse the scale of the weights after the requantization happened at the first loading.
|
||||
if weight_scale.format_ue8m0 and weight_scale.dtype == torch.uint8:
|
||||
if is_ue8m0_uint8:
|
||||
weight_scale = (weight_scale.to(torch.int32) << 23).view(
|
||||
torch.float32
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user