Agentic GRPO improvements: sampler-IS correction, eval fix, flash attn

Key changes:
- Fix PeftTrainer early exit when is_managed_externally=True (prevents
  spurious max_steps-triggered break inside the externally-managed
  agentic training loop)
- Fix eval deduplication in agentic_rl_learner: eval at a step boundary
  fired grad_accum_steps times instead of once; guard with
  _last_eval_train_step to skip repeat evals within the same train_step
- Add token-level truncated importance-sampling (TIS) correction in
  agentic GRPO to account for sampler-trainer log-probability drift in
  multi-turn rollouts (sampler_is='token', configurable threshold)
- Add sampler_is_weights field to TrainExample; apply in GRPO loss before
  aggregation
- Log sampler-trainer prob_diff and pearson correlation every training
  step for diagnosing numerical alignment between rollout and trainer
- Qwen3 model: thread segment_ids through flash-attention (splash kernel)
  forward pass so left-padded prompts do not contaminate attention output
- vllm_sampler: fix EOS token handling to prevent off-by-one in
  completion mask construction
- trajectory_collect_engine: include sampler logprobs in trajectory
  output so the GRPO learner can apply the TIS correction without a
  second forward pass; fix conversation mask alignment

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: colincai-mc

tunix/generate/vllm_sampler.py

@@ -349,10 +349,13 @@ def detokenize(
         input_strings, request_outputs
     ):
       for idx, single_output in enumerate(multi_sampling_output.outputs):
-        # vLLM still returns 1 eos id even if we ask it to stop at eos.
-        if single_output.token_ids[-1] == self.tokenizer.eos_id():
-          single_output.token_ids = single_output.token_ids[:-1]
-          single_output.logprobs = single_output.logprobs[:-1]
+        # KEEP the eos token in the returned token_ids — needed so multi-turn
+        # consumers (agentic engine) can reconstruct the exact sequence the
+        # next turn's prompt was rendered from. Combined with
+        # `include_stop_str_in_output=True`, vLLM emits one eos at the end of
+        # each generation. Stripping it (the previous behavior) made
+        # trainer-side concatenation miss `<|im_end|>` at every turn boundary
+        # and produced 30+ nat sampler-trainer logp diffs.
 
         out_tokens[idx].append(
             np.array(single_output.token_ids, dtype=np.int32)
@@ -461,6 +464,14 @@ def __call__(
         sampling_params.prompt_logprobs = 0
       sampling_params.stop_token_ids = [self.tokenizer.eos_id()]
       sampling_params.skip_special_tokens = True
+      # Keep the stop token in the returned ``token_ids`` so multi-turn
+      # consumers can reconstruct the exact sequence the model was sampled
+      # on. This makes the trainer-side concatenation align with what
+      # ``apply_chat_template`` produces for the next turn's prompt; without
+      # it, the trailing ``<|im_end|>`` (or equivalent eos token) is missing
+      # at every turn boundary in the recorded sequence, biasing logp
+      # recomputation against the model's actual sampling context.
+      sampling_params.include_stop_str_in_output = True
 
       if top_p is not None:
         sampling_params.top_p = top_p

tunix/models/qwen3/model.py

@@ -486,6 +486,7 @@ def block(
       segment_pos: jaxtyping.Array,
       cache: LayerCache | None,
       attn_mask: jaxtyping.Array | None,
+      segment_ids: jaxtyping.Array | None = None,
   ) -> tuple[LayerCache | None, jaxtyping.Array]:
     seq_len = x.shape[1]
 
@@ -571,19 +572,59 @@ def block(
           shd.NamedSharding(mesh, P(shd_n, shd_t))
       )
 
-      @partial(
-          shard_map,
-          mesh=mesh,
-          in_specs=(kernel_spec, shd_spec, unsharded_seq, unsharded_seq),
-          out_specs=shd_spec,
-          check_rep=False,
-      )
-      def sharded_splash_attn(kernel, q_block, k_block, v_block):
-        return jax.vmap(kernel)(q_block, k_block, v_block)
+      # Per-position segment ids let splash suppress cross-segment attention
+      # (e.g. real-token to pad-token, or sequence-packing cross-boundary).
+      # The pallas splash kernel only accepts a static causal mask kernel-side,
+      # so per-batch dynamic padding masks have to flow in via segment_ids.
+      if segment_ids is not None:
+        seg_spec = P(shd_b, shd_t)
+        unsharded_seg_spec = P(shd_b, None)
+
+        @partial(
+            shard_map,
+            mesh=mesh,
+            in_specs=(
+                kernel_spec,
+                shd_spec,
+                unsharded_seq,
+                unsharded_seq,
+                seg_spec,
+                unsharded_seg_spec,
+            ),
+            out_specs=shd_spec,
+            check_rep=False,
+        )
+        def sharded_splash_attn(
+            kernel, q_block, k_block, v_block, q_seg_block, kv_seg_block
+        ):
+          seg_ids = splash.SegmentIds(q=q_seg_block, kv=kv_seg_block)
+          return jax.vmap(kernel)(
+              q_block, k_block, v_block, segment_ids=seg_ids
+          )
+
+        qkv = sharded_splash_attn(
+            splash_attn_kernel,
+            query_proj,
+            key_proj,
+            value_proj,
+            segment_ids,
+            segment_ids,
+        )
+      else:
 
-      qkv = sharded_splash_attn(
-          splash_attn_kernel, query_proj, key_proj, value_proj
-      )
+        @partial(
+            shard_map,
+            mesh=mesh,
+            in_specs=(kernel_spec, shd_spec, unsharded_seq, unsharded_seq),
+            out_specs=shd_spec,
+            check_rep=False,
+        )
+        def sharded_splash_attn(kernel, q_block, k_block, v_block):
+          return jax.vmap(kernel)(q_block, k_block, v_block)
+
+        qkv = sharded_splash_attn(
+            splash_attn_kernel, query_proj, key_proj, value_proj
+        )
       qkv = qkv.transpose(0, 2, 1, 3)
     else:
       # GQA
@@ -621,6 +662,7 @@ def __call__(
       segment_pos: jaxtyping.Array,
       cache: LayerCache | None,
       attn_mask: jaxtyping.Array | None,
+      segment_ids: jaxtyping.Array | None = None,
   ) -> tuple[LayerCache | None, jaxtyping.Array]:
     if (
         self.config.remat_config == RematConfig.BLOCK
@@ -629,10 +671,10 @@ def __call__(
       # nnx.remat needs to be applied to the unbound function and take self
       # as the first argument.
       return nnx.remat(self.block.__func__)(
-          self, x, segment_pos, cache, attn_mask
+          self, x, segment_pos, cache, attn_mask, segment_ids
       )
     else:
-      return self.block(x, segment_pos, cache, attn_mask)
+      return self.block(x, segment_pos, cache, attn_mask, segment_ids=segment_ids)
 
   @property
   def head_dim(self):
@@ -1052,13 +1094,15 @@ def block(
       segment_pos: jaxtyping.Array,
       cache: LayerCache | None,
       attn_mask: jaxtyping.Array,
+      segment_ids: jaxtyping.Array | None = None,
   ) -> tuple[LayerCache | None, jaxtyping.Array]:
     inputs_normalized = self.input_layernorm(x)
     cache, attn_output = self.attn(
         inputs_normalized,
         segment_pos,
         cache,
         attn_mask,
+        segment_ids=segment_ids,
     )
     attn_output += x
     residual = attn_output
@@ -1073,14 +1117,19 @@ def __call__(
       segment_pos: jaxtyping.Array,
       cache: LayerCache | None,
       attn_mask: jaxtyping.Array,
+      segment_ids: jaxtyping.Array | None = None,
   ) -> tuple[LayerCache | None, jaxtyping.Array]:
     if (
         self.config.remat_config == RematConfig.DECODER
         or self.config.remat_config == RematConfig.DECODER.value
     ):
-      return nnx.remat(self.block.__func__)(self, x, segment_pos, cache, attn_mask)
+      return nnx.remat(self.block.__func__)(
+          self, x, segment_pos, cache, attn_mask, segment_ids
+      )
     else:
-      return self.block(x, segment_pos, cache, attn_mask)
+      return self.block(
+          x, segment_pos, cache, attn_mask, segment_ids=segment_ids
+      )
 
 
 class Qwen3(BackendMappingMixin, nnx.Module):
@@ -1146,6 +1195,7 @@ def __call__(
       cache: Cache | None,  # (sequence length L')
       attention_mask: jaxtyping.Array,  # [B, L, L']
       output_hidden_states: bool = False,
+      segment_ids: jaxtyping.Array | None = None,  # [B, L]
   ) -> tuple[jaxtyping.Array, Cache | None]:
     """Qwen3 model.
 
@@ -1155,6 +1205,11 @@ def __call__(
       cache: Attention KV cache or None.
       attention_mask: transformer input mask.
       output_hidden_states: whether to output the hidden states.
+      segment_ids: optional per-position segment identifiers, [B, L]. Used by
+        flash attention to suppress cross-segment attention (e.g. real-token
+        to pad-token, or sequence-packing across document boundaries). Pass a
+        1/0 mask to skip pad positions; pass increasing integer ids per packed
+        document for sequence packing.
 
     Returns:
       predicted_logits, new_cache
@@ -1173,6 +1228,7 @@ def __call__(
           positions,
           layer_cache,
           attention_mask,
+          segment_ids=segment_ids,
       )
       if cache is not None:
         new_cache[layer_name] = layer_cache  # pytype: disable=container-type-mismatch