Skip to content

[Bug fix] Fix lm head bias #3185

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jiang-Jia-Jun merged 1 commit into from
Aug 5, 2025
Merged

[Bug fix] Fix lm head bias #3185

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion fastdeploy/model_executor/layers/lm_head.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,13 @@ def __init__(
dtype=paddle.get_default_dtype(),
is_bias=False,
)
if self.bias_key is not None:
self.bias = self.create_parameter(

@yuanlehome yuanlehome Aug 5, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lm_head是否有bias不应该写在use_ep逻辑下吧？应该是通用的

@RichardWooSJTU RichardWooSJTU Aug 5, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TP下面用了ColumnLinear的类，类里面会判断。之所以EP不复用ColumnLinear，原因在于这里直接传入了从fleet获取的num_ranks信息。合并EP和TP可以用fd_config中的ep_rank和tp_rank去判断。限于人力这里没有做重构。

shape=[num_embeddings],
dtype=paddle.get_default_dtype(),
is_bias=True,
)

else:
if self.column_cut:
need_gather = True
Expand Down Expand Up @@ -107,6 +114,10 @@ def load_state_dict(self, state_dict: Dict[str, paddle.Tensor | np.ndarray]):

if self.use_ep:
self.weight.set_value(get_tensor(state_dict.pop(self.weight_key)).astype(paddle.get_default_dtype()))
if self.bias_key is not None:
self.bias.set_value(
get_tensor(state_dict.pop(self.linear_bias_key)).astype(paddle.get_default_dtype())
)
else:
if self.tie_word_embeddings:
self.linear.weight.set_value(
Expand Down Expand Up @@ -134,7 +145,10 @@ def forward(self, input: paddle.Tensor) -> paddle.Tensor:
"""
logits = input
if self.use_ep:
logits = paddle.matmul(logits, self.weight)
if self.linear_bias_key is None:
logits = paddle.matmul(logits, self.weight)
else:
logits = paddle.incubate.nn.functional.fused_linear(logits, self.weight, self.bias)

@heavengate heavengate Aug 4, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这两个底层调用的OP不一致么，会有性能diff么

@RichardWooSJTU RichardWooSJTU Aug 4, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

不一致，paddle.matmul调用cublas，paddle.incubate.nn.functional.fused_linear调用融合的cublasLt，可以将bias的加法融合到epilogue中，理论上比单独调用add算子性能更好

else:
logits = self.linear(logits)
return logits
Loading