[Bugfix] Use ReplicatedLinear for SequenceClassification head
#23836
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 fixes an accuracy issue with sequence classification models when using tensor parallelism. The change replaces RowParallelLinear with ReplicatedLinear for the classification head. This is the correct approach because the classification logic is executed only on a single rank, and using a tensor-parallel layer like RowParallelLinear would lead to hangs or incorrect results due to its communication primitives. The use of ReplicatedLinear ensures the layer is treated as a standard, non-distributed module, which aligns with the single-rank execution context. The change is sound and effectively resolves the bug.
|
Please set the threshold back to default |
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
DarkLight1337
added
the
ready
Problem solved! Big thanks！ |
…-project#23836) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
…-project#23836) Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
4 participants
Purpose
load_weights_using_from_2_way_softmax.cc @HelloZhang525, can you check if this PR can solve the accuracy issue at TP=4? I only have 2 T4 GPU to investigate this.
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.