[fix]: remove data type hardcoding from gptoss model implementation #23807
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the router in the GptOssMoE layer to use ReplicatedLinear, which enables quantization support and removes a hardcoded bfloat16 dtype. This is a good improvement for flexibility and performance. However, the change also modifies the router's architecture by removing the bias term, which could be a critical issue if the original model checkpoints expect a bias. Overall, the changes are well-targeted to address the quantization issue.
There was a problem hiding this comment.
The original implementation used torch.nn.Linear, which has bias=True by default. This change explicitly sets bias=False. This is a change in the model's architecture. If gpt-oss model checkpoints contain a bias term for the router, this change will cause weight loading to fail. If they do not contain a bias, then this change is a valuable correction, as the previous implementation would have used a randomly initialized bias, leading to incorrect behavior.
|
This pull request has merge conflicts that must be resolved before it can be |
There was a problem hiding this comment.
This is similar to glm4 moe, cannot use ReplicatedLinear
There was a problem hiding this comment.
Thanks for your review
There was a problem hiding this comment.
How should we handle the hardcoded dtype?
There was a problem hiding this comment.
Let's keep the hardcode for now, and then can we add acomment like in glm moe to indicate that we cannot use ReplicatedLinear?
There was a problem hiding this comment.
The hardcode causes a data type missmatch issue when we load model in fp32 or half. The in and out tensor data type conversion for this layer is not handled correctly.
There was a problem hiding this comment.
That makes sense, let's remove the hardcoded type
Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
nikhil-arm
force-pushed
the
gptoss_model_fix
branch
from
892226c to
bfbf330
Compare
nikhil-arm
changed the title
|
@jeejeelee can you please review and merge this PR. thanks! |
github-project-automation
bot
moved this from In progress
to Ready
in gpt-oss Issues & Enhancements
github-project-automation
bot
moved this from Ready
to Done
in gpt-oss Issues & Enhancements
…llm-project#23807) Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
…llm-project#23807) Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
…llm-project#23807) Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com> Signed-off-by: charlifu <charlifu@amd.com>
Purpose
Fix gptoss router layer to support quantization and remove hardcoded data type
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.