Spaces:
Running on Zero
Running on Zero
| """Causal-LM loading adapted for bidirectional denoising.""" | |
| from __future__ import annotations | |
| from collections import Counter, OrderedDict | |
| from typing import Any | |
| import torch | |
| import os | |
| from pathlib import Path | |
| from peft import LoraConfig, PeftModel, TaskType, get_peft_model, prepare_model_for_kbit_training | |
| from transformers import AutoModelForCausalLM | |
| _ATTENTION_MASK_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() | |
| _ATTENTION_MASK_CACHE_SIZE = 4 | |
| TRAINABLE_BASE_STATE_FILENAME = "trainable_base_state.pt" | |
| def bidirectional_attention_mask(padding_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: | |
| """A 4-D additive full-attention mask accepted unchanged by Transformers >=5. | |
| `padding_mask` is retained for loss/padding bookkeeping, but padding is | |
| intentionally visible to attention. Repeated EOS padding is part of the | |
| configured context-width signal: real and padded queries can attend to all | |
| positions, allowing the model to learn concise answers under wide contexts. | |
| """ | |
| # Every position is deliberately visible, so the additive mask is exactly | |
| # zero. Reuse the most recent shapes instead of allocating and clearing the | |
| # same dense tensor on every forward pass. | |
| shape = (padding_mask.shape[0], 1, padding_mask.shape[1], padding_mask.shape[1]) | |
| # Inference tensors cannot later be saved by autograd, so training and | |
| # inference-mode allocations must occupy separate cache entries. | |
| key = ( | |
| padding_mask.device.type, | |
| padding_mask.device.index, | |
| dtype, | |
| torch.is_inference_mode_enabled(), | |
| *shape, | |
| ) | |
| mask = _ATTENTION_MASK_CACHE.get(key) | |
| if mask is None: | |
| mask = torch.zeros(shape, device=padding_mask.device, dtype=dtype) | |
| _ATTENTION_MASK_CACHE[key] = mask | |
| if len(_ATTENTION_MASK_CACHE) > _ATTENTION_MASK_CACHE_SIZE: | |
| _ATTENTION_MASK_CACHE.popitem(last=False) | |
| else: | |
| _ATTENTION_MASK_CACHE.move_to_end(key) | |
| return mask | |
| def _target_modules(model: torch.nn.Module, requested: list[str]) -> list[str]: | |
| """Verify every requested LoRA projection exists in the loaded architecture.""" | |
| available = {name.rsplit(".", 1)[-1] for name, _ in model.named_modules()} | |
| missing = [name for name in requested if name not in available] | |
| if missing: | |
| raise ValueError(f"LoRA target modules missing from {model.config.model_type}: {missing}; available suffixes include {sorted(available)[:40]}") | |
| return requested | |
| def _transformer_layers(model: torch.nn.Module) -> torch.nn.ModuleList: | |
| """Return the ordered decoder blocks from a supported CausalLM/PEFT wrapper.""" | |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model | |
| candidates = ( | |
| ("model", "layers"), | |
| ("transformer", "h"), | |
| ("gpt_neox", "layers"), | |
| ("model", "decoder", "layers"), | |
| ) | |
| for path in candidates: | |
| value: Any = base | |
| for component in path: | |
| value = getattr(value, component, None) | |
| if value is None: | |
| break | |
| if isinstance(value, torch.nn.ModuleList): | |
| return value | |
| raise TypeError( | |
| f"Could not locate ordered transformer layers in {type(base).__name__}; " | |
| "train_last_n_layers is unsupported for this architecture" | |
| ) | |
| def _unfreeze_last_transformer_layers(model: torch.nn.Module, count: int) -> list[int]: | |
| """Unfreeze the final ``count`` complete decoder blocks.""" | |
| if count < 0: | |
| raise ValueError("train_last_n_layers must be non-negative") | |
| if count == 0: | |
| return [] | |
| layers = _transformer_layers(model) | |
| if count > len(layers): | |
| raise ValueError( | |
| f"train_last_n_layers={count} exceeds the model's {len(layers)} transformer layers" | |
| ) | |
| indices = list(range(len(layers) - count, len(layers))) | |
| for index in indices: | |
| for parameter in layers[index].parameters(): | |
| parameter.requires_grad = True | |
| return indices | |
| def _unfreeze_output_head(model: torch.nn.Module, enabled: bool) -> int: | |
| """Optionally unfreeze the untied language-model output projection.""" | |
| if not enabled: | |
| return 0 | |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model | |
| if not hasattr(base, "get_output_embeddings"): | |
| raise TypeError(f"{type(base).__name__} does not expose an output head") | |
| output_head = base.get_output_embeddings() | |
| input_embeddings = base.get_input_embeddings() if hasattr(base, "get_input_embeddings") else None | |
| if output_head is None: | |
| raise TypeError(f"{type(base).__name__} has no output head") | |
| input_parameter_ids = {id(parameter) for parameter in input_embeddings.parameters()} if input_embeddings else set() | |
| output_parameters = list(output_head.parameters()) | |
| if any(id(parameter) in input_parameter_ids for parameter in output_parameters): | |
| raise ValueError("train_lm_head requires untied input and output embeddings") | |
| for parameter in output_parameters: | |
| parameter.requires_grad = True | |
| return sum(parameter.numel() for parameter in output_parameters) | |
| def load_trainable_base_state(model: torch.nn.Module, adapter_path: str | Path) -> int: | |
| """Restore separately saved, non-LoRA trainable base parameters.""" | |
| state_path = Path(adapter_path) / TRAINABLE_BASE_STATE_FILENAME | |
| if not state_path.is_file(): | |
| return 0 | |
| state = torch.load(state_path, map_location="cpu", weights_only=True) | |
| parameters = dict(model.named_parameters()) | |
| missing = sorted(name for name in state if name not in parameters) | |
| mismatched = sorted( | |
| name for name, value in state.items() | |
| if name in parameters and parameters[name].shape != value.shape | |
| ) | |
| if missing or mismatched: | |
| raise ValueError( | |
| "The saved trainable base state does not match the adapter/base model: " | |
| f"missing={missing[:5]}, shape_mismatch={mismatched[:5]}" | |
| ) | |
| for name, value in state.items(): | |
| parameter = parameters[name] | |
| parameter.data.copy_(value.to(parameter.device, dtype=parameter.dtype)) | |
| return len(state) | |
| def parameter_audit( | |
| model: torch.nn.Module, | |
| train_last_n_layers: int = 0, | |
| train_lm_head: bool = False, | |
| full_finetuning: bool = False, | |
| ) -> dict[str, Any]: | |
| """Assert the intended trainable set and return parameter-count diagnostics.""" | |
| named = list(model.named_parameters()) | |
| trainable = [(name, p) for name, p in named if p.requires_grad] | |
| total = sum(p.numel() for _, p in named) | |
| if full_finetuning: | |
| frozen = [name for name, parameter in named if not parameter.requires_grad] | |
| if frozen: | |
| raise AssertionError({"unexpected_frozen_full_model_parameters": frozen[:20]}) | |
| return { | |
| "full_finetuning": True, | |
| "full_model_parameters": total, | |
| "lora_parameters": 0, | |
| "normalization_parameters": 0, | |
| "transformer_layer_parameters": 0, | |
| "lm_head_parameters": 0, | |
| "other_trainable_parameters": 0, | |
| "train_last_n_layers": 0, | |
| "train_lm_head": False, | |
| "trainable_transformer_layer_indices": [], | |
| "total_trainable_parameters": total, | |
| "total_model_parameters": total, | |
| "trainable_percentage": 100.0, | |
| "trainable_names": [name for name, _ in trainable], | |
| } | |
| selected_layers = list(_transformer_layers(model)[-train_last_n_layers:]) if train_last_n_layers else [] | |
| parameter_names = {id(parameter): name for name, parameter in named} | |
| selected_parameter_ids = { | |
| id(parameter) | |
| for layer in selected_layers | |
| for parameter in layer.parameters() | |
| if "lora_" not in parameter_names.get(id(parameter), "") | |
| } | |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model | |
| output_head = base.get_output_embeddings() if hasattr(base, "get_output_embeddings") else None | |
| output_head_parameter_ids = {id(parameter) for parameter in output_head.parameters()} if output_head else set() | |
| categories = Counter() | |
| unexpected = [] | |
| for name, p in trainable: | |
| if "lora_" in name: | |
| categories["lora"] += p.numel() | |
| elif id(p) in output_head_parameter_ids and train_lm_head: | |
| categories["lm_head"] += p.numel() | |
| elif id(p) in selected_parameter_ids: | |
| categories["transformer_layers"] += p.numel() | |
| elif "norm" in name.lower(): | |
| categories["norm"] += p.numel() | |
| else: | |
| categories["other"] += p.numel() | |
| unexpected.append(name) | |
| frozen_embedding = all(not p.requires_grad for n, p in named if any(x in n.lower() for x in ("embed_tokens", "embed_tokens", "wte"))) | |
| # An output head targeted by LoRA contains trainable adapter parameters even | |
| # when its original projection remains frozen. Only enforce the requested | |
| # state on non-LoRA parameters belonging to the head. | |
| non_lora_output_parameters = [ | |
| parameter | |
| for name, parameter in named | |
| if id(parameter) in output_head_parameter_ids and "lora_" not in name | |
| ] | |
| lm_head_trainability_correct = bool(non_lora_output_parameters) and all( | |
| parameter.requires_grad == train_lm_head | |
| for parameter in non_lora_output_parameters | |
| ) | |
| if not frozen_embedding or not lm_head_trainability_correct or unexpected: | |
| raise AssertionError({"embeddings_frozen": frozen_embedding, "lm_head_trainability_correct": lm_head_trainability_correct, "unexpected_trainable": unexpected}) | |
| frozen_selected = [ | |
| name for name, parameter in named | |
| if id(parameter) in selected_parameter_ids and not parameter.requires_grad | |
| ] | |
| if frozen_selected: | |
| raise AssertionError({"unexpected_frozen_last_layer_parameters": frozen_selected[:20]}) | |
| trainable_count = sum(p.numel() for _, p in trainable) | |
| layer_count = len(_transformer_layers(model)) if train_last_n_layers else 0 | |
| return {"lora_parameters": categories["lora"], "normalization_parameters": categories["norm"], "transformer_layer_parameters": categories["transformer_layers"], "lm_head_parameters": categories["lm_head"], "other_trainable_parameters": categories["other"], "train_last_n_layers": train_last_n_layers, "train_lm_head": train_lm_head, "trainable_transformer_layer_indices": list(range(layer_count - train_last_n_layers, layer_count)) if train_last_n_layers else [], "total_trainable_parameters": trainable_count, "total_model_parameters": total, "trainable_percentage": 100 * trainable_count / total, "trainable_names": [n for n, _ in trainable]} | |
| def load_denoising_model(config: dict[str, Any]) -> tuple[torch.nn.Module, dict[str, Any]]: | |
| """Load either a full model or a LoRA-adapted CausalLM and audit it.""" | |
| checkpoint = config["model_name_or_path"] | |
| precision = config.get("precision", "bf16") | |
| dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}.get(precision) | |
| if dtype is None: | |
| raise ValueError(f"Unknown precision {precision}") | |
| # No device_map: accelerate owns device placement in distributed runs. | |
| model_cache = Path(config.get("base_model_cache_dir", "base_models")); model_cache.mkdir(parents=True, exist_ok=True) | |
| quantization = str(config.get("quantization", "none")).lower() | |
| full_finetuning = bool(config.get("full_finetuning", False)) | |
| train_last_n_layers = int(config.get("train_last_n_layers", 0) or 0) | |
| train_lm_head = bool(config.get("train_lm_head", False)) | |
| if train_last_n_layers < 0: | |
| raise ValueError("train_last_n_layers must be non-negative") | |
| if (full_finetuning or train_last_n_layers or train_lm_head) and quantization not in {"none", "off", "false"}: | |
| raise ValueError("training full base-model layers or the LM head requires an unquantized base model") | |
| if full_finetuning and config.get("resume_from_adapter"): | |
| raise ValueError("full_finetuning starts from model_name_or_path and cannot use resume_from_adapter") | |
| if full_finetuning and (train_last_n_layers or train_lm_head): | |
| raise ValueError("train_last_n_layers/train_lm_head are redundant with full_finetuning") | |
| load_kwargs = dict(dtype=dtype, trust_remote_code=False, token=os.getenv("HF_TOKEN"), cache_dir=str(model_cache)) | |
| if quantization in {"4bit", "4-bit", "qlora"}: | |
| try: | |
| from transformers import BitsAndBytesConfig | |
| import bitsandbytes # noqa: F401 | |
| except ImportError as exc: | |
| raise ImportError("quantization=4bit requires CUDA bitsandbytes; install with `pip install -e '.[cuda]'`") from exc | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("4-bit bitsandbytes quantization requires an NVIDIA CUDA device") | |
| compute_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}.get(str(config.get("compute_dtype", precision)), dtype) | |
| load_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=str(config.get("quantization_type", "nf4")), bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=bool(config.get("double_quant", True))) | |
| elif quantization not in {"none", "off", "false"}: | |
| raise ValueError("quantization must be 'none' or '4bit'") | |
| model = AutoModelForCausalLM.from_pretrained(checkpoint, **load_kwargs) | |
| model.config.use_cache = False | |
| # PEFT's CAUSAL_LM task type only describes adapter integration; it does | |
| # not control attention direction. Make the base-model intent explicit as | |
| # well as supplying the prepared 4-D mask in forward_bidirectional(). | |
| model.config.is_causal = False | |
| if hasattr(model.config, "use_bidirectional_attention"): | |
| model.config.use_bidirectional_attention = True | |
| if full_finetuning: | |
| for parameter in model.parameters(): | |
| parameter.requires_grad = True | |
| if config.get("gradient_checkpointing", False): | |
| model.gradient_checkpointing_enable() | |
| model.enable_input_require_grads() | |
| audit = parameter_audit(model, full_finetuning=True) | |
| audit.update({"model_name": checkpoint, "resolved_lora_targets": []}) | |
| return model, audit | |
| if quantization in {"4bit", "4-bit", "qlora"}: | |
| model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=bool(config.get("gradient_checkpointing", False))) | |
| for parameter in model.parameters(): | |
| parameter.requires_grad = False | |
| targets = _target_modules(model, list(config.get("lora_targets", ["q_proj", "v_proj", "o_proj"]))) | |
| lora_config = LoraConfig(r=int(config.get("lora_r", 16)), lora_alpha=int(config.get("lora_alpha", 32)), lora_dropout=float(config.get("lora_dropout", 0.05)), target_modules=targets, bias="none", task_type=TaskType.CAUSAL_LM) | |
| resume_adapter = config.get("resume_from_adapter") | |
| if resume_adapter: | |
| adapter_path = Path(resume_adapter) | |
| if not (adapter_path / "adapter_config.json").is_file(): | |
| raise ValueError(f"resume_from_adapter is not a saved adapter directory: {adapter_path}") | |
| model = PeftModel.from_pretrained(model, adapter_path, is_trainable=True) | |
| load_trainable_base_state(model, adapter_path) | |
| norm_state_path = adapter_path / "normalization_state.pt" | |
| if norm_state_path.is_file(): | |
| norm_state = torch.load(norm_state_path, map_location="cpu", weights_only=True) | |
| named = dict(model.named_parameters()) | |
| for name, value in norm_state.items(): | |
| if name in named: | |
| named[name].data.copy_(value.to(named[name].device, dtype=named[name].dtype)) | |
| else: | |
| model = get_peft_model(model, lora_config) | |
| if config.get("train_normalization_layers", True): | |
| for name, parameter in model.named_parameters(): | |
| if "norm" in name.lower(): | |
| parameter.requires_grad = True | |
| # Accelerate's FP16 GradScaler cannot unscale FP16 gradients. | |
| # Keep trainable normalization parameters in FP32 so their | |
| # gradients are scaler-compatible; frozen base weights remain | |
| # in the configured compute dtype. | |
| if precision == "fp16" and parameter.dtype == torch.float16: | |
| parameter.data = parameter.data.float() | |
| trained_layer_indices = _unfreeze_last_transformer_layers(model, train_last_n_layers) | |
| _unfreeze_output_head(model, train_lm_head) | |
| if config.get("gradient_checkpointing", False): | |
| model.gradient_checkpointing_enable() | |
| model.enable_input_require_grads() | |
| audit = parameter_audit(model, train_last_n_layers=train_last_n_layers, train_lm_head=train_lm_head) | |
| audit.update({"model_name": checkpoint, "resolved_lora_targets": targets, "trainable_transformer_layer_indices": trained_layer_indices}) | |
| return model, audit | |
| def forward_bidirectional(model: torch.nn.Module, input_ids: torch.Tensor, padding_mask: torch.Tensor): | |
| """Run a CausalLM with the project’s explicit bidirectional padding mask.""" | |
| # 4-bit bitsandbytes weights are stored as uint8, which cannot represent | |
| # the floating additive attention mask. Use the first floating parameter | |
| # (normally a LoRA or normalization parameter) as the compute dtype. | |
| dtype = getattr(model, "_lad_attention_mask_dtype", None) | |
| if dtype is None: | |
| dtype = next((parameter.dtype for parameter in model.parameters() if parameter.is_floating_point()), torch.float32) | |
| model._lad_attention_mask_dtype = dtype | |
| return model(input_ids=input_ids, attention_mask=bidirectional_attention_mask(padding_mask, dtype), use_cache=False).logits | |
| def forward_bidirectional_selected( | |
| model: torch.nn.Module, | |
| input_ids: torch.Tensor, | |
| padding_mask: torch.Tensor, | |
| selection_mask: torch.Tensor, | |
| ): | |
| """Run the frozen LM head only at positions participating in the objective.""" | |
| dtype = getattr(model, "_lad_attention_mask_dtype", None) | |
| if dtype is None: | |
| dtype = next((parameter.dtype for parameter in model.parameters() if parameter.is_floating_point()), torch.float32) | |
| model._lad_attention_mask_dtype = dtype | |
| unwrapped = model.module if hasattr(model, "module") else model | |
| causal_lm = unwrapped.get_base_model() if hasattr(unwrapped, "get_base_model") else unwrapped | |
| backbone = getattr(causal_lm, "model", None) | |
| output_head = causal_lm.get_output_embeddings() if hasattr(causal_lm, "get_output_embeddings") else None | |
| if backbone is None or output_head is None: | |
| raise TypeError(f"Selected-logit optimization is unsupported for {type(causal_lm).__name__}") | |
| outputs = backbone( | |
| input_ids=input_ids, | |
| attention_mask=bidirectional_attention_mask(padding_mask, dtype), | |
| use_cache=False, | |
| ) | |
| example_ids, token_ids = selection_mask.nonzero(as_tuple=True) | |
| selected_logits = output_head(outputs.last_hidden_state[example_ids, token_ids]) | |
| return selected_logits, example_ids, token_ids | |