Code update

PiperOrigin-RevId: 874934348

Author: tianshub

tests/rl/common_test.py

@@ -123,8 +123,10 @@ def test_compute_per_token_logps(self):
     completion_tokens = jnp.array(
         [[10, 11, -1, 12], [10, 11, 12, 13], [10, 11, 12, -1]]
     )
+    graphdef, state = nnx.split(model)
     per_token_logps = common.compute_per_token_logps(
-        model,
+        graphdef,
+        state,
         prompt_tokens,
         completion_tokens,
         pad_id=0,
@@ -142,7 +144,8 @@ def test_compute_per_token_logps(self):
         rtol=1e-2,
     )
     _, logits = common.compute_per_token_logps(
-        model,
+        graphdef,
+        state,
         prompt_tokens,
         completion_tokens,
         pad_id=0,

tunix/generate/sampler.py

@@ -202,6 +202,10 @@ def __init__(
     self._compiled_decode_fn = jax.jit(self._decode_fn)
     self._compiled_prefill_fn = jax.jit(self._prefill_fn)
 
+  def model_def_and_state(self) -> tuple[graph.NodeDef, statelib.State]:
+    """Returns the transformer graphdef and state."""
+    return self._transformer_graphdef, self._flattened_transformer_state
+
   @property
   def transformer(self) -> nnx.Module:
     return nnx.merge(

tunix/rl/common.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Common RL helper classes and functions."""
 
+from functools import partial  # pylint: disable=g-importing-member
 from typing import Any, Iterable
 
 import flax
@@ -177,7 +178,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.
-@nnx.jit(static_argnames=("pad_id", "eos_id"))
+@partial(jax.jit, static_argnames=("pad_id", "eos_id"))
 def process_ids(
     prompt_tokens: jax.Array,
     completion_tokens: jax.Array,
@@ -202,9 +203,13 @@ def process_ids(
   return prompt_completion_ids, positions, attn_mask
 
 
-@nnx.jit(static_argnames=("pad_id", "eos_id", "stop_gradient", "return_logits"))
+@partial(
+    jax.jit,
+    static_argnames=("pad_id", "eos_id", "stop_gradient", "return_logits"),
+)
 def compute_per_token_logps(
-    model: nnx.Module,
+    graphdef,
+    state,
     prompt_tokens: jax.Array,
     completion_tokens: jax.Array,
     pad_id: int,
@@ -214,6 +219,7 @@ def compute_per_token_logps(
     return_logits: bool = False,
 ) -> jax.Array | tuple[jax.Array, jax.Array]:
   """Computes the per-token log probabilities."""
+  model = nnx.merge(graphdef, state)
   input_tokens, positions, attn_mask = process_ids(
       prompt_tokens, completion_tokens, pad_id, eos_id, completion_mask
   )

tunix/rl/experimental/agentic_grpo_learner.py

@@ -33,6 +33,7 @@
 from typing import Any, Dict, List, Sequence, Type, TypeVar
 
 from absl import logging
+from flax import nnx
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -345,7 +346,7 @@ def _process_results(
           completion_tokens=completion_ids,
           pad_id=pad_value,
           eos_id=eos_value,
-          micro_batch_size=1,
+          micro_batch_size=None,
       )
     else:
       ref_per_token_logps = None
@@ -390,7 +391,7 @@ def _process_results(
         rewards=rewards, num_generations=self.algo_config.num_generations
     )
 
-    policy_versions = jnp.array(policy_versions_list, dtype=jnp.int32)
+    policy_versions = np.array(policy_versions_list, dtype=np.int32)
 
     # Log completion lengths.
     agg_completion_mask = completion_mask.sum(axis=-1)
@@ -439,10 +440,7 @@ def _process_results(
         old_per_token_logps=old_per_token_logps,
         policy_version=policy_versions,
     )
-    return [
-        rl_utils.get_batch_slice(combined_batch, slice(i, i + 1))
-        for i in range(self.algo_config.num_generations)
-    ]
+    return [combined_batch]
 
 
 @function_registry.register_policy_loss_fn("agentic_grpo")
@@ -486,10 +484,11 @@ def grpo_loss_fn(
       train_example.completion_mask,
   )
 
-  # TODO(yangmu): trace this part as "actor_inference_and_training".
-  # with perf_tracer.span("...", list(completion_ids.devices())):
+  # TODO(tsbao): split can be avoided with updated peft_trainer model handling.
+  graphdef, state = nnx.split(model)
   per_token_logps = common.compute_per_token_logps(
-      model,
+      graphdef,
+      state,
       prompt_tokens=train_example.prompt_ids,
       completion_tokens=completion_ids,
       pad_id=pad_id,

tunix/rl/experimental/agentic_rl_learner.py

@@ -58,7 +58,7 @@
 
 @flax.struct.dataclass(frozen=True)
 class TrainExample(common.TrainExample):
-  policy_version: jax.Array | None = None
+  policy_version: np.ndarray | None = None
 
 
 @dataclasses.dataclass(slots=True, kw_only=True)
@@ -705,7 +705,7 @@ def train(
 
     # 2. Consume training examples and train.
     train_data_gen = self._data_consumer_batch_generator(
-        train_data_queue, train_micro_batch_size * self._num_generations()
+        train_data_queue, train_micro_batch_size
     )
     micro_batches_since_last_sync = 0
     micro_batches_per_full_batch = full_batch_size // train_micro_batch_size
@@ -720,13 +720,14 @@ def train(
         break
       self._iter_steps += 1
 
+      # TODO(tsbao): Re-enable this once off-policy filtering is needed.
       # Filter out examples that are too old (off-policy).
-      filtered_train_micro_batch = self._filter_outdated_offpolicy_examples(
-          train_micro_batch
-      )
-      if not filtered_train_micro_batch:
-        continue
-      train_micro_batch = filtered_train_micro_batch
+      # filtered_train_micro_batch = self._filter_outdated_offpolicy_examples(
+      #     train_micro_batch
+      # )
+      # if not filtered_train_micro_batch:
+      #   continue
+      # train_micro_batch = filtered_train_micro_batch
 
       merged_train_micro_batch = jax.tree.map(
           lambda *xs: jnp.concatenate(xs, axis=0), *train_micro_batch
@@ -770,7 +771,7 @@ async def _eval_runner_async(current_eval_orchestrator):
       )
       if hasattr(self.rl_cluster, "critic_trainer"):
         self.rl_cluster.update_critic(
-            train_micro_batch, current_eval_dataset, skip_jit
+            [merged_train_micro_batch], current_eval_dataset, skip_jit
         )
 
       # --- Weight Sync Logic ---

tunix/rl/grpo/grpo_learner.py

@@ -20,6 +20,7 @@
 from typing import Iterable, List, Sequence, TypeVar
 
 import flax
+from flax import nnx
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -455,8 +456,10 @@ def grpo_loss_fn(
 
   # TODO(yangmu): trace this part as "actor_inference_and_training".
   # with perf_tracer.span("...", list(completion_ids.devices())):
+  graphdef, state = nnx.split(model)
   per_token_logps = common.compute_per_token_logps(
-      model,
+      graphdef,
+      state,
       prompt_tokens=train_example.prompt_ids,
       completion_tokens=completion_ids,
       pad_id=pad_id,

tunix/rl/inference/inference_worker.py

@@ -16,7 +16,6 @@
 
 from flax import nnx
 import jax
-import jaxtyping
 from tunix.rl import common
 
 
@@ -31,6 +30,9 @@ def __init__(self, models: dict[str, nnx.Module]):
             " reference and reward."
         )
     self._models = models
+    self._model_states = {}
+    for k, m in models.items():
+      self._model_states[k] = nnx.split(m)
     # TODO(tsbao): support multiple reward models.
 
   def get_rewards(
@@ -55,11 +57,12 @@ def get_ref_per_token_logps(
       eos_id: int,
       completion_mask: jax.Array | None = None,
   ) -> jax.Array:
-    ref_model = self._models.get("reference")
-    if ref_model is None:
+    graphdef, state = self._model_states.get("reference")
+    if graphdef is None:
       raise ValueError("Reference model is not available.")
     return common.compute_per_token_logps(
-        ref_model,
+        graphdef,
+        state,
         prompt_tokens=prompt_tokens,
         completion_tokens=completion_tokens,
         pad_id=pad_id,
@@ -77,7 +80,8 @@ def get_values(
       eos_id: int,
       completion_mask: jax.Array | None = None,
   ) -> jax.Array:
-    critic_model = self._models.get("critic")
+    graphdef, state = self._model_states.get("critic")
+    critic_model = nnx.merge(graphdef, state)
     if critic_model is None:
       raise ValueError("Critic model is not available.")
     return common.compute_score(
@@ -93,8 +97,3 @@ def get_model(self, role: str) -> nnx.Module:
     if role not in self._models:
       raise ValueError(f"Model role {role} is not available.")
     return self._models[role]
-
-  def update_model(self, role: str, params: jaxtyping.PyTree):
-    if role not in self._models:
-      raise ValueError(f"Model role {role} is not available.")
-    nnx.update(self._models[role], params)

tunix/rl/ppo/ppo_learner.py

@@ -586,8 +586,10 @@ def ppo_policy_loss_fn(
   use_dual_clip_ppo = epsilon_c is not None
 
   # Get log probs.
+  graphdef, state = nnx.split(model)
   per_token_logps, logits = common.compute_per_token_logps(
-      model,
+      graphdef,
+      state,
       prompt_tokens=prompt_ids,
       completion_tokens=completion_ids,
       pad_id=pad_id,

tunix/rl/rollout/vanilla_rollout.py

@@ -78,8 +78,10 @@ def get_per_token_logps(
       completion_mask: jax.Array | None = None,
   ) -> jax.Array:
     """Returns per-token log probabilities from the rollout policy."""
+    graphdef, state = self._sampler.model_def_and_state()
     return common.compute_per_token_logps(
-        self.model(),
+        graphdef,
+        state,
         prompt_tokens=prompt_tokens,
         completion_tokens=completion_tokens,
         pad_id=self.pad_id(),