A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[Model] Use sigmoid for single-label classification (#18313)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
May 18, 2025 at 07:00 UTC
908733aca703f90476f91f400a77cbe31fc81d3f
1 file changed
+9
-2
vllm/model_executor/layers/pooler.py
+9
-2
@@ -242,9 +242,16 @@ class PoolerHead(nn.Module):
242
243
if self.softmax:
244
if isinstance(pooled_data, list):
245
- pooled_data = [F.softmax(data, dim=-1) for data in pooled_data]
245
+ pooled_data = [
246
+ F.softmax(data, dim=-1)
247
+ if data.shape[-1] >= 2 else F.sigmoid(data)
248
+ for data in pooled_data
249
+ ]
250
else:
247
- pooled_data = F.softmax(pooled_data, dim=-1)
251
+ if pooled_data.shape[-1] >= 2:
252
+ pooled_data = F.softmax(pooled_data, dim=-1)
253
+ else:
254
+ pooled_data = F.sigmoid(pooled_data)
255
256
return pooled_data
257