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() diff --git a/tunix/rl/common.py b/tunix/rl/common.py index ac8fc1330..4515b0f35 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,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)