From bae8c81e69f4e2f1d5025e1d26a8a83373121534 Mon Sep 17 00:00:00 2001 From: Yassir Atlas Date: Mon, 29 Jun 2026 22:03:13 +0000 Subject: [PATCH 1/3] Check that model supports segment ids before passing them when computing score --- tunix/rl/common.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tunix/rl/common.py b/tunix/rl/common.py index e84c290c4..ddb3fa16e 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -568,12 +568,25 @@ def compute_score( segment_positions, ) + try: + if hasattr(model, "transformer"): + target_call = model.transformer.__call__ + else: + target_call = model.__call__ + + sig = inspect.signature(target_call) + has_segment_ids = ("segment_ids" in sig.parameters) or any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + except Exception: + has_segment_ids = False + model_kwargs = {"positions": calculated_positions, "cache": None} - if segment_ids is not None: + if has_segment_ids and segment_ids is not None: model_kwargs["segment_ids"] = segment_ids else: model_kwargs["attention_mask"] = attn_mask - if input_seg_ids is not None: + if has_segment_ids and input_seg_ids is not None: model_kwargs["segment_ids"] = input_seg_ids out = model(prompt_completion_ids, **model_kwargs) From 4e90b4365b9763c218440f40e7062f25042dd0dc Mon Sep 17 00:00:00 2001 From: Yassir Atlas Date: Wed, 1 Jul 2026 18:34:28 +0000 Subject: [PATCH 2/3] Cache segment-id check --- tunix/rl/common.py | 48 ++++++++++++++++++++++------------------------ 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/tunix/rl/common.py b/tunix/rl/common.py index ddb3fa16e..77774b93e 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -13,7 +13,7 @@ # limitations under the License. """Common RL helper classes and functions.""" -from functools import partial # pylint: disable=g-importing-member +import functools import inspect from typing import Any, Iterable @@ -212,7 +212,7 @@ def get_per_token_logps( # TODO(abheesht): This is computed 4 times - twice in `compute_per_token_logps` # and twice in `compute_score`. We can factor this out and compute it just once. -@partial(jax.jit, static_argnames=("pad_id", "eos_id")) +@functools.partial(jax.jit, static_argnames=("pad_id", "eos_id")) def process_ids( prompt_tokens: jax.Array, completion_tokens: jax.Array, @@ -272,8 +272,26 @@ def process_ids( return prompt_completion_ids, positions, attn_mask, input_seg_ids +@functools.cache +def _call_contains_by_type(target_cls: type, target_arg: str) -> bool: + """Determines if a class' call function contains a target argument and caches the result""" -@partial( + try: + sig = inspect.signature(target_cls.__call__) + return (target_arg in sig.parameters) or any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + except Exception: + return False + +def model_call_contains(model, target_arg: str) -> bool: + """Determines if a model's call function contains a target argument""" + + target_obj = model.transformer if hasattr(model, "transformer") else model + + return _call_contains_by_type(type(target_obj), target_arg) + +@functools.partial( jax.jit, static_argnames=( "pad_id", @@ -365,15 +383,7 @@ def compute_per_token_logps( # precedence; otherwise we pass the per-position non-pad mask derived in # ``process_ids`` so flash-attention variants that lack a separate # padding-mask input still skip pad positions. - try: - sig = inspect.signature(model.__call__) - has_segment_ids = ("segment_ids" in sig.parameters) or any( - p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() - ) - except Exception: - has_segment_ids = False - - if has_segment_ids: + if model_call_contains(model, "segment_ids"): if segment_ids is not None: model_kwargs["segment_ids"] = segment_ids elif input_seg_ids is not None: @@ -568,19 +578,7 @@ def compute_score( segment_positions, ) - try: - if hasattr(model, "transformer"): - target_call = model.transformer.__call__ - else: - target_call = model.__call__ - - sig = inspect.signature(target_call) - has_segment_ids = ("segment_ids" in sig.parameters) or any( - p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() - ) - except Exception: - has_segment_ids = False - + has_segment_ids = model_call_contains(model, "segment_ids") model_kwargs = {"positions": calculated_positions, "cache": None} if has_segment_ids and segment_ids is not None: model_kwargs["segment_ids"] = segment_ids From af9f790d1e726871b1a52bab93431c61973aa6a6 Mon Sep 17 00:00:00 2001 From: Yassir Atlas Date: Fri, 10 Jul 2026 23:32:05 +0000 Subject: [PATCH 3/3] Unit tests --- tests/rl/common_test.py | 89 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/rl/common_test.py b/tests/rl/common_test.py index ffabd3f59..af3dce3ee 100644 --- a/tests/rl/common_test.py +++ b/tests/rl/common_test.py @@ -713,6 +713,95 @@ def test_aggregate_loss_bf16(self): self.assertEqual(loss.dtype, jnp.float32) self.assertAlmostEqual(loss, 1.5, places=5) + class ModelSupportingSegmentIds(tc.ToyTransformer): + """Model where __call__ explicitly accepts `segment_ids`.""" + + def __call__( + self, + x, + positions, + cache=None, + attention_mask=None, + output_hidden_states=False, + images=None, + segment_ids: jax.Array | None = None, + skip_lm_head: bool = False, + ): + if segment_ids is None: + raise AssertionError("Expected segment_ids to be passed, but got None.") + return super().__call__( + x, + positions, + cache=cache, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + images=images, + segment_ids=segment_ids, + skip_lm_head=skip_lm_head, + ) + + class ModelWithoutSegmentIds(tc.ToyTransformer): + """Model where __call__ does NOT accept `segment_ids` or `**kwargs`.""" + + def __call__( + self, + x, + positions, + cache=None, + attention_mask=None, + output_hidden_states=False, + images=None, + skip_lm_head: bool = False, + ): + # If common.py incorrectly passed `segment_ids=...` to this method, + # Python would automatically raise: + # TypeError: ModelWithoutSegmentIds.__call__() got an unexpected keyword argument 'segment_ids' + return super().__call__( + x, + positions, + cache=cache, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + images=images, + segment_ids=None, + skip_lm_head=skip_lm_head, + ) + + @parameterized.named_parameters( + dict( + testcase_name="explicit_parameter", + model_cls=ModelSupportingSegmentIds, + ), + dict( + testcase_name="unsupported_ignored_without_type_error", + model_cls=ModelWithoutSegmentIds, + ), + ) + def test_segment_ids_passed_only_when_supported(self, model_cls): + model = model_cls(config=tc.ModelConfig(), rngs=nnx.Rngs(0)) + graphdef, state = nnx.split(model) + + prompt_tokens = jnp.zeros((1, 0), dtype=jnp.int32) + completion_tokens = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + segment_ids = jnp.ones((1, 4), dtype=jnp.int32) + segment_positions = jnp.arange(4, dtype=jnp.int32) + + # 1. For ModelSupportingSegmentIds & ModelWithVarKeyword, our custom assertion inside + # __call__ verifies that segment_ids is passed and not None. + # 2. For ModelWithoutSegmentIds, if common.py passed segment_ids when unsupported, + # Python would raise a TypeError: got an unexpected keyword argument 'segment_ids'. + # Completing cleanly proves it was correctly filtered out. + logps = common.compute_per_token_logps( + graphdef, + state, + prompt_tokens, + completion_tokens, + pad_id=0, + eos_id=-1, + segment_ids=segment_ids, + segment_positions=segment_positions, + ) + self.assertEqual(logps.shape, (1, 4)) if __name__ == "__main__": absltest.main()