Skip to content

Adjust reward model's score module and pooler module order for reducing computation#1956

Merged
merrymercy merged 9 commits intosgl-project:mainfrom
aqweteddy:gemma2-rm
Nov 8, 2024
Merged

Adjust reward model's score module and pooler module order for reducing computation#1956
merrymercy merged 9 commits intosgl-project:mainfrom
aqweteddy:gemma2-rm

Conversation

@aqweteddy
Copy link
Copy Markdown
Contributor

Motivation

Modifications

  • Adjust the order of the pooler and score modules in class LlamaForSequenceClassification and class Gemma2ForSequenceClassification to reduce computation.
  • Remove redundant model-loading code in gemma2 reward model.

Checklist

  • [V] Format your code according to the Contributor Guide.
  • [V] Add unit tests as outlined in the Contributor Guide.
  • [V] Update documentation as needed, including docstrings or example tutorials.

return self.pooler(scores, forward_batch)
return EmbeddingPoolerOutput(scores)

def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

can you also simplify the weight loader of LlamaForSequenceClassification?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Done & verified.
image

@merrymercy
Copy link
Copy Markdown
Contributor

Can you fix the lint error?

Comment thread python/sglang/srt/models/gemma2_reward.py Outdated
Comment thread .pre-commit-config.yaml Outdated
@merrymercy merrymercy merged commit 4ade15d into sgl-project:main Nov 8, 2024
@merrymercy
Copy link
Copy Markdown
Contributor

@aqweteddy Thanks for the contribution. It is merged.

@aqweteddy aqweteddy deleted the gemma2-rm branch November 8, 2024 08:18
timethink pushed a commit to timethink/sglang that referenced this pull request Mar 9, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants