# How to Apply Different Bit Widths to Different Layers Using a Bits Map in Needle

> Learn how to apply different bit widths to specific layers in Needle using a bits map for mixed-precision quantization. Assign custom bit-widths to layers by name prefix.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-09-09

---

**Needle supports mixed-precision quantization through a bits map configuration that assigns distinct weight bit-widths to individual layers based on canonical name prefix matching.**

Needle is an open-source framework for model quantization that enables fine-grained control over compression strategies. When you apply different bit widths to different layers using a bits map, you override global quantization settings with a textual specification that maps layer name prefixes to specific precision levels. This mechanism allows you to aggressively compress insensitive components while preserving accuracy in critical layers.

## Understanding the Bits Map Syntax

A bits map is a comma-separated specification stored under the `weight_bits` key in your model configuration. According to the source code in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py), the `parse_bits_map()` function (line 301) expects a mandatory `default=<bits>` entry followed by optional prefix-specific assignments:

```text
default=4,attention=2,ffn=3

```

- **`default=<bits>`** – The fallback bit-width applied to any weight whose name does not match a more specific prefix.
- **`<prefix>=<bits>`** – Any string compared case-insensitively to the beginning of a weight’s canonical name.

### Longest-Prefix Matching Logic

The `_bits_for()` function (lines 218-225) implements the matching algorithm using longest-prefix-wins semantics. When processing a weight named `attention_head_0`, the system evaluates all configured prefixes and selects the bit width associated with the longest matching prefix. This ensures that specific rules override general ones.

## How Mixed-Precision Quantization Works Under the Hood

During the quantization pipeline, `cq_mixed_params()` (lines 321-327) traverses the parameter tree and invokes `_bits_for()` for each leaf node. The function examines each weight's canonical name, determines the appropriate bit width from the map, and quantizes that specific parameter accordingly. If a layer's name begins with *attention*, the function applies the 2-bit setting from the map; otherwise it falls back to the default of 4 bits.

In [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) (around line 389), the system distinguishes between global and mixed-precision modes through the assignment `bits_map, default_bits = parsed_bits_map`. When a valid bits map is present, Needle bypasses the global `--bits` flag and routes all quantization decisions through the per-layer mapping system.

## Three Ways to Specify a Bits Map

### Command-Line Interface

Pass the bits map directly to the `needle build` command using the `--bits` flag:

```bash
needle build \
    --ckpt tiny-checkpoint \
    --out quantised.ckpt \
    --bits "default=4,attention=2,ffn=3"

```

This configuration assigns 4-bit weights to all layers except those whose canonical names start with `attention` (2 bits) and `ffn` (3 bits).

### Configuration File

Add the specification to your [`config.json`](https://github.com/cactus-compute/needle/blob/main/config.json) or YAML model configuration:

```json
{
    "name": "tiny-model",
    "weight_bits": "default=4,attention=2"
}

```

When running `needle build`, the builder reads the configuration, parses the bits map, and applies 2-bit quantization to all attention-related weights while using 4 bits for everything else.

### Programmatic API

For custom training scripts, pass the specification string to `configure_qat()` or `configure_deploy()`, which forward the configuration to the underlying quantization functions:

```python
from needle.model.quantize import parse_bits_map, cq_ste_mixed_params

bits_spec = "default=4,attention=2"
bits_map, default = parse_bits_map(bits_spec)

# Apply mixed-precision STE quantisation to parameters

quantised_params = cq_ste_mixed_params(params, bits_map, default)

```

## Complete Implementation Example

The following tests demonstrate bits-map usage in the Needle test suite:

- **[`tests/test_finetune.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_finetune.py)** (lines 88-103): Verifies that setting `weight_bits` in the configuration correctly applies per-layer quantization during fine-tuning.
- **[`tests/test_build.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_build.py)** (lines 9-34): Validates CLI-level parsing of the `--bits` flag with a custom bits map.

## Summary

- A **bits map** uses comma-separated `prefix=bits` syntax with a mandatory `default=<bits>` entry to enable mixed-precision quantization.
- The **`parse_bits_map()`** function in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) (line 301) validates and parses the configuration string into a lookup dictionary.
- **Longest-prefix matching** in `_bits_for()` (lines 218-225) determines the appropriate bit width for each weight based on its canonical name.
- You can specify the map via **CLI flag**, **configuration file**, or **programmatically** through `configure_qat()` or `configure_deploy()`.
- The **`cq_mixed_params()`** function (lines 321-327) executes the actual quantization by walking the parameter tree and applying per-layer specifications.

## Frequently Asked Questions

### What is the exact syntax for a bits map in Needle?

The syntax requires a comma-separated list of `prefix=bits` pairs with a mandatory `default=<integer>` entry. For example: `default=4,attention=2,ffn=3`. The parser validates this format and returns a dictionary mapping prefixes to bit widths along with the default value. All prefix matching is performed case-insensitively against the beginning of each weight's canonical name.

### How does Needle resolve conflicts when multiple prefixes match a layer name?

The `_bits_for()` function implements longest-prefix-wins semantics. If you define both `att=3` and `attention=2`, a weight named `attention_head` receives 2-bit quantization because "attention" (9 characters) is longer than "att" (3 characters). This prioritization ensures that specific rules override general ones during the quantization walk.

### Can I use a bits map together with the global --bits flag?

No. According to the logic in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) (line 389), when a valid bits map is present, Needle ignores the global `--bits` flag entirely. The system unpacks the parsed bits map into `bits_map` and `default_bits` variables, routing all quantization decisions exclusively through the per-layer mapping system.

### Which layers benefit most from specific bit width assignments?

According to the Needle source analysis, attention mechanisms typically require higher precision (more bits) due to their sensitivity to quantization errors, while feed-forward networks (FFN) can tolerate aggressive compression. The bits map allows you to fine-tune the accuracy-efficiency trade-off by assigning specific bit widths to sensitive layers while compressing others.