inductive-mlxrl 0.2.0__tar.gz → 0.2.1__tar.gz

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to this project are documented here.
 The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project uses semantic versioning once releases are cut.
 
+## [0.2.1] - 2026-06-25
+
+### Fixed
+
+- Keep model-backed agentic/GiGPO decode tokens rank-2 during KV-cache decode,
+  fixing real Qwen3 model rollouts.
+- Use valid repo-relative output paths in the committed Qwen smoke configs.
+- Preserve eval-mode old-policy/reference logprob prep while restoring caller
+  train mode afterward.
+
 ## [0.2.0] - 2026-06-24
 
 ### Added

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inductive-mlxrl
-Version: 0.2.0
+Version: 0.2.1
 Summary: Small single-process RL post-training for LLMs on Apple Silicon with MLX.
 Project-URL: Homepage, https://github.com/inductiveML/mlxrl
 Project-URL: Repository, https://github.com/inductiveML/mlxrl

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/examples/qwen35_9b_g2_checkpoint.toml

@@ -13,7 +13,7 @@ gradient_checkpointing = true
 checkpoint_granularity = "per-layer"
 iogpu_wired_limit_mb = 49152
 micro_batch_size = 0
-output = "/tmp/mlxrl_qwen35_9b_g2_checkpoint.npz"
+output = "reference_outputs/mlxrl_qwen35_9b_g2_checkpoint.npz"
 
 [optimizer]
 name = "adam"

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/examples/qwen3_0_6b_grpo.toml

@@ -13,7 +13,7 @@ gradient_checkpointing = false
 checkpoint_granularity = "per-layer"
 iogpu_wired_limit_mb = 49152
 micro_batch_size = 0
-output = "/tmp/mlxrl_config_smoke.npz"
+output = "reference_outputs/mlxrl_config_smoke.npz"
 
 [optimizer]
 name = "adam"
@@ -27,4 +27,3 @@ top_k = 0
 
 [algorithm]
 name = "grpo"
-

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/__init__.py

@@ -2,4 +2,4 @@
 
 __all__ = ["__version__"]
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/algo/grpo.py

@@ -494,17 +494,28 @@ def filter_zero_advantage_groups(batch: Any, group_size: int) -> Any:
         raise ValueError("DAPO dynamic sampling dropped every group.")
 
     indices = mx.array(keep_indices, dtype=mx.int32)
-    return type(batch)(
-        prompt_token_ids=tuple(batch.prompt_token_ids[index] for index in keep_indices),
-        completion_token_ids=tuple(
+    filtered = {
+        "prompt_token_ids": tuple(batch.prompt_token_ids[index] for index in keep_indices),
+        "completion_token_ids": tuple(
             batch.completion_token_ids[index] for index in keep_indices
         ),
-        rewards=batch.rewards[indices],
-        advantages=batch.advantages[indices],
-        old_policy_logprobs=batch.old_policy_logprobs[indices],
-        reference_logprobs=batch.reference_logprobs[indices],
-        mask=batch.mask[indices],
-    )
+        "rewards": batch.rewards[indices],
+        "advantages": batch.advantages[indices],
+        "old_policy_logprobs": batch.old_policy_logprobs[indices],
+        "reference_logprobs": batch.reference_logprobs[indices],
+        "mask": batch.mask[indices],
+    }
+    if hasattr(batch, "reference_is_policy"):
+        filtered["reference_is_policy"] = batch.reference_is_policy
+    logprob_inputs = getattr(batch, "logprob_inputs", None)
+    if logprob_inputs is not None:
+        filtered["logprob_inputs"] = type(logprob_inputs)(
+            input_ids=logprob_inputs.input_ids[indices],
+            target_ids=logprob_inputs.target_ids[indices],
+            gather_indices=logprob_inputs.gather_indices[indices],
+            mask=logprob_inputs.mask[indices],
+        )
+    return type(batch)(**filtered)
 
 
 def _token_loss_metrics(
@@ -516,17 +527,26 @@ def _token_loss_metrics(
     beta: float,
     clipped_ratio: mx.array | None = None,
 ) -> AlgorithmLossMetrics:
-    token_kl = approximate_kl(policy_logprobs, reference_logprobs)
-    loss = masked_mean(token_policy_gradient + beta * token_kl, mask)
+    mask = mask.astype(mx.float32)
+    denominator = mx.maximum(mx.sum(mask), mx.array(1.0, dtype=mx.float32))
+    policy_gradient_loss = mx.sum(token_policy_gradient * mask) / denominator
+    if beta == 0.0 and reference_logprobs is policy_logprobs:
+        kl = mx.array(0.0, dtype=mx.float32)
+    else:
+        token_kl = approximate_kl(policy_logprobs, reference_logprobs)
+        kl = mx.sum(token_kl * mask) / denominator
+    mean_ratio = mx.sum(ratio * mask) / denominator
     if clipped_ratio is None:
         clip_fraction = mx.array(0.0, dtype=mx.float32)
     else:
-        clip_fraction = masked_mean(clip_indicator(ratio, clipped_ratio), mask)
+        clip_fraction = (
+            mx.sum(clip_indicator(ratio, clipped_ratio) * mask) / denominator
+        )
     return AlgorithmLossMetrics(
-        loss=loss,
-        policy_gradient_loss=masked_mean(token_policy_gradient, mask),
-        kl=masked_mean(token_kl, mask),
-        mean_ratio=masked_mean(ratio, mask),
+        loss=policy_gradient_loss + beta * kl,
+        policy_gradient_loss=policy_gradient_loss,
+        kl=kl,
+        mean_ratio=mean_ratio,
         clip_fraction=clip_fraction,
     )
 

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/policy/logprobs.py

@@ -27,6 +27,16 @@ class CompletionLogprobs:
     mask: mx.array
 
 
+@dataclass(frozen=True)
+class CompletionLogprobInputs:
+    """Prepared padded tensors for a completion full-forward logprob pass."""
+
+    input_ids: mx.array
+    target_ids: mx.array
+    gather_indices: mx.array
+    mask: mx.array
+
+
 @dataclass(frozen=True)
 class DualLogprobs:
     """Policy and reference logprobs for the same completion tokens."""
@@ -61,6 +71,25 @@ def completion_logprobs(
 ) -> CompletionLogprobs:
     """Gather logprobs assigned to completion tokens in a full forward pass."""
 
+    inputs = prepare_completion_logprob_inputs(
+        prompt_token_ids,
+        completion_token_ids,
+        pad_token_id=pad_token_id,
+    )
+    return completion_logprobs_from_inputs(
+        model,
+        inputs,
+        use_checkpoint=use_checkpoint,
+    )
+
+
+def prepare_completion_logprob_inputs(
+    prompt_token_ids: Sequence[Sequence[int]],
+    completion_token_ids: Sequence[Sequence[int]],
+    pad_token_id: int = 0,
+) -> CompletionLogprobInputs:
+    """Build reusable padded tensors for completion logprob gathering."""
+
     if len(prompt_token_ids) != len(completion_token_ids):
         raise ValueError("prompt_token_ids and completion_token_ids must have the same length.")
     if not prompt_token_ids:
@@ -81,8 +110,7 @@ def completion_logprobs(
     input_rows: list[list[int]] = []
     target_rows: list[list[int]] = []
     completion_mask_rows: list[list[float]] = []
-    completion_starts: list[int] = []
-    completion_lengths: list[int] = []
+    gather_index_rows: list[list[int]] = []
 
     for prompt, completion, sequence in zip(
         prompt_token_ids,
@@ -100,38 +128,50 @@ def completion_logprobs(
         mask_row = [0.0] * max_completion_len
         for index in range(len(completion)):
             mask_row[index] = 1.0
+        gather_index_rows.append(
+            [
+                completion_start + index if index < len(completion) else 0
+                for index in range(max_completion_len)
+            ]
+        )
 
         input_rows.append(input_row)
         target_rows.append(target_row)
         completion_mask_rows.append(mask_row)
-        completion_starts.append(completion_start)
-        completion_lengths.append(len(completion))
 
     input_ids = mx.array(input_rows, dtype=mx.int32)
     target_ids = mx.array(target_rows, dtype=mx.int32)
-    logits = _completion_forward(model, input_ids, use_checkpoint=use_checkpoint)
-    all_logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-    target_logprobs = mx.take_along_axis(all_logprobs, target_ids[..., None], axis=-1)
-    target_logprobs = mx.squeeze(target_logprobs, axis=-1)
-
-    rows: list[mx.array] = []
-    for row_index, (start, length) in enumerate(
-        zip(completion_starts, completion_lengths, strict=True)
-    ):
-        row = target_logprobs[row_index, start : start + length]
-        if length < max_completion_len:
-            row = mx.concatenate(
-                [
-                    row,
-                    mx.zeros((max_completion_len - length,), dtype=row.dtype),
-                ],
-                axis=0,
-            )
-        rows.append(row)
+    gather_indices = mx.array(gather_index_rows, dtype=mx.int32)
+    mask = mx.array(completion_mask_rows, dtype=mx.float32)
+    return CompletionLogprobInputs(
+        input_ids=input_ids,
+        target_ids=target_ids,
+        gather_indices=gather_indices,
+        mask=mask,
+    )
+
+
+def completion_logprobs_from_inputs(
+    model: nn.Module,
+    inputs: CompletionLogprobInputs,
+    use_checkpoint: bool = False,
+) -> CompletionLogprobs:
+    """Gather completion-token logprobs from prebuilt full-forward tensors."""
+
+    logits = _completion_forward(model, inputs.input_ids, use_checkpoint=use_checkpoint)
+    target_logprobs = target_logprobs_from_logits(logits, inputs.target_ids)
+    completion_logprob_values = mx.take_along_axis(
+        target_logprobs,
+        inputs.gather_indices,
+        axis=1,
+    )
+    completion_logprob_values = completion_logprob_values * inputs.mask.astype(
+        completion_logprob_values.dtype
+    )
 
     return CompletionLogprobs(
-        logprobs=mx.stack(rows, axis=0),
-        mask=mx.array(completion_mask_rows, dtype=mx.float32),
+        logprobs=completion_logprob_values,
+        mask=inputs.mask,
     )
 
 
@@ -186,16 +226,12 @@ def prefix_cached_completion_logprobs(
             dtype=mx.int32,
         )
         logits = model(chunks, cache=batch_cache)
-        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-        grouped_logprobs = mx.squeeze(
-            mx.take_along_axis(logprobs, targets[..., None], axis=-1),
-            axis=-1,
-        )
+        grouped_logprobs = target_logprobs_from_logits(logits, targets)
         group_mask = mx.array(
             [
                 [1.0] * len(completion_token_ids[index])
                 + [0.0] * (max_completion_len - len(completion_token_ids[index]))
-                for index in row_indices
+            for index in row_indices
             ],
             dtype=grouped_logprobs.dtype,
         )
@@ -216,16 +252,24 @@ def dual_logprobs(
     pad_token_id: int = 0,
     use_checkpoint: bool = False,
     compute_reference: bool = True,
+    prepared_inputs: CompletionLogprobInputs | None = None,
 ) -> DualLogprobs:
     """Compute policy logprobs and reference logprobs using one model object."""
 
-    policy = completion_logprobs(
-        model,
-        prompt_token_ids,
-        completion_token_ids,
-        pad_token_id,
-        use_checkpoint=use_checkpoint,
-    )
+    if prepared_inputs is None:
+        policy = completion_logprobs(
+            model,
+            prompt_token_ids,
+            completion_token_ids,
+            pad_token_id,
+            use_checkpoint=use_checkpoint,
+        )
+    else:
+        policy = completion_logprobs_from_inputs(
+            model,
+            prepared_inputs,
+            use_checkpoint=use_checkpoint,
+        )
     if not compute_reference:
         return DualLogprobs(
             policy=policy.logprobs,
@@ -287,6 +331,22 @@ def _completion_forward(
     return model(input_ids)
 
 
+def target_logprobs_from_logits(logits: mx.array, target_ids: mx.array) -> mx.array:
+    """Gather target-token logprobs without materializing full log-softmax.
+
+    This is algebraically equivalent to gathering from
+    ``logits - logsumexp(logits)`` but avoids a large [batch, seq, vocab]
+    intermediate on training forwards.
+    """
+
+    target_logits = mx.squeeze(
+        mx.take_along_axis(logits, target_ids[..., None], axis=-1),
+        axis=-1,
+    )
+    normalizer = mx.squeeze(mx.logsumexp(logits, axis=-1, keepdims=True), axis=-1)
+    return target_logits - normalizer
+
+
 def pad_token_id_from_tokenizer(tokenizer: Any) -> int:
     """Resolve a safe right-padding token id from a tokenizer."""
 

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/policy/model.py

@@ -6,7 +6,8 @@ per-layer checkpoint helpers (MIT, Copyright Apple Inc.).
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any, cast
 
@@ -152,6 +153,20 @@ def enable_grad_checkpointing(model: nn.Module) -> bool:
     return True
 
 
+@contextmanager
+def temporary_eval_mode(model: nn.Module) -> Iterator[None]:
+    """Run non-gradient work in eval mode, then restore the caller's mode."""
+
+    was_training = bool(getattr(model, "training", False))
+    if was_training:
+        model.eval()
+    try:
+        yield
+    finally:
+        if was_training:
+            model.train()
+
+
 def inject_lora_adapters(model: nn.Module, config: LoRAConfig) -> tuple[str, ...]:
     """Inject MLX-LM LoRA layers into all attention and MLP projections."""
 

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/policy/trajectory_logprobs.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 import mlx.core as mx
 import mlx.nn as nn
 
-from mlxrl.policy.logprobs import adapters_disabled
+from mlxrl.policy.logprobs import adapters_disabled, target_logprobs_from_logits
 from mlxrl.trajectory import Trajectory
 
 
@@ -47,7 +47,7 @@ def trajectory_action_logprobs(
 
     input_rows: list[list[int]] = []
     target_rows: list[list[int]] = []
-    action_target_indices: list[list[int]] = []
+    action_target_index_rows: list[list[int]] = []
     mask_rows: list[list[float]] = []
     target_width = max_sequence_len - 1
     for trajectory in trajectories:
@@ -59,35 +59,27 @@ def trajectory_action_logprobs(
             indices.extend(range(span.start - 1, span.end - 1))
         input_rows.append(input_row)
         target_rows.append(target_row)
-        action_target_indices.append(indices)
+        action_target_index_rows.append(
+            [indices[index] if index < len(indices) else 0 for index in range(max_action_len)]
+        )
         mask_rows.append([1.0] * len(indices) + [0.0] * (max_action_len - len(indices)))
 
     input_ids = mx.array(input_rows, dtype=mx.int32)
     target_ids = mx.array(target_rows, dtype=mx.int32)
     logits = model(input_ids)
-    all_logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-    target_logprobs = mx.squeeze(
-        mx.take_along_axis(all_logprobs, target_ids[..., None], axis=-1),
-        axis=-1,
+    target_logprobs = target_logprobs_from_logits(logits, target_ids)
+    action_indices = mx.array(action_target_index_rows, dtype=mx.int32)
+    mask = mx.array(mask_rows, dtype=mx.float32)
+    action_logprobs = mx.take_along_axis(
+        target_logprobs,
+        action_indices,
+        axis=1,
     )
-
-    rows: list[mx.array] = []
-    for row_index, indices in enumerate(action_target_indices):
-        row_values = [target_logprobs[row_index, index] for index in indices]
-        row = mx.stack(row_values, axis=0)
-        if len(indices) < max_action_len:
-            row = mx.concatenate(
-                [
-                    row,
-                    mx.zeros((max_action_len - len(indices),), dtype=row.dtype),
-                ],
-                axis=0,
-            )
-        rows.append(row)
+    action_logprobs = action_logprobs * mask.astype(action_logprobs.dtype)
 
     return TrajectoryLogprobs(
-        logprobs=mx.stack(rows, axis=0),
-        mask=mx.array(mask_rows, dtype=mx.float32),
+        logprobs=action_logprobs,
+        mask=mask,
     )
 
 

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/rollout/agentic.py

@@ -295,7 +295,7 @@ def _generate_model_action(
         if index > 0:
             if current is None:
                 raise RuntimeError("Internal error: missing previous decode token.")
-            logits = model(current, cache=state.cache)
+            logits = model(current[None], cache=state.cache)
             logprobs = logits[:, -1, :] - mx.logsumexp(logits[:, -1, :], axis=-1, keepdims=True)
         next_token = sampler(logprobs)
         sampled_logprob = sampled_token_logprobs(logprobs, next_token)

{inductive_mlxrl-0.2.0 → inductive_mlxrl-0.2.1}/mlxrl/rollout/optimized.py

@@ -264,13 +264,6 @@ def generate_from_prefix_cache(
         min_p=config.min_p,
         top_k=config.top_k,
     )
-    cache: list[Any]
-    if compile_decode_step:
-        cache = fixed_decode_cache_from_prefix(prefix.cache, config.max_tokens)
-        decode_logprobs = _compiled_decode_logprobs(model, cache)
-    else:
-        cache = clone_prompt_cache(prefix.cache)
-        decode_logprobs = _decode_logprobs(model, cache)
     completion: list[int] = []
     old_policy_logprobs: list[float] = []
 
@@ -290,6 +283,14 @@ def generate_from_prefix_cache(
             decode_completion(tokenizer, completion),
         )
 
+    cache: list[Any]
+    if compile_decode_step:
+        cache = fixed_decode_cache_from_prefix(prefix.cache, config.max_tokens)
+        decode_logprobs = _compiled_decode_logprobs(model, cache)
+    else:
+        cache = clone_prompt_cache(prefix.cache)
+        decode_logprobs = _decode_logprobs(model, cache)
+
     current = next_token
     for _ in range(1, config.max_tokens):
         logprobs = decode_logprobs(current)
@@ -364,16 +365,6 @@ def generate_prompt_set_from_prefix_caches(
             max_tokens=config.max_tokens,
         )
 
-    cache: list[Any] = fixed_decode_cache_from_prefixes(
-        prefixes,
-        group_size=group_size,
-        max_tokens=config.max_tokens,
-    )
-    decode_logprobs = (
-        _compiled_decode_logprobs(model, cache)
-        if compile_decode_step
-        else _decode_logprobs(model, cache)
-    )
     first_logprobs = mx.concatenate(
         [
             mx.broadcast_to(
@@ -409,6 +400,31 @@ def generate_prompt_set_from_prefix_caches(
             active[row_index] = False
             early_eos = config.max_tokens > 1
 
+    if early_eos and draw_keys is not None:
+        _set_random_state(start_random_state)
+        return _generate_prompt_set_sequentially(
+            model=model,
+            tokenizer=tokenizer,
+            prefixes=prefixes,
+            group_size=group_size,
+            config=config,
+            eos_token_ids=eos_token_ids,
+            compile_decode_step=compile_decode_step,
+        )
+    if not any(active):
+        return _rollout_rows_to_outputs(tokenizer, rows, old_policy_logprob_rows)
+
+    cache: list[Any] = fixed_decode_cache_from_prefixes(
+        prefixes,
+        group_size=group_size,
+        max_tokens=config.max_tokens,
+    )
+    decode_logprobs = (
+        _compiled_decode_logprobs(model, cache)
+        if compile_decode_step
+        else _decode_logprobs(model, cache)
+    )
+
     for step in range(1, config.max_tokens):
         logprobs = decode_logprobs(current)
         next_token = _sample_batched_logprobs(
@@ -432,29 +448,57 @@ def generate_prompt_set_from_prefix_caches(
                 active[row_index] = False
                 early_eos = step < config.max_tokens - 1
         current = next_token
+        if early_eos and draw_keys is not None:
+            _set_random_state(start_random_state)
+            return _generate_prompt_set_sequentially(
+                model=model,
+                tokenizer=tokenizer,
+                prefixes=prefixes,
+                group_size=group_size,
+                config=config,
+                eos_token_ids=eos_token_ids,
+                compile_decode_step=compile_decode_step,
+            )
         if not any(active):
             break
 
-    if early_eos:
-        _set_random_state(start_random_state)
-        outputs: list[tuple[tuple[int, ...], tuple[float, ...], str]] = []
-        for prefix in prefixes:
-            outputs.extend(
-                generate_from_prefix_cache(
-                    model=model,
-                    tokenizer=tokenizer,
-                    prefix=prefix,
-                    config=config,
-                    eos_token_ids=eos_token_ids,
-                    compile_decode_step=compile_decode_step,
-                )
-                for _ in range(group_size)
-            )
-        return tuple(outputs)
-
     if final_random_state is not None:
         _set_random_state(final_random_state)
 
+    return _rollout_rows_to_outputs(tokenizer, rows, old_policy_logprob_rows)
+
+
+def _generate_prompt_set_sequentially(
+    *,
+    model: nn.Module,
+    tokenizer: Any,
+    prefixes: Sequence[PrefixCache],
+    group_size: int,
+    config: SamplingConfig,
+    eos_token_ids: frozenset[int],
+    compile_decode_step: bool,
+) -> tuple[tuple[tuple[int, ...], tuple[float, ...], str], ...]:
+    outputs: list[tuple[tuple[int, ...], tuple[float, ...], str]] = []
+    for prefix in prefixes:
+        outputs.extend(
+            generate_from_prefix_cache(
+                model=model,
+                tokenizer=tokenizer,
+                prefix=prefix,
+                config=config,
+                eos_token_ids=eos_token_ids,
+                compile_decode_step=compile_decode_step,
+            )
+            for _ in range(group_size)
+        )
+    return tuple(outputs)
+
+
+def _rollout_rows_to_outputs(
+    tokenizer: Any,
+    rows: Sequence[Sequence[int]],
+    old_policy_logprob_rows: Sequence[Sequence[float]],
+) -> tuple[tuple[tuple[int, ...], tuple[float, ...], str], ...]:
     return tuple(
         (
             tuple(row),