inductive-mlxrl 0.1.0__py3-none-any.whl

inductive_mlxrl-0.1.0.dist-info/METADATA

@@ -0,0 +1,388 @@
+Metadata-Version: 2.4
+Name: inductive-mlxrl
+Version: 0.1.0
+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
+Project-URL: Issues, https://github.com/inductiveML/mlxrl/issues
+Project-URL: Changelog, https://github.com/inductiveML/mlxrl/blob/main/CHANGELOG.md
+License-Expression: MIT
+License-File: LICENSE
+License-File: THIRD_PARTY_LICENSES.md
+Keywords: apple-silicon,grpo,mlx,qlora,rl
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: mlx-lm>=0.31.0
+Requires-Dist: mlx>=0.31.0
+Description-Content-Type: text/markdown
+
+# mlxrl
+
+Fast on-policy MLX RL for Apple Silicon; not a general RL framework, not
+preference tuning, and not distributed training.
+
+`mlxrl` is a small, single-process RL post-training library for LLMs on Apple
+Silicon. It is built around one idea: GRPO on MLX should be a fast batched
+rollout path with a thin loss and optimizer step on top, not a framework.
+
+The current implementation targets QLoRA GRPO on local 4-bit MLX models. It
+reuses `mlx-lm` model loading, LoRA layers, KV caches, and sampling utilities,
+and keeps generation and training in one Python process with one model object.
+
+`mlxrl` is pre-1.0. The correctness gates are stable, but import APIs and config
+fields may change before a 1.0 release.
+
+## Quickstart
+
+```bash
+git clone https://github.com/inductiveML/mlxrl.git
+cd mlxrl
+UV_CACHE_DIR=.uv-cache uv sync --all-groups
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen3_0_6b_grpo.toml \
+  --available-memory-gb 48
+```
+
+For the measured 9B-on-48GB shape, use the checkpointed G=2 config:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen35_9b_g2_checkpoint.toml \
+  --available-memory-gb 48 \
+  --dry-run
+```
+
+## What Works
+
+- Batched group rollouts with MLX-LM KV caches and sampling.
+- Full-forward old-policy logprob recompute for training-time `pi_old`.
+- Adapter-disabled reference policy on the same model object.
+- GRPO, Dr. GRPO, DAPO, and GSPO loss variants.
+- RLOO (REINFORCE Leave-One-Out) as a critic-free rollout objective.
+- QLoRA injection on dense and heterogeneous/hybrid layer stacks.
+- Qwen3.5-style hybrid support via MLX-LM auto LoRA targeting, including
+  DeltaNet `linear_attn.in_proj_*` and dense attention `q/k/v/o_proj`.
+- Per-layer gradient checkpointing through `mlx_lm.tuner.trainer.grad_checkpoint`
+  for linear-attention/DeltaNet backward memory.
+- Micro-batched gradient accumulation for token-mean policy losses.
+- `beta == 0` reference-forward skip.
+- Phase 4 benchmark harness for `mlxrl`, `mlx-tune`, `mlx-lm-lora`, and `mlx-lm`.
+
+## Install
+
+After the first tagged release, install the PyPI distribution:
+
+```bash
+pip install inductive-mlxrl
+```
+
+The Python import package and CLI command are still `mlxrl`:
+
+```bash
+mlxrl --help
+```
+
+Source installs are also supported:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv sync --all-groups
+```
+
+Run commands through the local environment:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl --help
+```
+
+Python 3.11+ is required. Runtime dependencies are intentionally small:
+`mlx` and `mlx-lm`. Development dependencies include `pytest`, `ruff`,
+`pyright`, `mlx-tune`, and `mlx-lm-lora` for comparison benchmarks.
+The PyPI distribution name is `inductive-mlxrl`; the import package and console
+script remain `mlxrl`.
+
+## Quick Smoke Tests
+
+Dense Qwen3 0.6B:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase0-smoke \
+  --model mlx-community/Qwen3-0.6B-4bit \
+  --prompt "What is 2+2?"
+```
+
+Hybrid Qwen3.5 9B with rank-16 LoRA:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase0-smoke \
+  --model mlx-community/Qwen3.5-9B-MLX-4bit \
+  --rank 16 \
+  --scale 2.0 \
+  --prompt "What is 2+2?"
+```
+
+The smoke gate prints the model id, layer count, LoRA target keys, per-layer
+LoRA module counts, total/trainable parameter counts, and logits shape. It
+fails if any trainable leaf is not `lora_a` or `lora_b`.
+
+## Training Commands
+
+Toy hand-computed GRPO math gate:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-toy-gate
+```
+
+Small built-in GSM8K-style run:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-gsm8k \
+  --model mlx-community/Qwen3-0.6B-4bit \
+  --steps 20 \
+  --group-size 4 \
+  --max-tokens 64
+```
+
+Config-driven run:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen3_0_6b_grpo.toml \
+  --available-memory-gb 48
+```
+
+The config schema validates model id, quant bits, group size, completion/prompt
+lengths, checkpointing granularity, `iogpu.wired_limit_mb`, optimizer settings,
+algorithm hyperparameters, KL beta, and seed before a model is loaded. CLI
+overrides such as `--steps`, `--group-size`, `--max-tokens`, `--algorithm`,
+`--beta`, and `--seed` apply on top of the file.
+
+For DeltaNet / linear-attention models, enable per-layer checkpointing:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase1-gsm8k \
+  --model mlx-community/Qwen3.5-9B-MLX-4bit \
+  --rank 16 \
+  --scale 2.0 \
+  --checkpoint-completion-forward \
+  --steps 1 \
+  --group-size 2 \
+  --max-tokens 256
+```
+
+Despite the historical CLI name, `--checkpoint-completion-forward` now enables
+per-transformer-block checkpointing at model setup. The old whole-model
+`mx.checkpoint(...)` wrapper was removed because it does not cap DeltaNet's
+per-layer scan memory.
+
+Phase 2 rollout equivalence check:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl phase2-equivalence \
+  --model mlx-community/Qwen3-0.6B-4bit \
+  --group-size 4 \
+  --max-tokens 32 \
+  --compile-decode-step \
+  --batch-groups
+```
+
+## Import API
+
+Minimal model setup:
+
+```python
+from mlxrl.policy import LoRAConfig, load_policy_with_lora
+
+model, tokenizer, report = load_policy_with_lora(
+    model_id="mlx-community/Qwen3.5-9B-MLX-4bit",
+    config=LoRAConfig(
+        rank=16,
+        scale=2.0,
+        dropout=0.0,
+        grad_checkpoint=True,
+    ),
+)
+```
+
+One optimizer step:
+
+```python
+import mlx.optimizers as optim
+
+from mlxrl.algo import GRPOAlgorithm
+from mlxrl.train import batch_from_rollouts, optimizer_step
+
+optimizer = optim.Adam(learning_rate=1e-5)
+algorithm = GRPOAlgorithm()
+batch = batch_from_rollouts(
+    model=model,
+    completions=completions,
+    rewards=rewards,
+    group_size=4,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+    compute_reference=beta != 0.0,
+)
+metrics = optimizer_step(
+    model=model,
+    optimizer=optimizer,
+    batch=batch,
+    beta=beta,
+    pad_token_id=pad_token_id,
+    algorithm=algorithm,
+    use_checkpoint=True,
+    micro_batch_size=2,
+)
+```
+
+`micro_batch_size=0` keeps the original whole-batch path. Micro-batching is
+currently exact for token-mean policy losses: base GRPO, DAPO, GSPO token mode,
+RLOO, and Dr. GRPO with `loss_reduction="token_mean"`. Sequence-reduced losses
+should keep `micro_batch_size=0`.
+
+## Policy Semantics
+
+- The base model is frozen before LoRA injection.
+- Only LoRA adapter leaves are trainable.
+- Reference logprobs are computed by temporarily disabling adapters on the same
+  model object; there is no second reference model in memory.
+- Old-policy logprobs are recomputed with a full forward for the training batch.
+  Rollout-time logprobs are captured for inspection, but 4-bit sequential decode
+  and full-forward prefill are not numerically identical on hybrid/quantized
+  models, so recompute remains the default training semantics.
+- When `beta == 0`, the reference forward is skipped and the policy logprobs are
+  used as a zero-KL placeholder.
+- PPO, DPO, and ORPO are intentionally out of scope. PPO needs a separate critic
+  and value forward; DPO/ORPO are offline preference objectives with no rollout
+  phase. `mlxrl` is critic-free, on-policy, and rollout-based by design.
+
+## Algorithms
+
+Concrete algorithms implement the small `Algorithm` protocol: compute
+advantages, optionally filter a prepared batch, then compute a loss from policy,
+old-policy, and reference logprobs. `rollout/`, `policy/`, and `train/` do not
+import concrete algorithm implementations.
+
+| algorithm | defining behavior |
+| --- | --- |
+| GRPO | group-normalized rewards, token-level importance ratio |
+| Dr. GRPO | centered or normalized rewards with decoupled length reduction |
+| DAPO | asymmetric low/high clipping plus optional dynamic zero-advantage group filtering |
+| GSPO | sequence-level, length-normalized importance ratio and clipping |
+| RLOO | leave-one-out group baseline, no critic, no std-normalized advantage |
+
+## Memory Preflight
+
+`mlxrl train` can estimate memory before loading the model:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run mlxrl train \
+  --config examples/qwen3_0_6b_grpo.toml \
+  --available-memory-gb 48 \
+  --dry-run
+```
+
+The estimator is calibrated to measured anchors: `6.245 GB` for
+Qwen3-0.6B/G4/prompt≈19/T256, `25.9 GB` for
+Qwen3.5-9B/G2/seq609/per-layer-checkpointed, `45.9 GB` for
+Qwen3.5-9B/G4/seq609/per-layer-checkpointed, and `36 GB` for
+Qwen3.5-9B/G2/seq128/no-checkpoint. For hybrid 9B no-checkpoint long-sequence
+configs, it reports an OOM-risk lower bound rather than a fake precise peak.
+For an obviously too-large Qwen3.5-9B/G8/prompt97/T512/no-checkpoint config on
+48 GB, it flags the run and suggests the measured-boundary fallback around
+G4/T512/checkpointed.
+
+## Benchmarks
+
+Local M4 Max Phase 4 snapshot:
+
+- `454` rollout tok/s on Qwen3-0.6B GRPO with G=4 and 256-token completions.
+- `0.283` end-to-end it/s with full `mlxrl` training semantics.
+- `3.2x` faster rollout and `2.2x` higher end-to-end it/s than `mlx-tune`
+  v0.5.1 on the same run shape.
+- `1.3x` faster rollout than sequential `mlx-lm` generation at G=4.
+
+These are the two-pass means from
+`benchmarks/results/gate5_full_reconciled.md`, run with MLX 0.31.2,
+MLX-LM 0.31.3, `mlx-community/Qwen3-0.6B-4bit`, 100 measured steps with
+5 warmup steps discarded:
+
+| target | comparison | rollout tok/s | grad s/step | samples/s | it/s | peak GB |
+| --- | --- | ---: | ---: | ---: | ---: | ---: |
+| `mlxrl` | apples-to-apples GRPO | 454.1 | 1.282 | 1.133 | 0.283 | 6.25 |
+| `mlx-lm` | generation-only, G=1 | 347.0 | - | 1.355 | - | 0.52 |
+| `mlx-lm-g4` | generation-only, sequential G=4 | 349.7 | - | 1.366 | - | 0.52 |
+| `mlx-tune` | package-speed reference | 142.2 | 0.502 | 0.519 | 0.130 | 6.16 |
+| `mlx-lm-lora` | package-speed reference | 557.9 | 0.592 | 1.648 | 0.412 | 5.32 |
+
+`mlx-lm-lora` reports higher raw package-speed throughput in this snapshot, but
+its benchmarked path is not the same training problem as `mlxrl`'s live
+old-policy/reference semantics and completion-loss masking. That is the honest
+case where `mlxrl` is not faster; the apples-to-apples comparison label is
+reserved for `mlxrl`'s own semantic path. On the 9B Noether real workload, the
+checkpointed MLX path measured about 6x faster than the previous torch-MPS path;
+that workload is separate from the public Phase 4 package-speed harness.
+
+Run the Phase 4 harness:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run python benchmarks/run_phase4.py run \
+  --targets mlxrl,mlx-lm,mlx-tune,mlx-lm-lora \
+  --model mlx-community/Qwen3-0.6B-4bit \
+  --steps 100 \
+  --warmup-steps 5 \
+  --group-size 4 \
+  --max-tokens 256 \
+  --passes 2 \
+  --output benchmarks/results/phase4.jsonl \
+  --summary benchmarks/results/phase4.md \
+  --allow-missing-baselines
+```
+
+The harness reports synchronized rollout tok/s, gradient seconds per step,
+samples/s, it/s, and peak MLX memory. `mlx-lm` targets are generation-only;
+external package targets are useful speed references but may not match `mlxrl`
+training semantics.
+
+## Development
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) and [DESIGN.md](DESIGN.md) before adding
+algorithms or changing rollout/logprob semantics.
+
+Run the quality gates:
+
+```bash
+UV_CACHE_DIR=.uv-cache uv run pytest
+UV_CACHE_DIR=.uv-cache uv run ruff check .
+UV_CACHE_DIR=.uv-cache uv run pyright
+```
+
+MLX lazy evaluation matters. Any `mx.eval(...)` or `mx.synchronize()` in this
+repo should mark a real boundary: sampled token append/EOS checks, logprob
+freezing before adapter mutation, per-micro-batch graph release, optimizer
+updates, or benchmark timing boundaries.
+
+## Layout
+
+```text
+mlxrl/
+  rollout/   # batched group generation
+  policy/    # model loading, LoRA setup, logprob passes
+  algo/      # GRPO-family advantages and losses
+  train/     # value_and_grad and optimizer integration
+  data/      # toy GSM8K data and rewards
+  cli.py
+tests/
+benchmarks/
+```
+
+## Non-Goals
+
+- No inference server or second model copy.
+- No CUDA or torch fallback.
+- No distributed training.
+- No broad RL framework abstractions beyond the small algorithm interface.

inductive_mlxrl-0.1.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+mlxrl/__init__.py,sha256=OmZqyDOcKBeRwSNfiu_wdm_zZOvPdooMOaYTnzoHLtU,122
+mlxrl/algorithm.py,sha256=hBhTlIVOnL69_CoE7pxkt_zCk63ImxIJHi2VG7KSAYE,1778
+mlxrl/cli.py,sha256=6aovbx1Sv4r02RLrmexTDsbQGxXpcAy8rv8LDuIqrxo,29934
+mlxrl/config.py,sha256=RiWu-pWl8u8OGJrI8Rr2ozIiA-ecu5fwntM3G9cQZ1g,14675
+mlxrl/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+mlxrl/algo/__init__.py,sha256=VblutTkTpF4PQ-qKiCYUQ9EizJ3cUPGixqsQsWbbpPM,942
+mlxrl/algo/grpo.py,sha256=SQfY38IY_zr33S2Ph49xS4HN3LylWQYh2qUejw1jJMg,18824
+mlxrl/data/__init__.py,sha256=udtJ0TQfbEEj1LELw5NynXpsOv2_afFft54dKBgZCZ0,589
+mlxrl/data/gsm8k.py,sha256=SBc6SIRq1TlPOBe9KhIJP0KXxn8Gmay2f1yEFpyDoE8,2075
+mlxrl/data/rewards.py,sha256=IT6N07DkXgQidMI9BJbGG9KkiprGaa6QQpnIaLxDWKo,2988
+mlxrl/policy/__init__.py,sha256=ieu_rA9aq6BQdDXohixt02i1Ka3wqeTT7JLeBWHFy9s,1025
+mlxrl/policy/logprobs.py,sha256=nAFjIE7MUWUAQ_7faFdF3okOjNlLT1eagJ4S8BA7IzI,10411
+mlxrl/policy/model.py,sha256=DA_MWaugxsIl_s0QMD3bzwEwN9rxmZG6yoc7PbKAX6M,10015
+mlxrl/rollout/__init__.py,sha256=MCSESkZR7FDIZwj50J_ThA_RInquCQHB2Jh6OnZM9NM,847
+mlxrl/rollout/naive.py,sha256=Jo3B7RChEJI3ty0GC3mRHyVVzCgCWLiFEPRpnrQWlQs,5843
+mlxrl/rollout/optimized.py,sha256=MDaDVBHdqF4FeXGp6r4N8umDnOFFii-hc3JIJXd2p-g,21274
+mlxrl/train/__init__.py,sha256=oSs0IZBzMjw_eJgUBlyfd0XisIMz0SCvJWXjupE5q6M,219
+mlxrl/train/grpo.py,sha256=rgWEq16D07jfFKjLmJszlS016ViUmLoc5Ece0gszxxk,10486
+inductive_mlxrl-0.1.0.dist-info/METADATA,sha256=a6PV1SFsPyB98ZZTyBLiVuRXLDLBmnbsReLO8QQb0Ec,12932
+inductive_mlxrl-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+inductive_mlxrl-0.1.0.dist-info/entry_points.txt,sha256=IyD-wJ1EUMhrd40AjCGTEmDyMMb9M3lNyRuzB3Vd63w,41
+inductive_mlxrl-0.1.0.dist-info/licenses/LICENSE,sha256=Q2qW7SWCM5UjB3Rxs2rmhMAhhk5IbsI2sYuhmn5mgUo,1075
+inductive_mlxrl-0.1.0.dist-info/licenses/THIRD_PARTY_LICENSES.md,sha256=X5OXyGhte6HfWf3B1B4VwjoL8e6OoXpFEfHF5v9CjrE,2067
+inductive_mlxrl-0.1.0.dist-info/RECORD,,

inductive_mlxrl-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

inductive_mlxrl-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+mlxrl = mlxrl.cli:main

inductive_mlxrl-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 mlxrl contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

inductive_mlxrl-0.1.0.dist-info/licenses/THIRD_PARTY_LICENSES.md

@@ -0,0 +1,46 @@
+# Third-Party Notices
+
+`mlxrl` reuses and adapts APIs, cache semantics, and sampling utilities from
+upstream MLX projects. The repository keeps local attribution headers on files
+that adapt those patterns.
+
+## mlx-lm
+
+- Source: https://github.com/ml-explore/mlx-lm
+- License: MIT
+- Used for: model loading, LoRA utilities, gradient checkpoint helper,
+  prompt/KV-cache construction, cache state conventions, and sampling filters.
+
+Upstream notice:
+
+```text
+MIT License Copyright © 2023 Apple Inc.
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
+Software, and to permit persons to whom the Software is furnished to do so,
+subject to the following conditions: The above copyright notice and this
+permission notice shall be included in all copies or substantial portions of the
+Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+```
+
+## mlx-tune
+
+- Source: https://github.com/ARahim3/mlx-tune
+- License: Apache-2.0 per installed package metadata.
+- Used for: benchmark configuration inspection and external baseline adapter
+  patterns only. No `mlx-tune` implementation code is vendored in `mlxrl`.
+
+## mlx-lm-lora
+
+- Source: https://github.com/Goekdeniz-Guelmez/mlx-lm-lora
+- License: MIT per installed package metadata.
+- Used for: benchmark configuration inspection and external baseline adapter
+  patterns only. No `mlx-lm-lora` implementation code is vendored in `mlxrl`.

mlxrl/__init__.py

@@ -0,0 +1,6 @@
+"""Small single-process RL post-training for LLMs on Apple Silicon."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.1.0"
+

mlxrl/algo/__init__.py

@@ -0,0 +1,43 @@
+"""Algorithm-specific advantage and loss code."""
+
+from mlxrl.algo.grpo import (
+    AlgorithmLossMetrics,
+    DAPOAlgorithm,
+    DrGRPOAlgorithm,
+    GRPOAlgorithm,
+    GRPOLossMetrics,
+    GSPOAlgorithm,
+    PolicyAlgorithm,
+    RLOOAlgorithm,
+    algorithm_by_name,
+    approximate_kl,
+    clip_ratio,
+    filter_zero_advantage_groups,
+    group_center_rewards,
+    group_normalize_rewards,
+    grpo_loss,
+    masked_mean,
+    sequence_importance_ratio,
+    token_importance_ratio,
+)
+
+__all__ = [
+    "AlgorithmLossMetrics",
+    "DAPOAlgorithm",
+    "DrGRPOAlgorithm",
+    "GRPOAlgorithm",
+    "GRPOLossMetrics",
+    "GSPOAlgorithm",
+    "PolicyAlgorithm",
+    "RLOOAlgorithm",
+    "algorithm_by_name",
+    "approximate_kl",
+    "clip_ratio",
+    "filter_zero_advantage_groups",
+    "group_center_rewards",
+    "group_normalize_rewards",
+    "grpo_loss",
+    "masked_mean",
+    "sequence_importance_ratio",
+    "token_importance_ratio",
+]