diff --git a/src/heretic/model.py b/src/heretic/model.py index 4aa813ec..574bd19b 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -2,9 +2,9 @@ # Copyright (C) 2025-2026 Philipp Emanuel Weidmann + contributors import math -from contextlib import suppress +from contextlib import contextmanager, suppress from dataclasses import dataclass -from typing import Any, Type, cast +from typing import Any, Generator, Type, cast import bitsandbytes as bnb import torch @@ -31,6 +31,9 @@ GenerateDecoderOnlyOutput, # ty:ignore[possibly-missing-import] ) +with suppress(ImportError): + from transformers.models.diffusion_gemma import DiffusionGemmaForBlockDiffusion # type: ignore[assignment] + from .config import QuantizationMethod, RowNormalization, Settings from .system import empty_cache from .utils import Prompt, batchify, format_exception, print @@ -41,7 +44,17 @@ def get_model_class( ) -> Type[AutoModelForImageTextToText] | Type[AutoModelForCausalLM]: configs = PretrainedConfig.get_config_dict(model) - if any([("vision_config" in config) for config in configs]): + for config in configs: + if isinstance(config, dict) and config.get("model_type") == "diffusion_gemma": + if "DiffusionGemmaForBlockDiffusion" not in globals(): + raise ImportError( + "DiffusionGemma support requires a newer version of the transformers library." + ) + return DiffusionGemmaForBlockDiffusion # type: ignore[return-value] + + if any( + [("vision_config" in config) for config in configs if isinstance(config, dict)] + ): return AutoModelForImageTextToText else: return AutoModelForCausalLM @@ -61,7 +74,6 @@ class Model: # Set for multimodal models, None for text-only ones. processor: ProcessorMixin | None peft_config: LoraConfig - dtype: torch.dtype def __init__(self, settings: Settings): self.settings = settings @@ -76,6 +88,7 @@ def __init__(self, settings: Settings): self.tokenizer = AutoTokenizer.from_pretrained( settings.model, + trust_remote_code=settings.trust_remote_code, **self.revision_kwargs, ) @@ -84,6 +97,7 @@ def __init__(self, settings: Settings): if get_model_class(settings.model) == AutoModelForImageTextToText: self.processor = AutoProcessor.from_pretrained( settings.model, + trust_remote_code=settings.trust_remote_code, **self.revision_kwargs, ) @@ -102,8 +116,10 @@ def __init__(self, settings: Settings): if settings.max_memory else None ) + self.trusted_models = {settings.model: settings.trust_remote_code} - self.trusted_models = set() + if self.settings.evaluate_model is not None: + self.trusted_models[settings.evaluate_model] = settings.trust_remote_code for dtype in settings.dtypes: print(f"* Trying dtype [bold]{dtype}[/]...") @@ -122,18 +138,15 @@ def __init__(self, settings: Settings): dtype=dtype, device_map=settings.device_map, max_memory=self.max_memory, - trust_remote_code=True - if settings.model in self.trusted_models - else None, + trust_remote_code=self.trusted_models.get(settings.model), **self.revision_kwargs, **extra_kwargs, ) - self.dtype = self.model.dtype # If we reach this point and the model requires trust_remote_code, - # the user must have agreed when prompted to execute remote code, - # because from_pretrained raises an exception otherwise. - self.trusted_models.add(settings.model) + # either the user accepted, or settings.trust_remote_code is True. + if self.trusted_models.get(settings.model) is None: + self.trusted_models[settings.model] = True # A test run can reveal dtype-related problems such as the infamous # "RuntimeError: probability tensor contains either `inf`, `nan` or element < 0" @@ -170,6 +183,9 @@ def __init__(self, settings: Settings): # LoRA B matrices are initialized to zero by default in PEFT, # so we don't need to do anything manually. + if self._is_diffusion_gemma(): + self._save_dg_expert_weights() + print(f"* Transformer model with [bold]{len(self.get_layers())}[/] layers") all_components = {} @@ -215,6 +231,11 @@ def _apply_lora(self): # Row magnitude preservation introduces nonlinear effects. lora_rank = self.settings.full_normalization_lora_rank + # DiffusionGemmaForBlockDiffusion uses its own generation mixin and does not + # implement prepare_inputs_for_generation, so CAUSAL_LM task type makes PEFT + # fail. Use FEATURE_EXTRACTION instead (no generation hooks required). + task_type = "FEATURE_EXTRACTION" if self._is_diffusion_gemma() else "CAUSAL_LM" + self.peft_config = LoraConfig( r=lora_rank, target_modules=target_modules, @@ -223,7 +244,7 @@ def _apply_lora(self): bias="none", # Even if we're using AutoModelForImageTextToText, this is still correct, # as VL models are typically just causal LMs with an added image encoder. - task_type="CAUSAL_LM", + task_type=task_type, ) # self.peft_config is a LoraConfig object rather than a dictionary, @@ -281,9 +302,7 @@ def get_merged_model(self) -> PreTrainedModel: self.settings.model, torch_dtype=self.model.dtype, device_map="cpu", - trust_remote_code=True - if self.settings.model in self.trusted_models - else None, + trust_remote_code=self.trusted_models.get(self.settings.model), **self.revision_kwargs, ) @@ -319,25 +338,24 @@ def reset_model(self): - Slow path: If switching models or after merge_and_unload(), performs full model reload with quantization config. """ - # If a prior model load was interrupted/cancelled mid-process, self.model will be None. - current_model = None - if self.model is not None: - current_model = getattr(self.model.config, "name_or_path", None) - + current_model = getattr(self.model.config, "name_or_path", None) if current_model == self.settings.model and not self.needs_reload: # Reset LoRA adapters to zero (identity transformation). for name, module in self.model.named_modules(): if "lora_B" in name and hasattr(module, "weight"): torch.nn.init.zeros_(module.weight) + # Restore expert weights that were modified in-place by EGA + if self._is_diffusion_gemma(): + self._restore_dg_expert_weights() return + dtype = self.model.dtype + # Purge existing model object from memory to make space. self.model = None # ty:ignore[invalid-assignment] empty_cache() - quantization_config = self._get_quantization_config( - str(self.dtype).split(".")[-1] - ) + quantization_config = self._get_quantization_config(str(dtype).split(".")[-1]) # Build kwargs, only include quantization_config if it's not None. extra_kwargs = {} @@ -346,20 +364,170 @@ def reset_model(self): self.model = get_model_class(self.settings.model).from_pretrained( self.settings.model, - dtype=self.dtype, + dtype=dtype, device_map=self.settings.device_map, max_memory=self.max_memory, - trust_remote_code=True - if self.settings.model in self.trusted_models - else None, + trust_remote_code=self.trusted_models.get(self.settings.model), **self.revision_kwargs, **extra_kwargs, ) self._apply_lora() + # On full reload the expert weights are restored from disk; just re-save + # them so future _restore_dg_expert_weights() calls have fresh copies. + if self._is_diffusion_gemma(): + self._save_dg_expert_weights() + self.needs_reload = False + @contextmanager + def _dg_lora_merged(self) -> Generator[None, None, None]: + """Temporarily merge LoRA adapters into the shared encoder/decoder weight tensors. + + DiffusionGemma ties encoder and decoder weights (same data_ptr). PEFT wraps + only the encoder layers, so the decoder's forward passes won't see the LoRA + abliteration. We temporarily fold the delta (lora_B @ lora_A) into the shared + base weight tensor so that the decoder-driven diffusion generation also reflects + the current abliteration state, then restore the originals when done. + """ + saved: dict[int, tuple[Tensor, Tensor]] = {} + + for module in self.model.modules(): + if ( + isinstance(module, Linear) + and hasattr(module, "lora_A") + and hasattr(module, "lora_B") + ): + base_w = module.base_layer.weight.data + ptr = base_w.data_ptr() + if ptr in saved: + continue # Already handled this shared tensor. + if base_w.dtype == torch.uint8 or hasattr(base_w, "quant_state"): + raise RuntimeError( + "DiffusionGemma LoRA merge is not supported with 4-bit quantization." + ) + lora_A_w = cast(Tensor, module.lora_A["default"].weight.data) + lora_B_w = cast(Tensor, module.lora_B["default"].weight.data) + delta = (lora_B_w @ lora_A_w).to(base_w.dtype) + saved[ptr] = (base_w, delta) + base_w.add_(delta) + + try: + yield + finally: + for base_w, delta in saved.values(): + base_w.sub_(delta) + + def _save_dg_expert_weights(self) -> None: + """Save CPU copies of expert down_proj tensors for fast trial reset.""" + self._dg_expert_saved: dict[int, Tensor] = {} + layers = self.get_layers() + for layer_idx, layer in enumerate(layers): + if hasattr(layer, "experts") and hasattr(layer.experts, "down_proj"): + self._dg_expert_saved[layer_idx] = ( + layer.experts.down_proj.data.cpu().clone() + ) + + def _restore_dg_expert_weights(self) -> None: + """Restore expert weights modified in-place by EGA, ready for the next trial.""" + if not hasattr(self, "_dg_expert_saved"): + return + layers = self.get_layers() + for layer_idx, layer in enumerate(layers): + if layer_idx in self._dg_expert_saved: + layer.experts.down_proj.data.copy_( + self._dg_expert_saved[layer_idx].to(layer.experts.down_proj.device) + ) + + def _abliterate_dg_experts( + self, + refusal_directions: Tensor, + direction_index: float | None, + parameters: dict[str, AbliterationParameters], + ) -> None: + """EGA: abliterate each expert's down_proj slice in-place. + + Expert weights [n_experts, hidden, moe_inter] are nn.Parameters, not nn.Linear, + so LoRA can't wrap them. We apply norm-preserving biprojected ablation to every + expert slice using the same refusal direction and kernel as mlp.down_proj. + Because encoder and decoder share these tensors (confirmed same data_ptr), + abliterating here automatically fixes the decoder too. + """ + layers = self.get_layers() + for layer_idx, layer in enumerate(layers): + if not (hasattr(layer, "experts") and hasattr(layer.experts, "down_proj")): + continue + + if "mlp.down_proj" in parameters: + params = parameters["mlp.down_proj"] + distance = cast(float, abs(layer_idx - params.max_weight_position)) + if distance > params.min_weight_distance: + continue + weight_scale = params.max_weight + ( + distance / params.min_weight_distance + ) * (params.min_weight - params.max_weight) + else: + weight_scale = 1.0 + + # Select refusal direction for this layer (same logic as abliterate()) + if direction_index is None: + layer_refusal_direction = refusal_directions[layer_idx + 1] + else: + w_frac, idx = math.modf(direction_index + 1) + layer_refusal_direction = F.normalize( + refusal_directions[int(idx)].lerp( + refusal_directions[int(idx) + 1], w_frac + ), + p=2, + dim=0, + ) + + expert_down = layer.experts.down_proj # [n_experts, hidden, moe_inter] + v = F.normalize(layer_refusal_direction.float(), dim=0).to( + expert_down.device + ) + + for expert_idx in range(expert_down.shape[0]): + # W: [hidden, moe_inter] — out_features=hidden, in_features=moe_inter + W = expert_down.data[expert_idx].float() + W_norms = W.norm(dim=1, keepdim=True) # [hidden, 1] + W_dirs = F.normalize(W, dim=1) # Row-normalised. + + # Projection 1: remove refusal component from each column of W + refusal_comp = v @ W_dirs # [moe_inter] + W_dirs = F.normalize( + W_dirs - weight_scale * v.unsqueeze(1) * refusal_comp.unsqueeze(0), + dim=1, + ) + + # Projection 2: biprojection to catch residual leakage + refusal_comp2 = v @ W_dirs + W_dirs = F.normalize( + W_dirs - v.unsqueeze(1) * refusal_comp2.unsqueeze(0), dim=1 + ) + + expert_down.data[expert_idx] = (W_norms * W_dirs).to(expert_down.dtype) + + def _is_diffusion_gemma(self) -> bool: + model = self.model + if isinstance(model, PeftModel): + model = model.base_model.model + return "DiffusionGemmaForBlockDiffusion" in type(model).__name__ + + def _get_dg_encoder(self) -> Module: + """Return the DiffusionGemmaEncoderModel (not the text sub-model).""" + model = self.model + if isinstance(model, PeftModel): + model = model.base_model.model + return model.model.encoder + + def _get_dg_lm_head(self) -> Module: + model = self.model + if isinstance(model, PeftModel): + model = model.base_model.model + return model.lm_head + def get_layers(self) -> ModuleList: model = self.model @@ -367,6 +535,11 @@ def get_layers(self) -> ModuleList: if isinstance(model, PeftModel): model = model.base_model.model + # DiffusionGemma encoder-decoder (encoder layers are abliterated; + # decoder shares weights via tied_weights_keys so it gets them for free). + with suppress(Exception): + return model.model.encoder.language_model.layers + # Most multimodal models. with suppress(Exception): return model.model.language_model.layers @@ -576,10 +749,8 @@ def abliterate( W = W - W_org # Use a low-rank SVD to get an approximation of the matrix. r = self.peft_config.r - # svd_lowrank is randomized: - # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109 - # Reseed immediately before the call so restoring a trial is independent of RNG history. - torch.manual_seed(self.settings.seed) + if self.settings.seed is not None: + torch.manual_seed(self.settings.seed) U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6) # Truncate it to the part we want to store in the LoRA adapter. # Note: svd_lowrank actually returns V, so transpose it to get Vh. @@ -601,11 +772,17 @@ def abliterate( weight_A.data = lora_A.to(weight_A.dtype) weight_B.data = lora_B.to(weight_B.dtype) + # For DiffusionGemma, also abliterate the batched MoE expert parameters in-place. + # LoRA can't wrap nn.Parameter tensors, so EGA modifies them directly. The + # decoder shares these tensors (same data_ptr), so it's abliterated for free. + if self._is_diffusion_gemma(): + self._abliterate_dg_experts(refusal_directions, direction_index, parameters) + def generate( self, prompts: list[Prompt], **kwargs: Any, - ) -> tuple[BatchEncoding, GenerateDecoderOnlyOutput | LongTensor]: + ) -> tuple[BatchEncoding, GenerateDecoderOnlyOutput | LongTensor | Any]: chats = [ [ {"role": "system", "content": prompt.system}, @@ -639,6 +816,28 @@ def generate( return_token_type_ids=False, ).to(self.model.device) + if self._is_diffusion_gemma(): + # DiffusionGemmaGenerationMixin.generate() does not support do_sample, + # pad_token_id, output_hidden_states, output_logits, return_dict_in_generate, + # or use_cache — these are handled via the diffusion sampler config instead. + dg_allowed = { + "max_new_tokens", + "streamer", + "generation_config", + "logits_processor", + "stopping_criteria", + "max_length", + } + dg_kwargs = {k: v for k, v in kwargs.items() if k in dg_allowed} + # Merge LoRA into shared encoder/decoder tensors so the decoder-driven + # diffusion generation sees the current abliteration state. + with self._dg_lora_merged(): + outputs = self.model.generate( # ty:ignore[call-non-callable] + input_ids=cast(Tensor, inputs["input_ids"]), + **dg_kwargs, + ) + return inputs, outputs + # FIXME: The type checker has been disabled here because of the extremely complex # interplay between different generate() signatures and dynamic delegation. outputs = self.model.generate( @@ -660,11 +859,16 @@ def get_responses( max_new_tokens=self.settings.max_response_length, ) + # DiffusionGemmaGenerationOutput stores the full sequence in .sequences. + sequences: LongTensor = ( + outputs.sequences if hasattr(outputs, "sequences") else outputs + ) # ty:ignore[assignment] + return self.tokenizer.batch_decode( # Extract the newly generated part. # This cast is valid because the input_ids property is a Tensor # if the tokenizer is invoked with return_tensors="pt", as above. - outputs[:, cast(Tensor, inputs["input_ids"]).shape[1] :], + sequences[:, cast(Tensor, inputs["input_ids"]).shape[1] :], skip_special_tokens=skip_special_tokens, ) @@ -684,7 +888,101 @@ def get_responses_batched( return responses + def _get_residuals_dg(self, prompts: list[Prompt]) -> Tensor: + """Get per-layer residuals for DiffusionGemma by hooking the encoder layers. + + DiffusionGemma's block-diffusion generate() doesn't expose per-layer hidden + states, so we run a causal encoder forward pass instead and capture each + layer's output via forward hooks. + """ + chats = [ + [ + {"role": "system", "content": p.system}, + {"role": "user", "content": p.user}, + ] + for p in prompts + ] + chat_prompts = cast( + list[str], + self.tokenizer.apply_chat_template( + chats, add_generation_prompt=True, tokenize=False + ), + ) + if self.settings.response_prefix: + chat_prompts = [cp + self.settings.response_prefix for cp in chat_prompts] + + inputs = self.tokenizer( + chat_prompts, return_tensors="pt", padding=True, return_token_type_ids=False + ).to(self.model.device) + + encoder = self._get_dg_encoder() + text_model = encoder.language_model + + # layer_idx -> (batch, hidden) tensor captured at the last prompt token + captured: dict[int, Tensor] = {} + hooks = [] + + # Capture the embed_tokens output as "layer 0" (embedding layer). + def make_embed_hook(): + def hook(module: Module, input: Any, output: Tensor) -> None: + captured[-1] = output.detach() + + return hook + + hooks.append(text_model.embed_tokens.register_forward_hook(make_embed_hook())) + + for idx, layer in enumerate(text_model.layers): + + def make_layer_hook(i: int): + def hook(module: Module, input: Any, output: Tensor) -> None: + # Encoder layers return a plain Tensor, not a tuple. + captured[i] = ( + output.detach() + if isinstance(output, Tensor) + else output[0].detach() + ) + + return hook + + hooks.append(layer.register_forward_hook(make_layer_hook(idx))) + + try: + with torch.no_grad(): + encoder( + input_ids=cast(Tensor, inputs["input_ids"]), + attention_mask=inputs.get("attention_mask"), # type: ignore[arg-type] + ) + finally: + for h in hooks: + h.remove() + + # Build (prompt, layer, component) tensor; heretic convention: index 0 = embeddings. + layer_outputs = [] + if -1 in captured: + layer_outputs.append(captured[-1][:, -1, :]) # Last token position. + for idx in range(len(text_model.layers)): + if idx in captured: + layer_outputs.append(captured[idx][:, -1, :]) + + residuals = torch.stack(layer_outputs, dim=1).to(torch.float32) + + if 0 <= self.settings.winsorization_quantile < 1: + abs_residuals = torch.abs(residuals) + thresholds = torch.quantile( + abs_residuals, self.settings.winsorization_quantile, dim=2, keepdim=True + ) + residuals = torch.clamp(residuals, -thresholds, thresholds) + + if self.settings.offload_outputs_to_cpu: + residuals = residuals.cpu() + empty_cache() + + return residuals + def get_residuals(self, prompts: list[Prompt]) -> Tensor: + if self._is_diffusion_gemma(): + return self._get_residuals_dg(prompts) + # We only generate one token, and we return the residual vectors # at that token position, for each prompt and layer. _, outputs = self.generate( @@ -768,9 +1066,58 @@ def get_residuals_mean(self, prompts: list[Prompt]) -> Tensor: return (running_sum / total_count).to(torch.float32) + def _get_logprobs_dg(self, prompts: list[Prompt]) -> Tensor: + """Get next-token logprobs for DiffusionGemma. + + Runs the encoder in causal mode and applies lm_head to the last-position + hidden state, giving the autoregressive prediction for the first new token. + This is the appropriate signal for KL divergence: it captures how the + abliteration has shifted the encoder's output distribution. + """ + chats = [ + [ + {"role": "system", "content": p.system}, + {"role": "user", "content": p.user}, + ] + for p in prompts + ] + chat_prompts = cast( + list[str], + self.tokenizer.apply_chat_template( + chats, add_generation_prompt=True, tokenize=False + ), + ) + if self.settings.response_prefix: + chat_prompts = [cp + self.settings.response_prefix for cp in chat_prompts] + + inputs = self.tokenizer( + chat_prompts, return_tensors="pt", padding=True, return_token_type_ids=False + ).to(self.model.device) + + with torch.no_grad(): + encoder = self._get_dg_encoder() + enc_out = encoder( + input_ids=cast(Tensor, inputs["input_ids"]), + attention_mask=inputs.get("attention_mask"), # type: ignore[arg-type] + ) + last_hidden = enc_out.last_hidden_state[:, -1, :] # (batch, hidden) + lm_head = self._get_dg_lm_head() + logits = lm_head(last_hidden.to(lm_head.weight.dtype)) # (batch, vocab) + + logprobs = F.log_softmax(logits.float(), dim=-1) + + if self.settings.offload_outputs_to_cpu: + logprobs = logprobs.cpu() + empty_cache() + + return logprobs + # We work with logprobs rather than probabilities for numerical stability # when computing the KL divergence. def get_logprobs(self, prompts: list[Prompt]) -> Tensor: + if self._is_diffusion_gemma(): + return self._get_logprobs_dg(prompts) + # We only generate one token, and we return the (log) probability distributions # over the vocabulary at that token position, for each prompt. _, outputs = self.generate( @@ -834,6 +1181,24 @@ def stream_chat_response(self, chat: list[dict[str, str]]) -> str: skip_special_tokens=True, ) + if self._is_diffusion_gemma(): + with self._dg_lora_merged(): + outputs = self.model.generate( # ty:ignore[call-non-callable] + input_ids=cast(Tensor, inputs["input_ids"]), + streamer=streamer, + max_new_tokens=4096, + ) + sequences: LongTensor = ( + outputs.sequences if hasattr(outputs, "sequences") else outputs + ) # ty:ignore[assignment] + return cast( + str, + self.tokenizer.decode( + sequences[0, inputs["input_ids"].shape[1] :], + skip_special_tokens=True, + ), + ) + # FIXME: The type checker has been disabled here because of the extremely complex # interplay between different generate() signatures and dynamic delegation. outputs = self.model.generate(