chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Let reward model take text inputs instead of message lists (#1907)

Browse Source 
Co-authored-by: Kyle Corbitt <kyle@corbt.com>
This commit is contained in:
Lianmin Zheng
2024-11-03 13:27:12 -08:00
committed by GitHub
parent 793b79dbe9
commit 2ce32db6fb
12 changed files with 43 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/models/llama_embedding.py
View File

@@ -36,9 +36,7 @@ class LlamaEmbeddingModel(nn.Module):
hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
return self.pooler(hidden_states, forward_batch)
def load_weights(
self, weights: Iterable[Tuple[str, torch.Tensor]], name=None, loaded_weight=None
):
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
@@ -49,7 +47,7 @@ class LlamaEmbeddingModel(nn.Module):
]
params_dict = dict(self.model.named_parameters())
def load_weights_per_param(name, loaded_weight):
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name or "projector" in name:
return
if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
@@ -78,12 +76,6 @@ class LlamaEmbeddingModel(nn.Module):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
if name is None or loaded_weight is None:
for name, loaded_weight in weights:
load_weights_per_param(name, loaded_weight)
else:
load_weights_per_param(name, loaded_weight)
class MistralModel(LlamaEmbeddingModel):
pass