Verify model supports segment-ids

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()

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,12 +578,13 @@ def compute_score(
       segment_positions,
   )
 
+  has_segment_ids = model_call_contains(model, "segment_ids")
   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)