circuitkit 0.2.0__tar.gz → 0.2.1__tar.gz

{circuitkit-0.2.0/src/circuitkit.egg-info → circuitkit-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: circuitkit
-Version: 0.2.0
+Version: 0.2.1
 Summary: Unified Discover, Evaluate, Intervene toolkit for mechanistic interpretability: 13 circuit discovery algorithms, 6-pillar faithfulness evaluation, and pruning / quantization / knowledge editing / unlearning / activation steering / contrastive weight steering interventions behind one API.
 Author: Hem Gosalia, Aditya Kasliwal, Vinay Kumar Sankarapu
 Author-email: Pratinav Seth <pratinav.seth@lexsi.ai>

{circuitkit-0.2.0 → circuitkit-0.2.1}/docs/INSTALLATION.md

@@ -39,7 +39,7 @@ transformers>=5.15,<6, ordered-set, pyyaml, networkx,
 matplotlib, ipywidgets, scikit-learn, evaluate,
 sacrebleu, rouge_score, bert-score, accelerate,
 sentence-transformers, click, rich, torch-pruning>=1.0,
-transformer-lens==3.8.0 (CircuitKIT applies its Gemma-4 / Sarvam-MoE / Cohere patches at import)
+transformer-lens==3.8.0 (CircuitKIT applies its Gemma-4 / Sarvam-MoE / Cohere patches and the Pythia / GPT-NeoX loading fix at import)
 ```
 
 Install core only:

{circuitkit-0.2.0 → circuitkit-0.2.1}/docs/advanced/experimental-models.md

@@ -42,11 +42,17 @@ repository — no `HF_TOKEN` needed.
     [Tests](#tests) below).
 
 !!! note "Local checkpoint folders"
-    TransformerLens 3.8 uses `llama` and `gemma` substrings in a local checkpoint
-    path for loader selection before it reads `config.json`. CircuitKIT rejects
-    ambiguous local paths with a rename/symlink suggestion; use a neutral folder
-    name such as `./checkpoint`. Local SmolLM3 checkpoint folders are not
-    supported by the current port; load `HuggingFaceTB/SmolLM3-3B` by Hub ID.
+    While CircuitKIT's in-memory TransformerLens patch series is active
+    (`transformer-lens==3.8.0`, with `import circuitkit` before the first
+    `import transformer_lens`), a local checkpoint folder loads from its
+    `config.json` whatever it is called (a fine-tune in
+    `meta-llama-Llama-3.2-1B_sft/` is not mistaken for base Llama-3.2-1B).
+    Without the series, a folder whose path contains "llama" or "gemma" is
+    refused. Local Gemma checkpoints, fine-tunes included, are refused until
+    CircuitKIT moves to TransformerLens 4: TransformerLens 3.8 has Gemma configs
+    for Hub names only, and loading the Hub ID gives the base weights. Local SmolLM3
+    checkpoint folders are not supported by the current port; load
+    `HuggingFaceTB/SmolLM3-3B` by Hub ID.
 
 ## Gemma-4 and Sarvam-MoE: discovery-only TransformerLens ports
 

{circuitkit-0.2.0 → circuitkit-0.2.1}/docs/user-guide/visualization.md

@@ -27,7 +27,7 @@ pipe.visualize(mode="graph", output="./circuit.html")
 
 ## Mode 1: Graph view
 
-Renders the circuit as a directed graph: nodes are attention heads and MLP layers, edges represent information flow weighted by attribution score.
+Renders the circuit as a directed graph: nodes are attention heads and MLP layers, edges join circuit nodes of adjacent layers, weighted by node importance. They are not measured connections.
 
 ```python
 fig = ck.visualize_circuit(circuit, mode="graph")
@@ -35,7 +35,8 @@ fig = ck.visualize_circuit(circuit, mode="graph")
 
 - Nodes coloured by type (blue = attention head, amber = MLP)
 - Node size proportional to importance score
-- Edges only for components at or above sparsity threshold
+- Edges are drawn only between components in the circuit (pruned components have none)
+- A graph with more than 20,000 edges opens with the threshold raised so it draws at most 20,000 (the page says how many are shown); move the threshold to 0 to show all
 - Hover tooltips with layer, head index, and importance score
 
 Graph mode delegates to `Circuit.plot`, which takes only an optional output path.

{circuitkit-0.2.0 → circuitkit-0.2.1}/examples/finetuning/finetune-llama.py

@@ -385,8 +385,8 @@ def _build_parser() -> argparse.ArgumentParser:
     mdl = p.add_argument_group("model loading")
     mdl.add_argument(
         "--dtype",
-        default="float16",
-        choices=["float16", "bfloat16", "float32"],
+        default="bfloat16",
+        choices=["bfloat16", "float32"],
         help="Model dtype for inference and training.",
     )
     mdl.add_argument(

{circuitkit-0.2.0 → circuitkit-0.2.1}/examples/finetuning/finetune-qwen.py

@@ -367,7 +367,7 @@ def _build_parser() -> argparse.ArgumentParser:
     mdl.add_argument(
         "--dtype",
         default="bfloat16",
-        choices=["float16", "bfloat16", "float32"],
+        choices=["bfloat16", "float32"],
         help="Model dtype. Qwen models default to bfloat16.",
     )
     mdl.add_argument(

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/__init__.py

@@ -8,7 +8,7 @@ from ._tl_patches import install as _install_tl_patches
 _install_tl_patches()
 
 # The release pipeline (scripts/release.py set-version) rewrites this line.
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 __author__ = "Pratinav Seth, Hem Gosalia, Aditya Kasliwal, Vinay Kumar Sankarapu"
 __description__ = "Unified Discover, Evaluate, Intervene toolkit for mechanistic interpretability"
 _API_EXPORTS = {"discover_circuit", "evaluate_circuit", "load_circuit"}

circuitkit-0.2.1/src/circuitkit/_tl_patches/0007-local-checkpoint-config.patch

@@ -0,0 +1,128 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: hemgo-lexsi <hem.gosalia@lexsi.ai>
+Date: Thu, 29 Oct 2026 13:36:47 +0530
+Subject: [PATCH 7/7] local checkpoints: config and model class from
+ config.json, not the path
+
+convert_hf_model_config chose the architecture, and for 29 branches the
+whole config, from substrings of the model name. A local checkpoint's
+directory name is arbitrary: a fine-tuned "meta-llama-Llama-3.2-1B_sft"
+folder loaded with base Llama-3.2-1B's hard-coded config and never read
+its own config.json, and other Llama-named folders crashed on
+hf_config=None. For a local directory with a config.json, the chain now
+sees an empty name, so the architecture comes from config.json and only
+architecture-keyed branches match (Llama: the generic LlamaForCausalLM
+branch). get_pretrained_state_dict picks a local checkpoint's model class
+from its architecture instead of "bert" / "t5" / "hubert" / "wav2vec2" in
+the path. Hub names are unchanged.
+---
+ transformer_lens/loading_from_pretrained.py | 34 ++++++++++++++++-----
+ 1 file changed, 27 insertions(+), 7 deletions(-)
+
+diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py
+index 1f80721..bd33dd8 100644
+--- a/transformer_lens/loading_from_pretrained.py
++++ b/transformer_lens/loading_from_pretrained.py
+@@ -151,6 +151,11 @@ def get_official_model_name(model_name: str) -> str:
+     return official_model_name
+ 
+ 
++# Patch 0007: local checkpoints take their config and model class from
++# config.json, never from their path name. Read by circuitkit.load_model.
++LOCAL_CHECKPOINTS_FROM_CONFIG = True
++
++
+ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]:
+     """
+     Returns the model config for a HuggingFace model, converted to a dictionary
+@@ -159,12 +164,21 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]:
+     Takes the official_model_name as an input.
+     """
+     # In case the user passed in an alias
+-    if (Path(model_name) / "config.json").exists():
++    is_local = (Path(model_name) / "config.json").exists()
++    if is_local:
+         logging.info("Loading model config from local directory")
+         official_model_name = model_name
+     else:
+         official_model_name = get_official_model_name(model_name)
+ 
++    # A local checkpoint is described by its config.json alone. Its directory
++    # name says nothing reliable (a fine-tuned "meta-llama-Llama-3.2-1B_sft"
++    # folder is not base Llama-3.2-1B), so no branch below may match on it:
++    # they see an empty name until the chain is done.
++    checkpoint_name = official_model_name
++    if is_local:
++        official_model_name = ""
++
+     # Load HuggingFace model config. Stays None on the name-based branches
+     # below, which infer the architecture from the model name and never fetch.
+     hf_config: Any = None
+@@ -194,7 +208,7 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]:
+     else:
+         huggingface_token = os.environ.get("HF_TOKEN", "")
+         hf_config = AutoConfig.from_pretrained(
+-            official_model_name,
++            checkpoint_name,
+             token=huggingface_token if len(huggingface_token) > 0 else None,
+             **kwargs,
+         )
+@@ -1841,6 +1855,7 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]:
+         }
+     else:
+         raise NotImplementedError(f"{architecture} is not currently supported.")
++    official_model_name = checkpoint_name
+     # All of these models use LayerNorm
+     cfg_dict["original_architecture"] = architecture
+     # Carried on the cfg so the loader can act on the quantization without a
+@@ -2213,11 +2228,16 @@ def get_pretrained_state_dict(
+     if "torch_dtype" in kwargs:
+         dtype = kwargs["torch_dtype"]
+         del kwargs["torch_dtype"]
+-    if Path(official_model_name).exists():
++    is_local = Path(official_model_name).exists()
++    if is_local:
+         official_model_name = str(Path(official_model_name).resolve())
+         logging.info(f"Loading model from local path {official_model_name}")
+     else:
+         official_model_name = get_official_model_name(official_model_name)
++    # The model classes below are chosen by name for Hub models; a local
++    # checkpoint's path can contain "bert" or "t5" by chance, so use its
++    # architecture (e.g. "T5ForConditionalGeneration") instead.
++    class_key = (cfg.original_architecture or "").lower() if is_local else official_model_name
+     if official_model_name.startswith(NEED_REMOTE_CODE_MODELS) and not kwargs.get(
+         "trust_remote_code", False
+     ):
+@@ -2275,28 +2295,28 @@ def get_pretrained_state_dict(
+             huggingface_token = os.environ.get("HF_TOKEN", "")
+             if official_model_name in NON_HF_HOSTED_MODEL_NAMES:
+                 raise NotImplementedError("Model not hosted on HuggingFace, must pass in hf_model")
+-            elif "hubert" in official_model_name:
++            elif "hubert" in class_key:
+                 hf_model = HubertModel.from_pretrained(
+                     official_model_name,
+                     dtype=dtype,
+                     token=huggingface_token if len(huggingface_token) > 0 else None,
+                     **kwargs,
+                 )
+-            elif "wav2vec2" in official_model_name:
++            elif "wav2vec2" in class_key:
+                 hf_model = Wav2Vec2Model.from_pretrained(
+                     official_model_name,
+                     dtype=dtype,
+                     token=huggingface_token if len(huggingface_token) > 0 else None,
+                     **kwargs,
+                 )
+-            elif "bert" in official_model_name:
++            elif "bert" in class_key:
+                 hf_model = BertForPreTraining.from_pretrained(
+                     official_model_name,
+                     dtype=dtype,
+                     token=huggingface_token if len(huggingface_token) > 0 else None,
+                     **kwargs,
+                 )
+-            elif "t5" in official_model_name:
++            elif "t5" in class_key:
+                 hf_model = T5ForConditionalGeneration.from_pretrained(
+                     official_model_name,
+                     dtype=dtype,

circuitkit-0.2.1/src/circuitkit/_tl_patches/0008-neox-lm-head.patch

@@ -0,0 +1,27 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: Pratinav Seth <pratinav.seth@lexsi.ai>
+Date: Sat, 10 Oct 2026 12:00:00 +0530
+Subject: [PATCH] neox: read the unembedding from lm_head when embed_out is gone
+
+transformers renamed GPTNeoXForCausalLM.embed_out to lm_head (TransformerLens
+reports the change at 5.13 in one place and 5.14 in another). 3.8.0 reads
+embed_out only, so every GPT-NeoX / Pythia checkpoint fails in
+convert_neox_weights with AttributeError on the transformers versions CircuitKIT
+requires (>=5.15). Same change as the HookedTransformer half of the upstream fix
+released in 3.9.0; hasattr-based, so it does not depend on the exact version.
+
+diff --git a/transformer_lens/pretrained/weight_conversions/neox.py b/transformer_lens/pretrained/weight_conversions/neox.py
+index 70d8a6b..76827ed 100644
+--- a/transformer_lens/pretrained/weight_conversions/neox.py
++++ b/transformer_lens/pretrained/weight_conversions/neox.py
+@@ -54,6 +54,9 @@ def convert_neox_weights(neox, cfg: HookedTransformerConfig):
+     state_dict["ln_final.w"] = neox.gpt_neox.final_layer_norm.weight
+     state_dict["ln_final.b"] = neox.gpt_neox.final_layer_norm.bias
+ 
+-    state_dict["unembed.W_U"] = neox.embed_out.weight.T
++    # transformers renamed GPTNeoXForCausalLM.embed_out to lm_head (5.13/5.14);
++    # read whichever the installed transformers exposes.
++    unembed = neox.lm_head if hasattr(neox, "lm_head") else neox.embed_out
++    state_dict["unembed.W_U"] = unembed.weight.T
+     state_dict["unembed.b_U"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype)
+     return state_dict

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/_tl_patches/__init__.py

@@ -1,7 +1,8 @@
 """Apply CircuitKIT's TransformerLens patch series at import time.
 
 The ``*.patch`` files in this directory add Gemma-4, Sarvam-MoE and Cohere support
-to ``transformer-lens==3.8.0`` (see README.md here). :func:`install` puts an import
+to ``transformer-lens==3.8.0`` and fix local-checkpoint and GPT-NeoX / Pythia loading
+(see README.md here). :func:`install` puts an import
 hook on ``sys.meta_path`` that serves each patched ``transformer_lens`` module from
 its installed source with the series applied in memory. Nothing on disk changes,
 so ``pip install circuitkit`` needs no extra step. ``import circuitkit`` installs

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/api.py

@@ -6,7 +6,7 @@ from contextlib import ExitStack, contextmanager
 from datetime import datetime
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import torch as t
 
@@ -878,6 +878,64 @@ def _build_artifact_stem(config: Dict[str, Any]) -> str:
     return "_".join(str(p) for p in parts)
 
 
+def _flatten_ib_neuron_scores(
+    neuron_scores: Dict[str, Any],
+) -> Tuple[List[Tuple[float, Tuple[str, int, Any, int]]], Dict[str, float], Dict[str, int]]:
+    """Flatten IBCircuit per-neuron scores for pruning, node scores and graph metadata.
+
+    One pass over the scores yields everything the IBCircuit neuron branch of
+    :func:`discover_circuit` needs, so no tensor is read twice.
+
+    Args:
+        neuron_scores: ``{component name: 1-D score tensor}`` (values are read with
+            ``len`` and iteration, so any 1-D sequence of tensors works) with names
+            ``A{layer}.{head}`` and ``MLP {layer}``, as returned by
+            ``IBHookedTransformer.extract_neuron_scores``. Other names are skipped.
+
+    Returns:
+        A tuple ``(entries, node_scores, neurons_per_kind)``:
+
+        * ``entries``: one ``(abs score, (kind, layer, head, neuron index))`` per
+          neuron, with ``kind`` ``"attn"`` or ``"mlp"`` and ``head`` ``None`` for MLPs.
+        * ``node_scores``: the largest absolute neuron score per component, the
+          same shape node-level discovery writes, so ``Circuit.from_artifact``
+          and ``Circuit.plot`` can use it.
+        * ``neurons_per_kind``: ``neurons_per_head`` / ``neurons_per_mlp`` when
+          every component of that kind has the same neuron count. A kind whose
+          sizes differ (or that is absent) gets no key, so the graph treats no
+          component of it as pruned rather than guessing.
+    """
+    entries: List[Tuple[float, Tuple[str, int, Any, int]]] = []
+    node_scores: Dict[str, float] = {}
+    sizes: Dict[str, set] = {"neurons_per_head": set(), "neurons_per_mlp": set()}
+
+    for ib_name, score_tensor in neuron_scores.items():
+        attn_match = re.match(r"A(\d+)\.(\d+)$", ib_name)
+        mlp_match = re.match(r"MLP (\d+)$", ib_name)
+        if attn_match:
+            entry_key: Tuple[str, int, Any] = (
+                "attn",
+                int(attn_match.group(1)),
+                int(attn_match.group(2)),
+            )
+            sizes["neurons_per_head"].add(len(score_tensor))
+        elif mlp_match:
+            entry_key = ("mlp", int(mlp_match.group(1)), None)
+            sizes["neurons_per_mlp"].add(len(score_tensor))
+        else:
+            continue
+        top = float("-inf")
+        for neuron_idx, score in enumerate(score_tensor):
+            abs_score = abs(score.item())
+            entries.append((abs_score, (*entry_key, neuron_idx)))
+            top = max(top, abs_score)
+        if top != float("-inf"):
+            node_scores[ib_name] = top
+
+    neurons_per_kind = {k: next(iter(v)) for k, v in sizes.items() if len(v) == 1}
+    return entries, node_scores, neurons_per_kind
+
+
 def _save_artifact(
     data: Any, output_path: str, suffix: str, logger, config: Dict = None
 ) -> Optional[str]:
@@ -1522,8 +1580,17 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                     pruned_attn_neurons = defaultdict(list)
 
                     all_neuron_scores = []
+                    # Neurons per component, for the circuit graph: a component
+                    # counts as pruned only when all of them were. Constant per
+                    # kind (set by the model config), so one number each.
+                    neurons_per_mlp: Optional[int] = None
+                    neurons_per_head: Optional[int] = None
                     for node in tqdm(graph.nodes.values(), desc="Extracting neuron scores"):
                         if isinstance(node, (MLPNode, AttentionNode)):
+                            if isinstance(node, MLPNode):
+                                neurons_per_mlp = node.d_neuron
+                            else:
+                                neurons_per_head = node.d_neuron
                             # Filter by scope
                             if effective_scope == "mlp" and not isinstance(node, MLPNode):
                                 continue
@@ -1575,13 +1642,18 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                                 (int(attn_match.group(1)), int(attn_match.group(2)))
                             ].append(neuron_idx)
 
+                    result_meta: Dict[str, Any] = {
+                        "mlp_hook": discovery_cfg.get("mlp_hook", "mlp_out"),
+                        "heads_hook": "attn.hook_result",  # EAP uses attn.hook_result for heads
+                    }
+                    if neurons_per_mlp is not None:
+                        result_meta["neurons_per_mlp"] = neurons_per_mlp
+                    if neurons_per_head is not None:
+                        result_meta["neurons_per_head"] = neurons_per_head
                     result = {
                         "mlp": dict(pruned_mlp_neurons),
                         "heads": dict(pruned_attn_neurons),
-                        "_meta": {
-                            "mlp_hook": discovery_cfg.get("mlp_hook", "mlp_out"),
-                            "heads_hook": "attn.hook_result",  # EAP uses attn.hook_result for heads
-                        },
+                        "_meta": result_meta,
                     }
                     if config.get("output_path"):
                         os.makedirs(os.path.dirname(config["output_path"]), exist_ok=True)
@@ -1590,6 +1662,10 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                     _save_artifact(
                         {
                             "algo": algo,
+                            # Task and model, so Circuit.from_artifact can label a
+                            # neuron-level graph (there is no scores JSON here).
+                            "task": discovery_cfg["task"],
+                            "model": model_cfg["name"],
                             "level": "neuron",
                             "node_scores": node_scores,
                             "neurons_scores": graph.neurons_scores.cpu(),
@@ -1690,40 +1766,18 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                 # save in the same format as EAP neuron so evaluate_circuit works.
                 pruned_mlp_neurons = defaultdict(list)
                 pruned_attn_neurons = defaultdict(list)
-                all_neuron_scores = []
-
-                for ib_name, score_tensor in node_scores.items():
-                    attn_match = re.match(r"A(\d+)\.(\d+)$", ib_name)
-                    mlp_match = re.match(r"MLP (\d+)$", ib_name)
-                    for neuron_idx, score in enumerate(score_tensor):
-                        if attn_match:
-                            all_neuron_scores.append(
-                                (
-                                    abs(score.item()),
-                                    (
-                                        "attn",
-                                        int(attn_match.group(1)),
-                                        int(attn_match.group(2)),
-                                        neuron_idx,
-                                    ),
-                                )
-                            )
-                        elif mlp_match:
-                            all_neuron_scores.append(
-                                (
-                                    abs(score.item()),
-                                    ("mlp", int(mlp_match.group(1)), None, neuron_idx),
-                                )
-                            )
+                ib_neuron_entries, ib_node_scores, neurons_per_kind = _flatten_ib_neuron_scores(
+                    node_scores
+                )
 
-                all_neuron_scores.sort(key=lambda x: x[0])  # ascending: lowest = least important
-                n_to_prune = int(len(all_neuron_scores) * pruning_cfg["target_sparsity"])
+                ib_neuron_entries.sort(key=lambda x: x[0])  # ascending: lowest = least important
+                n_to_prune = int(len(ib_neuron_entries) * pruning_cfg["target_sparsity"])
 
                 logger.info(
-                    f"IBCircuit neuron discovery: {len(all_neuron_scores)} total neurons, pruning {n_to_prune}"
+                    f"IBCircuit neuron discovery: {len(ib_neuron_entries)} total neurons, pruning {n_to_prune}"
                 )
 
-                for _, (kind, layer, head, neuron_idx) in all_neuron_scores[:n_to_prune]:
+                for _, (kind, layer, head, neuron_idx) in ib_neuron_entries[:n_to_prune]:
                     if kind == "mlp":
                         pruned_mlp_neurons[layer].append(neuron_idx)
                     else:
@@ -1732,15 +1786,22 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                 result = {
                     "mlp": dict(pruned_mlp_neurons),
                     "heads": dict(pruned_attn_neurons),
-                    "_meta": {"mlp_hook": ib_config["mlp_hook"]},
+                    "_meta": {"mlp_hook": ib_config["mlp_hook"], **neurons_per_kind},
                 }
 
                 _save_artifact(
                     {
                         "algo": algo,
+                        "task": discovery_cfg["task"],
+                        "model": model_cfg["name"],
                         "level": "neuron",
+                        # Largest absolute neuron score per component, as the EAP
+                        # neuron branch writes, so Circuit.from_artifact gets
+                        # non-empty scores and the circuit graph can be drawn.
+                        "node_scores": ib_node_scores,
                         "neurons_scores": node_scores,
-                        "total_neurons": len(all_neuron_scores),
+                        "mlp_hook": ib_config["mlp_hook"],
+                        "total_neurons": len(ib_neuron_entries),
                     },
                     config.get("output_path"),
                     "_scores",

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/applications/selective_finetuning/examples/finetune_llama.py

@@ -383,8 +383,8 @@ def _build_parser() -> argparse.ArgumentParser:
     mdl = p.add_argument_group("model loading")
     mdl.add_argument(
         "--dtype",
-        default="float16",
-        choices=["float16", "bfloat16", "float32"],
+        default="bfloat16",
+        choices=["bfloat16", "float32"],
         help="Model dtype for inference and training.",
     )
     mdl.add_argument(

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/applications/selective_finetuning/examples/finetune_qwen.py

@@ -365,7 +365,7 @@ def _build_parser() -> argparse.ArgumentParser:
     mdl.add_argument(
         "--dtype",
         default="bfloat16",
-        choices=["float16", "bfloat16", "float32"],
+        choices=["bfloat16", "float32"],
         help="Model dtype. Qwen models default to bfloat16.",
     )
     mdl.add_argument(

{circuitkit-0.2.0 → circuitkit-0.2.1}/src/circuitkit/applications/selective_finetuning/finetune_utils.py

@@ -511,13 +511,15 @@ def run_finetuning(
     eval_dataloader: Optional[DataLoader] = None,
     eval_fn: Optional[Any] = None,
     log_every: int = 10,
+    weight_decay: float = 0.01,
 ) -> Tuple[nn.Module, List[Dict[str, float]]]:
     """
     Run selective gradient-masked finetuning.
 
     Parameters
     ----------
-    model               : HuggingFace CausalLM on `device`.
+    model               : HuggingFace CausalLM on `device`, in float32 or
+                          bfloat16. float16 raises ValueError (see below).
     selection           : SelectionResult controlling which params are trained.
     finetune_dataloader : DataLoader from build_finetune_dataloader.
     device              : Training device.
@@ -530,6 +532,9 @@ def run_finetuning(
     eval_fn             : Optional callable(model, eval_dataloader) → Dict[str, float].
                           Called at the end of every epoch if provided.
     log_every           : Print training loss every N steps.
+    weight_decay        : AdamW decoupled weight decay (0.01, the AdamW default).
+                          For a masked parameter it is applied to the selected
+                          entries only; masked-out entries are never changed.
 
     Returns
     -------
@@ -540,10 +545,38 @@ def run_finetuning(
     -----
     Hooks are always removed in a finally block.  If training is interrupted
     the model is left in its partially-trained state with requires_grad reset.
+
+    The weights are trained in their own dtype, with no float32 master copy.
+    AdamW cannot do that in float16: eps=1e-8 and the squared gradient of any
+    gradient below about 5e-3 round to 0, so the update is 0/0 or x/0 and those
+    weights become NaN/inf after the first step. A float16 model is therefore
+    rejected before anything is changed.
     """
+    if any(p.dtype == torch.float16 for p in model.parameters()):
+        raise ValueError(
+            "run_finetuning trains the weights in place with AdamW, which turns "
+            "float16 weights into NaN/inf after one step. Load the model in "
+            "torch.bfloat16 or torch.float32."
+        )
+
     hook_handles, trainable_params, masks_dict = setup_selective_training(model, selection, device)
 
-    optimiser = torch.optim.AdamW(trainable_params, lr=lr)
+    # AdamW's decoupled weight decay shrinks every entry of a parameter, including
+    # the masked-out entries whose gradient is zero. Masked parameters therefore
+    # get no decay from the optimiser; it is applied to their selected entries
+    # just before each step instead. The factor is kept in float32, as AdamW's
+    # scalar factor is, so half-precision weights decay by the same amount.
+    masked_params = [p for p in trainable_params if str(id(p)) in masks_dict]
+    decay_factors = [
+        1.0 - lr * weight_decay * masks_dict[str(id(p))].float() for p in masked_params
+    ]
+    param_groups = [
+        {"params": [p for p in trainable_params if str(id(p)) not in masks_dict]},
+        {"params": masked_params, "weight_decay": 0.0},
+    ]
+    optimiser = torch.optim.AdamW(
+        [g for g in param_groups if g["params"]], lr=lr, weight_decay=weight_decay
+    )
     epoch_logs: List[Dict[str, float]] = []
     verified = False
 
@@ -605,6 +638,11 @@ def run_finetuning(
                 # during backward, so masks are applied before clipping acts
                 # on the masked gradients.
                 torch.nn.utils.clip_grad_norm_(trainable_params, max_grad_norm)
+                if weight_decay:
+                    with torch.no_grad():
+                        for param, factor in zip(masked_params, decay_factors):
+                            if param.grad is not None:
+                                param.mul_(factor)
                 optimiser.step()
 
                 epoch_loss += loss.item()