libthx 0.1.8__tar.gz → 0.2.1__tar.gz

{libthx-0.1.8 → libthx-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libthx
-Version: 0.1.8
+Version: 0.2.1
 Summary: Architecture experimentation and training infrastructure.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{libthx-0.1.8 → libthx-0.2.1}/libthx.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libthx
-Version: 0.1.8
+Version: 0.2.1
 Summary: Architecture experimentation and training infrastructure.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{libthx-0.1.8 → libthx-0.2.1}/libthx.egg-info/SOURCES.txt

@@ -7,6 +7,15 @@ libthx.egg-info/dependency_links.txt
 libthx.egg-info/entry_points.txt
 libthx.egg-info/requires.txt
 libthx.egg-info/top_level.txt
+tests/test_contrastive_roundtrip.py
+tests/test_datasets.py
+tests/test_eval_padding.py
+tests/test_gpu_availability.py
+tests/test_hardware_dispatch.py
+tests/test_kv_cache.py
+tests/test_lora.py
+tests/test_mamba.py
+tests/test_registries.py
 theseus/__init__.py
 theseus/cli.py
 theseus/config.py
@@ -31,6 +40,7 @@ theseus/data/datasets/ccaligned.py
 theseus/data/datasets/cfq.py
 theseus/data/datasets/clutrr.py
 theseus/data/datasets/dataset.py
+theseus/data/datasets/dictlearn.py
 theseus/data/datasets/fever.py
 theseus/data/datasets/fineweb.py
 theseus/data/datasets/harmfulqa.py
@@ -43,6 +53,7 @@ theseus/data/datasets/pes2o.py
 theseus/data/datasets/pg19.py
 theseus/data/datasets/pile.py
 theseus/data/datasets/pile_detoxify.py
+theseus/data/datasets/pile_injected.py
 theseus/data/datasets/qqp.py
 theseus/data/datasets/siqa.py
 theseus/data/datasets/squad.py
@@ -67,11 +78,14 @@ theseus/evaluation/__init__.py
 theseus/evaluation/base.py
 theseus/evaluation/huggingface.py
 theseus/evaluation/datasets/__init__.py
+theseus/evaluation/datasets/alpaca.py
+theseus/evaluation/datasets/arithmetic.py
 theseus/evaluation/datasets/bbq.py
 theseus/evaluation/datasets/blimp.py
 theseus/evaluation/datasets/ccaligned.py
 theseus/evaluation/datasets/cfq.py
 theseus/evaluation/datasets/clutrr.py
+theseus/evaluation/datasets/dictlearn.py
 theseus/evaluation/datasets/fever.py
 theseus/evaluation/datasets/longbench.py
 theseus/evaluation/datasets/longhealth.py
@@ -81,7 +95,9 @@ theseus/evaluation/datasets/mtob.py
 theseus/evaluation/datasets/perplexity_evals.py
 theseus/evaluation/datasets/pes2o.py
 theseus/evaluation/datasets/pg19.py
+theseus/evaluation/datasets/pg19_lengthgen.py
 theseus/evaluation/datasets/pile.py
+theseus/evaluation/datasets/pile_injected.py
 theseus/evaluation/datasets/qqp.py
 theseus/evaluation/datasets/siqa.py
 theseus/evaluation/datasets/squad.py
@@ -91,12 +107,17 @@ theseus/evaluation/datasets/winogrande.py
 theseus/experiments/__init__.py
 theseus/experiments/continual/__init__.py
 theseus/experiments/continual/abcd.py
+theseus/experiments/continual/benchmark.py
 theseus/experiments/models/__init__.py
 theseus/experiments/models/forking.py
 theseus/experiments/models/gpt.py
 theseus/experiments/models/gpt_neox.py
 theseus/experiments/models/llama.py
+theseus/experiments/models/moe.py
 theseus/experiments/models/qwen.py
+theseus/experiments/mok/__init__.py
+theseus/experiments/mok/reward.py
+theseus/experiments/mok/smoke.py
 theseus/experiments/redcodegen/__init__.py
 theseus/experiments/redcodegen/hardening.py
 theseus/inference/__init__.py
@@ -120,6 +141,8 @@ theseus/model/block/block.py
 theseus/model/block/forking.py
 theseus/model/block/gpt_neox.py
 theseus/model/block/llama.py
+theseus/model/block/mamba.py
+theseus/model/block/moe.py
 theseus/model/block/qwen.py
 theseus/model/block/scratching.py
 theseus/model/layers/__init__.py
@@ -129,6 +152,9 @@ theseus/model/layers/rmsnorm.py
 theseus/model/layers/rope.py
 theseus/model/models/__init__.py
 theseus/model/models/base.py
+theseus/model/models/hybrid.py
+theseus/model/models/mamba.py
+theseus/model/models/moe.py
 theseus/model/models/scratchbubbles.py
 theseus/model/models/thoughtbubbles.py
 theseus/model/models/contrib/__init__.py
@@ -136,12 +162,18 @@ theseus/model/models/contrib/gpt_neox.py
 theseus/model/models/contrib/llama.py
 theseus/model/models/contrib/marin.py
 theseus/model/models/contrib/qwen.py
+theseus/model/moe/__init__.py
+theseus/model/moe/base.py
+theseus/model/moe/bias_balanced.py
 theseus/training/__init__.py
 theseus/training/backbone.py
 theseus/training/base.py
 theseus/training/contrastive.py
+theseus/training/grpo.py
 theseus/training/huggingface.py
 theseus/training/kl_divergence.py
+theseus/training/lora.py
+theseus/training/ppo.py
 theseus/training/utils.py
 theseus/training/flywheel/__init__.py
 theseus/training/flywheel/contrastive.py
@@ -152,6 +184,7 @@ theseus/training/optimizers/__init__.py
 theseus/training/optimizers/adamw.py
 theseus/training/optimizers/muon.py
 theseus/training/schedules/__init__.py
+theseus/training/schedules/cosine_rewarm.py
 theseus/training/schedules/wsd.py
 theseus/training/schedules/wsds.py
 theseus/web/__init__.py

{libthx-0.1.8 → libthx-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "libthx"
-version = "0.1.8"
+version = "0.2.1"
 description = "Architecture experimentation and training infrastructure."
 readme = "README.md"
 requires-python = ">=3.11"
@@ -229,7 +229,7 @@ plugins = ['pydantic.mypy']
 python_version = "3.12"
 strict = true
 ignore_missing_imports = true
-exclude = ["tests", "examples"]
+exclude = ["tests", "examples", "theseus/dispatch/bootstrap.py"]
 
 [[tool.mypy.overrides]]
 module = "huggingface_hub.*"

libthx-0.2.1/tests/test_contrastive_roundtrip.py

@@ -0,0 +1,143 @@
+"""Round-trip check for contrastive tokenization + loader + async strategy.
+
+Builds a tiny in-memory contrastive dataset, writes memmaps in a temp dir,
+loads via ContrastivePaddedDataset and Strategy(Async), and verifies shapes.
+
+Migrated from scripts/test_contrastive_roundtrip.py.
+"""
+
+import json
+import random
+import tempfile
+from pathlib import Path
+from typing import List, Tuple
+
+import numpy as np
+import pytest
+
+from theseus.data.tokenize import (
+    TokenizeContrastiveDatasetConfig,
+    _build_padded_arrays,
+    _encode_dataset_item,
+)
+from theseus.data.tokenizer import TokenizerConfig, get_tokenizer
+from theseus.training.flywheel.contrastive import ContrastivePaddedDataset
+from theseus.training.flywheel.strategy import DatasetStyle, Sampling, Strategy
+from theseus.base.job import ExecutionSpec
+
+
+PAIRS: List[Tuple[str, str]] = [
+    ("the quick brown fox", "the slow red fox"),
+    ("hello world", "hello mars"),
+    ("good code is readable", "good code is obscure"),
+    ("unit tests prevent bugs", "unit tests waste time"),
+]
+
+
+class TinyContrastive:
+    def __init__(self, data: List[Tuple[str, str]]):
+        self.data = data
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> Tuple[str, str]:
+        return self.data[idx]
+
+
+def _write_contrastive_data(
+    td_path: Path, tokenizer_cfg: TokenizerConfig, label: str
+) -> str:
+    """Write contrastive memmaps and return dataset name."""
+    name = f"tmpcontrast_{label}"
+    out = td_path / "data" / name
+    out.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = get_tokenizer(tokenizer_cfg)
+    args = TokenizeContrastiveDatasetConfig(
+        name="fever", block_size=32, pad_token=0, val_pct=0.25, seed=42,
+    )
+
+    dataset = TinyContrastive(PAIRS)
+    indices = list(range(len(dataset)))
+    random.seed(args.seed)
+    random.shuffle(indices)
+    val_size = int(len(dataset) * args.val_pct)
+    splits = {"train": indices[val_size:], "val": indices[:val_size]}
+    shapes = {}
+    dtype = np.uint32
+
+    for split_name, split_indices in splits.items():
+        num_samples = len(split_indices)
+        pos_tokens = np.memmap(
+            out / f"{split_name}.pos.bin", dtype=dtype, mode="w+",
+            shape=(num_samples, args.block_size),
+        )
+        pos_mask = np.memmap(
+            out / f"{split_name}.pos.bin.mask", dtype=np.bool_, mode="w+",
+            shape=(num_samples, args.block_size),
+        )
+        neg_tokens = np.memmap(
+            out / f"{split_name}.neg.bin", dtype=dtype, mode="w+",
+            shape=(num_samples, args.block_size),
+        )
+        neg_mask = np.memmap(
+            out / f"{split_name}.neg.bin.mask", dtype=np.bool_, mode="w+",
+            shape=(num_samples, args.block_size),
+        )
+
+        for arr_idx, didx in enumerate(split_indices):
+            pos_str, neg_str = dataset[didx]
+            for s, target_tokens, target_mask in [
+                (pos_str, pos_tokens, pos_mask),
+                (neg_str, neg_tokens, neg_mask),
+            ]:
+                ids, mask_list = _encode_dataset_item(s, False, tokenizer, args)
+                t, m, *_ = _build_padded_arrays(
+                    ids, mask_list, args.block_size, args.pad_token, dtype
+                )
+                target_tokens[arr_idx] = t
+                target_mask[arr_idx] = m
+
+        pos_tokens.flush()
+        pos_mask.flush()
+        neg_tokens.flush()
+        neg_mask.flush()
+        shapes[split_name] = {
+            "pos": [num_samples, args.block_size],
+            "neg": [num_samples, args.block_size],
+        }
+
+    with open(out / "shape.json", "w") as f:
+        json.dump(shapes, f)
+    with open(out / "config.json", "w") as f:
+        json.dump({}, f)
+
+    return name
+
+
+class TestContrastiveRoundtrip:
+    def test_tiktoken_roundtrip(self):
+        with tempfile.TemporaryDirectory() as td:
+            td_path = Path(td)
+            tok_cfg = TokenizerConfig(backend="tiktoken", name="cl100k_base")
+            name = _write_contrastive_data(td_path, tok_cfg, "tiktoken")
+
+            spec = ExecutionSpec.local(root_dir=str(td_path))
+            ds = ContrastivePaddedDataset(spec, block_size=32, name=name, suffix="")
+            batch = ds.get_batch(batch_size=2, split="train")
+
+            assert batch["padding_mask_pos"].shape == (2, 32)
+            assert batch["padding_mask_neg"].shape == (2, 32)
+
+            # Async loader
+            strat = Strategy(
+                spec, block_size=32,
+                mixture=[Sampling(name=name, rate=1.0, style=DatasetStyle.CONTRASTIVE)],
+            )
+            async_loader = strat.get_async_batches(batch_size=2, split="train")
+            async_batch = async_loader.get_batch()
+            async_loader.close()
+
+            assert async_batch["padding_mask_pos"].shape == (2, 32)
+            assert async_batch["padding_mask_neg"].shape == (2, 32)

libthx-0.2.1/tests/test_datasets.py

@@ -0,0 +1,78 @@
+"""Tests for new benchmark datasets and evaluations."""
+
+import pytest
+
+
+class TestPileInjected:
+    """Tests for the injected Pile dataset."""
+
+    def test_injected_texts_reproducible(self):
+        """Same seed always produces the same injected sequences."""
+        from theseus.data.datasets.pile_injected import (
+            _generate_injected_texts,
+            INJECTED_TEXTS,
+        )
+
+        texts_a = _generate_injected_texts(n_sequences=100, seed=42)
+        texts_b = _generate_injected_texts(n_sequences=100, seed=42)
+        assert texts_a == texts_b
+        assert texts_a == list(INJECTED_TEXTS)
+
+    def test_injected_texts_count(self):
+        from theseus.data.datasets.pile_injected import INJECTED_TEXTS
+
+        assert len(INJECTED_TEXTS) == 100
+
+    def test_injected_texts_nonempty(self):
+        from theseus.data.datasets.pile_injected import INJECTED_TEXTS
+
+        for text in INJECTED_TEXTS:
+            assert len(text) > 100  # each should be substantial
+
+    def test_injection_positions_sorted(self):
+        from theseus.data.datasets.pile_injected import INJECTION_POSITIONS
+
+        assert INJECTION_POSITIONS == sorted(INJECTION_POSITIONS)
+        assert len(INJECTION_POSITIONS) == 100
+
+    def test_different_seeds_different_texts(self):
+        from theseus.data.datasets.pile_injected import _generate_injected_texts
+
+        texts_a = _generate_injected_texts(n_sequences=10, seed=42)
+        texts_b = _generate_injected_texts(n_sequences=10, seed=99)
+        assert texts_a != texts_b
+
+
+class TestPileInjectedEval:
+    """Tests for the injected sequence memorization evaluation."""
+
+    def test_eval_uses_same_texts(self):
+        from theseus.data.datasets.pile_injected import INJECTED_TEXTS
+        from theseus.evaluation.datasets.pile_injected import PileInjectedEval
+
+        ev = PileInjectedEval()
+        assert len(ev) == len(INJECTED_TEXTS)
+        for i in range(len(ev)):
+            assert ev.get(i) == INJECTED_TEXTS[i]
+
+    def test_eval_name(self):
+        from theseus.evaluation.datasets.pile_injected import PileInjectedEval
+
+        ev = PileInjectedEval()
+        assert ev.name == "pile_injected_ppl"
+
+
+class TestPG19LengthGen:
+    """Tests for variable-length PG-19 evaluations."""
+
+    def test_eval_names_registered(self):
+        from theseus.registry import EVALUATIONS
+
+        for name in [
+            "pg19_2k_ppl",
+            "pg19_4k_ppl",
+            "pg19_8k_ppl",
+            "pg19_16k_ppl",
+            "pg19_32k_ppl",
+        ]:
+            assert name in EVALUATIONS, f"{name} not registered"

libthx-0.2.1/tests/test_eval_padding.py

@@ -0,0 +1,184 @@
+"""Evaluation padding tests — verifies evaluations work when dataset
+size is not a multiple of the batch unit.
+
+Migrated from scripts/test_eval_padding.py.
+"""
+
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
+
+from theseus.base import Axis
+from theseus.evaluation import (
+    EncodingEvaluation,
+    PerplexityComparisonEvaluation,
+    PerplexityEvaluation,
+    RolloutEvaluation,
+)
+
+
+class ToyTokenizer:
+    pad_token = 0
+
+    def encode(self, text: str) -> list[int]:
+        return [ord(ch) + 1 for ch in text]
+
+    def encode_batch(
+        self, text_list: list[str], allowed_special: str | None = None
+    ) -> list[list[int]]:
+        del allowed_special
+        return [self.encode(text) for text in text_list]
+
+    def decode_batch(self, seqs: list[list[int]]) -> list[str]:
+        return ["".join(chr(tok - 1) for tok in seq if tok > 0) for seq in seqs]
+
+
+@flax.struct.dataclass
+class DummyState:
+    params: jax.Array
+
+
+class DummyInference:
+    def __init__(self, batch_unit: int = 8, block_size: int = 8, rollout_token: str = "!"):
+        self.replicas = 1
+        self.local_replicas = 1
+        self.per_device_batch_size = batch_unit
+        self.block_size = block_size
+        self.mesh = Mesh(np.array(jax.devices()).reshape((1,)), (Axis.BATCH,))
+        scalar_sharding = NamedSharding(self.mesh, P())
+        self.state = DummyState(
+            params=jax.device_put(jnp.array(0, dtype=jnp.int32), scalar_sharding)
+        )
+        self.state_sharding = DummyState(params=scalar_sharding)
+        self.key = jax.random.PRNGKey(0)
+        self.vocab_size = 256
+        self.rollout_token_id = ord(rollout_token) + 1
+
+    @staticmethod
+    def pad(seqs: list[list[int]], pad_token: int = 0) -> tuple[jax.Array, jax.Array]:
+        max_len = max(len(seq) for seq in seqs)
+        padded = [([pad_token] * (max_len - len(seq))) + seq for seq in seqs]
+        masks = [([False] * (max_len - len(seq))) + ([True] * len(seq)) for seq in seqs]
+        return jnp.array(padded, dtype=jnp.int32), jnp.array(masks, dtype=jnp.bool_)
+
+    def _autoregress(self, state, key, input, input_mask, num_tokens, temperature, top_p):
+        del state, key, input_mask, temperature, top_p
+        extra = num_tokens - input.shape[-1]
+        if extra <= 0:
+            return input[:, :num_tokens]
+        generated = jnp.full((input.shape[0], extra), self.rollout_token_id, dtype=jnp.int32)
+        return jnp.concatenate([input, generated], axis=-1)
+
+    def forward(self, state, params, batch, key, deterministic):
+        del state, params, key, deterministic
+        x_batch, _, mask_batch = batch
+        next_tokens = jnp.roll(x_batch, -1, axis=-1)
+        next_tokens = next_tokens.at[:, -1].set(0)
+        next_tokens = jnp.where(mask_batch, next_tokens, 0)
+        logits = jax.nn.one_hot(next_tokens, self.vocab_size, dtype=jnp.float32) * 20.0
+        return logits, None, None
+
+
+class ToyRolloutEval(RolloutEvaluation):
+    def __init__(self):
+        self.items = [("aa", "!"), ("bb", "!"), ("cc", "!")]
+
+    @property
+    def name(self):
+        return "toy_rollout"
+
+    def __len__(self):
+        return len(self.items)
+
+    def clean(self, y_hat):
+        return y_hat[-1:]
+
+    def get(self, indx):
+        return self.items[indx]
+
+    def max_new_tokens(self, inference):
+        return 1
+
+    def check(self, y, y_hat):
+        return y == y_hat
+
+
+class ToyEncodingEval(EncodingEvaluation):
+    def __init__(self):
+        self.items = ["abc", "def", "ghi"]
+
+    @property
+    def name(self):
+        return "toy_encoding"
+
+    def __len__(self):
+        return len(self.items)
+
+    def clean(self, y_hat):
+        return y_hat
+
+    def get(self, indx):
+        return self.items[indx]
+
+    def check(self, x, y_hat):
+        return y_hat == x[1:]
+
+
+class ToyPerplexityEval(PerplexityEvaluation):
+    def __init__(self):
+        self.items = ["abc", "def", "ghi"]
+
+    @property
+    def name(self):
+        return "toy_ppl"
+
+    def __len__(self):
+        return len(self.items)
+
+    def get(self, indx):
+        return self.items[indx]
+
+
+class ToyComparisonEval(PerplexityComparisonEvaluation):
+    def __init__(self):
+        self.items = [("a", ["bc", "zz"], 0), ("d", ["ef", "yy"], 0), ("g", ["hi", "xx"], 0)]
+
+    @property
+    def name(self):
+        return "toy_compare"
+
+    def __len__(self):
+        return len(self.items)
+
+    def get(self, indx):
+        return self.items[indx]
+
+
+class TestEvalPadding:
+    """Test that evaluations handle padding correctly when dataset size < batch unit."""
+
+    def test_rollout_eval(self):
+        tokenizer = ToyTokenizer()
+        inference = DummyInference(batch_unit=8, block_size=8)
+        score = ToyRolloutEval()(inference, tokenizer)
+        assert score == 1.0
+
+    def test_encoding_eval(self):
+        tokenizer = ToyTokenizer()
+        inference = DummyInference(batch_unit=8, block_size=8)
+        score = ToyEncodingEval()(inference, tokenizer)
+        assert score == 1.0
+
+    def test_perplexity_eval(self):
+        tokenizer = ToyTokenizer()
+        inference = DummyInference(batch_unit=8, block_size=8)
+        score = ToyPerplexityEval()(inference, tokenizer)
+        assert 0.99 < score <= 1.0
+
+    def test_comparison_eval(self):
+        tokenizer = ToyTokenizer()
+        inference = DummyInference(batch_unit=8, block_size=8)
+        score = ToyComparisonEval()(inference, tokenizer)
+        assert score == 1.0

libthx-0.2.1/tests/test_gpu_availability.py

@@ -0,0 +1,71 @@
+"""Tests for _check_plain_host_availability GPU memory parsing logic.
+
+Migrated from scripts/test_gpu_availability.py.
+"""
+
+from unittest.mock import MagicMock, patch
+import pytest
+
+
+def make_run_result(stdout: str, ok: bool = True) -> MagicMock:
+    r = MagicMock()
+    r.ok = ok
+    r.stdout = stdout
+    r.stderr = ""
+    return r
+
+
+def check(nvidia_smi_output: str, configured: int = 4) -> int:
+    from theseus.dispatch.solve import _check_plain_host_availability
+
+    with patch(
+        "theseus.dispatch.ssh.run", return_value=make_run_result(nvidia_smi_output)
+    ):
+        return _check_plain_host_availability("fake-host", configured, timeout=5.0)
+
+
+class TestGPUAvailability:
+    def test_no_processes(self):
+        assert check("") == 4
+
+    def test_xorg_only(self):
+        output = "GPU-abc, 3374, /usr/lib/xorg/Xorg, 4\nGPU-def, 3374, /usr/lib/xorg/Xorg, 4\n"
+        assert check(output) == 4
+
+    def test_zero_memory(self):
+        output = "GPU-abc, 1234, some-daemon, 0\n"
+        assert check(output) == 4
+
+    def test_na_memory(self):
+        output = "GPU-abc, 1234, some-process, N/A\n"
+        assert check(output) == 4
+
+    def test_missing_memory_field(self):
+        output = "GPU-abc, 1234, some-process\n"
+        assert check(output) == 4
+
+    def test_real_training_job(self):
+        output = "GPU-abc, 9999, python, 38000\n"
+        assert check(output) == 0
+
+    def test_mixed_noise_and_real(self):
+        output = "GPU-abc, 3374, /usr/lib/xorg/Xorg, 4\nGPU-def, 9999, python, 38000\n"
+        assert check(output) == 0
+
+    def test_exactly_at_threshold(self):
+        output = "GPU-abc, 9999, python, 100\n"
+        assert check(output) == 0
+
+    def test_just_below_threshold(self):
+        output = "GPU-abc, 9999, some-process, 99\n"
+        assert check(output) == 4
+
+    def test_nvidia_smi_failure(self):
+        from theseus.dispatch.solve import _check_plain_host_availability
+
+        with patch(
+            "theseus.dispatch.ssh.run",
+            return_value=make_run_result("", ok=False),
+        ):
+            result = _check_plain_host_availability("fake-host", 4, timeout=5.0)
+        assert result == 0