From 457473d39f665b78daea7db22962d47eea498741 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:54:13 -0700 Subject: [PATCH 1/4] Add Differential K-Means implementation --- changelog.d/+dkm.added | 1 + docs/src/examples/toy_models.md | 1 + docs/src/palettization/config.md | 111 +++ docs/src/palettization/overview.md | 40 + .../tutorials/mnist_dkm_palettization.ipynb | 798 ++++++++++++++++++ .../palettization/kmeans/__init__.py | 1 + .../palettization/kmeans/dkm/__init__.py | 18 + .../palettization/kmeans/dkm/_dispatch.py | 73 ++ .../palettization/kmeans/dkm/_forward.py | 85 ++ .../kmeans/dkm/training_strategy.py | 151 ++++ tests/palettization/test_dkm_config.py | 94 +++ tests/palettization/test_dkm_dispatch.py | 66 ++ tests/palettization/test_dkm_forward.py | 106 +++ .../test_dkm_training_strategy.py | 232 +++++ 14 files changed, 1777 insertions(+) create mode 100644 changelog.d/+dkm.added create mode 100644 docs/src/tutorials/mnist_dkm_palettization.ipynb create mode 100644 src/coreai_opt/palettization/kmeans/dkm/__init__.py create mode 100644 src/coreai_opt/palettization/kmeans/dkm/_dispatch.py create mode 100644 src/coreai_opt/palettization/kmeans/dkm/_forward.py create mode 100644 src/coreai_opt/palettization/kmeans/dkm/training_strategy.py create mode 100644 tests/palettization/test_dkm_config.py create mode 100644 tests/palettization/test_dkm_dispatch.py create mode 100644 tests/palettization/test_dkm_forward.py create mode 100644 tests/palettization/test_dkm_training_strategy.py diff --git a/changelog.d/+dkm.added b/changelog.d/+dkm.added new file mode 100644 index 00000000..5731c1fe --- /dev/null +++ b/changelog.d/+dkm.added @@ -0,0 +1 @@ +Added Differentiable K-Means (DKM) training-time palettization (`coreai_opt.palettization.kmeans.dkm.DKMTrainingSpec`). diff --git a/docs/src/examples/toy_models.md b/docs/src/examples/toy_models.md index eefc1097..1a9ee46b 100644 --- a/docs/src/examples/toy_models.md +++ b/docs/src/examples/toy_models.md @@ -7,6 +7,7 @@ These tutorials demonstrate model compression techniques using simple toy models /tutorials/mnist_quantization /tutorials/mnist_palettization +/tutorials/mnist_dkm_palettization /tutorials/mnist_palettization_and_activation_quantization /tutorials/mnist_pruning ``` diff --git a/docs/src/palettization/config.md b/docs/src/palettization/config.md index 42cecaa3..feb03338 100644 --- a/docs/src/palettization/config.md +++ b/docs/src/palettization/config.md @@ -346,3 +346,114 @@ kmeans_palettization_config: n_init: 1 max_iter: 1 ``` + +## Training strategy: DKM (`training_strategy_spec`) + +Each {class}`~coreai_opt.palettization.spec.PalettizationSpec` carries a `training_strategy_spec` field that selects the training-time behavior for that weight. It defaults to vanilla (frozen) k-means, where the codebook and index assignments computed at `prepare()` stay fixed during training. Setting it to a `DKMTrainingSpec` instead enables [Differentiable K-Means](overview.md#differentiable-k-means-dkm-api), so the codebook and assignments are learned through gradient-based fine-tuning. + +`DKMTrainingSpec` exposes the following settings: + +- `tau`: Softmax temperature for the soft centroid assignment (default `1e-4`). Smaller values sharpen toward hard k-means; larger values soften it. The loss-minimizing `tau` is model-dependent — **smaller models and lower bit-widths tend to prefer a larger `tau`**. A good starting point is `1e-5` for LLMs and the default `1e-4` for smaller models. +- `kmeans_max_iter`: Number of soft-assignment EM iterations per DKM training forward pass (default `3`). +- `enforce_zero`: If `True`, pins the centroid nearest the origin to exactly zero (default `False`). Useful for joint pruning and palettization. +- `zero_threshold`: Numerical-stability floor applied to the softmax assignment weights and per-centroid mass at small `tau` (default `1e-7`). + +### Apply DKM to all modules + +Set `training_strategy_spec=DKMTrainingSpec()` on the `global_config` spec to enable DKM for every palettized weight in the model: + +```python +from coreai_opt.palettization import ( + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, + PalettizationSpec, +) +from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec + +config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={ + "weight": PalettizationSpec( + n_bits=4, + training_strategy_spec=DKMTrainingSpec(), + ) + }, + ) +) +``` + +```yaml +# yaml +kmeans_palettization_config: + global_config: + op_state_spec: + weight: + n_bits: 4 + training_strategy_spec: + type: dkm +``` + +### Apply DKM to a subset of modules, vanilla k-means to the rest + +Because `training_strategy_spec` lives on the per-weight `PalettizationSpec`, DKM can be enabled for selected modules while the rest keep the default vanilla k-means. Below, all layers are palettized 4-bit; two named modules train their palettization with DKM, and every other module (covered by `global_config`) uses vanilla k-means simply by not setting a `training_strategy_spec`. + +```python +from coreai_opt.palettization import ( + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, + PalettizationSpec, +) +from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec + +config = KMeansPalettizerConfig( + # Default for all layers: vanilla 4-bit k-means (no training_strategy_spec). + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={"weight": PalettizationSpec(n_bits=4)}, + ), + module_name_configs={ + "encoder.layer.0.attention": ModuleKMeansPalettizerConfig( + op_state_spec={ + "weight": PalettizationSpec( + n_bits=4, + training_strategy_spec=DKMTrainingSpec(tau=1e-4, kmeans_max_iter=3), + ) + }, + ), + "encoder.layer.1.attention": ModuleKMeansPalettizerConfig( + op_state_spec={ + "weight": PalettizationSpec( + n_bits=4, + training_strategy_spec=DKMTrainingSpec(tau=5e-5), + ) + }, + ), + }, +) +``` + +```yaml +# yaml +kmeans_palettization_config: + global_config: + op_state_spec: + weight: + n_bits: 4 + module_name_configs: + encoder.layer.0.attention: + op_state_spec: + weight: + n_bits: 4 + training_strategy_spec: + type: dkm + tau: 1.0e-4 + kmeans_max_iter: 3 + encoder.layer.1.attention: + op_state_spec: + weight: + n_bits: 4 + training_strategy_spec: + type: dkm + tau: 5.0e-5 +``` + +For a DKM spec with the defaults, the `default_dkm_weight_palettization_spec()` helper (importable from `coreai_opt.palettization.kmeans`) returns a ready-made 4-bit, per-tensor `PalettizationSpec` with `training_strategy_spec=DKMTrainingSpec()`. diff --git a/docs/src/palettization/overview.md b/docs/src/palettization/overview.md index ce747975..4584fdd0 100644 --- a/docs/src/palettization/overview.md +++ b/docs/src/palettization/overview.md @@ -196,3 +196,43 @@ with palettizer.training_mode(): Each time `palettizer.step()` is called, the palettizer checks each module's `PATSchedule` to see whether any palettizer should be enabled. With the above config settings of `enable_fake_palettize` = 2 and 4 for `module1` and `module2` respectively, this results in both palettizers being disabled to start, `module1`'s palettizer enabled after `palettizer.step()` has been called twice, and `module2`'s palettizer enabled after `palettizer.step()` has been called 4 times in total. Note that `palettizer.step()` can be called whenever the user desires, whether it be per-batch, per-epoch, or on some other cadence. + +## Differentiable K-Means (DKM) API + +DKM palettization allows the palettization codebook and cluster assignments to be learned through training, removing the inability for already palettized weights to update their codebook and index assignments. This is achieved through calculation of attention scores between all weights and all centroids, along with utilizing a soft-assignment instead of hard-assignment during training forward, allowing gradients to propagate through palettizers to the original weight parameters. + +For more details on DKM, refer to [DKM (Cho et al., 2022)](https://arxiv.org/abs/2108.12659). + +Enabling DKM simply requires setting `training_strategy_spec` of `PalettizationSpec` to the `DKMTrainingSpec` class. The same `KMeansPalettizer` class is used for palettization. For example, the below code shows enabling DKM for all weights in the model: + +```python +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, + PalettizationSpec, +) +from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec + +kmeans_palettizer_config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + op_state_spec={ + "weight": PalettizationSpec( + ..., training_strategy_spec=DKMTrainingSpec(...) + ) + } + ) +) + +palettizer = KMeansPalettizer(model, kmeans_palettizer_config) +prepared_model = palettizer.prepare(example_inputs) + +with palettizer.training_mode(): + for i in range(...): + # Perform whatever training forward pass / loss calculation / optimizer gradient application is desired + ... + +finalized_model = palettizer.finalize(...) +``` + +As `training_strategy_spec` is part of `PalettizationSpec`, which in turn can be configured on a per-op or module basis in `KMeansPalettizerConfig`, users may choose to enable DKM for a subset of palettized weights in the model while leaving another portion of palettized weights using non-DKM `Vanilla K-means`. diff --git a/docs/src/tutorials/mnist_dkm_palettization.ipynb b/docs/src/tutorials/mnist_dkm_palettization.ipynb new file mode 100644 index 00000000..b32e65bb --- /dev/null +++ b/docs/src/tutorials/mnist_dkm_palettization.ipynb @@ -0,0 +1,798 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3b8a7a96", + "metadata": {}, + "source": [ + "# Applying DKM palettization to an MNIST model" + ] + }, + { + "cell_type": "markdown", + "id": "ca89f9d5", + "metadata": {}, + "source": [ + "This tutorial shows how to palettize an MNIST model with **Differentiable K-Means (DKM)**, which learns the palettization codebook and cluster assignments during training rather than fixing them from a single one-shot K-means fit.\n", + "\n", + "By the end of this tutorial you should be familiar with:\n", + "1. [How to configure DKM palettization](#DKM-Palettization)\n", + "2. [How to fine-tune a DKM-palettized model](#Fine-tuning-with-DKM)\n", + "3. [How to export the palettized model to Core AI](#Export-to-Core-AI)\n", + "\n", + "**Table of Contents:**\n", + "\n", + "- [Setup](#Setup)\n", + " - [MNIST Dataset download](#MNIST-Dataset-download)\n", + " - [Model definition](#Model-definition)\n", + " - [Training and Evaluation](#Training-and-Evaluation)\n", + " - [Train baseline model](#Train-baseline-model)\n", + "- [DKM Palettization](#DKM-Palettization)\n", + " - [Fine-tuning with DKM](#Fine-tuning-with-DKM)\n", + "- [Export to Core AI](#Export-to-Core-AI)" + ] + }, + { + "cell_type": "markdown", + "id": "4b4fcd7a", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "markdown", + "id": "e1545b8a", + "metadata": {}, + "source": [ + "We train a small CNN on MNIST to get a baseline, then palettize it to 4-bit with DKM and fine-tune so the codebooks are learned jointly with the model." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "6eb782d9", + "metadata": {}, + "outputs": [], + "source": [ + "import logging\n", + "import random\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import torch\n", + "from torch import nn\n", + "from torch.utils.data import DataLoader\n", + "from torchinfo import summary\n", + "from torchvision import datasets, transforms\n", + "from tqdm import tqdm\n", + "\n", + "logging.basicConfig(level=logging.INFO)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8e7abd23", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Using device cpu\n" + ] + } + ], + "source": [ + "# Use CUDA when available; otherwise fall back to CPU.\n", + "if torch.cuda.is_available():\n", + " DEVICE = torch.device(\"cuda\")\n", + "else:\n", + " DEVICE = torch.device(\"cpu\")\n", + "\n", + "print(f\"Using device {DEVICE}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "8a9cfac5", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "SEED = 1976\n", + "\n", + "random.seed(SEED)\n", + "np.random.seed(SEED)\n", + "torch.manual_seed(SEED)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "97f74827", + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [], + "source": [ + "# Used to save intermediate results and datasets\n", + "SAVE_DIRECTORY = \".\"" + ] + }, + { + "cell_type": "markdown", + "id": "c08e532d", + "metadata": {}, + "source": [ + "### MNIST Dataset download\n", + "\n", + "Helper to download the MNIST dataset with standard normalization applied." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "6201e35e", + "metadata": {}, + "outputs": [], + "source": [ + "def mnist_transforms() -> transforms.Compose:\n", + " return transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n", + "\n", + "\n", + "def download_mnist_dataset(\n", + " download_path: Path, transform: transforms.Compose | None = None\n", + ") -> tuple[datasets.MNIST, datasets.MNIST]:\n", + " if transform is None:\n", + " transform = mnist_transforms()\n", + " train = datasets.MNIST(download_path, train=True, download=True, transform=transform)\n", + " test = datasets.MNIST(download_path, train=False, download=True, transform=transform)\n", + " return train, test" + ] + }, + { + "cell_type": "markdown", + "id": "e779d29e", + "metadata": {}, + "source": [ + "### Model definition\n", + "\n", + "A simple CNN with a single Conv2d → ReLU → MaxPool block, followed by Flatten and a Linear classifier." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "9c4f555f", + "metadata": {}, + "outputs": [], + "source": [ + "class MnistNetwork(nn.Module):\n", + " def __init__(self, num_classes: int = 10, state_dict: dict | None = None) -> None:\n", + " super().__init__()\n", + " self.model = nn.Sequential(\n", + " nn.Conv2d(1, 12, 3, padding=1),\n", + " nn.ReLU(),\n", + " nn.MaxPool2d(2, stride=2, padding=0),\n", + " nn.Flatten(),\n", + " nn.Linear(2352, num_classes),\n", + " )\n", + " if state_dict is not None:\n", + " self.load_state_dict(state_dict)\n", + "\n", + " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", + " return self.model(x)" + ] + }, + { + "cell_type": "markdown", + "id": "98b11457", + "metadata": {}, + "source": [ + "### Training and Evaluation\n", + "\n", + "Standard PyTorch training loop and evaluation function that computes accuracy." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "9bd3e929", + "metadata": {}, + "outputs": [], + "source": [ + "def train_step(model, optimizer, loss_fn, inputs, ground_truth) -> float:\n", + " model.train()\n", + " device = next(model.parameters()).device\n", + " inputs = inputs.to(device)\n", + " ground_truth = ground_truth.to(device)\n", + " optimizer.zero_grad()\n", + " predictions = model(inputs)\n", + " loss = loss_fn(predictions, ground_truth)\n", + " loss.backward()\n", + " optimizer.step()\n", + " return loss.item()\n", + "\n", + "\n", + "def train_epoch(model, train_loader, optimizer, loss_fn) -> float:\n", + " total_loss = 0.0\n", + " for inputs, ground_truth in train_loader:\n", + " loss = train_step(model, optimizer, loss_fn, inputs, ground_truth)\n", + " total_loss += loss\n", + " return total_loss / len(train_loader)\n", + "\n", + "\n", + "def create_adam_optimizer(model: nn.Module, lr: float = 1e-3) -> torch.optim.Adam:\n", + " return torch.optim.Adam(model.parameters(), lr=lr)\n", + "\n", + "\n", + "def eval_model(model: nn.Module, test_dataloader: DataLoader) -> float:\n", + " model.eval()\n", + " device = next(model.parameters()).device\n", + " num_correct = 0\n", + " total = 0\n", + " with torch.no_grad():\n", + " for inputs, ground_truth in test_dataloader:\n", + " inputs = inputs.to(device)\n", + " ground_truth = ground_truth.to(device)\n", + " predictions = model(inputs)\n", + " _, predicted = torch.max(predictions.data, 1)\n", + " total += ground_truth.size(0)\n", + " num_correct += (predicted == ground_truth).sum().item()\n", + " return num_correct / total" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "fe7cefb8", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9.91M/9.91M [00:00<00:00, 13.9MB/s]\n", + "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 28.9k/28.9k [00:00<00:00, 441kB/s]\n", + "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.65M/1.65M [00:00<00:00, 4.67MB/s]\n", + "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4.54k/4.54k [00:00<00:00, 1.83MB/s]\n" + ] + } + ], + "source": [ + "# Download and instantiate datasets\n", + "DOWNLOAD_PATH = Path(SAVE_DIRECTORY) / \".mnist_dataset\"\n", + "\n", + "train_dataset, test_dataset = download_mnist_dataset(download_path=DOWNLOAD_PATH)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "69c18e59", + "metadata": {}, + "outputs": [], + "source": [ + "BATCH_SIZE = 128\n", + "\n", + "# Instantiate dataloaders\n", + "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n", + "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE)" + ] + }, + { + "cell_type": "markdown", + "id": "ab5805ae", + "metadata": {}, + "source": [ + "The CNN model used for this tutorial contains a single Conv2d, ReLU, MaxPool, Flatten, and Linear layer. Here's the structure:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "d6c7f0e6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "==========================================================================================\n", + "Layer (type:depth-idx) Output Shape Param #\n", + "==========================================================================================\n", + "MnistNetwork [1, 10] --\n", + "├─Sequential: 1-1 [1, 10] --\n", + "│ └─Conv2d: 2-1 [1, 12, 28, 28] 120\n", + "│ └─ReLU: 2-2 [1, 12, 28, 28] --\n", + "│ └─MaxPool2d: 2-3 [1, 12, 14, 14] --\n", + "│ └─Flatten: 2-4 [1, 2352] --\n", + "│ └─Linear: 2-5 [1, 10] 23,530\n", + "==========================================================================================\n", + "Total params: 23,650\n", + "Trainable params: 23,650\n", + "Non-trainable params: 0\n", + "Total mult-adds (Units.MEGABYTES): 0.12\n", + "==========================================================================================\n", + "Input size (MB): 0.00\n", + "Forward/backward pass size (MB): 0.08\n", + "Params size (MB): 0.09\n", + "Estimated Total Size (MB): 0.17\n", + "==========================================================================================" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "basic_cnn_model = MnistNetwork(num_classes=10)\n", + "\n", + "# Print summary of model\n", + "summary(basic_cnn_model, input_size=(1, 1, 28, 28))" + ] + }, + { + "cell_type": "markdown", + "id": "fc95cd52", + "metadata": {}, + "source": [ + "### Train baseline model" + ] + }, + { + "cell_type": "markdown", + "id": "0ba457d2", + "metadata": {}, + "source": [ + "Let's train this model so we can get a baseline accuracy. We save the trained weights so we can reload them for the palettization experiment." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "3343708f", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [01:00<00:00, 6.05s/it]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " Epoch 1: loss=0.3126\n", + " Epoch 2: loss=0.1170\n", + " Epoch 3: loss=0.0821\n", + " Epoch 4: loss=0.0669\n", + " Epoch 5: loss=0.0590\n", + " Epoch 6: loss=0.0522\n", + " Epoch 7: loss=0.0481\n", + " Epoch 8: loss=0.0452\n", + " Epoch 9: loss=0.0415\n", + " Epoch 10: loss=0.0379\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "EPOCHS = 10\n", + "\n", + "loss_fn = torch.nn.CrossEntropyLoss()\n", + "optimizer = create_adam_optimizer(basic_cnn_model)\n", + "\n", + "basic_cnn_model = basic_cnn_model.to(DEVICE)\n", + "\n", + "epoch_results = []\n", + "for epoch in tqdm(range(EPOCHS)):\n", + " epoch_avg_loss = train_epoch(\n", + " model=basic_cnn_model, train_loader=train_loader, optimizer=optimizer, loss_fn=loss_fn\n", + " )\n", + " epoch_results.append(f\" Epoch {epoch + 1}: loss={epoch_avg_loss:.4f}\")\n", + "\n", + "basic_cnn_model = basic_cnn_model.cpu().eval()\n", + "print(\"\\n\".join(epoch_results))\n", + "\n", + "# Save trained weights for reuse in the palettization section\n", + "pretrained_state_dict = basic_cnn_model.state_dict()" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "2dcb4fab", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Baseline accuracy: 0.9798\n" + ] + } + ], + "source": [ + "accuracy = eval_model(basic_cnn_model, test_loader)\n", + "print(f\"Baseline accuracy: {accuracy:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "3992c86e", + "metadata": {}, + "source": [ + "## DKM Palettization" + ] + }, + { + "cell_type": "markdown", + "id": "6fe61ea9", + "metadata": {}, + "source": [ + "Palettization compresses a model's weights by clustering them into a small look-up table (LUT) of centroids, and replacing each weight with an index into that table.\n", + "\n", + "Plain K-means palettization computes the LUT once, post-training, and freezes it. DKM instead makes the clustering step differentiable: during training it uses a soft assignment of weights to centroids, so gradients flow back into the palettized weights and their codebooks. This lets the codebook and assignments be learned jointly with the rest of the model, which can recover accuracy that a one-shot fit leaves behind.\n", + "\n", + "DKM is enabled by setting the `training_strategy_spec` field of a `PalettizationSpec` to a `DKMTrainingSpec`. The same `KMeansPalettizer` drives it." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "707d627e", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "WARNING:torchao.kernel.intmm:Warning: Detected no triton, on systems without Triton certain kernels will not work\n", + "W1030 11:30:48.699000 8891 torch/distributed/elastic/multiprocessing/redirects.py:35] NOTE: Redirects are currently not supported in MacOs.\n", + "W1030 11:30:48.728000 8891 torch/utils/_pytree.py:630] is an Enum subclass and is now natively supported by torch.compile as an opaque value type. Calling register_constant() on Enum subclasses is deprecated and will be an error in a future release.\n" + ] + } + ], + "source": [ + "from coreai_opt.palettization import (\n", + " KMeansPalettizer,\n", + " KMeansPalettizerConfig,\n", + " ModuleKMeansPalettizerConfig,\n", + " PalettizationSpec,\n", + ")\n", + "from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec\n", + "from coreai_opt.palettization.spec import PerGroupedChannelGranularity\n", + "\n", + "example_inputs = (torch.randn(1, 1, 28, 28),)" + ] + }, + { + "cell_type": "markdown", + "id": "16956cd8", + "metadata": {}, + "source": [ + "We use 4-bit **per-grouped-channel** palettization (each group of 2 output channels gets its own LUT of 2⁴ = 16 centroids), and attach a `DKMTrainingSpec` so the palettization is trained with DKM.\n", + "\n", + "`DKMTrainingSpec` exposes a few hyperparameters (all optional): `tau` (softmax temperature for the soft assignment), `kmeans_max_iter` (soft-assignment iterations per forward), `enforce_zero`, and `zero_threshold`. Here we keep the defaults." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "7563c363", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO:coreai_opt.palettization.kmeans.palettizer:Preparing model for palettization\n", + "Palettizing layers (num_workers=1): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00, 1.16it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Prepared DKM palettization with n_bits=4\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "dkm_model = MnistNetwork(num_classes=10, state_dict=pretrained_state_dict)\n", + "dkm_model.eval()\n", + "\n", + "weight_spec = PalettizationSpec(\n", + " n_bits=4,\n", + " granularity=PerGroupedChannelGranularity(axis=0, group_size=2),\n", + " training_strategy_spec=DKMTrainingSpec(),\n", + ")\n", + "\n", + "dkm_config = KMeansPalettizerConfig(\n", + " global_config=ModuleKMeansPalettizerConfig(\n", + " op_state_spec={\"weight\": weight_spec},\n", + " )\n", + ")\n", + "\n", + "dkm_palettizer = KMeansPalettizer(dkm_model, dkm_config)\n", + "dkm_prepared = dkm_palettizer.prepare(example_inputs)\n", + "print(\"Prepared DKM palettization with n_bits=4\")\n", + "\n", + "# Move the prepared model onto the target device so DKM's training-time forward runs there.\n", + "# prepare() itself runs on CPU: k-means centroid initialization uses a CPU extension.\n", + "dkm_prepared = dkm_prepared.to(DEVICE)" + ] + }, + { + "cell_type": "markdown", + "id": "f7c65646", + "metadata": {}, + "source": [ + "After `prepare()`, the model has fake-palettization modules whose codebooks are initialized from one-shot K-means. Measure the accuracy of the palettized model before DKM fine-tuning:" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "1adaeadf", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DKM (before fine-tuning) accuracy: 0.9789\n" + ] + } + ], + "source": [ + "dkm_initial_accuracy = eval_model(dkm_prepared, test_loader)\n", + "print(f\"DKM (before fine-tuning) accuracy: {dkm_initial_accuracy:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "efb0c387", + "metadata": {}, + "source": [ + "### Fine-tuning with DKM" + ] + }, + { + "cell_type": "markdown", + "id": "7aad8b57", + "metadata": {}, + "source": [ + "We now fine-tune inside `dkm_palettizer.training_mode()`. Within this context each forward pass runs DKM's differentiable soft assignment, so gradients flow into the palettized weights and their codebooks are updated by the optimizer.\n", + "\n", + "`training_mode()` also restores the model's original train/eval state on exit, so we can evaluate immediately afterwards." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "57b17b18", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + " 0%| | 0/469 [00:00 0.0175\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "import torch.nn.functional as F\n", + "\n", + "DKM_FINETUNE_STEPS = 200\n", + "optimizer = create_adam_optimizer(dkm_prepared, lr=1e-4)\n", + "\n", + "losses = []\n", + "step = 0\n", + "with dkm_palettizer.training_mode():\n", + " while step < DKM_FINETUNE_STEPS:\n", + " for inputs, targets in tqdm(train_loader):\n", + " inputs = inputs.to(DEVICE)\n", + " targets = targets.to(DEVICE)\n", + " optimizer.zero_grad()\n", + " loss = F.cross_entropy(dkm_prepared(inputs), targets)\n", + " loss.backward()\n", + " optimizer.step()\n", + " losses.append(loss.item())\n", + " step += 1\n", + " if step >= DKM_FINETUNE_STEPS:\n", + " break\n", + "\n", + "print(f\"Fine-tuned DKM for {step} steps; loss {losses[0]:.4f} -> {losses[-1]:.4f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "df0af64d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DKM (after fine-tuning) accuracy: 0.9827\n" + ] + } + ], + "source": [ + "dkm_accuracy = eval_model(dkm_prepared, test_loader)\n", + "print(f\"DKM (after fine-tuning) accuracy: {dkm_accuracy:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "38d95fe8", + "metadata": {}, + "source": [ + "## Export to Core AI" + ] + }, + { + "cell_type": "markdown", + "id": "2e943aff", + "metadata": {}, + "source": [ + "Once the palettized model is ready, call `finalize()` to convert the fake-palettization modules into the real LUT-based representation for deployment.\n", + "\n", + "Pass `ExportBackend.CoreAI` to `finalize(backend=...)` to target the `.aimodel` format produced by `coreai-torch`." + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "8f0c0e9b", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO:coreai_opt.palettization.kmeans.palettizer:Finalizing model for backend: coreai\n", + "/Volumes/Data/coreai-opt/external/src/coreai_opt/palettization/kmeans/_prepare_for_export.py:302: UserWarning: coreai-torch has only been validated with torch<=2.13.0; found torch 2.14.0. Some functionality may not work as expected.\n", + " from coreai_torch._compression.custom_layers import ( # noqa: PLC0415\n" + ] + } + ], + "source": [ + "from coreai_opt import ExportBackend\n", + "\n", + "# Export (torch.export + Core AI conversion) runs on CPU and traces with the CPU\n", + "# example_inputs, so move the model back off the GPU first.\n", + "dkm_prepared = dkm_prepared.to(\"cpu\")\n", + "coreai_model = dkm_palettizer.finalize(backend=ExportBackend.CoreAI)" + ] + }, + { + "cell_type": "markdown", + "id": "917d5ed5", + "metadata": {}, + "source": [ + "The export proceeds in three steps:\n", + "\n", + "- Trace the model with `torch.export.export()` to obtain a graph representation.\n", + "- Apply `cast_to_16_bit_precision()` to cast remaining FP32 parameters to FP16 for optimal on-device performance.\n", + "- Convert the exported program to Core AI format using `coreai-torch.TorchConverter`." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "9da6672e", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
coreai-torch 0.0.0.dev202610011705+745fc1b: converting 1 program(s) to Core AI\n",
+       "
\n" + ], + "text/plain": [ + "\u001b[1;36mcoreai-torch\u001b[0m \u001b[1;2;36m0.0\u001b[0m\u001b[2m.\u001b[0m\u001b[1;2;36m0.\u001b[0m\u001b[2mdev202610011705+745fc1b\u001b[0m: converting \u001b[1;36m1\u001b[0m \u001b[1;35mprogram\u001b[0m\u001b[1m(\u001b[0ms\u001b[1m)\u001b[0m to Core AI\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Exported: exported_model.aimodel\n" + ] + } + ], + "source": [ + "import shutil\n", + "\n", + "from coreai_opt.casting import cast_to_16_bit_precision\n", + "from coreai_torch import TorchConverter, get_decomp_table\n", + "\n", + "exported_program = torch.export.export(coreai_model, example_inputs, strict=False)\n", + "exported_program = exported_program.run_decompositions(get_decomp_table())\n", + "cast_to_16_bit_precision(exported_program)\n", + "\n", + "coreai_program = TorchConverter().add_exported_program(exported_program).to_coreai()\n", + "\n", + "output_path = Path(SAVE_DIRECTORY) / \"exported_model.aimodel\"\n", + "if output_path.exists():\n", + " shutil.rmtree(output_path)\n", + "coreai_program.save_asset(output_path)\n", + "print(f\"Exported: {output_path}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/coreai_opt/palettization/kmeans/__init__.py b/src/coreai_opt/palettization/kmeans/__init__.py index e1926163..182ad64f 100644 --- a/src/coreai_opt/palettization/kmeans/__init__.py +++ b/src/coreai_opt/palettization/kmeans/__init__.py @@ -5,6 +5,7 @@ """K-means palettization compressor.""" +from . import dkm # Registers the "dkm" training strategy. from .palettizer import KMeansPalettizer __all__ = [ diff --git a/src/coreai_opt/palettization/kmeans/dkm/__init__.py b/src/coreai_opt/palettization/kmeans/dkm/__init__.py new file mode 100644 index 00000000..e81e5252 --- /dev/null +++ b/src/coreai_opt/palettization/kmeans/dkm/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""DKM (Differentiable K-Means) training-time palettization.""" + +from .training_strategy import ( + SUPPORTED_CLUSTER_DIMS, + DKMTrainingSpec, + default_dkm_weight_palettization_spec, +) + +__all__ = [ + "SUPPORTED_CLUSTER_DIMS", + "DKMTrainingSpec", + "default_dkm_weight_palettization_spec", +] diff --git a/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py new file mode 100644 index 00000000..fe7bc19c --- /dev/null +++ b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py @@ -0,0 +1,73 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Resolve the DKM forward function for a device. + +Providers registered under the ``coreai_opt.dkm_providers`` entry-point group may supply +the forward. +""" + +import logging +from collections.abc import Callable +from functools import cache +from importlib.metadata import entry_points +from typing import NamedTuple, Protocol + +import torch + +__all__ = ["Forward", "Provider", "get_forward"] + +logger = logging.getLogger(__name__) + +ENTRY_POINT_GROUP = "coreai_opt.dkm_providers" + + +class Forward(Protocol): + """Signature of a DKM forward pass, as called by the DKM training strategy.""" + + def __call__( + self, + weights: torch.Tensor, + centroids: torch.Tensor, + tau: float, + kmeans_max_iter: int, + zero_threshold: float, + enforce_zero: bool, + quantize_fn: Callable[[torch.Tensor], torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(reconstructed_weights, centroids)`` for ``weights`` of shape ``(P, N, D)``.""" + ... + + +class Provider(NamedTuple): + """A named provider; ``load(device)`` returns its forward, or ``None`` if it does not apply.""" + + name: str + load: Callable[[torch.device], Forward | None] + + +@cache +def _providers() -> tuple[Provider, ...]: + providers: list[Provider] = [] + for entry_point in entry_points(group=ENTRY_POINT_GROUP): + try: + providers = entry_point.load()(providers) + except Exception: + logger.warning( + "DKM: provider %s failed to load; skipping.", entry_point.name, exc_info=True + ) + return tuple(providers) + + +@cache +def get_forward(device: torch.device) -> Forward | None: + """Return the forward supplied by a provider for ``device``, or ``None``.""" + for provider in _providers(): + forward = provider.load(device) + if forward is not None: + logger.info("DKM: using the %s provider on %s", provider.name, device) + return forward + logger.info("DKM: using the PyTorch implementation on %s", device) + return None diff --git a/src/coreai_opt/palettization/kmeans/dkm/_forward.py b/src/coreai_opt/palettization/kmeans/dkm/_forward.py new file mode 100644 index 00000000..e8e27777 --- /dev/null +++ b/src/coreai_opt/palettization/kmeans/dkm/_forward.py @@ -0,0 +1,85 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""PyTorch implementation of DKM.""" + +from collections.abc import Callable + +import torch +import torch.nn.functional as F + + +def _pin_zero_centroid(centroids: torch.Tensor) -> torch.Tensor: + """Pin each partition's nearest-origin centroid to exactly zero, keeping grad on the rest.""" + centroids = centroids.clone() + p_idx = torch.arange(centroids.shape[0], device=centroids.device) + zero_idx = centroids.detach().norm(dim=-1).argmin(dim=1) # (P,) nearest-origin code + centroids[p_idx, zero_idx, :] = centroids[p_idx, zero_idx, :].detach().zero_() + return centroids + + +def dkm_forward( + weights: torch.Tensor, + centroids: torch.Tensor, + tau: float, + kmeans_max_iter: int, + zero_threshold: float = 1e-7, + enforce_zero: bool = False, + quantize_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DKM soft k-means and return the reconstructed weights and updated centroids. + + Args: + weights: Weight tensor of shape ``(P, N, D)``. + centroids: Initial centroid tensor of shape ``(P, K, D)``. + tau: Softmax temperature. + kmeans_max_iter: Number of soft-assignment iterations. + zero_threshold: Floor applied to the softmax assignment weights and the + per-centroid mass (denominator) for numerical stability at small + ``tau``. + enforce_zero: If True, pins each partition's nearest-origin centroid to + exactly zero after every update step, keeping an exact-zero codebook + entry for joint pruning + palettization. + quantize_fn: If provided, fake-quantizes centroids after every update + step (after ``enforce_zero``'s pin, if both are set), keeping the + EM loop quantization-aware at every iteration rather than only at + the initial seed. + + Returns: + ``(reconstructed_weights, centroids_new)`` — both of shape matching ``weights`` + and ``centroids`` respectively. ``reconstructed_weights`` is differentiable + w.r.t. the input ``weights`` and ``centroids``. + """ + # Work in float32 for numerical stability regardless of input dtype + weights_f = weights.float() + centroids_f = centroids.float() + + attn = None + for _ in range(kmeans_max_iter): + # (P, N, K) — squared L2 distance from each weight to each centroid + dist = (weights_f.unsqueeze(-2) - centroids_f.unsqueeze(-3)).pow(2).sum(-1) + + # Soft assignment: (P, N, K) + attn = F.softmax(-dist / tau, dim=-1).clamp(min=zero_threshold) + + # Weighted centroid update: (P, K, D) + # attn.mT is (P, K, N), weights_f is (P, N, D) → matmul gives (P, K, D) + attn_sum = attn.sum(dim=1, keepdim=True).mT.clamp(min=zero_threshold) + centroids_new = attn.mT @ weights_f / attn_sum + + # Pin each partition's nearest-origin centroid to exactly zero so the + # codebook keeps an exact-zero entry (block-sparse code) that never drifts. + if enforce_zero: + centroids_new = _pin_zero_centroid(centroids_new) + + if quantize_fn is not None: + centroids_new = quantize_fn(centroids_new) + + centroids_f = centroids_new + + # Reconstructed weights: (P, N, D) + reconstructed_weights = attn @ centroids_f + + return reconstructed_weights.to(weights.dtype), centroids_f.to(centroids.dtype) diff --git a/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py b/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py new file mode 100644 index 00000000..84be7c91 --- /dev/null +++ b/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py @@ -0,0 +1,151 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""DKM (Differentiable K-Means) training strategy and its specification.""" + +from functools import partial + +import torch +from pydantic import Field, PositiveInt + +from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize +from coreai_opt.palettization.spec import PalettizationSpec, PerTensorGranularity +from coreai_opt.palettization.spec.training_strategy import TrainingStrategy, TrainingStrategySpec + +from ._dispatch import get_forward as _get_forward +from ._forward import dkm_forward as _dkm_forward + +# Cluster dim values which are supported by DKM +SUPPORTED_CLUSTER_DIMS = {1, 2, 4} + + +def _quantize_centroids(module: _KMeansFakePalettize, centroids: torch.Tensor) -> torch.Tensor: + """Bridge DKM's ``(P, K, D)`` centroid shape to ``module.quantize_lut()``'s + ``(num_blocks, num_clusters[, cluster_dim])`` convention, which drops the + trailing ``D`` axis when ``cluster_dim == 1`` (same as ``module._raw_lut``). + """ + if module.cluster_dim == 1: + # (P, K, 1) -> (P, K) for quantize_lut, then restore (P, K, 1) + return module.quantize_lut(centroids.squeeze(-1)).unsqueeze(-1) + return module.quantize_lut(centroids) + + +class _DKMTrainingStrategy(TrainingStrategy): + """Differentiable soft k-means training step. Mutates ``module.centroids`` + in place every training step via the EM loop. + + For details on DKM, refer to: + Cho et al., "DKM: Differentiable K-Means Clustering Layer for Neural Network Compression", + ICLR 2022. arXiv:2108.12659 + """ + + def __init__( + self, + tau: float = 1e-4, + kmeans_max_iter: int = 3, + enforce_zero: bool = False, + zero_threshold: float = 1e-7, + ): + self.tau = tau + self.kmeans_max_iter = kmeans_max_iter + self.enforce_zero = enforce_zero + self.zero_threshold = zero_threshold + + def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> torch.Tensor: + """Run one DKM EM step and reconstruct the weight from the updated centroids.""" + if not isinstance(module, _KMeansFakePalettize): + raise TypeError( + "DKM training strategy only supports k-means palettization " + f"(_KMeansFakePalettize), got {type(module).__name__}." + ) + if module.cluster_dim not in SUPPORTED_CLUSTER_DIMS: + raise ValueError( + f"DKM supports cluster_dim in {SUPPORTED_CLUSTER_DIMS}, got {module.cluster_dim}." + ) + + # Shapes: P=#blocks, N=vectors/block, D=cluster_dim, K=#clusters. + + # Scale (if enabled), reshape, block-split, and vectorize -> (P, N, D). + weights, vectorization = module.vectorize(weight) + + quantize_fn = partial(_quantize_centroids, module) + + # Clone is needed to ensure "centroids" and "module.centroids" do not share the same storage + # (occurs in the case of no LUT quantization). + # DKM requires gradients to flow through "centroids", however if they share the same + # storage, module.centroids.copy_() further below will invalidate gradients. + centroids = quantize_fn(module.centroids).clone() # (P, K, D) + + assert centroids.shape[0] == weights.shape[0], ( + f"partition count mismatch: centroids P={centroids.shape[0]} " + f"vs weight P={weights.shape[0]}" + ) + + forward = _get_forward(weight.device) + + if forward is not None: + reconstructed_weights, centroids_new = forward( + weights, + centroids, + self.tau, + kmeans_max_iter=self.kmeans_max_iter, + zero_threshold=self.zero_threshold, + enforce_zero=self.enforce_zero, + quantize_fn=quantize_fn, + ) + else: + reconstructed_weights, centroids_new = _dkm_forward( + weights, + centroids, + self.tau, + self.kmeans_max_iter, + zero_threshold=self.zero_threshold, + enforce_zero=self.enforce_zero, + quantize_fn=quantize_fn, + ) + + module.centroids.copy_(centroids_new.detach()) + module._indices_stale = True + + # reconstructed_weights is (P, N, D): undo vectorize + block-split (+ unscale) + # to rebuild the weight. + return module.devectorize(reconstructed_weights, vectorization) + + +@TrainingStrategySpec.register("dkm") +class DKMTrainingSpec(TrainingStrategySpec): + """DKM (Differentiable K-Means) training-strategy hyperparameters. + + Attributes: + tau: Softmax temperature for DKM's soft centroid assignment. Smaller + values sharpen toward hard k-means; larger values soften it. + Typical range: 1e-5 to 1e-3. + kmeans_max_iter: Number of soft-assignment EM iterations per DKM + training forward pass. + enforce_zero: If True, pins the centroid nearest the origin to + exactly zero. Useful for joint pruning + palettization. + zero_threshold: Floor applied to the softmax assignment weights and the + per-centroid mass (denominator) for numerical stability at small + ``tau``. + """ + + _strategy_cls = _DKMTrainingStrategy + + tau: float = Field(default=1e-4, gt=0) + kmeans_max_iter: PositiveInt = 3 + enforce_zero: bool = False + zero_threshold: float = Field(default=1e-7, gt=0) + + +def default_dkm_weight_palettization_spec() -> PalettizationSpec: + """Return the default DKM weight-palettization spec (4-bit, per-tensor, no LUT quantization).""" + return PalettizationSpec( + n_bits=4, + lut_qspec=None, + granularity=PerTensorGranularity(), + cluster_dim=1, + enable_per_channel_scale=False, + training_strategy_spec=DKMTrainingSpec(), + ) diff --git a/tests/palettization/test_dkm_config.py b/tests/palettization/test_dkm_config.py new file mode 100644 index 00000000..3153fe72 --- /dev/null +++ b/tests/palettization/test_dkm_config.py @@ -0,0 +1,94 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""CPU tests for DKMTrainingConfig: defaults, validation, registration, and build_strategy().""" + +import subprocess +import sys + +import pytest +from pydantic import ValidationError + +from coreai_opt.palettization.kmeans.dkm.training_strategy import ( + DKMTrainingSpec, + _DKMTrainingStrategy, + default_dkm_weight_palettization_spec, +) +from coreai_opt.palettization.spec import PerTensorGranularity +from coreai_opt.palettization.spec.training_strategy import TrainingStrategySpec + + +class TestDKMTrainingConfig: + def test_defaults(self): + c = DKMTrainingSpec() + assert c.tau == 1e-4 + assert c.kmeans_max_iter == 3 + assert c.enforce_zero is False + assert c.zero_threshold == 1e-7 + + @pytest.mark.parametrize( + "kwargs", + [ + {"tau": 0.0}, + {"tau": -1.0}, + {"zero_threshold": 0.0}, + {"zero_threshold": -1e-7}, + {"kmeans_max_iter": 0}, + {"kmeans_max_iter": -1}, + ], + ) + def test_invalid_values_rejected(self, kwargs): + with pytest.raises(ValidationError): + DKMTrainingSpec(**kwargs) + + def test_unknown_field_rejected(self): + with pytest.raises(ValidationError): + DKMTrainingSpec(epsilon=1e-4) + + def test_frozen(self): + c = DKMTrainingSpec() + with pytest.raises(ValidationError): + c.tau = 1e-3 + + def test_registered_under_dkm(self): + assert "dkm" in TrainingStrategySpec.REGISTRY + rebuilt = TrainingStrategySpec.maybe_build_from_dict( + {"type": "dkm", "tau": 2e-4, "kmeans_max_iter": 5} + ) + assert isinstance(rebuilt, DKMTrainingSpec) + assert rebuilt.tau == 2e-4 + assert rebuilt.kmeans_max_iter == 5 + + def test_registered_on_palettization_import_alone(self): + """A fresh process that imports only ``coreai_opt.palettization`` can rebuild a DKM spec.""" + code = ( + "import coreai_opt.palettization\n" + "from coreai_opt.palettization.spec.training_strategy import TrainingStrategySpec\n" + "spec = TrainingStrategySpec.maybe_build_from_dict({'type': 'dkm', 'tau': 2e-4})\n" + "assert type(spec).__name__ == 'DKMTrainingSpec' and spec.tau == 2e-4\n" + ) + subprocess.run([sys.executable, "-c", code], check=True) + + def test_serialization_round_trip(self): + c = DKMTrainingSpec(tau=3e-4, enforce_zero=True, zero_threshold=1e-6) + data = c.model_dump() + assert data["type"] == "dkm" + assert TrainingStrategySpec.maybe_build_from_dict(data) == c + + def test_build_strategy_threads_fields(self): + c = DKMTrainingSpec(tau=2e-4, kmeans_max_iter=7, enforce_zero=True, zero_threshold=1e-6) + s = c.build_strategy() + assert isinstance(s, _DKMTrainingStrategy) + assert s.tau == 2e-4 + assert s.kmeans_max_iter == 7 + assert s.enforce_zero is True + assert s.zero_threshold == 1e-6 + + def test_default_dkm_spec(self): + spec = default_dkm_weight_palettization_spec() + assert isinstance(spec.training_strategy_spec, DKMTrainingSpec) + assert spec.n_bits == 4 + assert isinstance(spec.granularity, PerTensorGranularity) + assert spec.cluster_dim == 1 diff --git a/tests/palettization/test_dkm_dispatch.py b/tests/palettization/test_dkm_dispatch.py new file mode 100644 index 00000000..80ad4f7d --- /dev/null +++ b/tests/palettization/test_dkm_dispatch.py @@ -0,0 +1,66 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +from unittest import mock + +import pytest +import torch + +from coreai_opt.palettization.kmeans.dkm import _dispatch + + +@pytest.fixture(autouse=True) +def _clear_caches(): + _dispatch._providers.cache_clear() + _dispatch.get_forward.cache_clear() + yield + _dispatch._providers.cache_clear() + _dispatch.get_forward.cache_clear() + + +def _plugin_entry_point(plugin): + entry_point = mock.Mock() + entry_point.load.return_value = plugin + return entry_point + + +def test_no_providers_returns_none_and_logs(caplog): + with ( + mock.patch.object(_dispatch, "entry_points", return_value=[]), + caplog.at_level("INFO", logger=_dispatch.logger.name), + ): + assert _dispatch.get_forward(torch.device("cpu")) is None + assert "DKM: using the PyTorch implementation on cpu" in caplog.text + + +def test_providers_are_chained_and_first_applicable_wins(caplog): + def forward(*args, **kwargs): + raise NotImplementedError + + def plugin(providers): + return [ + *providers, + _dispatch.Provider("unusable", lambda device: None), + _dispatch.Provider("good", lambda device: forward), + ] + + with ( + mock.patch.object(_dispatch, "entry_points", return_value=[_plugin_entry_point(plugin)]), + caplog.at_level("INFO", logger=_dispatch.logger.name), + ): + assert _dispatch.get_forward(torch.device("cpu")) is forward + assert "DKM: using the good provider on cpu" in caplog.text + + +def test_failing_plugin_is_skipped_with_warning(caplog): + def plugin(providers): + raise ImportError("boom") + + with ( + mock.patch.object(_dispatch, "entry_points", return_value=[_plugin_entry_point(plugin)]), + caplog.at_level("WARNING", logger=_dispatch.logger.name), + ): + assert _dispatch.get_forward(torch.device("cpu")) is None + assert "failed to load; skipping" in caplog.text diff --git a/tests/palettization/test_dkm_forward.py b/tests/palettization/test_dkm_forward.py new file mode 100644 index 00000000..e28b7bd0 --- /dev/null +++ b/tests/palettization/test_dkm_forward.py @@ -0,0 +1,106 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""CPU tests for ``dkm_forward``: forward, knobs, and backward.""" + +import pytest +import torch + +from coreai_opt.palettization.kmeans.dkm._forward import dkm_forward + + +def _inputs(P=2, N=64, K=4, D=1, seed=0, requires_grad=False): + g = torch.Generator().manual_seed(seed) + W = torch.randn(P, N, D, generator=g, requires_grad=requires_grad) + C = torch.randn(P, K, D, generator=g, requires_grad=requires_grad) + return W, C + + +class TestDKMForward: + @pytest.mark.parametrize("D", [1, 2, 4]) + @pytest.mark.parametrize("kmeans_max_iter", [1, 3]) + def test_shapes_and_finite(self, D, kmeans_max_iter): + W, C = _inputs(D=D) + W_hat, C_new = dkm_forward(W, C, tau=1e-4, kmeans_max_iter=kmeans_max_iter) + assert W_hat.shape == W.shape + assert C_new.shape == C.shape + assert torch.isfinite(W_hat).all() + assert torch.isfinite(C_new).all() + + def test_reconstructs_exact_k_valued_input(self): + # Weight takes exactly K distinct values; with centroids seeded at those + # values and a sharp softmax, the soft assignment is effectively hard, so + # the reconstruction reproduces the input. + vals = torch.tensor([-2.0, -1.0, 1.0, 2.0]) + W = vals.repeat(16).reshape(1, 64, 1) + C = vals.reshape(1, 4, 1) + W_hat, _ = dkm_forward(W, C, tau=1e-4, kmeans_max_iter=1) + torch.testing.assert_close(W_hat, W, atol=1e-3, rtol=0) + + def test_enforce_zero_reserves_exact_zero_code(self): + W, C = _inputs() + _, C_new = dkm_forward(W, C, tau=1e-4, kmeans_max_iter=3, enforce_zero=True) + # Each partition keeps an exact-zero centroid (the nearest-origin one is pinned). + assert torch.all(C_new.norm(dim=-1).amin(dim=1) == 0) + + def test_zero_threshold_changes_output(self): + # zero_threshold floors the softmax weights; a large floor visibly changes + # the centroid update, proving the knob is live (not hardcoded). + W, C = _inputs() + _, C_small = dkm_forward(W, C, tau=1e-4, kmeans_max_iter=3, zero_threshold=1e-7) + _, C_large = dkm_forward(W, C, tau=1e-4, kmeans_max_iter=3, zero_threshold=1e-1) + assert not torch.allclose(C_small, C_large) + + def test_quantize_fn_applied_each_step(self): + # A quantize_fn that snaps centroids to zero forces C_new to zero. + W, C = _inputs() + _, C_new = dkm_forward( + W, C, tau=1e-4, kmeans_max_iter=2, quantize_fn=lambda c: torch.zeros_like(c) + ) + assert torch.equal(C_new, torch.zeros_like(C_new)) + + +class TestDKMBackward: + def test_grad_flows_to_inputs(self): + W, C = _inputs(requires_grad=True) + W_hat, C_new = dkm_forward(W, C, tau=1e-2, kmeans_max_iter=2) + (W_hat.pow(2).sum() + C_new.sum()).backward() + for g in (W.grad, C.grad): + assert g is not None + assert torch.isfinite(g).all() + # Gradients are non-trivial (the soft assignment is genuinely differentiable). + assert W.grad.abs().sum() > 0 + assert C.grad.abs().sum() > 0 + + def test_backward_matches_finite_difference(self): + # dkm_forward casts internally to fp32, so use a smooth tau and a loose + # tolerance for the central-difference check. + W = torch.randn(1, 6, 1, requires_grad=True) + C = torch.randn(1, 3, 1) + tau = 1.0 + + def loss(w): + W_hat, _ = dkm_forward(w, C, tau=tau, kmeans_max_iter=1) + return W_hat.pow(2).sum() + + loss(W).backward() + analytic = W.grad.clone() + + h = 1e-2 + with torch.no_grad(): + for idx in [(0, 0, 0), (0, 3, 0), (0, 5, 0)]: + wp = W.detach().clone() + wp[idx] += h + wm = W.detach().clone() + wm[idx] -= h + numeric = (loss(wp) - loss(wm)) / (2 * h) + assert abs(numeric.item() - analytic[idx].item()) < 1e-1 + + def test_enforce_zero_backward_is_finite(self): + W, C = _inputs(requires_grad=True) + W_hat, _ = dkm_forward(W, C, tau=1e-2, kmeans_max_iter=2, enforce_zero=True) + W_hat.pow(2).sum().backward() + assert torch.isfinite(W.grad).all() + assert torch.isfinite(C.grad).all() diff --git a/tests/palettization/test_dkm_training_strategy.py b/tests/palettization/test_dkm_training_strategy.py new file mode 100644 index 00000000..45644b04 --- /dev/null +++ b/tests/palettization/test_dkm_training_strategy.py @@ -0,0 +1,232 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""CPU tests for the DKM training strategy and its export path. + +The palettizer-level invariant tests are parametrized over ``_SCENARIOS``, a +curated list of ``PalettizationSpec`` combinations covering n_bits, granularity, +cluster_dim, lut_qspec, and enable_per_channel_scale. It is a hand-picked +representative set, not a full Cartesian product: many field combinations are +invalid (e.g. ``enable_fast_kmeans_mode`` forbids ``cluster_dim > 1``), so the +list crosses the settings pairwise around the interesting corners. +""" + +import pytest +import torch +import torch.nn as nn + +from coreai_opt import ExportBackend +from coreai_opt.palettization import ( + KMeansPalettizer, + KMeansPalettizerConfig, + ModuleKMeansPalettizerConfig, +) +from coreai_opt.palettization.kmeans.dkm.training_strategy import DKMTrainingSpec +from coreai_opt.palettization.kmeans.kmeans_fake_palettize import _KMeansFakePalettize +from coreai_opt.palettization.spec import ( + PalettizationSpec, + PerGroupedChannelGranularity, + PerTensorGranularity, +) +from coreai_opt.quantization.spec import QuantizationScheme, QuantizationSpec + +# Weight shape shared by every scenario: nn.Linear(_IN, _OUT) -> weight (_OUT, _IN). +# Sized so even n_bits=8 (256 centroids) has enough vectors to seed at the +# finest granularity/cluster_dim used below. +_IN, _OUT = 64, 32 + +_INT8_LUT = QuantizationSpec(dtype=torch.int8, qscheme=QuantizationScheme.SYMMETRIC) + +# Each scenario is the set of PalettizationSpec kwargs that differ from the +# defaults (n_bits=4, per-tensor, cluster_dim=1, no LUT quant, no per-channel scale). +_SCENARIOS = [ + pytest.param({}, id="n4-pt"), + pytest.param({"n_bits": 2, "lut_qspec": _INT8_LUT}, id="n2-pt-int8lut"), + pytest.param({"n_bits": 8}, id="n8-pt"), + pytest.param({"n_bits": 3, "enable_per_channel_scale": True}, id="n3-pt-pcs"), + pytest.param({"n_bits": 6, "cluster_dim": 2}, id="n6-pt-cd2"), + pytest.param( + {"granularity": PerGroupedChannelGranularity(axis=0, group_size=8)}, + id="n4-g8-cd1", + ), + pytest.param( + { + "granularity": PerGroupedChannelGranularity(axis=0, group_size=8), + "cluster_dim": 2, + "lut_qspec": _INT8_LUT, + }, + id="n4-g8-cd2-int8lut", + ), + pytest.param( + {"granularity": PerGroupedChannelGranularity(axis=0, group_size=16), "cluster_dim": 4}, + id="n4-g16-cd4", + ), + pytest.param( + { + "n_bits": 2, + "granularity": PerGroupedChannelGranularity(axis=0, group_size=8), + "enable_per_channel_scale": True, + }, + id="n2-g8-pcs", + ), + pytest.param({"cluster_dim": 2, "enable_per_channel_scale": True}, id="n4-pt-cd2-pcs"), +] + + +def _dkm_spec( + *, + n_bits=4, + granularity=None, + cluster_dim=1, + lut_qspec=None, + enable_per_channel_scale=False, + **dkm_kwargs, +): + return PalettizationSpec( + n_bits=n_bits, + granularity=granularity or PerTensorGranularity(), + cluster_dim=cluster_dim, + enable_per_channel_scale=enable_per_channel_scale, + lut_qspec=lut_qspec, + training_strategy_spec=DKMTrainingSpec(**dkm_kwargs), + ) + + +def _fp_module(spec): + return _KMeansFakePalettize(**spec.__dict__) + + +def _build_and_train(spec_kwargs, *, steps=3): + """Prepare a single-Linear model for DKM and run ``steps`` training forwards. + + Returns ``(palettizer, prepared, x)`` with the model left in eval mode. + """ + spec = _dkm_spec(**spec_kwargs) + model = nn.Sequential(nn.Linear(_IN, _OUT, bias=False)) + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + enable_fast_kmeans_mode=(spec.cluster_dim == 1), + op_state_spec={"weight": spec}, + ) + ) + palettizer = KMeansPalettizer(model, config) + prepared = palettizer.prepare((torch.randn(2, _IN),)) + x = torch.randn(2, _IN) + if steps: + with palettizer.training_mode(): + for _ in range(steps): + prepared(x) + prepared.eval() + return palettizer, prepared, x + + +class TestDKMSpecMatrix: + """Invariants that hold for every valid spec combination, at the palettizer level.""" + + @pytest.mark.parametrize("spec_kwargs", _SCENARIOS) + def test_training_forward_reconstructs_and_flows_grad(self, spec_kwargs): + spec = _dkm_spec(**spec_kwargs) + model = nn.Sequential(nn.Linear(_IN, _OUT, bias=False)) + config = KMeansPalettizerConfig( + global_config=ModuleKMeansPalettizerConfig( + enable_fast_kmeans_mode=(spec.cluster_dim == 1), + op_state_spec={"weight": spec}, + ) + ) + palettizer = KMeansPalettizer(model, config) + prepared = palettizer.prepare((torch.randn(2, _IN),)) + x = torch.randn(2, _IN) + with palettizer.training_mode(): + out = prepared(x) # DKM soft k-means runs here + out.sum().backward() + + assert out.shape == (2, _OUT) + assert torch.isfinite(out).all() + # Unlike the default (frozen) strategy, DKM's soft assignment is + # differentiable, so the palettized weight receives a finite gradient. + weight = prepared[0].parametrizations.weight.original + assert weight.grad is not None and torch.isfinite(weight.grad).all() + + @pytest.mark.parametrize("spec_kwargs", _SCENARIOS) + def test_train_eval_consistency(self, spec_kwargs): + """Train- and eval-mode reconstructions of the same weight must agree at tiny ``tau``. + + A sparse tail can differ: train applies ``attn`` (assigned against the pre-update + centroids) to the post-update quantized centroids, while eval hard-assigns against the + final centroids. With in-loop LUT quant the last rounding can shift a centroid across + the int8 grid and flip a few boundary weights by ~a centroid spacing (up to ~0.6% of + weights in quant scenarios, ~0% otherwise). Bound the 99th percentile, not the max. + """ + m = _fp_module(_dkm_spec(**spec_kwargs, tau=1e-6)) + w = torch.randn(_OUT, _IN) + m._initialize(w) + + m.train() + for _ in range(3): + train_out = m.forward(w) + + m.eval() + eval_out = m.forward(w) + assert (train_out - eval_out).abs().quantile(0.99) < 1e-4 + + @pytest.mark.parametrize("spec_kwargs", _SCENARIOS) + def test_state_dict_roundtrip_preserves_eval_output(self, spec_kwargs): + _, prepared, x = _build_and_train(spec_kwargs) + ref = prepared(x) + + _, reloaded, _ = _build_and_train(spec_kwargs, steps=0) + reloaded.load_state_dict(prepared.state_dict()) + reloaded.eval() + assert torch.equal(reloaded(x), ref) + + @pytest.mark.parametrize("spec_kwargs", _SCENARIOS) + def test_finalize_torch_roundtrip(self, spec_kwargs): + palettizer, prepared, x = _build_and_train(spec_kwargs) + eval_out = prepared(x) + finalized = palettizer.finalize(backend=ExportBackend.Torch) + assert torch.equal(eval_out, finalized(x)) + + def test_finalize_coreai_backend_succeeds(self): + palettizer, _, _ = _build_and_train({}) + finalized = palettizer.finalize(backend=ExportBackend.CoreAI) + assert finalized is not None + + +class TestDKMStrategyInternals: + """Module-level checks on behavior the palettizer workflow does not expose.""" + + def test_guard_rejects_non_kmeans_module(self): + strategy = DKMTrainingSpec().build_strategy() + with pytest.raises(TypeError): + strategy.train_forward(object(), torch.randn(4, 8)) + + @pytest.mark.parametrize("cluster_dim", [3, 8]) + def test_guard_rejects_unsupported_cluster_dim(self, cluster_dim): + strategy = DKMTrainingSpec().build_strategy() + module = _fp_module(_dkm_spec(cluster_dim=cluster_dim)) + with pytest.raises(ValueError, match="cluster_dim"): + strategy.train_forward(module, torch.randn(_OUT, _IN)) + + def test_train_forward_mutates_centroids_and_marks_indices_stale(self): + m = _fp_module(_dkm_spec()) + w = torch.randn(_OUT, _IN) + m._initialize(w) # seed centroids without a DKM step + before = m.centroids.clone() + + m.train() + out = m.forward(w) + + assert out.shape == w.shape + assert not torch.equal(m.centroids, before) # DKM mutated the centroids + assert m._indices_stale is True + + def test_enforce_zero_reserves_exact_zero_code(self): + m = _fp_module(_dkm_spec(enforce_zero=True)) + w = torch.randn(_OUT, _IN) + m._initialize(w) + m.train() + m.forward(w) + # the nearest-origin centroid is pinned, so an exact-zero code exists + assert torch.all(m.centroids.norm(dim=-1).amin(dim=1) == 0) From b4a3f133e06ae35d1a7e9bf47a35bb46b8eeed47 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 3 Nov 2026 15:52:04 -0800 Subject: [PATCH 2/4] Use registry for dkm dispatch --- .../palettization/kmeans/dkm/_dispatch.py | 42 ++++++---------- tests/palettization/test_dkm_dispatch.py | 50 ++++--------------- 2 files changed, 25 insertions(+), 67 deletions(-) diff --git a/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py index fe7bc19c..3a7e9133 100644 --- a/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py +++ b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py @@ -5,23 +5,21 @@ """Resolve the DKM forward function for a device. -Providers registered under the ``coreai_opt.dkm_providers`` entry-point group may supply -the forward. +Backends registered with :class:`DKMBackend` may supply the forward. """ import logging from collections.abc import Callable from functools import cache -from importlib.metadata import entry_points -from typing import NamedTuple, Protocol +from typing import Protocol import torch -__all__ = ["Forward", "Provider", "get_forward"] +from coreai_opt._utils.registry_utils import FunctionRegistryMixin -logger = logging.getLogger(__name__) +__all__ = ["DKMBackend", "Forward", "get_forward"] -ENTRY_POINT_GROUP = "coreai_opt.dkm_providers" +logger = logging.getLogger(__name__) class Forward(Protocol): @@ -41,33 +39,21 @@ def __call__( ... -class Provider(NamedTuple): - """A named provider; ``load(device)`` returns its forward, or ``None`` if it does not apply.""" - - name: str - load: Callable[[torch.device], Forward | None] +class DKMBackend(FunctionRegistryMixin): + """Registry of DKM backends, tried in registration order. - -@cache -def _providers() -> tuple[Provider, ...]: - providers: list[Provider] = [] - for entry_point in entry_points(group=ENTRY_POINT_GROUP): - try: - providers = entry_point.load()(providers) - except Exception: - logger.warning( - "DKM: provider %s failed to load; skipping.", entry_point.name, exc_info=True - ) - return tuple(providers) + Each entry is a loader ``(device) -> Forward | None`` that returns its forward, or ``None`` + if it does not apply to the device. + """ @cache def get_forward(device: torch.device) -> Forward | None: - """Return the forward supplied by a provider for ``device``, or ``None``.""" - for provider in _providers(): - forward = provider.load(device) + """Return the forward supplied by a registered backend for ``device``, or ``None``.""" + for name, load in DKMBackend.REGISTRY.items(): + forward = load(device) if forward is not None: - logger.info("DKM: using the %s provider on %s", provider.name, device) + logger.info("DKM: using the %s backend on %s", name, device) return forward logger.info("DKM: using the PyTorch implementation on %s", device) return None diff --git a/tests/palettization/test_dkm_dispatch.py b/tests/palettization/test_dkm_dispatch.py index 80ad4f7d..ae61063f 100644 --- a/tests/palettization/test_dkm_dispatch.py +++ b/tests/palettization/test_dkm_dispatch.py @@ -3,7 +3,7 @@ # Use of this source code is governed by a BSD-3-Clause license that can # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause -from unittest import mock +from collections import OrderedDict import pytest import torch @@ -12,55 +12,27 @@ @pytest.fixture(autouse=True) -def _clear_caches(): - _dispatch._providers.cache_clear() +def _isolated_registry(monkeypatch): + monkeypatch.setattr(_dispatch.DKMBackend, "REGISTRY", OrderedDict()) _dispatch.get_forward.cache_clear() yield - _dispatch._providers.cache_clear() _dispatch.get_forward.cache_clear() -def _plugin_entry_point(plugin): - entry_point = mock.Mock() - entry_point.load.return_value = plugin - return entry_point - - -def test_no_providers_returns_none_and_logs(caplog): - with ( - mock.patch.object(_dispatch, "entry_points", return_value=[]), - caplog.at_level("INFO", logger=_dispatch.logger.name), - ): +def test_no_backends_returns_none_and_logs(caplog): + with caplog.at_level("INFO", logger=_dispatch.logger.name): assert _dispatch.get_forward(torch.device("cpu")) is None assert "DKM: using the PyTorch implementation on cpu" in caplog.text -def test_providers_are_chained_and_first_applicable_wins(caplog): +def test_first_applicable_backend_wins(caplog): def forward(*args, **kwargs): raise NotImplementedError - def plugin(providers): - return [ - *providers, - _dispatch.Provider("unusable", lambda device: None), - _dispatch.Provider("good", lambda device: forward), - ] + _dispatch.DKMBackend.register("unusable")(lambda device: None) + _dispatch.DKMBackend.register("good")(lambda device: forward) + _dispatch.DKMBackend.register("later")(lambda device: pytest.fail("loaded after a match")) - with ( - mock.patch.object(_dispatch, "entry_points", return_value=[_plugin_entry_point(plugin)]), - caplog.at_level("INFO", logger=_dispatch.logger.name), - ): + with caplog.at_level("INFO", logger=_dispatch.logger.name): assert _dispatch.get_forward(torch.device("cpu")) is forward - assert "DKM: using the good provider on cpu" in caplog.text - - -def test_failing_plugin_is_skipped_with_warning(caplog): - def plugin(providers): - raise ImportError("boom") - - with ( - mock.patch.object(_dispatch, "entry_points", return_value=[_plugin_entry_point(plugin)]), - caplog.at_level("WARNING", logger=_dispatch.logger.name), - ): - assert _dispatch.get_forward(torch.device("cpu")) is None - assert "failed to load; skipping" in caplog.text + assert "DKM: using the good backend on cpu" in caplog.text From 8319e466907032249ca36764d9931748c9712c40 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:32:43 -0800 Subject: [PATCH 3/4] Update PATSchedule and DKM wording --- docs/src/palettization/overview.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/src/palettization/overview.md b/docs/src/palettization/overview.md index 4584fdd0..fbd7383d 100644 --- a/docs/src/palettization/overview.md +++ b/docs/src/palettization/overview.md @@ -140,7 +140,7 @@ with palettizer.training_mode(): ### Using PATSchedule -Users may also choose to enable palettization of parameters at a future point during training instead of right from the start. Weights which are marked for palettization, but have not yet had their palettizers enabled, will still be able to update during training. When a weight's palettizer is switched from disabled to enabled, its codebook and index assignments are recomputed using the most current weight values. +Users may also choose to enable palettization of parameters at a future point during training instead of right from the start. Weights which are marked for palettization, but have not yet had their palettizers enabled, will still be able to update during training. When a weight's palettizer is switched from disabled to enabled, its codebook and index assignments are recomputed using the most current weight values. `PATSchedule` applies to any training strategy, including [DKM](#differentiable-k-means-dkm-api). Since palettizer enabling can be configured on a module by module basis, this gives users more flexibility on training to be palettization-aware by allowing weights to update in response to palettization noise from other palettized weights before themselves being palettized. @@ -149,7 +149,7 @@ For example, users could choose to apply the following workflow: 1. Begin training with all parameters marked to be palettized but with all palettizers disabled 2. Fine-tune the model with no palettization noise for a few epochs 3. Enable a subset of palettizers while keeping the rest of the palettizers disabled -4. Fine-tune the model with palettization noise from the subset of enabled palettizers (any weights with enabled palettizers will not update further, but any weights with palettizers not yet enabled will continue training) +4. Fine-tune the model with palettization noise from the subset of enabled palettizers (with the default strategy, any weights with enabled palettizers will not update further, but any weights with palettizers not yet enabled will continue training) 5. Enable the rest of palettizers in the model and evaluate the final accuracy Steps 3 and 4 can be repeated as many times as desired with more and more palettizers being enabled. @@ -235,4 +235,6 @@ with palettizer.training_mode(): finalized_model = palettizer.finalize(...) ``` +DKM works with `PATSchedule` as well. A schedule decides when a module's palettizer becomes active, while the training strategy decides what that palettizer's forward pass does once it is. For example, you can leave DKM modules unpalettized for the first few hundred steps and then turn on soft-assignment training. To do this, set `pat_schedule` on the module config and call `palettizer.step()` in the training loop as described in [Using PATSchedule](#using-patschedule). Unlike with the default strategy, a DKM weight keeps updating after its palettizer is enabled, because gradients flow through the soft assignment. + As `training_strategy_spec` is part of `PalettizationSpec`, which in turn can be configured on a per-op or module basis in `KMeansPalettizerConfig`, users may choose to enable DKM for a subset of palettized weights in the model while leaving another portion of palettized weights using non-DKM `Vanilla K-means`. From bd9db144f75018b6b5be0e8ad4be833e132bdb59 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Wed, 4 Nov 2026 09:13:06 -0800 Subject: [PATCH 4/4] Address review comments --- docs/src/palettization/config.md | 2 +- .../palettization/kmeans/dkm/_dispatch.py | 11 +++++-- .../kmeans/dkm/training_strategy.py | 30 ++++++------------- tests/palettization/test_dkm_dispatch.py | 6 ++-- 4 files changed, 21 insertions(+), 28 deletions(-) diff --git a/docs/src/palettization/config.md b/docs/src/palettization/config.md index feb03338..f0f9b041 100644 --- a/docs/src/palettization/config.md +++ b/docs/src/palettization/config.md @@ -456,4 +456,4 @@ kmeans_palettization_config: tau: 5.0e-5 ``` -For a DKM spec with the defaults, the `default_dkm_weight_palettization_spec()` helper (importable from `coreai_opt.palettization.kmeans`) returns a ready-made 4-bit, per-tensor `PalettizationSpec` with `training_strategy_spec=DKMTrainingSpec()`. +For a DKM spec with the defaults, the `default_dkm_weight_palettization_spec()` helper (importable from `coreai_opt.palettization.kmeans.dkm`) returns a ready-made 4-bit, per-tensor `PalettizationSpec` with `training_strategy_spec=DKMTrainingSpec()`. diff --git a/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py index 3a7e9133..afb6d925 100644 --- a/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py +++ b/src/coreai_opt/palettization/kmeans/dkm/_dispatch.py @@ -17,6 +17,8 @@ from coreai_opt._utils.registry_utils import FunctionRegistryMixin +from ._forward import dkm_forward + __all__ = ["DKMBackend", "Forward", "get_forward"] logger = logging.getLogger(__name__) @@ -48,12 +50,15 @@ class DKMBackend(FunctionRegistryMixin): @cache -def get_forward(device: torch.device) -> Forward | None: - """Return the forward supplied by a registered backend for ``device``, or ``None``.""" +def get_forward(device: torch.device) -> Forward: + """ + Return the forward from the first applicable registered backend, else the PyTorch + ``dkm_forward``. + """ for name, load in DKMBackend.REGISTRY.items(): forward = load(device) if forward is not None: logger.info("DKM: using the %s backend on %s", name, device) return forward logger.info("DKM: using the PyTorch implementation on %s", device) - return None + return dkm_forward diff --git a/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py b/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py index 84be7c91..d93e649a 100644 --- a/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py +++ b/src/coreai_opt/palettization/kmeans/dkm/training_strategy.py @@ -15,7 +15,6 @@ from coreai_opt.palettization.spec.training_strategy import TrainingStrategy, TrainingStrategySpec from ._dispatch import get_forward as _get_forward -from ._forward import dkm_forward as _dkm_forward # Cluster dim values which are supported by DKM SUPPORTED_CLUSTER_DIMS = {1, 2, 4} @@ -85,26 +84,15 @@ def train_forward(self, module: _KMeansFakePalettize, weight: torch.Tensor) -> t forward = _get_forward(weight.device) - if forward is not None: - reconstructed_weights, centroids_new = forward( - weights, - centroids, - self.tau, - kmeans_max_iter=self.kmeans_max_iter, - zero_threshold=self.zero_threshold, - enforce_zero=self.enforce_zero, - quantize_fn=quantize_fn, - ) - else: - reconstructed_weights, centroids_new = _dkm_forward( - weights, - centroids, - self.tau, - self.kmeans_max_iter, - zero_threshold=self.zero_threshold, - enforce_zero=self.enforce_zero, - quantize_fn=quantize_fn, - ) + reconstructed_weights, centroids_new = forward( + weights, + centroids, + self.tau, + kmeans_max_iter=self.kmeans_max_iter, + zero_threshold=self.zero_threshold, + enforce_zero=self.enforce_zero, + quantize_fn=quantize_fn, + ) module.centroids.copy_(centroids_new.detach()) module._indices_stale = True diff --git a/tests/palettization/test_dkm_dispatch.py b/tests/palettization/test_dkm_dispatch.py index ae61063f..3bce1aeb 100644 --- a/tests/palettization/test_dkm_dispatch.py +++ b/tests/palettization/test_dkm_dispatch.py @@ -8,7 +8,7 @@ import pytest import torch -from coreai_opt.palettization.kmeans.dkm import _dispatch +from coreai_opt.palettization.kmeans.dkm import _dispatch, _forward @pytest.fixture(autouse=True) @@ -19,9 +19,9 @@ def _isolated_registry(monkeypatch): _dispatch.get_forward.cache_clear() -def test_no_backends_returns_none_and_logs(caplog): +def test_no_backends_returns_pytorch_forward_and_logs(caplog): with caplog.at_level("INFO", logger=_dispatch.logger.name): - assert _dispatch.get_forward(torch.device("cpu")) is None + assert _dispatch.get_forward(torch.device("cpu")) is _forward.dkm_forward assert "DKM: using the PyTorch implementation on cpu" in caplog.text