tensorless-pytorch 0.2.0__tar.gz → 0.3.0__tar.gz

{tensorless_pytorch-0.2.0/tensorless_pytorch.egg-info → tensorless_pytorch-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorless-pytorch
-Version: 0.2.0
+Version: 0.3.0
 Summary: Automatic PyTorch training with portable models and minimal setup.
 Author: Tensorless PyTorch Contributors
 License: MIT
@@ -41,10 +41,18 @@ tokenizer; use `tokenizer="char"` for a character-level model. Tensorless PyTorc
 derives model size, batch size, epochs, validation, device, and BPE vocabulary
 size from the data, while every setting can be overridden.
 
-Long text is tokenized lazily and fed through PyTorch in fixed-size batches.
-CUDA training automatically uses `fp16` or `bf16` when supported, including
-gradient scaling and checkpointed scaler state. Reduce `batch_size` if memory
-is limited.
+Text is tokenized once up front (not re-tokenized every epoch) and streamed
+through PyTorch in fixed-size batches. Model size is auto-scaled with corpus
+size across four tiers -- small, lower-mid, upper-mid, large -- capped at
+"upper-mid" (roughly 50M-350M parameters) unless you pass explicit
+`d_model=`/`layers=`/etc. yourself. New training runs use the `architecture="v2"`
+backbone (RoPE + RMSNorm + SwiGLU + KV-cached generation); older `.tl`
+checkpoints keep loading and running on the original `"v1"` backbone
+automatically. CUDA training automatically uses `fp16` or `bf16` when
+supported (with gradient scaling and checkpointed scaler state), TPU (XLA)
+training is supported via `device="tpu"`, and large auto-sized models enable
+gradient checkpointing automatically to fit in memory. Reduce `batch_size` if
+memory is limited.
 
 ## English starter pretraining
 

{tensorless_pytorch-0.2.0 → tensorless_pytorch-0.3.0}/README.md

@@ -27,10 +27,18 @@ tokenizer; use `tokenizer="char"` for a character-level model. Tensorless PyTorc
 derives model size, batch size, epochs, validation, device, and BPE vocabulary
 size from the data, while every setting can be overridden.
 
-Long text is tokenized lazily and fed through PyTorch in fixed-size batches.
-CUDA training automatically uses `fp16` or `bf16` when supported, including
-gradient scaling and checkpointed scaler state. Reduce `batch_size` if memory
-is limited.
+Text is tokenized once up front (not re-tokenized every epoch) and streamed
+through PyTorch in fixed-size batches. Model size is auto-scaled with corpus
+size across four tiers -- small, lower-mid, upper-mid, large -- capped at
+"upper-mid" (roughly 50M-350M parameters) unless you pass explicit
+`d_model=`/`layers=`/etc. yourself. New training runs use the `architecture="v2"`
+backbone (RoPE + RMSNorm + SwiGLU + KV-cached generation); older `.tl`
+checkpoints keep loading and running on the original `"v1"` backbone
+automatically. CUDA training automatically uses `fp16` or `bf16` when
+supported (with gradient scaling and checkpointed scaler state), TPU (XLA)
+training is supported via `device="tpu"`, and large auto-sized models enable
+gradient checkpointing automatically to fit in memory. Reduce `batch_size` if
+memory is limited.
 
 ## English starter pretraining
 

{tensorless_pytorch-0.2.0 → tensorless_pytorch-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorless-pytorch"
-version = "0.2.0"
+version = "0.3.0"
 description = "Automatic PyTorch training with portable models and minimal setup."
 readme = "README.md"
 requires-python = ">=3.9"

tensorless_pytorch-0.3.0/tensorless/auto/config.py

@@ -0,0 +1,242 @@
+"""Automatic configuration.
+
+Turns a `Dataset` + a user-supplied `TrainConfig` (with mostly-None fields)
+into a fully-resolved `ResolvedConfig`. Every automatic choice is driven by
+simple, explainable heuristics based on dataset size and available
+hardware -- this is not meant to be state-of-the-art NAS, it's meant to
+produce a *sane, working* default so `tl.train("./data")` just works,
+scaling all the way from tiny CPU-only experiments up to "upper-mid"
+scale models (roughly 100M-350M parameters for transformers) when the
+data and hardware support it.
+
+Model-size tiers (approximate parameter counts, transformer LM):
+    small       < 10M
+    lower-mid   10M  - 50M
+    upper-mid   50M  - 350M   <- Tensorless auto-config tops out here
+    large       350M - 3B     (never auto-selected; pass explicit
+                                d_model=/layers=/etc. yourself)
+    extreme     3B+           (out of scope for this project)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Tuple
+
+from ..config import TrainConfig, ResolvedConfig
+from ..data.loader import Dataset
+from ..devices.device import auto_select_device, get_torch_device, recommend_max_params
+from .detector import detect_task
+
+
+# ---------------------------------------------------------------------------
+# Architecture sizing
+# ---------------------------------------------------------------------------
+
+def _text_model_tier(n_examples: int) -> Tuple[int, int, int, int]:
+    """Return (d_model, layers, heads, ff_mult), scaled up to the
+    "upper-mid" ceiling as effective dataset size grows. `ff_mult` is
+    interpreted by the v2 (SwiGLU) architecture as a multiplier on the
+    gated hidden width, same as v1's plain MLP hidden width.
+    """
+    if n_examples < 200:
+        return 64, 2, 2, 2          # tiny / smoke-test scale
+    elif n_examples < 5_000:
+        return 128, 4, 4, 3         # small
+    elif n_examples < 50_000:
+        return 256, 6, 8, 4         # small / lower-mid boundary
+    elif n_examples < 200_000:
+        return 384, 8, 8, 4         # lower-mid
+    elif n_examples < 1_000_000:
+        return 640, 12, 10, 4       # lower-mid / upper-mid boundary (~80M params)
+    else:
+        return 1024, 16, 16, 4      # upper-mid ceiling (~250-300M params)
+
+
+def _tabular_model_tier(n_examples: int) -> Tuple[int, int, int, int]:
+    if n_examples < 500:
+        return 32, 2, 1, 2
+    elif n_examples < 20_000:
+        return 64, 3, 1, 2
+    elif n_examples < 200_000:
+        return 128, 5, 1, 2
+    else:
+        return 256, 8, 1, 2         # upper-mid ceiling for tabular data
+
+
+def _estimate_transformer_params(vocab_size: int, d_model: int, layers: int, ff_mult: int) -> int:
+    swiglu_hidden = max(int(d_model * ff_mult * (2 / 3)), d_model)
+    per_layer = 4 * d_model * d_model + 3 * d_model * swiglu_hidden  # attn proj + swiglu gate/up/down
+    return vocab_size * d_model + layers * per_layer
+
+
+def _fit_within_param_budget(
+    vocab_size: int, d_model: int, layers: int, heads: int, ff_mult: int, max_params: int
+) -> Tuple[int, int, int, int]:
+    """If the tier-selected architecture would exceed what the detected
+    hardware can comfortably hold, scale it back down -- first by
+    trimming depth, then width -- until it fits. Keeps `d_model`
+    divisible by `heads` throughout.
+    """
+    while _estimate_transformer_params(vocab_size, d_model, layers, ff_mult) > max_params and layers > 2:
+        layers -= 1
+    while _estimate_transformer_params(vocab_size, d_model, layers, ff_mult) > max_params and d_model > 64:
+        d_model = max(64, d_model - 64)
+        while d_model % heads != 0 and heads > 1:
+            heads -= 1
+    return d_model, layers, heads, ff_mult
+
+
+def _auto_batch_size(n_examples: int, max_seq_len: int, device: str) -> int:
+    # GPUs/TPUs benefit from a larger token budget per step than CPU.
+    token_budget = 32768 if device in ("cuda", "tpu") else 8192
+    sequence_batch = max(1, token_budget // max_seq_len)
+    if n_examples < 200:
+        return min(8, sequence_batch)
+    elif n_examples < 2000:
+        return min(16, sequence_batch)
+    elif n_examples < 20000:
+        return min(32, sequence_batch)
+    elif n_examples < 200000:
+        return min(64, sequence_batch)
+    else:
+        return min(128, sequence_batch)
+
+
+def _auto_epochs(n_examples: int) -> int:
+    if n_examples < 200:
+        return 40
+    elif n_examples < 2000:
+        return 20
+    elif n_examples < 20000:
+        return 10
+    else:
+        return 5
+
+
+def _effective_text_size(ds: Dataset) -> int:
+    """Estimate useful training examples for raw corpora."""
+    if ds.kind in ("text", "text_labeled"):
+        return max(len(ds), sum(len(text) for text in ds.texts) // 200)
+    return len(ds)
+
+
+def _auto_vocab_size(ds: Dataset, n_examples: int) -> int:
+    if ds.kind not in ("text", "text_labeled"):
+        return 1000
+    unique_chars = len(set("".join(ds.texts)))
+    # Bigger corpora justify (and benefit from) a bigger BPE vocabulary --
+    # fewer tokens per sequence, more of max_seq_len spent on real context.
+    if n_examples < 5_000:
+        ceiling = 2048
+    elif n_examples < 200_000:
+        ceiling = 8192
+    else:
+        ceiling = 32000
+    return min(ceiling, max(64, unique_chars * 8))
+
+
+def _auto_num_workers(device: str) -> int:
+    import os
+
+    cpu_count = os.cpu_count() or 1
+    if device in ("cuda", "tpu"):
+        # Leave headroom for the main process / accelerator driver threads.
+        return max(0, min(4, cpu_count - 1))
+    return 0  # CPU training: extra worker processes just add overhead
+
+
+def resolve_config(ds: Dataset, user: TrainConfig) -> ResolvedConfig:
+    n = _effective_text_size(ds)
+    task = user.task or detect_task(ds)
+    model_type = user.model_type or (
+        "transformer" if task in ("text-generation", "text-classification") else "mlp"
+    )
+    tokenizer = user.tokenizer or "bpe"
+    if tokenizer not in ("char", "bpe"):
+        raise ValueError("tokenizer must be 'char' or 'bpe'")
+
+    architecture = user.architecture or "v2"
+    if architecture not in ("v1", "v2"):
+        raise ValueError("architecture must be 'v1' or 'v2'")
+
+    device, precision = auto_select_device(user.device, user.precision)
+    torch_device = get_torch_device(device)
+
+    if model_type == "transformer":
+        d_model, layers, heads, ff_mult = _text_model_tier(n)
+    else:
+        d_model, layers, heads, ff_mult = _tabular_model_tier(n)
+
+    max_seq_len = user.max_seq_len or (256 if ds.kind in ("text", "text_labeled") else 1)
+    bpe_vocab_size = (
+        user.bpe_vocab_size if user.bpe_vocab_size is not None else _auto_vocab_size(ds, n)
+    )
+
+    # Memory-aware cap: only clamp automatically-chosen sizes, never a
+    # user's explicit d_model=/layers=/heads=/ff_mult= override.
+    user_set_arch = any(
+        v is not None for v in (user.d_model, user.layers, user.heads, user.ff_mult)
+    )
+    if model_type == "transformer" and not user_set_arch:
+        max_params = recommend_max_params(torch_device)
+        d_model, layers, heads, ff_mult = _fit_within_param_budget(
+            bpe_vocab_size if tokenizer == "bpe" else 256, d_model, layers, heads, ff_mult, max_params
+        )
+
+    gradient_checkpointing = user.gradient_checkpointing
+    if gradient_checkpointing is None:
+        # Auto-enable on GPU/TPU once the model is large enough that
+        # activation memory (not parameter memory) becomes the bottleneck.
+        approx_params = _estimate_transformer_params(
+            bpe_vocab_size if tokenizer == "bpe" else 256, d_model, layers, ff_mult
+        ) if model_type == "transformer" else 0
+        gradient_checkpointing = device in ("cuda", "tpu") and approx_params > 60_000_000
+
+    compile_model = user.compile
+    if compile_model is None:
+        compile_model = device == "cuda"
+
+    num_workers = user.num_workers if user.num_workers is not None else _auto_num_workers(device)
+
+    out = user.out or "model.tl"
+    checkpoint_dir = user.checkpoint_dir or (out + ".ckpt")
+
+    resolved = ResolvedConfig(
+        out=out,
+        force=bool(user.force),
+        resume=user.resume,
+        ask_on_data_change=bool(user.ask_on_data_change),
+        task=task,
+        model_type=model_type,
+        architecture=architecture,
+        d_model=user.d_model or d_model,
+        layers=user.layers or layers,
+        heads=user.heads or heads,
+        ff_mult=user.ff_mult or ff_mult,
+        dropout=user.dropout if user.dropout is not None else 0.1,
+        max_seq_len=max_seq_len,
+        tokenizer=tokenizer,
+        bpe_vocab_size=bpe_vocab_size,
+        optimizer=user.optimizer or "adamw",
+        learning_rate=user.learning_rate or (3e-4 if model_type == "transformer" else 1e-3),
+        weight_decay=user.weight_decay if user.weight_decay is not None else 0.01,
+        batch_size=user.batch_size if user.batch_size is not None else _auto_batch_size(n, max_seq_len, device),
+        epochs=user.epochs or _auto_epochs(n),
+        max_steps=user.max_steps,
+        gradient_accumulation_steps=user.gradient_accumulation_steps or 1,
+        grad_clip=user.grad_clip if user.grad_clip is not None else 1.0,
+        warmup_steps=user.warmup_steps if user.warmup_steps is not None else min(100, max(1, n // 10)),
+        val_split=user.val_split if user.val_split is not None else (0.1 if n >= 50 else 0.0),
+        patience=user.patience if user.patience is not None else 3,
+        min_delta=user.min_delta if user.min_delta is not None else 1e-4,
+        device=device,
+        precision=precision,
+        checkpoint_every=user.checkpoint_every or 50,
+        checkpoint_dir=checkpoint_dir,
+        gradient_checkpointing=gradient_checkpointing,
+        compile=compile_model,
+        num_workers=num_workers,
+        seed=user.seed,
+        verbose=user.verbose,
+    )
+    return resolved

{tensorless_pytorch-0.2.0 → tensorless_pytorch-0.3.0}/tensorless/config.py

@@ -27,6 +27,7 @@ class TrainConfig:
     # --- task / architecture ---
     task: Optional[str] = None                  # "text-generation", "classification", "regression"
     model_type: Optional[str] = None            # "transformer", "mlp"
+    architecture: Optional[str] = None          # "v1" (legacy) or "v2" (RoPE/RMSNorm/SwiGLU); default "v2"
     d_model: Optional[int] = None
     layers: Optional[int] = None
     heads: Optional[int] = None
@@ -35,6 +36,9 @@ class TrainConfig:
     max_seq_len: Optional[int] = None
     tokenizer: Optional[str] = None           # "char" or "bpe"
     bpe_vocab_size: Optional[int] = None
+    gradient_checkpointing: Optional[bool] = None  # trade compute for memory on big models
+    compile: Optional[bool] = None                  # torch.compile() on CUDA when available
+    num_workers: Optional[int] = None               # DataLoader worker processes
 
     # --- optimization ---
     optimizer: Optional[str] = None             # "adamw", "adam", "sgd"
@@ -86,6 +90,7 @@ class ResolvedConfig:
 
     task: str
     model_type: str
+    architecture: str
     d_model: int
     layers: int
     heads: int
@@ -94,6 +99,9 @@ class ResolvedConfig:
     max_seq_len: int
     tokenizer: str
     bpe_vocab_size: int
+    gradient_checkpointing: bool
+    compile: bool
+    num_workers: int
 
     optimizer: str
     learning_rate: float

tensorless_pytorch-0.3.0/tensorless/devices/__init__.py

@@ -0,0 +1,17 @@
+from .device import (
+    auto_select_device,
+    get_torch_device,
+    get_device_info,
+    mark_step,
+    recommend_max_params,
+    tpu_core_count,
+)
+
+__all__ = [
+    "auto_select_device",
+    "get_torch_device",
+    "get_device_info",
+    "mark_step",
+    "recommend_max_params",
+    "tpu_core_count",
+]

{tensorless_pytorch-0.2.0 → tensorless_pytorch-0.3.0}/tensorless/devices/device.py

@@ -9,7 +9,8 @@ trainer (see `training/trainer.py`) will catch that and fall back too.
 
 from __future__ import annotations
 
-from typing import Optional, Tuple
+import os
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 
@@ -85,6 +86,85 @@ def auto_select_device(user_device: Optional[str], user_precision: Optional[str]
     return device, precision
 
 
+def tpu_core_count() -> int:
+    """Number of TPU cores visible to this process (1 if not on TPU)."""
+    try:
+        import torch_xla.core.xla_model as xm
+
+        return xm.xrt_world_size()
+    except Exception:
+        try:
+            import torch_xla.runtime as xr
+
+            return xr.world_size()
+        except Exception:
+            return 1
+
+
+def mark_step(device: torch.device) -> None:
+    """XLA/TPU tensors are lazily evaluated -- without periodically calling
+    `xm.mark_step()` the compute graph never actually executes and a
+    'TPU-enabled' training run will just silently build up an ever-larger
+    graph (or hang) instead of doing real work. This is a no-op on every
+    other device, so it's safe to call unconditionally from the trainer.
+    """
+    if device.type == "xla":
+        try:
+            import torch_xla.core.xla_model as xm
+
+            xm.mark_step()
+        except Exception:
+            pass
+
+
+def get_device_info(device: torch.device) -> Dict[str, Any]:
+    """Human-readable + programmatic info about the active compute device,
+    used both for the startup log line and for memory-aware auto-sizing.
+    """
+    info: Dict[str, Any] = {"type": device.type, "name": device.type, "memory_gb": None, "num_devices": 1}
+    try:
+        if device.type == "cuda":
+            idx = device.index if device.index is not None else 0
+            props = torch.cuda.get_device_properties(idx)
+            info["name"] = props.name
+            info["memory_gb"] = round(props.total_memory / (1024 ** 3), 2)
+            info["num_devices"] = torch.cuda.device_count()
+            info["bf16"] = _cuda_supports_bf16()
+        elif device.type == "xla":
+            info["name"] = "TPU"
+            info["num_devices"] = tpu_core_count()
+        elif device.type == "mps":
+            info["name"] = "Apple Silicon (MPS)"
+        else:
+            info["name"] = "CPU"
+            info["num_devices"] = os.cpu_count() or 1
+    except Exception:
+        pass
+    return info
+
+
+def recommend_max_params(device: torch.device) -> int:
+    """Rough ceiling on trainable parameter count so an auto-configured
+    model has a reasonable chance of fitting in memory alongside AdamW's
+    optimizer state (~2x params) and activations. Deliberately
+    conservative -- users can always override architecture args by hand.
+    """
+    info = get_device_info(device)
+    if device.type == "cuda" and info.get("memory_gb"):
+        # Rule of thumb for AdamW + fp16/bf16 mixed precision + activations:
+        # budget roughly 1M params per 15MB of GPU memory, capped well
+        # inside the "upper-mid" tier regardless of how much memory exists.
+        budget = int(info["memory_gb"] * 1024 * (1_000_000 / 15))
+        return max(2_000_000, min(budget, 350_000_000))
+    if device.type == "xla":
+        # TPU v3/v4 cores typically expose >=16GB HBM each.
+        return 350_000_000
+    if device.type == "mps":
+        return 60_000_000
+    # CPU: keep it trainable in a reasonable amount of wall-clock time.
+    return 40_000_000
+
+
 def get_torch_device(device: str) -> torch.device:
     """Convert our string device name into a torch.device, with a
     runtime fallback to CPU if the requested backend is unavailable.

tensorless_pytorch-0.3.0/tensorless/models/__init__.py

@@ -0,0 +1,13 @@
+from .registry import build_model
+from .transformer import TinyTransformer, TinyTransformerV1, TinyTransformerV2
+from .mlp import TabularMLP, TabularMLPV1, TabularMLPV2
+
+__all__ = [
+    "build_model",
+    "TinyTransformer",
+    "TinyTransformerV1",
+    "TinyTransformerV2",
+    "TabularMLP",
+    "TabularMLPV1",
+    "TabularMLPV2",
+]

tensorless_pytorch-0.3.0/tensorless/models/mlp.py

@@ -0,0 +1,142 @@
+"""MLP model for tabular data (classification / regression).
+
+Numeric columns feed directly into the network; each categorical column
+gets its own small embedding table, and all features are concatenated
+before the hidden layers.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _embedding_dim(vocab_size: int) -> int:
+    # Common rule of thumb, capped to keep tiny-data models tiny.
+    return max(2, min(32, round(1.6 * (vocab_size ** 0.56))))
+
+
+class TabularMLPV1(nn.Module):
+    def __init__(
+        self,
+        n_numeric: int,
+        categorical_vocab_sizes: List[int],
+        d_model: int,
+        layers: int,
+        dropout: float,
+        task: str,
+        n_classes: int = 0,
+    ):
+        super().__init__()
+        self.task = task
+        self.n_numeric = n_numeric
+        self.categorical_vocab_sizes = categorical_vocab_sizes
+
+        self.embeddings = nn.ModuleList(
+            [nn.Embedding(v, _embedding_dim(v)) for v in categorical_vocab_sizes]
+        )
+        cat_dim = sum(_embedding_dim(v) for v in categorical_vocab_sizes)
+        in_dim = n_numeric + cat_dim
+
+        dims = [in_dim] + [d_model] * layers
+        modules = []
+        for i in range(len(dims) - 1):
+            modules += [nn.Linear(dims[i], dims[i + 1]), nn.ReLU(), nn.Dropout(dropout)]
+        self.backbone = nn.Sequential(*modules)
+
+        out_dim = n_classes if task == "classification" else 1
+        self.head = nn.Linear(dims[-1], out_dim)
+
+    def forward(self, numeric: torch.Tensor, categorical: torch.Tensor) -> torch.Tensor:
+        parts = []
+        if self.n_numeric > 0:
+            parts.append(numeric)
+        for i, emb in enumerate(self.embeddings):
+            parts.append(emb(categorical[:, i]))
+        x = torch.cat(parts, dim=1) if parts else numeric
+        x = self.backbone(x)
+        out = self.head(x)
+        if self.task == "regression":
+            return out.squeeze(-1)
+        return out
+
+
+class _ResidualBlock(nn.Module):
+    """Pre-norm residual MLP block (LayerNorm -> Linear -> GELU -> Linear
+    -> residual add). Deeper `TabularMLPV2` stacks of these train more
+    stably than plain stacked Linear+ReLU at larger widths/depths.
+    """
+
+    def __init__(self, d_model: int, dropout: float):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        self.fc1 = nn.Linear(d_model, d_model)
+        self.fc2 = nn.Linear(d_model, d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = F.gelu(self.fc1(h))
+        h = self.drop(h)
+        h = self.fc2(h)
+        return x + self.drop(h)
+
+
+class TabularMLPV2(nn.Module):
+    """Residual, pre-norm tabular MLP -- same public interface as
+    `TabularMLPV1`, but scales more gracefully to the wider/deeper configs
+    Tensorless now auto-selects for larger tabular datasets.
+    """
+
+    def __init__(
+        self,
+        n_numeric: int,
+        categorical_vocab_sizes: List[int],
+        d_model: int,
+        layers: int,
+        dropout: float,
+        task: str,
+        n_classes: int = 0,
+    ):
+        super().__init__()
+        self.task = task
+        self.n_numeric = n_numeric
+        self.categorical_vocab_sizes = categorical_vocab_sizes
+
+        self.embeddings = nn.ModuleList(
+            [nn.Embedding(v, _embedding_dim(v)) for v in categorical_vocab_sizes]
+        )
+        cat_dim = sum(_embedding_dim(v) for v in categorical_vocab_sizes)
+        in_dim = n_numeric + cat_dim
+
+        self.in_proj = nn.Linear(in_dim, d_model)
+        self.blocks = nn.ModuleList([_ResidualBlock(d_model, dropout) for _ in range(layers)])
+        self.norm_f = nn.LayerNorm(d_model)
+
+        out_dim = n_classes if task == "classification" else 1
+        self.head = nn.Linear(d_model, out_dim)
+
+    def forward(self, numeric: torch.Tensor, categorical: torch.Tensor) -> torch.Tensor:
+        parts = []
+        if self.n_numeric > 0:
+            parts.append(numeric)
+        for i, emb in enumerate(self.embeddings):
+            parts.append(emb(categorical[:, i]))
+        x = torch.cat(parts, dim=1) if parts else numeric
+        x = self.in_proj(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm_f(x)
+        out = self.head(x)
+        if self.task == "regression":
+            return out.squeeze(-1)
+        return out
+
+
+# New training runs default to the residual v2 architecture; v1 stays
+# available (and is what old checkpoints without an "architecture" key
+# in their config resolve to) for backward compatibility.
+TabularMLP = TabularMLPV2

{tensorless_pytorch-0.2.0 → tensorless_pytorch-0.3.0}/tensorless/models/registry.py

@@ -5,6 +5,14 @@ config + the dataset-derived metadata (vocab size, n_classes, etc). Kept
 separate from `auto/config.py` so new model types/backends can be added
 here without touching auto-configuration logic, per the "extensible
 architecture" requirement.
+
+Also owns the v1/v2 architecture switch: `cfg.get("architecture", "v1")`
+decides which concrete class gets built. Checkpoints saved before the v2
+upgrade have no "architecture" key in their stored config at all, so
+`.get(..., "v1")` transparently resolves them to the original
+implementation -- old `.tl` files keep loading and running correctly.
+New training runs get "v2" (better quality, KV-cache generation, scales
+to "upper-mid" sizes) unless the user explicitly pins `architecture="v1"`.
 """
 
 from __future__ import annotations
@@ -13,8 +21,8 @@ from typing import Any, Dict
 
 import torch.nn as nn
 
-from .transformer import TinyTransformer
-from .mlp import TabularMLP
+from .transformer import TinyTransformerV1, TinyTransformerV2
+from .mlp import TabularMLPV1, TabularMLPV2
 from ..errors import ModelError
 
 
@@ -26,8 +34,10 @@ def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str,
       - text tasks: {"vocab_size": int, "pad_id": int, "n_classes": int}
       - tabular tasks: {"n_numeric": int, "categorical_vocab_sizes": [...], "n_classes": int}
     """
+    architecture = cfg.get("architecture", "v1")
+
     if model_type == "transformer":
-        return TinyTransformer(
+        common = dict(
             vocab_size=meta["vocab_size"],
             d_model=cfg["d_model"],
             layers=cfg["layers"],
@@ -39,8 +49,15 @@ def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str,
             n_classes=meta.get("n_classes", 0),
             pad_id=meta.get("pad_id", 0),
         )
+        if architecture == "v2":
+            return TinyTransformerV2(
+                gradient_checkpointing=bool(cfg.get("gradient_checkpointing", False)),
+                **common,
+            )
+        return TinyTransformerV1(**common)
+
     elif model_type == "mlp":
-        return TabularMLP(
+        common = dict(
             n_numeric=meta["n_numeric"],
             categorical_vocab_sizes=meta["categorical_vocab_sizes"],
             d_model=cfg["d_model"],
@@ -49,5 +66,9 @@ def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str,
             task=task,
             n_classes=meta.get("n_classes", 0),
         )
+        if architecture == "v2":
+            return TabularMLPV2(**common)
+        return TabularMLPV1(**common)
+
     else:
         raise ModelError(f"Unknown model_type '{model_type}'.")