gradientmail 0.0.1__py3-none-any.whl

gradientmail/__init__.py

@@ -0,0 +1,59 @@
+"""GradientMail: distributed training through asynchronous message exchange.
+
+``import gradientmail`` needs only the standard library and safetensors. ``Session`` wraps a
+training loop; every session also takes part in closing rounds, so a run has no separate
+process to operate. What a node reports while it runs is in ``gradientmail.events``, handed
+to the session's ``callbacks``. ``S3Mailbox`` is resolved on first access and raises
+``MissingExtra`` naming the ``s3`` extra when boto3 is absent. The framework backends live in
+``gradientmail.adapters.torch`` and ``gradientmail.adapters.mlx``.
+"""
+
+from typing import TYPE_CHECKING
+
+from gradientmail.adapters import Adapter, Backend, MemoryAdapter, MissingExtra
+from gradientmail.algorithms import Algorithm
+from gradientmail.algorithms.diloco import ByBytes, DiLoCo, Rda, Strided, WeightedMean
+from gradientmail.events import Callback, Event, log_events
+from gradientmail.mailbox import Mailbox, MailboxUnavailable
+from gradientmail.mailbox.memory import MemoryMailbox
+from gradientmail.protocol import RunConfigMismatch
+from gradientmail.serialization import Schema, SchemaMismatch, TensorBytes
+from gradientmail.session import Session, align
+
+if TYPE_CHECKING:
+    from gradientmail.mailbox.s3 import S3Mailbox as S3Mailbox
+
+__all__ = [
+    "Adapter",
+    "Algorithm",
+    "Backend",
+    "ByBytes",
+    "Callback",
+    "DiLoCo",
+    "Event",
+    "Mailbox",
+    "MailboxUnavailable",
+    "MemoryAdapter",
+    "MemoryMailbox",
+    "MissingExtra",
+    "Rda",
+    "RunConfigMismatch",
+    "Schema",
+    "SchemaMismatch",
+    "Session",
+    "Strided",
+    "TensorBytes",
+    "WeightedMean",
+    "align",
+    "log_events",
+]
+
+
+def __getattr__(name: str) -> object:
+    if name == "S3Mailbox":
+        try:
+            from gradientmail.mailbox.s3 import S3Mailbox
+        except ImportError as exc:
+            raise MissingExtra("s3") from exc
+        return S3Mailbox
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

gradientmail/adapters/__init__.py

@@ -0,0 +1,174 @@
+"""Framework backends: how a model's parameters become bytes, and how bytes become arithmetic.
+
+A ``Backend`` is everything GradientMail knows about one framework. On a learner it wraps a
+model in an ``Adapter``, the three-method byte boundary (``schema``, ``snapshot``, ``apply``);
+on a syncer it loads ``TensorBytes`` as ``Value``s, the handful of fp32 operations the outer
+step is made of, which serialize themselves back. ``TorchBackend`` and ``MlxBackend`` live
+in their own modules and are imported only when asked for, so this package needs nothing
+beyond the standard library. Validation against the schema is the
+caller's job (``Learner`` before every ``apply``, ``Syncer`` before merging), so adapters
+stay thin.
+"""
+
+from collections.abc import Iterable, Mapping
+from importlib.util import find_spec
+from typing import Protocol, Self, runtime_checkable
+
+from gradientmail.serialization import (
+    DType,
+    ParameterSpec,
+    Quantized,
+    Schema,
+    SchemaMismatch,
+    TensorBytes,
+)
+
+__all__ = [
+    "Adapter",
+    "Backend",
+    "MemoryAdapter",
+    "MissingExtra",
+    "Value",
+    "adapter_for",
+    "backend_for",
+    "default_backend",
+]
+
+
+class MissingExtra(ImportError):
+    """An optional dependency is not installed; ``extra`` names the extra that provides it."""
+
+    extra: str
+
+    def __init__(self, extra: str) -> None:
+        self.extra = extra
+        super().__init__(f"this needs the {extra!r} extra: pip install 'gradientmail[{extra}]'")
+
+
+@runtime_checkable
+class Adapter(Protocol):
+    def schema(self) -> Schema[DType]:
+        """Canonical trainable parameters, sorted by name."""
+        ...
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        """Owned copies of ``names``, consistent at call time."""
+        ...
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        """Overwrite parameters in place; never rebinds them, so optimizer state survives."""
+        ...
+
+
+@runtime_checkable
+class Value(Protocol):
+    """An fp32 tensor held by a backend, with the arithmetic the outer step is made of.
+
+    Every operation returns a new value of the same backend. ``norm`` is the Euclidean norm
+    over all elements and ``dot`` the sum of the elementwise products with ``other``.
+    ``to_bytes`` rounds to ``dtype`` (nearest even) and serializes; ``quantize`` serializes to
+    a ``Quantized`` format as ``(data, scales)`` in the layout ``serialization.quantized_specs``
+    describes, each scale a power of two that puts the largest magnitude of its tensor or
+    block at the top of the format, so dequantizing is exact in fp32.
+    """
+
+    def __add__(self, other: Self) -> Self: ...
+
+    def __sub__(self, other: Self) -> Self: ...
+
+    def __mul__(self, scalar: float) -> Self: ...
+
+    def norm(self) -> float: ...
+
+    def dot(self, other: Self) -> float: ...
+
+    def to_bytes(self, dtype: DType) -> TensorBytes: ...
+
+    def quantize(self, dtype: Quantized) -> tuple[TensorBytes, TensorBytes]: ...
+
+
+@runtime_checkable
+class Backend(Protocol):
+    """One framework: it adapts that framework's models and does the syncer's arithmetic."""
+
+    def adapter(self, model: object) -> Adapter:
+        """An Adapter over ``model``; ``TypeError`` when it is not this framework's model."""
+        ...
+
+    def load(self, tensor: TensorBytes) -> Value:
+        """``tensor`` as an fp32 value of its shape; ``ValueError`` for a quantized dtype,
+        whose bytes mean nothing without their scales."""
+        ...
+
+    def dequantize(self, data: TensorBytes, scale: TensorBytes, shape: tuple[int, ...]) -> Value:
+        """The fp32 value of ``shape`` whose ``quantize`` gave ``(data, scale)``, up to the
+        format's rounding; the same floats on every backend."""
+        ...
+
+
+class MemoryAdapter(Adapter):
+    """An Adapter over a dict of ``TensorBytes``: the Phase 2 learner and the fake in tests."""
+
+    def __init__(self, tensors: Mapping[str, TensorBytes]) -> None:
+        self._tensors = dict(tensors)
+
+    @property
+    def tensors(self) -> Mapping[str, TensorBytes]:
+        return self._tensors
+
+    def schema(self) -> Schema[DType]:
+        def spec(name: str, t: TensorBytes) -> ParameterSpec[DType]:
+            match t.dtype:
+                case "F32" | "F16" | "BF16" as dtype:
+                    return ParameterSpec(name, dtype, t.shape)
+                case _:
+                    raise SchemaMismatch([name], "parameter dtype is not F32, F16 or BF16")
+
+        return Schema(tuple(spec(name, t) for name, t in sorted(self._tensors.items())))
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        return {name: self._tensors[name] for name in names}
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        self._tensors.update(tensors)
+
+
+def backend_for(model: object) -> Backend:
+    """The backend of ``model``'s framework, read from the modules of ``type(model)``'s MRO.
+
+    A user's ``class MLP(nn.Module)`` lives in the user's module, so the whole MRO is searched;
+    no framework is imported to decide. Anything else is a ``TypeError``.
+    """
+    roots = {cls.__module__.partition(".")[0] for cls in type(model).__mro__}
+    if "torch" in roots:
+        from gradientmail.adapters.torch import TorchBackend
+
+        return TorchBackend()
+    if "mlx" in roots:
+        from gradientmail.adapters.mlx import MlxBackend
+
+        return MlxBackend()
+    raise TypeError(
+        f"no backend for {type(model).__module__}.{type(model).__qualname__}; "
+        "pass a torch or MLX module, or an Adapter"
+    )
+
+
+def adapter_for(model: object) -> Adapter:
+    """``model`` itself when it is an Adapter; otherwise ``backend_for(model).adapter(model)``."""
+    if isinstance(model, Adapter):
+        return model
+    return backend_for(model).adapter(model)
+
+
+def default_backend() -> Backend:
+    """The installed backend: torch, then MLX. ``MissingExtra("torch")`` when neither is."""
+    if find_spec("torch") is not None:
+        from gradientmail.adapters.torch import TorchBackend
+
+        return TorchBackend()
+    if find_spec("mlx") is not None:
+        from gradientmail.adapters.mlx import MlxBackend
+
+        return MlxBackend()
+    raise MissingExtra("torch")

gradientmail/adapters/mlx.py

@@ -0,0 +1,252 @@
+"""MLX backend: the trainable arrays of an ``mlx.nn.Module`` as ``TensorBytes`` and back, and
+the outer step on ``mx.array``.
+
+Parameter names are the dotted paths of ``model.trainable_parameters()`` (``fc1.weight``,
+``layers.0.bias``), the names ``mlx.utils.tree_flatten`` and ``load_weights`` use, so a
+torch model with the same attribute structure produces the same schema. ``apply`` hands
+the payload to MLX's own safetensors reader and sets the arrays with ``Module.update``;
+MLX optimizer state is keyed by path rather than by array identity, so it survives. bf16 and
+f16 are native, so ``MlxValue`` converts with ``astype``; MLX flushes subnormals to zero on the
+way, where NumPy and torch round them. E4M3 goes through ``mx.to_fp8`` and ``mx.from_fp8``,
+which agree with torch's casts bit for bit; E2M1 nibbles and E8M0 scales are packed and
+unpacked by hand, as in the torch backend, so both produce the same bytes.
+
+Only ``mlx.core`` ships type information a checker accepts, so this module imports nothing
+from ``mlx.nn`` and types the model structurally through ``Module``.
+"""
+
+import io
+import json
+import math
+from collections import Counter
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass, replace
+from typing import Protocol, Self, assert_never, runtime_checkable
+
+import mlx.core as mx
+from mlx.utils import tree_unflatten
+
+from gradientmail.adapters import Adapter, Backend, Value
+from gradientmail.serialization import (
+    F4_BLOCK,
+    DType,
+    ParameterSpec,
+    Quantized,
+    Schema,
+    SchemaMismatch,
+    TensorBytes,
+    pack,
+    validate,
+)
+
+__all__ = ["MLX_DTYPE", "MlxAdapter", "MlxBackend", "MlxValue", "Module", "ParameterTree"]
+
+type ParameterTree = mx.array | Mapping[str, ParameterTree] | Sequence[ParameterTree]
+
+MLX_DTYPE: Mapping[DType, mx.Dtype] = {
+    "F32": mx.float32,
+    "F16": mx.float16,
+    "BF16": mx.bfloat16,
+}
+
+
+@runtime_checkable
+class Module(Protocol):
+    """What the adapter needs from a model; every ``mlx.nn.Module`` satisfies it."""
+
+    def trainable_parameters(self) -> Mapping[str, ParameterTree]: ...
+
+    def update(self, parameters: dict[str, ParameterTree], strict: bool = True) -> object: ...
+
+
+class MlxBackend(Backend):
+    def adapter(self, model: object) -> Adapter:
+        if not isinstance(model, Module):
+            raise TypeError(f"not an mlx.nn.Module: {type(model).__qualname__}")
+        return MlxAdapter(model)
+
+    def load(self, tensor: TensorBytes) -> Value:
+        match tensor.dtype:
+            case "F32" | "F16" | "BF16":
+                return MlxValue(_as_array(tensor).astype(mx.float32))
+            case other:
+                raise ValueError(f"{other} bytes need their scale: use dequantize")
+
+    def dequantize(self, data: TensorBytes, scale: TensorBytes, shape: tuple[int, ...]) -> Value:
+        match data.dtype:
+            case "F8_E4M3":
+                values = mx.from_fp8(_as_uint8(data.data), mx.float32) * _as_array(scale)
+            case "F4":
+                codes = _as_uint8(data.data)
+                nibbles = mx.stack([codes & 0xF, codes >> 4], axis=1).reshape(-1, F4_BLOCK)
+                magnitudes = mx.take(_F4_VALUES, nibbles & 7)
+                signed = mx.where((nibbles & 8) != 0, -magnitudes, magnitudes)
+                scales = _pow2(_as_uint8(scale.data).astype(mx.int32) - 127)
+                values = (signed * scales[:, None]).reshape(-1)[: math.prod(shape)]
+            case other:
+                raise ValueError(f"{other} is not a quantized dtype")
+        return MlxValue(values.reshape(shape))
+
+
+class MlxAdapter(Adapter):
+    """Adapter over an MLX module.
+
+    ``schema`` rejects two paths that hold the same array object: MLX ties are
+    identity-only and break at the first ``update``, so weights must be tied
+    structurally (``embed.as_linear``), which lists them once.
+    """
+
+    def __init__(self, model: Module) -> None:
+        self._model = model
+
+    def schema(self) -> Schema[DType]:
+        parameters = self._parameters()
+        holders = Counter(id(array) for array in parameters.values())
+        tied = tuple(name for name, array in parameters.items() if holders[id(array)] > 1)
+        if tied:
+            raise SchemaMismatch(tied, "parameters tied by identity; tie them structurally")
+        unsupported = tuple(
+            name for name, array in parameters.items() if array.dtype not in _DTYPE_CODE
+        )
+        if unsupported:
+            raise SchemaMismatch(unsupported, "dtype is not float32, float16 or bfloat16")
+        return Schema(
+            tuple(
+                ParameterSpec(name, _DTYPE_CODE[parameters[name].dtype], parameters[name].shape)
+                for name in sorted(parameters)
+            )
+        )
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        parameters = self._parameters()
+        return {name: _tensor_bytes(parameters[name]) for name in names}
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        validate(self.schema(), tensors, tensors)
+        loaded = mx.load(io.BytesIO(pack(tensors)), format="safetensors")
+        assert isinstance(loaded, dict)  # the stub's union; safetensors always decodes to a dict
+        self._model.update(tree_unflatten(list(loaded.items())))
+
+    def _parameters(self) -> Mapping[str, mx.array]:
+        def flattened(tree: ParameterTree, prefix: str) -> Iterator[tuple[str, mx.array]]:
+            match tree:
+                case mx.array():
+                    yield prefix.removeprefix("."), tree
+                case Mapping():
+                    for key, child in tree.items():
+                        yield from flattened(child, f"{prefix}.{key}")
+                case Sequence():
+                    for index, child in enumerate(tree):
+                        yield from flattened(child, f"{prefix}.{index}")
+                case _:
+                    assert_never(tree)
+
+        return dict(flattened(self._model.trainable_parameters(), ""))
+
+
+@dataclass(frozen=True, eq=False)
+class MlxValue(Value):
+    data: mx.array
+
+    def __add__(self, other: Self) -> Self:
+        return replace(self, data=self.data + other.data)
+
+    def __sub__(self, other: Self) -> Self:
+        return replace(self, data=self.data - other.data)
+
+    def __mul__(self, scalar: float) -> Self:
+        return replace(self, data=self.data * scalar)
+
+    def norm(self) -> float:
+        return float(mx.linalg.norm(mx.reshape(self.data, (-1,))))
+
+    def dot(self, other: Self) -> float:
+        return float(mx.sum(self.data * other.data))
+
+    def to_bytes(self, dtype: DType) -> TensorBytes:
+        return _tensor_bytes(self.data.astype(MLX_DTYPE[dtype]))
+
+    def quantize(self, dtype: Quantized) -> tuple[TensorBytes, TensorBytes]:
+        flat = mx.reshape(self.data, (-1,))
+        match dtype:
+            case "F8_E4M3":
+                largest = mx.max(mx.abs(flat)).reshape(1) if flat.size else mx.zeros((1,))
+                _, scale = _scales(largest, _E4M3_HEADROOM)
+                codes = mx.to_fp8(self.data / scale.reshape(()))
+                return (
+                    TensorBytes(dtype, self.data.shape, _bytes(codes)),
+                    TensorBytes("F32", (), _bytes(scale.reshape(()))),
+                )
+            case "F4":
+                blocks = -(-flat.size // F4_BLOCK)
+                blocked = mx.pad(flat, (0, blocks * F4_BLOCK - flat.size)).reshape(blocks, F4_BLOCK)
+                exponents, scales = _scales(mx.max(mx.abs(blocked), axis=1), _E2M1_HEADROOM)
+                q = mx.minimum(mx.abs(blocked) / scales[:, None], 6.0)
+                # Round to nearest even on the E2M1 grid 0, .5, 1, 1.5, 2, 3, 4, 6, whose codes
+                # are 0 to 7 in that order.
+                codes = mx.where(
+                    q < 2, mx.round(q * 2), mx.where(q < 4, mx.round(q) + 2, mx.round(q / 2) + 4)
+                ).astype(mx.uint8)
+                codes = codes | mx.where(blocked < 0, 8, 0).astype(mx.uint8)
+                pairs = codes.reshape(-1, 2)
+                packed = pairs[:, 0] | (pairs[:, 1] << 4)
+                return (
+                    TensorBytes(dtype, (blocks * F4_BLOCK,), _bytes(packed)),
+                    TensorBytes("F8_E8M0", (blocks,), _bytes((exponents + 127).astype(mx.uint8))),
+                )
+            case _:
+                raise ValueError(f"{dtype} is not a quantized dtype")
+
+
+_DTYPE_CODE: Mapping[mx.Dtype, DType] = {dtype: code for code, dtype in MLX_DTYPE.items()}
+_F4_VALUES = mx.array([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
+# Scales put the largest magnitude of a tensor or block in the binade just under the format's
+# top: below 2^8 for E4M3 (whose maximum is 448), in [4, 8) for E2M1 (whose maximum is 6).
+_E4M3_HEADROOM = 7
+_E2M1_HEADROOM = 2
+
+
+def _scales(largest: mx.array, headroom: int) -> tuple[mx.array, mx.array]:
+    """Per entry of ``largest``: the exponent ``k`` with ``largest < 2^(k + headroom + 1)``
+    (zero gets ``k = -headroom``, a scale of one), clamped to what E8M0 holds, and ``2^k``."""
+    exponents = (
+        mx.where(largest > 0, (largest.view(mx.uint32) >> 23).astype(mx.int32) - 127, headroom)
+        - headroom
+    )
+    exponents = mx.clip(exponents, -126, 127)
+    return exponents, _pow2(exponents)
+
+
+def _pow2(exponents: mx.array) -> mx.array:
+    """``2^k`` for int32 ``k`` in ``[-126, 127]``, built from the fp32 exponent bits."""
+    return ((exponents + 127).astype(mx.uint32) << 23).view(mx.float32)
+
+
+def _as_array(t: TensorBytes) -> mx.array:
+    # The only typed way from bytes to an array is MLX's own safetensors reader.
+    return _loaded(pack({"tensor": t}))
+
+
+def _as_uint8(data: bytes) -> mx.array:
+    """``data`` as a flat uint8 array, through a safetensors blob written here: the reader is the
+    typed way in, and ``U8`` is not a payload dtype ``pack`` knows."""
+    header = json.dumps(
+        {"tensor": {"dtype": "U8", "shape": [len(data)], "data_offsets": [0, len(data)]}}
+    ).encode()
+    return _loaded(len(header).to_bytes(8, "little") + header + data)
+
+
+def _loaded(blob: bytes) -> mx.array:
+    loaded = mx.load(io.BytesIO(blob), format="safetensors")
+    assert isinstance(loaded, dict)  # the stub's union; safetensors always decodes to a dict
+    return loaded["tensor"]
+
+
+def _bytes(array: mx.array) -> bytes:
+    # Flattening first makes a 0-d array viewable as uint8 and gives a transposed array its
+    # logical row-major bytes; the buffer protocol evaluates the array.
+    return memoryview(mx.reshape(array, (-1,)).view(mx.uint8)).tobytes()
+
+
+def _tensor_bytes(array: mx.array) -> TensorBytes:
+    return TensorBytes(_DTYPE_CODE[array.dtype], array.shape, _bytes(array))

gradientmail/adapters/torch.py

@@ -0,0 +1,205 @@
+"""PyTorch backend: ``nn.Module`` parameters in and out as bytes, and the outer step on tensors.
+
+``TorchAdapter.snapshot`` reads host memory through ``ctypes`` because ``memoryview(tensor)``
+is not supported. ``apply`` writes with ``Tensor.copy_`` under ``no_grad``: it moves the data
+to the parameter's device, writes through shared storage (tied weights) and never rebinds the
+``nn.Parameter``, so inner-optimizer state keyed by it survives. bf16, f16 and E4M3 are
+native, so ``TorchValue`` converts with ``Tensor.to``; E2M1 nibbles and E8M0 scales are
+packed and unpacked by hand.
+"""
+
+import ctypes
+import math
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass, replace
+from typing import Self
+
+import torch
+from torch import nn
+
+from gradientmail.adapters import Adapter, Backend, Value
+from gradientmail.serialization import (
+    F4_BLOCK,
+    DType,
+    ParameterSpec,
+    Quantized,
+    Schema,
+    SchemaMismatch,
+    TensorBytes,
+    validate,
+)
+
+__all__ = ["TORCH_DTYPE", "TorchAdapter", "TorchBackend", "TorchValue"]
+
+TORCH_DTYPE: Mapping[DType, torch.dtype] = {
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+}
+
+
+class TorchBackend(Backend):
+    """Adapts ``nn.Module``s; loads values on ``device`` (``"cpu"`` unless told otherwise)."""
+
+    def __init__(self, device: str = "cpu") -> None:
+        self._device = device
+
+    def adapter(self, model: object) -> Adapter:
+        if not isinstance(model, nn.Module):
+            raise TypeError(f"not a torch.nn.Module: {type(model).__qualname__}")
+        return TorchAdapter(model)
+
+    def load(self, tensor: TensorBytes) -> Value:
+        return TorchValue(_as_tensor(tensor).to(self._device, torch.float32))
+
+    def dequantize(self, data: TensorBytes, scale: TensorBytes, shape: tuple[int, ...]) -> Value:
+        match data.dtype:
+            case "F8_E4M3":
+                values = _raw(data).view(torch.float8_e4m3fn).to(torch.float32) * _as_tensor(scale)
+            case "F4":
+                codes = _raw(data)
+                nibbles = torch.stack((codes & 0xF, codes >> 4), dim=1).reshape(-1, F4_BLOCK)
+                magnitudes = _F4_VALUES[(nibbles & 7).long()]
+                signed = torch.where((nibbles & 8) != 0, -magnitudes, magnitudes)
+                scales = _pow2(_raw(scale).to(torch.int32) - 127)
+                values = (signed * scales[:, None]).reshape(-1)[: math.prod(shape)]
+            case other:
+                raise ValueError(f"{other} is not a quantized dtype")
+        return TorchValue(values.reshape(shape).to(self._device))
+
+
+class TorchAdapter(Adapter):
+    """Adapter over ``model.named_parameters()``, where a tied parameter appears once."""
+
+    def __init__(self, model: nn.Module) -> None:
+        self._model = model
+
+    def schema(self) -> Schema[DType]:
+        def spec(name: str) -> ParameterSpec[DType]:
+            p = parameters[name]
+            return ParameterSpec(name, _DTYPE_CODE[p.dtype], tuple(p.shape))
+
+        parameters = dict(self._model.named_parameters())
+        unsupported = tuple(name for name, p in parameters.items() if p.dtype not in _DTYPE_CODE)
+        if unsupported:
+            raise SchemaMismatch(unsupported, "parameter dtype is not float32, float16 or bfloat16")
+        return Schema(tuple(spec(name) for name in sorted(parameters)))
+
+    def snapshot(self, names: Iterable[str]) -> Mapping[str, TensorBytes]:
+        parameters = dict(self._model.named_parameters())
+        return {name: _tensor_bytes(parameters[name].detach()) for name in names}
+
+    def apply(self, tensors: Mapping[str, TensorBytes]) -> None:
+        # copy_ would cast a wrong dtype and broadcast a compatible wrong shape silently.
+        validate(self.schema(), tensors, tensors)
+        parameters = dict(self._model.named_parameters())
+        with torch.no_grad():
+            for name, t in tensors.items():
+                parameters[name].copy_(_as_tensor(t))
+
+
+@dataclass(frozen=True, eq=False)
+class TorchValue(Value):
+    data: torch.Tensor
+
+    def __add__(self, other: Self) -> Self:
+        return replace(self, data=self.data + other.data)
+
+    def __sub__(self, other: Self) -> Self:
+        return replace(self, data=self.data - other.data)
+
+    def __mul__(self, scalar: float) -> Self:
+        return replace(self, data=self.data * scalar)
+
+    def norm(self) -> float:
+        return float(self.data.square().sum().sqrt())
+
+    def dot(self, other: Self) -> float:
+        return float((self.data * other.data).sum())
+
+    def to_bytes(self, dtype: DType) -> TensorBytes:
+        return _tensor_bytes(self.data.to(TORCH_DTYPE[dtype]))
+
+    def quantize(self, dtype: Quantized) -> tuple[TensorBytes, TensorBytes]:
+        flat = self.data.reshape(-1)
+        match dtype:
+            case "F8_E4M3":
+                largest = flat.abs().max().reshape(1) if flat.numel() else torch.zeros(1)
+                _, scale = _scales(largest, _E4M3_HEADROOM)
+                codes = (self.data / scale.reshape(())).to(torch.float8_e4m3fn)
+                return (
+                    TensorBytes(dtype, tuple(self.data.shape), _bytes(codes)),
+                    TensorBytes("F32", (), _bytes(scale.reshape(()))),
+                )
+            case "F4":
+                blocks = -(-flat.numel() // F4_BLOCK)
+                padding = torch.zeros(blocks * F4_BLOCK - flat.numel(), dtype=flat.dtype)
+                blocked = torch.cat((flat, padding)).reshape(blocks, F4_BLOCK)
+                exponents, scales = _scales(blocked.abs().amax(dim=1), _E2M1_HEADROOM)
+                q = (blocked.abs() / scales[:, None]).clamp(max=6.0)
+                # Round to nearest even on the E2M1 grid 0, .5, 1, 1.5, 2, 3, 4, 6, whose codes
+                # are 0 to 7 in that order.
+                codes = torch.where(
+                    q < 2,
+                    torch.round(q * 2),
+                    torch.where(q < 4, torch.round(q) + 2, torch.round(q / 2) + 4),
+                ).to(torch.uint8)
+                codes = codes | torch.where(blocked < 0, 8, 0).to(torch.uint8)
+                pairs = codes.reshape(-1, 2)
+                packed = pairs[:, 0] | (pairs[:, 1] << 4)
+                return (
+                    TensorBytes(dtype, (blocks * F4_BLOCK,), _bytes(packed)),
+                    TensorBytes("F8_E8M0", (blocks,), _bytes((exponents + 127).to(torch.uint8))),
+                )
+            case _:
+                raise ValueError(f"{dtype} is not a quantized dtype")
+
+
+_DTYPE_CODE: Mapping[torch.dtype, DType] = {dtype: code for code, dtype in TORCH_DTYPE.items()}
+_F4_VALUES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
+# Scales put the largest magnitude of a tensor or block in the binade just under the format's
+# top: below 2^8 for E4M3 (whose maximum is 448), in [4, 8) for E2M1 (whose maximum is 6).
+_E4M3_HEADROOM = 7
+_E2M1_HEADROOM = 2
+
+
+def _scales(largest: torch.Tensor, headroom: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per entry of ``largest``: the exponent ``k`` with ``largest < 2^(k + headroom + 1)``
+    (zero gets ``k = -headroom``, a scale of one), clamped to what E8M0 holds, and ``2^k``."""
+    exponents = (
+        torch.where(largest > 0, (largest.view(torch.int32) >> 23) - 127, headroom) - headroom
+    )
+    exponents = exponents.clamp(-126, 127)
+    return exponents, _pow2(exponents)
+
+
+def _pow2(exponents: torch.Tensor) -> torch.Tensor:
+    """``2^k`` for int32 ``k`` in ``[-126, 127]``, built from the fp32 exponent bits."""
+    return ((exponents + 127) << 23).view(torch.float32)
+
+
+def _as_tensor(t: TensorBytes) -> torch.Tensor:
+    match t.dtype:
+        case "F32" | "F16" | "BF16" as dtype:
+            # frombuffer rejects an empty buffer and warns on a read-only one.
+            if not t.data:
+                return torch.empty(t.shape, dtype=TORCH_DTYPE[dtype])
+            return torch.frombuffer(bytearray(t.data), dtype=TORCH_DTYPE[dtype]).reshape(t.shape)
+        case other:
+            raise ValueError(f"{other} bytes need their scale: use dequantize")
+
+
+def _raw(t: TensorBytes) -> torch.Tensor:
+    """``t``'s bytes as a flat uint8 tensor."""
+    if not t.data:
+        return torch.empty(0, dtype=torch.uint8)
+    return torch.frombuffer(bytearray(t.data), dtype=torch.uint8)
+
+
+def _bytes(tensor: torch.Tensor) -> bytes:
+    t = tensor.to("cpu").contiguous()
+    return ctypes.string_at(t.data_ptr(), t.numel() * t.element_size())
+
+
+def _tensor_bytes(tensor: torch.Tensor) -> TensorBytes:
+    return TensorBytes(_DTYPE_CODE[tensor.dtype], tuple(tensor.shape), _bytes(tensor))