diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index 606279c01a8a..e12e039299d0 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -118,6 +118,15 @@ class Model(msgspec.Struct): bool, "Whether to use a CausalLM as an embedding model.", ] = False + disable_normalize_embedding: A[ + bool, + "Disable the L2 normalization applied by the embedding Pooler after pooling. " + "By default, models that construct their Pooler with normalize=True L2-normalize " + "the pooled embedding; this flag turns that off at serve time without modifying " + "the model code. Models that already pool without normalization are unaffected, " + "and so are models with a custom pooler that does not go through Pooler (for " + "example the cross-encoder, BERT-pooler and vision-pooler paths).", + ] = False revision: A[ Optional[str], "The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version.", diff --git a/python/sglang/srt/layers/pooler.py b/python/sglang/srt/layers/pooler.py index 4557bc5c94c1..cfa91ef89c81 100644 --- a/python/sglang/srt/layers/pooler.py +++ b/python/sglang/srt/layers/pooler.py @@ -12,6 +12,7 @@ from transformers import PretrainedConfig from sglang.srt.layers.activation import get_cross_encoder_activation_function +from sglang.srt.runtime_context import get_model, is_config_published if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -165,6 +166,20 @@ def score_and_pool( ) +def _disable_normalize_embedding() -> bool: + """Whether ``--disable-normalize-embedding`` is set. + + ``Pooler`` is constructible outside a published server context (offline / + standalone model construction), while config bags fail closed until + ``publish`` has run. Absent a published model namespace there is no flag to + honor, so fall back to the default of keeping normalization. Production + model init runs after ``publish``, so serving always sees the real value. + """ + if not is_config_published("model"): + return False + return get_model().disable_normalize_embedding + + class Pooler(nn.Module): """A layer that pools specific information from hidden states. This layer does the following: @@ -179,6 +194,8 @@ class Pooler(nn.Module): def __init__(self, pooling_type: PoolingType, normalize: bool): super().__init__() self.pooling_type = pooling_type + if normalize and _disable_normalize_embedding(): + normalize = False self.normalize = normalize def forward( diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 92c7889f5ace..34f03d145f92 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -1069,6 +1069,20 @@ def set_server_args(self, server_args: ServerArgs) -> None: self._overrides_log = [] self._publish_role = None + def is_config_published(self, name: str) -> bool: + """Whether ``name``'s config bag has been projected. + + For the rare reader that must tolerate an unpublished context (a layer + constructible offline, outside any published server context) and has a + well-defined pre-publish default. Regular reads must keep using + ``config_bag`` and fail closed. Deliberately does not run the role + namespace check: this answers "is it there", not "may I read it", so a + role violation still raises from the subsequent ``config_bag`` call + rather than being reported as unpublished. + """ + bags = self._config_bags + return bool(bags) and name in bags + def config_bag(self, name: str) -> _ConfigBag: """Return the top-level config namespace bag (``device`` / ``model`` / ``exec`` / ``schedule`` / ``memory`` / ``spec`` / ``lora`` / ``mm`` / @@ -1374,6 +1388,15 @@ def get_device() -> _ConfigBag: return _CONTEXT.config_bag("device") +def is_config_published(name: str) -> bool: + """Whether the ``name`` config namespace has been projected by ``publish``. + + Only for readers that must stay constructible without a published context; + everything else reads the bag directly and fails closed. + """ + return _CONTEXT.is_config_published(name) + + def get_model() -> _ConfigBag: return _CONTEXT.config_bag("model") diff --git a/test/registered/unit/layers/test_pooler_score_and_pool.py b/test/registered/unit/layers/test_pooler_score_and_pool.py index 4b2ce14a4b95..90d37934d6d9 100644 --- a/test/registered/unit/layers/test_pooler_score_and_pool.py +++ b/test/registered/unit/layers/test_pooler_score_and_pool.py @@ -16,6 +16,7 @@ PoolingType, score_and_pool, ) +from sglang.srt.runtime_context import get_context, reset_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -187,5 +188,68 @@ def test_empty_delimiter_indices(self): self.assertEqual(out.embeddings[0].shape, (0, self.num_labels)) +class TestDisableNormalizeEmbedding(CustomTestCase): + """--disable-normalize-embedding turns off the Pooler's L2 normalization. + + ``Pooler.__init__`` reads the resolved flag from the ``model`` config bag, + so these tests publish a context via ``override_server_args`` rather than + faking the accessor. Construction against an *unpublished* context stays + supported and keeps normalization — see + ``test_unpublished_context_keeps_normalization``. + """ + + def setUp(self): + torch.manual_seed(42) + self.hidden = torch.randn(4, 8) + self.fb = _make_forward_batch(extend_seq_lens=[2, 2]) + + def _pooled(self, pooler): + return pooler(self.hidden, self.fb).embeddings + + def test_flag_disables_normalization(self): + with get_context().override_server_args(disable_normalize_embedding=True): + pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + self.assertFalse(pooler.normalize) + pooled = self._pooled(pooler) + + # Un-normalized rows keep the raw hidden-state norms. + expected = self.hidden[torch.tensor([1, 3])] + torch.testing.assert_close(pooled, expected) + + def test_default_keeps_normalization(self): + with get_context().override_server_args(disable_normalize_embedding=False): + pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + self.assertTrue(pooler.normalize) + pooled = self._pooled(pooler) + + torch.testing.assert_close( + pooled.norm(p=2, dim=-1), torch.ones(pooled.shape[0]) + ) + + def test_flag_does_not_force_enable_normalization(self): + """The override is one-way: normalize=False models stay un-normalized.""" + with get_context().override_server_args(disable_normalize_embedding=True): + self.assertFalse( + Pooler(pooling_type=PoolingType.LAST, normalize=False).normalize + ) + with get_context().override_server_args(disable_normalize_embedding=False): + self.assertFalse( + Pooler(pooling_type=PoolingType.LAST, normalize=False).normalize + ) + + def test_unpublished_context_keeps_normalization(self): + """Pooler stays constructible without a published server context. + + Config bags fail closed before ``publish``, so reading the flag + unconditionally would make ``Pooler(normalize=True)`` raise in any + standalone/offline construction path. With no published model + namespace there is no override to honor, so normalization is kept. + """ + self.addCleanup(reset_context) + reset_context() + pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + self.assertTrue(pooler.normalize) + + if __name__ == "__main__": unittest.main()