Module neuralops.core
Efficient Transformer Toolkit – Core module. Provides model‑agnostic architecture detection, attention injection, compression pipeline, and base classes.
Sub-modules
neuralops.core.adapters-
Model‑family adapters. Extracts projection weights and biases cleanly across all Transformer architectures.
neuralops.core.architecture-
Model‑agnostic architecture helpers. Provides unified access to layers, attention, FFN, embeddings, classifiers across all common HuggingFace …
neuralops.core.attention_mask-
Universal attention mask handler for standard and efficient attention modules …
neuralops.core.attention_utils-
Unified Attention Injection & Weight Adaptation Engine for Foundation Models …
neuralops.core.base-
Base classes for attention modules and compressors.
neuralops.core.common-
Unified shared primitives and utility layers for attention and compression modules. Centralizes repeatedly copied helpers (_cfg_get, repeat_kv, …
neuralops.core.compression_engine-
Compression pipeline – orchestrates multiple compressors sequentially.
neuralops.core.config-
Simple YAML config loader.
neuralops.core.gqa_utils-
Shared Grouped-Query Attention (GQA) & Multi-Head Attention (MHA) Utilities. Provides dynamic dimension extraction, head repetition, and reshaping for …
neuralops.core.legacy_ensemble-
DEPRECATED — Legacy ensemble modules …
neuralops.core.mask_utils-
core/mask_utils.py — Unified mask handling for all attention and compression modules. Handles 2D/4D, bool/float, additive/multiplicative masks …
neuralops.core.memory-
Universal GPU/MPS/CPU Memory Management & Optimization Utilities. Ensures 100% clean memory deallocation between training steps, epochs, and …
neuralops.core.metrics-
core/metrics.py — Dataset-specific metric computation for NeurIPS/ICLR publication. Supports: Accuracy, F1, MCC (CoLA), Pearson/Spearman (STS-B), …
neuralops.core.model_utils-
Generic helpers: calibration, weight‑copying, safe module replacement, activation collection.
neuralops.core.model_wrapper-
UnifiedTransformerModel – wraps any HuggingFace model with Hub support, attention replacement, and config serialization …
neuralops.core.multi_mode_benchmark-
Unified Multi-Mode Evaluation Runner for Efficient Transformer Toolkit …
neuralops.core.pruning_utils-
Pruning Utilities and Saliency Estimation for Structural Pruning and Layer Shedding …
neuralops.core.registry-
Central registry system for the Efficient Transformer Toolkit. All framework components (attentions, compressors, trainers, visualizers, etc.) are …
neuralops.core.replace_attention-
Model Attention Replacement and Dynamic Graph Surgery Utilities …
neuralops.core.snapshot_restore-
Model State Snapshot & Exact Bit-Identical Restorer …
neuralops.core.svd_utils-
Centralized SVD & Factorization Utilities for Attention and Compression Modules. Provides numerical stability (FP32 execution for half-precision …
neuralops.core.t5_compat-
T5 Architecture Compatibility Module. Provides robust support for T5 relative position bias, encoder-decoder cross-attention, and numerically stable …
neuralops.core.triton_kernels-
Shared Triton Kernels for Attention and Linear Architectures …
neuralops.core.utils-
General utilities (device, seeding, formatting).
Functions
def _apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]-
Expand source code
def _apply_rope( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """ Applies Rotary Position Embedding (RoPE) to Query and Key tensors. """ if cos.dim() < q.dim(): cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) def rotate_half(x: torch.Tensor) -> torch.Tensor: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed.type_as(q), k_embed.type_as(k)Applies Rotary Position Embedding (RoPE) to Query and Key tensors.
def _cfg_get(config: Any, key: str, default: Any = None) ‑> Any-
Expand source code
def _cfg_get(config: Any, key: str, default: Any = None) -> Any: """ Safely retrieves a configuration attribute from a dict, object, or None. Handles nested fallbacks without throwing AttributeErrors. """ if config is None: return default if isinstance(config, dict): return config.get(key, default) val = getattr(config, key, None) return val if val is not None else defaultSafely retrieves a configuration attribute from a dict, object, or None. Handles nested fallbacks without throwing AttributeErrors.
def _detect_decoder(config: Any = None,
is_decoder: bool | None = None,
old_attn: nn.Module | None = None) ‑> bool-
Expand source code
def _detect_decoder( config: Any = None, is_decoder: bool | None = None, old_attn: nn.Module | None = None, ) -> bool: """ Detects whether an attention layer should operate in causal decoder mode. """ if is_decoder is not None: return bool(is_decoder) if config is not None: if getattr(config, "is_decoder", False) or getattr(config, "causal", False): return True mt = str(getattr(config, "model_type", "")).lower() if mt.startswith( ( "gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt", "bloom", "falcon", "olmo", "mpt", "stablelm", ) ): return True arch = getattr(config, "architectures", None) if arch and any(("GPT" in str(a) or "CausalLM" in str(a)) for a in arch): return True if old_attn is not None: if getattr(old_attn, "is_decoder", False): return True if hasattr(old_attn, "c_attn"): return True if "gpt" in old_attn.__class__.__name__.lower(): return True return FalseDetects whether an attention layer should operate in causal decoder mode.
def activation_hooks(model: torch.nn.modules.module.Module,
target_module_type: type = torch.nn.modules.linear.Linear,
collect_input: bool = True,
collect_output: bool = False) ‑> Any-
Expand source code
@contextmanager def activation_hooks( model: nn.Module, target_module_type: type = nn.Linear, collect_input: bool = True, collect_output: bool = False, ) -> Any: """Context manager that collects activations for all modules of a given type.""" activations: dict[str, Any] = {} hooks = [] def make_hook(name) -> Any: """Make hook.""" def hook(module, input, output) -> None: """Hook.""" data = {} if collect_input: data["input"] = input[0].detach() if collect_output: data["output"] = ( output.detach() if isinstance(output, torch.Tensor) else output[0].detach() ) activations[name] = data return hook for name, module in model.named_modules(): if isinstance(module, target_module_type): hooks.append(module.register_forward_hook(make_hook(name))) try: yield activations finally: for hook in hooks: hook.remove()Context manager that collects activations for all modules of a given type.
def apply_attention_mask(scores: torch.Tensor,
attention_mask: torch.Tensor | None,
mask_value: float = -10000.0) ‑> torch.Tensor-
Expand source code
def apply_attention_mask( scores: torch.Tensor, attention_mask: torch.Tensor | None, mask_value: float = -10000.0, ) -> torch.Tensor: r"""Apply attention mask to attention logits in a numerically robust and universal manner. Args: scores (torch.Tensor): Attention logits of shape `[Batch, Heads, SeqLen_Q, SeqLen_K]`. attention_mask (Optional[torch.Tensor]): Mask tensor in arbitrary standard format: - None: No-op, returns `scores` directly. - 2D `[B, N_k]`: Binary (1=keep, 0=mask) or boolean. - 3D `[B, 1, N_k]` or `[B, N_q, N_k]`. - 4D `[B, 1, 1, N_k]` or `[B, 1, N_q, N_k]`: Pre-computed additive float mask or binary. mask_value (float): Large negative value to assign to masked positions. Returns: torch.Tensor: Masked attention logits of same shape and dtype as `scores`. """ if attention_mask is None: return scores if not isinstance(attention_mask, torch.Tensor): return scores dtype = scores.dtype device = scores.device # Ensure mask is on the same device mask = attention_mask.to(device=device) # Case 1: 2D binary / boolean mask [B, N_k] if mask.dim() == 2: if mask.dtype == torch.bool: # True = keep, False = mask additive = torch.zeros_like(mask, dtype=dtype) additive = additive.masked_fill(~mask, mask_value) elif torch.is_floating_point(mask): # Check if it's already negative additive or 0/1 binary if (mask < 0).any(): # Already additive additive = mask.to(dtype=dtype) else: # 1.0 = keep, 0.0 = mask additive = (1.0 - mask.to(dtype=dtype)) * mask_value else: # Integer binary mask (1=keep, 0=mask) additive = (1.0 - mask.to(dtype=dtype)) * mask_value # Broadcast to [B, 1, 1, N_k] return scores + additive.unsqueeze(1).unsqueeze(2) # Case 2: 3D mask [B, 1, N_k] or [B, N_q, N_k] elif mask.dim() == 3: if mask.dtype == torch.bool: additive = torch.zeros_like(mask, dtype=dtype) additive = additive.masked_fill(~mask, mask_value) return scores + additive.unsqueeze(1) elif torch.is_floating_point(mask) and (mask < 0).any(): return scores + mask.unsqueeze(1).to(dtype=dtype) else: additive = (1.0 - mask.to(dtype=dtype)) * mask_value return scores + additive.unsqueeze(1) # Case 3: 4D mask [B, 1, 1, N_k] or [B, 1, N_q, N_k] or [B, H, N_q, N_k] elif mask.dim() == 4: if mask.dtype == torch.bool: additive = torch.zeros_like(mask, dtype=dtype) additive = additive.masked_fill(~mask, mask_value) return scores + additive elif torch.is_floating_point(mask): # HuggingFace extended attention mask already contains 0.0 for unmasked and -10000.0 for masked if (mask < 0).any(): return scores + mask.to(dtype=dtype) else: # 4D binary mask return scores + (1.0 - mask.to(dtype=dtype)) * mask_value else: # 4D integer mask return scores + (1.0 - mask.to(dtype=dtype)) * mask_value return scoresApply attention mask to attention logits in a numerically robust and universal manner.
- Args
- -----=
scores:torch.Tensor- Attention logits of shape
[Batch, Heads, SeqLen_Q, SeqLen_K]. attention_mask:Optional[torch.Tensor]- Mask tensor in arbitrary standard format:
- None: No-op, returns
scoresdirectly. - 2D[B, N_k]: Binary (1=keep, 0=mask) or boolean. - 3D[B, 1, N_k]or[B, N_q, N_k]. - 4D[B, 1, 1, N_k]or[B, 1, N_q, N_k]: Pre-computed additive float mask or binary. mask_value:float- Large negative value to assign to masked positions.
- Returns
- -----=
torch.Tensor- Masked attention logits of same shape and dtype as
scores.
def available_attention_types() ‑> list[str]-
Expand source code
def available_attention_types() -> list[str]: """Return all registered attention mechanism identifiers.""" return sorted(list(ATTENTION_REGISTRY.keys()))Return all registered attention mechanism identifiers.
def cleanup_gpu_memory(force: bool = True) ‑> None-
Expand source code
def cleanup_gpu_memory(force: bool = True) -> None: """ Perform exhaustive GPU / MPS / CPU memory deallocation. Frees all cached blocks, calls cyclic garbage collector, and flushes allocators. """ # 1. Python garbage collection gc.collect() # 2. CUDA Memory Cleanup if torch.cuda.is_available(): torch.cuda.empty_cache() if hasattr(torch.cuda, "ipc_collect"): torch.cuda.ipc_collect() if force and hasattr(torch.cuda, "reset_peak_memory_stats"): torch.cuda.reset_peak_memory_stats() # 3. Apple Silicon (MPS) Memory Cleanup if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): try: torch.mps.empty_cache() except ( AttributeError, KeyError, RuntimeError, TypeError, ValueError, OSError, ImportError, ): passPerform exhaustive GPU / MPS / CPU memory deallocation. Frees all cached blocks, calls cyclic garbage collector, and flushes allocators.
def compress_model(model: torch.nn.modules.module.Module,
config: dict[str, typing.Any],
tokenizer=None) ‑> torch.nn.modules.module.Module-
Expand source code
def compress_model( model: nn.Module, config: dict[str, Any], tokenizer=None, ) -> nn.Module: """Convenience function to compress a model with a configuration dict.""" pipeline = CompressionPipeline(model, config, tokenizer) return pipeline.run()Convenience function to compress a model with a configuration dict.
def copy_attention_weights(src_module: torch.nn.modules.module.Module,
dst_module: torch.nn.modules.module.Module,
source_map: dict[str, list[str]] | None = None) ‑> None-
Expand source code
def copy_attention_weights( src_module: nn.Module, dst_module: nn.Module, source_map: dict[str, list[str]] | None = None, ) -> None: """ Copy pretrained attention weights from src to dst. Tries multiple possible attribute names for each projection. """ if source_map is None: source_map = ATTENTION_WEIGHT_MAP for target_attr, src_attrs in source_map.items(): src = None for attr in src_attrs: if hasattr(src_module, attr): src = getattr(src_module, attr) break if src is None: raise ValueError( f"Cannot find source for '{target_attr}' in " f"{src_module.__class__.__name__}. Tried: {src_attrs}" ) dst = getattr(dst_module, target_attr) dst.weight.data.copy_(src.weight.data) if src.bias is not None and dst.bias is not None: dst.bias.data.copy_(src.bias.data) elif src.bias is None and dst.bias is not None: nn.init.zeros_(dst.bias)Copy pretrained attention weights from src to dst. Tries multiple possible attribute names for each projection.
def count_params(model: torch.nn.modules.module.Module, trainable_only: bool = False) ‑> int-
Expand source code
def count_parameters(model: nn.Module, trainable_only: bool = False) -> int: """Count parameters.""" if trainable_only: return sum(p.numel() for p in model.parameters() if p.requires_grad) return sum(p.numel() for p in model.parameters())Count parameters.
def create_band_mask(seq_len: int, window_size: int, device: torch.device | None = None) ‑> torch.Tensor-
Expand source code
def create_band_mask( seq_len: int, window_size: int, device: torch.device | None = None ) -> torch.Tensor: r"""Create a sliding window band attention mask [SeqLen, SeqLen] where positions |i - j| <= window_size are True.""" idx = torch.arange(seq_len, device=device) diff = (idx.unsqueeze(0) - idx.unsqueeze(1)).abs() return diff <= window_sizeCreate a sliding window band attention mask [SeqLen, SeqLen] where positions |i - j| <= window_size are True.
def create_causal_mask(seq_len: int, device: torch.device | None = None) ‑> torch.Tensor-
Expand source code
def create_causal_mask(seq_len: int, device: torch.device | None = None) -> torch.Tensor: r"""Create a lower-triangular causal attention mask [SeqLen, SeqLen] where True = attend, False = mask.""" mask = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=device)) return maskCreate a lower-triangular causal attention mask [SeqLen, SeqLen] where True = attend, False = mask.
def detect_architecture(model: torch.nn.modules.module.Module) ‑> str-
Expand source code
def detect_architecture(model: nn.Module) -> str: """Return canonical architecture name.""" config = getattr(model, "config", None) model_type = getattr(config, "model_type", "") if config else "" name = model.__class__.__name__.lower() if "distilbert" in model_type or "distilbert" in name: return "distilbert" if "roberta" in model_type or "roberta" in name: return "roberta" if "bert" in model_type or "bert" in name: return "bert" if "gpt2" in model_type or "gpt2" in name: return "gpt2" if "qwen" in model_type or "qwen" in name: return "qwen" if "llama" in model_type: return "llama" if "mistral" in model_type: return "mistral" if "gemma" in model_type: return "gemma" if "phi" in model_type: return "phi" if "t5" in model_type: return "t5" return "unknown"Return canonical architecture name.
def format_number(num: int) ‑> str-
Expand source code
def format_number(num: int) -> str: """Format number.""" if num >= 1_000_000: return f"{num / 1_000_000:.1f}M" if num >= 1_000: return f"{num / 1_000:.0f}K" return str(num)Format number.
def get_adapter(arch: str) ‑> Any-
Expand source code
def get_adapter(arch: str) -> Any: """Return universal adapter for any architecture.""" return UniversalAdapterReturn universal adapter for any architecture.
def get_attention_modules(model: torch.nn.modules.module.Module) ‑> list[tuple[str, torch.nn.modules.module.Module]]-
Expand source code
def get_attention_modules(model: nn.Module) -> list[tuple[str, nn.Module]]: """Return (name, module) for all leaf attention sub-modules across any transformer architecture. Applies strict path-prefix deduplication so parent composite containers are never returned alongside their inner leaf attention blocks, preventing accidental deletion of LayerNorm, output projections, and residual streams. """ raw_candidates: list[tuple[str, nn.Module]] = [] seen_names = set() for name, module in model.named_modules(): if not name: continue if name in seen_names: continue class_name = module.__class__.__name__.lower() if "universalattentionwrapper" in class_name or "adaptivespan" in class_name: continue if name.endswith(".inner") or name.endswith("._wrapped_attention"): continue is_output_submodule = ( name.endswith(".output") or name.endswith("_output") or "selfoutput" in class_name ) if is_output_submodule: continue # Heuristic 1: Module class name indicates attention or linear-recurrent block if any( x in class_name for x in [ "attention", "selfattn", "selfattention", "crossattn", "deltanet", "gated_deltanet", "gatedattention", ] ): raw_candidates.append((name, module)) seen_names.add(name) continue # Heuristic 2: Attribute-based detection (modules holding projections) if ( hasattr(module, "q_proj") or hasattr(module, "k_proj") or hasattr(module, "q_lin") or hasattr(module, "query") or hasattr(module, "c_attn") or hasattr(module, "qk_proj") ): # Avoid leaf nn.Linear layers that might happen to match attributes if not isinstance(module, (nn.Linear, nn.LayerNorm, nn.Dropout, nn.Embedding)): raw_candidates.append((name, module)) seen_names.add(name) # Apply strict path-prefix deduplication: # If module A's path is a strict prefix of module B's path (e.g. A="...attention", B="...attention.self"), # A is an outer container holding B. We keep only the deepest/leaf modules (B) and drop A. all_cand_names = [name for name, _ in raw_candidates] leaf_modules: list[tuple[str, nn.Module]] = [] for name, module in raw_candidates: is_parent = any(other != name and other.startswith(name + ".") for other in all_cand_names) if not is_parent: leaf_modules.append((name, module)) # Expected count sanity validation based on model configuration cfg = getattr(model, "config", None) if cfg is not None: enc_layers = getattr( cfg, "num_hidden_layers", getattr(cfg, "num_layers", getattr(cfg, "n_layer", None)), ) is_enc_dec = getattr(cfg, "is_encoder_decoder", False) dec_layers = getattr(cfg, "num_decoder_layers", enc_layers) if is_enc_dec else 0 if enc_layers is not None: expected_min = (enc_layers + 2 * dec_layers) if is_enc_dec else enc_layers if len(leaf_modules) != expected_min: logger.warning( "get_attention_modules detected %d modules, expected approximately %d " "(layers=%s, is_encoder_decoder=%s). Detected: %s", len(leaf_modules), expected_min, enc_layers, is_enc_dec, [n for n, _ in leaf_modules], ) return leaf_modulesReturn (name, module) for all leaf attention sub-modules across any transformer architecture.
Applies strict path-prefix deduplication so parent composite containers are never returned alongside their inner leaf attention blocks, preventing accidental deletion of LayerNorm, output projections, and residual streams.
def get_calibration_tensors(model: nn.Module,
tokenizer: Any | None = None,
calibration_loader: Any | None = None,
calibration_texts: list[str] | None = None,
device: torch.device | None = None,
sample_size: int = 8,
seq_len: int = 32,
hidden_dim: int | None = None) ‑> Any-
Expand source code
def get_calibration_tensors( model: nn.Module, tokenizer: Any | None = None, calibration_loader: Any | None = None, calibration_texts: list[str] | None = None, device: torch.device | None = None, sample_size: int = 8, seq_len: int = 32, hidden_dim: int | None = None, ) -> Any: """ Generate or extract calibration inputs for a model. Handles tokenizers, DataLoaders, and synthetic tensor fallbacks. """ if device is None: try: device = next(model.parameters()).device except StopIteration: device = torch.device("cpu") if calibration_loader is not None: for batch in calibration_loader: if isinstance(batch, dict): return {k: v.to(device) if hasattr(v, "to") else v for k, v in batch.items()} elif isinstance(batch, (tuple, list)): return [b.to(device) if hasattr(b, "to") else b for b in batch] elif isinstance(batch, torch.Tensor): return batch.to(device) if tokenizer is not None: texts = list(calibration_texts or DEFAULT_CALIB_TEXTS) try: # H4: Do not artificially repeat texts; clamp sample_size to available unique texts actual_samples = min(sample_size, len(texts)) texts = texts[:actual_samples] encoded = tokenizer( texts, return_tensors="pt", padding=True, truncation=True, max_length=seq_len, ) return {k: v.to(device) for k, v in encoded.items()} except ( AttributeError, KeyError, RuntimeError, TypeError, ValueError, OSError, ImportError, ): pass # Synthetic fallback if hidden_dim is None: for m in model.modules(): if isinstance(m, nn.Linear): hidden_dim = m.in_features break if hidden_dim is None: hidden_dim = 64 # Try input_ids (for language models) or continuous hidden states (for sequential/vision) v_size = None for m in model.modules(): if isinstance(m, nn.Embedding): v_size = m.num_embeddings break if v_size is None: v_size = getattr(getattr(model, "config", None), "vocab_size", 1000) low_idx = min(1, max(0, v_size // 4)) high_idx = max(low_idx + 1, v_size - 1) input_ids = torch.randint( low_idx, high_idx, (sample_size, seq_len), device=device, dtype=torch.long ) eos_token_id = getattr(getattr(model, "config", None), "eos_token_id", 1) or 1 input_ids[:, -1] = eos_token_id return { "input_ids": input_ids, "attention_mask": torch.ones((sample_size, seq_len), device=device, dtype=torch.long), }Generate or extract calibration inputs for a model. Handles tokenizers, DataLoaders, and synthetic tensor fallbacks.
def get_classifier_head(model: torch.nn.modules.module.Module) ‑> torch.nn.modules.module.Module | None-
Expand source code
def get_classifier_head(model: nn.Module) -> nn.Module | None: """Return the final classification layer.""" if hasattr(model, "classifier"): return model.classifier if hasattr(model, "score"): return model.score if hasattr(model, "pre_classifier"): return model.pre_classifier, model.classifier return NoneReturn the final classification layer.
def get_device(device_str: str = 'auto') ‑> torch.device-
Expand source code
def get_device(device_str: str = "auto") -> torch.device: """Return the device.""" if device_str == "auto": if torch.cuda.is_available(): return torch.device("cuda") if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") device = torch.device(device_str) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA was requested, but no CUDA device is available") if device.type == "mps" and ( not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available() ): raise RuntimeError("MPS was requested, but MPS is not available") if device.type == "cuda" and device.index is not None: count = torch.cuda.device_count() if device.index >= count: raise RuntimeError( f"CUDA device index {device.index} is unavailable; {count} device(s) detected" ) return deviceReturn the device.
def get_embedding_module(model: torch.nn.modules.module.Module) ‑> torch.nn.modules.module.Module | None-
Expand source code
def get_embedding_module(model: nn.Module) -> nn.Module | None: """Return the input word embedding module.""" if hasattr(model, "get_input_embeddings"): return model.get_input_embeddings() if hasattr(model, "model") and hasattr(model.model, "embed_tokens"): return model.model.embed_tokens if hasattr(model, "transformer") and hasattr(model.transformer, "wte"): return model.transformer.wte for module in model.modules(): if isinstance(module, nn.Embedding): return module return NoneReturn the input word embedding module.
def get_ffn_modules(model: torch.nn.modules.module.Module) ‑> list[tuple[str, torch.nn.modules.module.Module]]-
Expand source code
def get_ffn_modules(model: nn.Module) -> list[tuple[str, nn.Module]]: """Return (name, module) for all feed‑forward blocks.""" ffn_modules = [] for name, module in model.named_modules(): class_name = module.__class__.__name__.lower() if any(x in class_name for x in ["mlp", "ffn", "feedforward"]): ffn_modules.append((name, module)) continue if ( hasattr(module, "gate_proj") or hasattr(module, "up_proj") or hasattr(module, "intermediate") ): ffn_modules.append((name, module)) return ffn_modulesReturn (name, module) for all feed‑forward blocks.
def get_gpu_memory_stats() ‑> dict[str, float]-
Expand source code
def get_gpu_memory_stats() -> dict[str, float]: """ Get detailed memory statistics in Megabytes (MB). """ stats: dict[str, float] = { "allocated_mb": 0.0, "reserved_mb": 0.0, "max_allocated_mb": 0.0, } if torch.cuda.is_available(): stats["allocated_mb"] = torch.cuda.memory_allocated() / (1024 * 1024) stats["reserved_mb"] = torch.cuda.memory_reserved() / (1024 * 1024) stats["max_allocated_mb"] = torch.cuda.max_memory_allocated() / (1024 * 1024) elif hasattr(torch, "mps") and hasattr(torch.mps, "current_allocated_memory"): try: stats["allocated_mb"] = torch.mps.current_allocated_memory() / (1024 * 1024) if hasattr(torch.mps, "driver_allocated_memory"): stats["reserved_mb"] = torch.mps.driver_allocated_memory() / (1024 * 1024) except ( AttributeError, KeyError, RuntimeError, TypeError, ValueError, OSError, ImportError, ): pass return statsGet detailed memory statistics in Megabytes (MB).
def get_parent_and_child_name(model: torch.nn.modules.module.Module, full_name: str) ‑> Any-
Expand source code
def get_parent_and_child_name(model: nn.Module, full_name: str) -> Any: """Return (parent_module, child_attribute_name), supporting ModuleList and ModuleDict.""" parts = [p for p in full_name.split(".") if p] parent = model for part in parts[:-1]: if isinstance(parent, (nn.ModuleList, list)): parent = parent[int(part)] elif isinstance(parent, (nn.ModuleDict, dict)): parent = parent[part] else: parent = getattr(parent, part) return parent, parts[-1]Return (parent_module, child_attribute_name), supporting ModuleList and ModuleDict.
def get_token_valid_mask(attention_mask: torch.Tensor | None,
batch_size: int,
seq_len: int,
device: torch.device) ‑> torch.Tensor-
Expand source code
def get_token_valid_mask( attention_mask: torch.Tensor | None, batch_size: int, seq_len: int, device: torch.device, ) -> torch.Tensor: r"""Extract a 1D boolean validity mask [Batch, SeqLen] where True = valid token, False = pad. Args: attention_mask (Optional[torch.Tensor]): Mask in any format. batch_size (int): Batch size B. seq_len (int): Sequence length N. device (torch.device): Device to place the mask on. Returns: torch.Tensor: Boolean tensor of shape `[Batch, SeqLen]`. """ if attention_mask is None: return torch.ones((batch_size, seq_len), dtype=torch.bool, device=device) mask = attention_mask.to(device=device) if mask.dim() == 2: if mask.dtype == torch.bool: return mask elif torch.is_floating_point(mask) and (mask < 0).any(): return mask >= -1.0 # unmasked is ~0.0 else: return mask > 0.5 elif mask.dim() == 4: # [B, 1, 1, N] or [B, 1, N, N] if mask.size(2) == 1: last_dim_mask = mask[:, 0, 0, :] if last_dim_mask.dtype == torch.bool: return last_dim_mask elif torch.is_floating_point(last_dim_mask) and (last_dim_mask < 0).any(): return last_dim_mask >= -1.0 else: return last_dim_mask > 0.5 elif mask.size(2) == seq_len: # 4D causal or 2D attention mask [B, 1, N, N] # Do NOT use mask[:, 0, 0, :] as row 0 only unmasks token 0 in causal masks! if mask.dtype == torch.bool: return mask.any(dim=2)[:, 0, :] elif torch.is_floating_point(mask) and (mask < 0).any(): return (mask >= -1000.0).any(dim=2)[:, 0, :] else: return (mask > 0.5).any(dim=2)[:, 0, :] return torch.ones((batch_size, seq_len), dtype=torch.bool, device=device)Extract a 1D boolean validity mask [Batch, SeqLen] where True = valid token, False = pad.
- Args
- -----=
attention_mask:Optional[torch.Tensor]- Mask in any format.
batch_size:int- Batch size B.
seq_len:int- Sequence length N.
device:torch.device- Device to place the mask on.
- Returns
- -----=
torch.Tensor- Boolean tensor of shape
[Batch, SeqLen].
def get_transformer_layers(model: torch.nn.modules.module.Module) ‑> list[torch.nn.modules.module.Module]-
Expand source code
def get_transformer_layers(model: nn.Module) -> list[nn.Module]: """Return a flat list of transformer blocks.""" # Llama / Gemma / Phi / Mistral / Falcon if hasattr(model, "model") and hasattr(model.model, "layers"): return list(model.model.layers) # GPT‑2 / OPT / GPT‑J if hasattr(model, "transformer") and hasattr(model.transformer, "h"): return list(model.transformer.h) # BERT / RoBERTa / DistilBert / ALBERT if hasattr(model, "encoder") and hasattr(model.encoder, "layer"): return list(model.encoder.layer) if hasattr(model, "bert") and hasattr(model.bert, "encoder"): return list(model.bert.encoder.layer) if hasattr(model, "roberta") and hasattr(model.roberta, "encoder"): return list(model.roberta.encoder.layer) if hasattr(model, "distilbert") and hasattr(model.distilbert, "transformer"): return list(model.distilbert.transformer.layer) # T5 encoder / decoder if hasattr(model, "encoder") and hasattr(model.encoder, "block"): return list(model.encoder.block) if hasattr(model, "decoder") and hasattr(model.decoder, "block"): return list(model.decoder.block) raise ValueError("Could not locate transformer layers automatically.")Return a flat list of transformer blocks.
def human_readable_time(seconds: float) ‑> str-
Expand source code
def human_readable_time(seconds: float) -> str: """Human readable time.""" hours = int(seconds // 3600) minutes = int((seconds % 3600) // 60) secs = int(seconds % 60) if hours > 0: return f"{hours}h {minutes}m {secs}s" if minutes > 0: return f"{minutes}m {secs}s" return f"{secs}s"Human readable time.
def inject_attention(model: nn.Module,
attention_type: str,
rank: int | None = None,
copy_weights: bool = True,
**kwargs: Any) ‑> torch.nn.modules.module.Module-
Expand source code
def inject_attention( model: nn.Module, attention_type: str, rank: int | None = None, copy_weights: bool = True, **kwargs: Any, ) -> nn.Module: """Universal attention layer replacement engine for HuggingFace Transformers.""" attn_class = ATTENTION_REGISTRY.resolve(attention_type) if callable(attn_class) and not isinstance(attn_class, type): try: res = attn_class(model=model, rank=rank, copy_weights=copy_weights, **kwargs) if res is not None: return res except Exception as inj_err: logger.debug( "Dedicated injector %s failed, falling back to layer replacement: %s", attention_type, inj_err, ) config = getattr(model, "config", None) if config is not None: cfg_source = getattr(config, "text_config", config) hidden_size = getattr( cfg_source, "hidden_size", getattr( cfg_source, "d_model", getattr(cfg_source, "n_embd", getattr(cfg_source, "dim", None)), ), ) num_heads = getattr( cfg_source, "num_attention_heads", getattr( cfg_source, "num_heads", getattr(cfg_source, "n_heads", getattr(cfg_source, "n_head", None)), ), ) head_dim = getattr( cfg_source, "head_dim", getattr( cfg_source, "d_kv", (hidden_size // max(1, num_heads) if (hidden_size and num_heads) else None), ), ) num_kv_heads = getattr(cfg_source, "num_key_value_heads", num_heads) if hidden_size is not None: config.hidden_size = hidden_size if num_heads is not None: config.num_attention_heads = num_heads if head_dim is not None: config.head_dim = head_dim if num_kv_heads is not None: config.num_key_value_heads = num_kv_heads if not hasattr(config, "dropout"): dropout_val = getattr( config, "hidden_dropout_prob", getattr(config, "attention_probs_dropout_prob", 0.1), ) config.dropout = dropout_val else: hidden_size = getattr(model, "hidden_size", None) num_heads = getattr(model, "num_attention_heads", None) head_dim = getattr(model, "head_dim", None) num_kv_heads = getattr(model, "num_key_value_heads", num_heads) arch = detect_architecture(model) adapter = get_adapter(arch) family = detect_family(model) family_adapter = ModelAdapter(model, family) is_dec = family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral") if config is not None: config.is_decoder = bool(is_dec or getattr(config, "is_decoder", False)) if not hasattr(config, "model_type") or not config.model_type: config.model_type = family replaced = 0 modules_list = list(get_attention_modules(model)) if not modules_list: logger.warning("No attention modules found to replace for architecture '%s'", arch) for name, old_attn in modules_list: # In hybrid architectures (like Qwen3.5), only replace linear_attn if candidate explicitly supports DeltaNet if "linear_attn" in name and not getattr(attn_class, "supports_delta_net", False): continue # Dynamically infer dimensions from family_adapter and old_attn if family_adapter is not None: if hidden_size is None: hidden_size = getattr(family_adapter, "hidden_size", None) if num_heads is None: num_heads = getattr(family_adapter, "num_heads", None) if num_kv_heads is None: num_kv_heads = getattr(family_adapter, "num_key_value_heads", None) if head_dim is None: head_dim = getattr(family_adapter, "head_dim", None) if hidden_size is None or num_heads is None: for _, child in old_attn.named_children(): if ( hasattr(child, "weight") and isinstance(child.weight, torch.Tensor) and child.weight.ndim == 2 ): if hidden_size is None: hidden_size = child.weight.shape[1] break if num_heads is None: num_heads = getattr( old_attn, "num_heads", getattr(old_attn, "num_attention_heads", None), ) if num_kv_heads is None: num_kv_heads = getattr(old_attn, "num_key_value_heads", num_heads) if head_dim is None and hidden_size is not None and num_heads is not None: head_dim = getattr(old_attn, "head_dim", hidden_size // max(1, num_heads)) # Ensure config exists or construct minimal namespace config if config is None: from types import SimpleNamespace is_decoder_family_bool = bool( is_dec or family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral") ) config = SimpleNamespace( hidden_size=hidden_size, num_attention_heads=num_heads, head_dim=head_dim, num_key_value_heads=num_kv_heads or num_heads, dropout=0.1, is_decoder=is_decoder_family_bool, is_causal=is_decoder_family_bool, model_type=family, ) else: if family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral") or is_dec: config.is_decoder = True config.is_causal = True # Ensure old_attn has canonical projection attributes for custom modules inspecting pretrained_attn for proj_key, alias_list in [ ("q", ["query", "q_proj", "q_lin", "Wq", "wq", "q"]), ("k", ["key", "k_proj", "k_lin", "Wk", "wk", "k"]), ("v", ["value", "v_proj", "v_lin", "Wv", "wv", "v"]), ("o", ["dense", "out_proj", "o_proj", "out_lin", "c_proj", "wo", "o"]), ]: try: existing_mod = next( (getattr(old_attn, a) for a in alias_list if hasattr(old_attn, a)), None, ) if existing_mod is None and hasattr(old_attn, "self"): existing_mod = next( ( getattr(old_attn.self, a) for a in alias_list if hasattr(old_attn.self, a) ), None, ) if existing_mod is not None: for a in alias_list: if not hasattr(old_attn, a): try: setattr(old_attn, a, existing_mod) except Exception: pass else: w, b = adapter.get_weight_bias(old_attn, proj_key) if w is not None: proxy = nn.Linear( w.shape[1], w.shape[0], bias=(b is not None), device=w.device, dtype=w.dtype, ) proxy.weight.data.copy_(w) if b is not None: proxy.bias.data.copy_(b) for a in alias_list: if not hasattr(old_attn, a): try: setattr(old_attn, a, proxy) except Exception: pass except Exception as ex: logger.debug( "Failed to set projection aliases for %s on %s: %s", proj_key, name, ex, ) effective_rank = rank if rank is not None else getattr(config, "rank", 32) # Standardized instantiation: try passing pretrained_attn=old_attn first so modules # can inherit pretrained weights during initialization, then fallback without it attempts = [ lambda: attn_class(config, rank=effective_rank, pretrained_attn=old_attn, **kwargs), lambda: attn_class(config, pretrained_attn=old_attn, **kwargs), lambda: attn_class(config=config, pretrained_attn=old_attn, **kwargs), lambda: attn_class(config, rank=effective_rank, **kwargs), lambda: attn_class(config, **kwargs), lambda: attn_class(rank=effective_rank, **kwargs), lambda: attn_class(**kwargs), ] new_attn = None for attempt in attempts: try: new_attn = attempt() break except TypeError: continue except Exception as e: logger.debug("Instantiation attempt failed: %s", e) continue if new_attn is None: raise RuntimeError(f"Could not instantiate attention class {attn_class.__name__}") # Do not let a legacy full-MHA replacement silently retain randomly # initialized K/V projections on GQA decoders. Such a module must # explicitly advertise the geometry it implements. is_gqa = ( family in ("qwen", "qwen3_5") and family_adapter.num_key_value_heads != family_adapter.num_heads ) if is_gqa and not getattr(new_attn, "supports_gqa", False): expected = { "q": family_adapter.q_proj_out, "k": family_adapter.kv_proj_out, "v": family_adapter.kv_proj_out, } for key, patterns in ( ("q", PROJ_REGEX_PATTERNS["q"]), ("k", PROJ_REGEX_PATTERNS["k"]), ("v", PROJ_REGEX_PATTERNS["v"]), ): target = _find_submodule_by_regex(new_attn, patterns) out_features = getattr(target, "out_features", None) if target is not None and out_features not in {expected[key], None}: if expected[key] is not None and out_features is not None: if ( out_features % expected[key] == 0 or expected[key] % out_features == 0 or _find_submodule_by_regex(new_attn, PROJ_REGEX_PATTERNS["o"]) is not None ): pass else: raise QwenIncompatibleError( f"{attention_type}: {key}_projection_out={out_features}, " f"requires={expected[key]} (GQA)" ) # Match device and dtype of old attention layer old_param = next(old_attn.parameters(), None) if old_param is not None: new_attn.to(device=old_param.device, dtype=old_param.dtype) if copy_weights: custom_copied = False if ( is_dec or hasattr(old_attn, "c_attn") or "gpt" in old_attn.__class__.__name__.lower() ): old_attn.is_decoder = True if hasattr(new_attn, "is_decoder"): new_attn.is_decoder = True if hasattr(new_attn, "copy_weights") and callable(new_attn.copy_weights): try: new_attn.copy_weights(old_attn) custom_copied = True if is_dec and hasattr(new_attn, "is_decoder"): new_attn.is_decoder = True if old_param is not None: new_attn.to(device=old_param.device, dtype=old_param.dtype) except (AttributeError, RuntimeError, ValueError) as e: logger.debug("Custom copy_weights failed on %s: %s", name, e) except Exception as e: logger.warning("Unexpected error in custom copy_weights on %s: %s", name, e) if not custom_copied: # Detect fused QKV projection modules via regex and slice parameters fused_mod = None for child_name, child_mod in old_attn.named_children(): if FUSED_QKV_REGEX.match(child_name) and hasattr(child_mod, "weight"): fused_mod = child_mod break if fused_mod is None: for candidate in [ "c_attn", "qkv_proj", "query_key_value", "in_proj", "Wqkv", ]: mod = getattr(old_attn, candidate, None) if mod is not None and hasattr(mod, "weight"): fused_mod = mod break if fused_mod is not None: # Sliced Fused QKV handling sliced = _slice_fused_qkv_weights( fused_mod.weight.data, ( fused_mod.bias.data if hasattr(fused_mod, "bias") and fused_mod.bias is not None else None ), num_heads=num_heads if num_heads is not None else 12, num_kv_heads=num_kv_heads or num_heads, head_dim=head_dim, source_mod=fused_mod, ) for proj_key in ["q", "k", "v"]: w, b = sliced[proj_key] target_proj = _find_submodule_by_regex( new_attn, PROJ_REGEX_PATTERNS[proj_key] ) if target_proj is not None: _copy_projection_weights(target_proj, w, b) else: lin1 = getattr(new_attn, f"{proj_key}_proj1", None) lin2 = getattr(new_attn, f"{proj_key}_proj2", None) if lin1 is not None and lin2 is not None: _copy_projection_weights(lin1, w, None) with torch.no_grad(): W_f = w.detach().float() if hasattr(lin1, "weight") and hasattr(lin2, "weight"): if ( W_f.shape[1] != lin1.weight.shape[1] and W_f.shape[0] == lin1.weight.shape[1] ): W_f = W_f.t() U, S, Vt = torch.linalg.svd(W_f, full_matrices=False) r = min( lin1.weight.shape[0], lin2.weight.shape[1], U.size(1), Vt.size(0), ) sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12)) U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to( device=w.device, dtype=w.dtype ) V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to( device=w.device, dtype=w.dtype ) lin1.weight.data[:r, :].copy_(V_b) lin2.weight.data[:, :r].copy_(U_b) if ( b is not None and getattr(lin2, "bias", None) is not None ): lin2.bias.data.copy_( b.to( lin2.bias.device, lin2.bias.dtype, ) ) # Output projection try: src_w_o, src_b_o = adapter.get_weight_bias(old_attn, "o") target_o = _find_submodule_by_regex(new_attn, PROJ_REGEX_PATTERNS["o"]) if target_o is not None: _copy_projection_weights(target_o, src_w_o, src_b_o) else: lin1_o = getattr( new_attn, "out_proj1", getattr(new_attn, "o_proj1", None), ) lin2_o = getattr( new_attn, "out_proj2", getattr(new_attn, "o_proj2", None), ) if lin1_o is not None and lin2_o is not None: with torch.no_grad(): W_f = src_w_o.detach().float() if ( W_f.shape[1] != lin1_o.weight.shape[1] and W_f.shape[0] == lin1_o.weight.shape[1] ): W_f = W_f.t() U, S, Vt = torch.linalg.svd(W_f, full_matrices=False) r = min( lin1_o.weight.shape[0], lin2_o.weight.shape[1], U.size(1), Vt.size(0), ) sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12)) U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to( device=src_w_o.device, dtype=src_w_o.dtype ) V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to( device=src_w_o.device, dtype=src_w_o.dtype ) min_out = min(lin2_o.weight.shape[0], U_b.shape[0]) min_in = min(lin1_o.weight.shape[1], V_b.shape[1]) lin1_o.weight.data[:r, :min_in].copy_(V_b[:, :min_in]) lin2_o.weight.data[:min_out, :r].copy_(U_b[:min_out, :]) if ( src_b_o is not None and getattr(lin2_o, "bias", None) is not None ): min_b = min(lin2_o.bias.shape[0], src_b_o.shape[0]) lin2_o.bias.data[:min_b].copy_( src_b_o[:min_b].to( lin2_o.bias.device, lin2_o.bias.dtype, ) ) except (ValueError, AttributeError) as e: logger.debug("Could not copy output projection for %s: %s", name, e) else: # Standard decoupled Q, K, V, O projections for proj_key in ["q", "k", "v", "o"]: src_weight, src_bias = None, None try: src_weight, src_bias = adapter.get_weight_bias(old_attn, proj_key) except (ValueError, AttributeError) as e: logger.debug( "Could not get %s projection weights for %s: %s", proj_key, name, e, ) target_proj = _find_submodule_by_regex( new_attn, PROJ_REGEX_PATTERNS[proj_key] ) if target_proj is not None: if src_weight is not None: _copy_projection_weights(target_proj, src_weight, src_bias) elif proj_key == "o": # Encoder self-attention (BERT/RoBERTa .self) has no internal out_proj # because the outer attention container already applies output.dense. # Initialize new_attn.out_proj to Identity so context passes through transparently. with torch.no_grad(): if ( hasattr(target_proj, "weight") and target_proj.weight.ndim == 2 ): if ( target_proj.weight.shape[0] == target_proj.weight.shape[1] ): nn.init.eye_(target_proj.weight) else: nn.init.zeros_(target_proj.weight) min_d = min( target_proj.weight.shape[0], target_proj.weight.shape[1], ) target_proj.weight.data[ :min_d, :min_d ].fill_diagonal_(1.0) elif hasattr(target_proj, "set_svd_factors") or ( hasattr(target_proj, "U") and hasattr(target_proj, "V") ): # Low-rank / factorized linear layer: set factors to truncated identity d_out = getattr( target_proj, "out_features", getattr( target_proj, "d_out", hidden_size or 768, ), ) d_in = getattr( target_proj, "in_features", getattr( target_proj, "d_in", hidden_size or 768, ), ) eye_mat = torch.eye( d_out, d_in, device=( target_proj.U.device if hasattr(target_proj, "U") else None ), ) _copy_projection_weights(target_proj, eye_mat, None) if ( hasattr(target_proj, "bias") and target_proj.bias is not None ): nn.init.zeros_(target_proj.bias) else: # Support indexed dual projections (e.g. aegis_hybrid, svd_linear_attention: q_proj1, q_proj2) lin1 = getattr( new_attn, f"{proj_key}_proj1", getattr(new_attn, f"{proj_key}1", None), ) lin2 = getattr( new_attn, f"{proj_key}_proj2", getattr(new_attn, f"{proj_key}2", None), ) if proj_key == "o": lin1 = lin1 or getattr( new_attn, "out_proj1", getattr(new_attn, "o_proj1", None), ) lin2 = lin2 or getattr( new_attn, "out_proj2", getattr(new_attn, "o_proj2", None), ) if lin1 is not None and lin2 is not None: if src_weight is not None: with torch.no_grad(): W_f = src_weight.detach().float() if hasattr(lin1, "weight") and hasattr(lin2, "weight"): if ( W_f.shape[1] != lin1.weight.shape[1] and W_f.shape[0] == lin1.weight.shape[1] ): W_f = W_f.t() U, S, Vt = torch.linalg.svd( W_f, full_matrices=False ) r = min( lin1.weight.shape[0], lin2.weight.shape[1], U.size(1), Vt.size(0), ) sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12)) U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to( device=src_weight.device, dtype=src_weight.dtype, ) V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to( device=src_weight.device, dtype=src_weight.dtype, ) min_out = min(lin2.weight.shape[0], U_b.shape[0]) min_in = min(lin1.weight.shape[1], V_b.shape[1]) lin1.weight.data[:r, :min_in].copy_(V_b[:, :min_in]) lin2.weight.data[:min_out, :r].copy_( U_b[:min_out, :] ) if ( src_bias is not None and getattr(lin2, "bias", None) is not None ): min_b = min( lin2.bias.shape[0], src_bias.shape[0], ) lin2.bias.data[:min_b].copy_( src_bias[:min_b].to( lin2.bias.device, lin2.bias.dtype, ) ) elif proj_key == "o": # Identity for dual out_proj with torch.no_grad(): if hasattr(lin1, "weight") and hasattr(lin2, "weight"): nn.init.zeros_(lin1.weight) nn.init.zeros_(lin2.weight) r = min( lin1.weight.shape[0], lin2.weight.shape[1], lin1.weight.shape[1], lin2.weight.shape[0], ) lin1.weight.data[:r, :r].fill_diagonal_(1.0) lin2.weight.data[:r, :r].fill_diagonal_(1.0) if getattr(lin2, "bias", None) is not None: nn.init.zeros_(lin2.bias) from neuralops.utils.replace_attention import UniversalAttentionWrapper return_tuple = ( "self_attn" in name.lower() or "attention" in name.lower() or getattr(old_attn, "is_decoder", False) or family == "gpt" or name.endswith(".attn") ) wrapped_attn = UniversalAttentionWrapper( new_attn, layer_name=name, return_tuple=return_tuple ) replace_module(model, name, wrapped_attn) replaced += 1 logger.info( "Successfully injected '%s' into %d transformer blocks.", attention_type, replaced, ) if not hasattr(model, "compression_stats"): model.compression_stats = {} model.compression_stats["num_compressed"] = replaced model.compression_stats["total_params_after"] = sum(p.numel() for p in model.parameters()) return modelUniversal attention layer replacement engine for HuggingFace Transformers.
def list_all_components() ‑> dict-
Expand source code
def list_all_components() -> dict: """Convenience function that returns all global registry keys.""" return { "attentions": ATTENTION_REGISTRY.keys(), "compressors": COMPRESSOR_REGISTRY.keys(), "methods": METHOD_REGISTRY.keys(), "auxiliary": AUX_REGISTRY.keys(), "models": MODEL_REGISTRY.keys(), "datasets": DATASET_REGISTRY.keys(), "losses": LOSS_REGISTRY.keys(), "trainers": TRAINER_REGISTRY.keys(), "visualizers": VISUALIZER_REGISTRY.keys(), }Convenience function that returns all global registry keys.
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) ‑> torch.Tensor-
Expand source code
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ Repeats Key/Value heads along the head dimension for Grouped-Query Attention (GQA). hidden_states: (batch, num_key_value_heads, seqlen, head_dim) Returns: (batch, num_attention_heads, seqlen, head_dim) """ if n_rep <= 1: return hidden_states b, h, s, d = hidden_states.shape return hidden_states[:, :, None, :, :].expand(b, h, n_rep, s, d).reshape(b, h * n_rep, s, d)Repeats Key/Value heads along the head dimension for Grouped-Query Attention (GQA). hidden_states: (batch, num_key_value_heads, seqlen, head_dim) Returns: (batch, num_attention_heads, seqlen, head_dim)
def replace_module(model: torch.nn.modules.module.Module,
full_name: str,
new_module: torch.nn.modules.module.Module) ‑> None-
Expand source code
def replace_module(model: nn.Module, full_name: str, new_module: nn.Module) -> None: """Replace a sub-module identified by its dotted name, supporting ModuleList and ModuleDict.""" parts = [p for p in full_name.split(".") if p] parent = model for part in parts[:-1]: if isinstance(parent, (nn.ModuleList, list)): parent = parent[int(part)] elif isinstance(parent, (nn.ModuleDict, dict)): parent = parent[part] else: parent = getattr(parent, part) last = parts[-1] if isinstance(parent, (nn.ModuleList, list)): parent[int(last)] = new_module elif isinstance(parent, (nn.ModuleDict, dict)): parent[last] = new_module else: setattr(parent, last, new_module)Replace a sub-module identified by its dotted name, supporting ModuleList and ModuleDict.
def run_calibration_forward(model: torch.nn.modules.module.Module,
tokenizer,
texts: list[str],
device: torch.device | None = None,
max_length: int = 128) ‑> torch.nn.modules.module.Module-
Expand source code
def run_calibration_forward( model: nn.Module, tokenizer, texts: list[str], device: torch.device | None = None, max_length: int = 128, ) -> nn.Module: """Run a single forward pass on calibration texts; return model (in eval).""" model.eval() inputs = tokenizer( texts, return_tensors="pt", padding=True, truncation=True, max_length=max_length ) if device is None: device = next(model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): model(**inputs) return modelRun a single forward pass on calibration texts; return model (in eval).
def selective_scan(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
z: torch.Tensor | None = None,
delta_bias: torch.Tensor | None = None,
delta_softplus: bool = True,
chunk_size: int = 32) ‑> torch.Tensor-
Expand source code
def selective_scan( u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, D: torch.Tensor | None = None, z: torch.Tensor | None = None, delta_bias: torch.Tensor | None = None, delta_softplus: bool = True, chunk_size: int = 32, ) -> torch.Tensor: """ Dispatch to CUDA kernel if available, else stable chunked fallback. """ if HAS_MAMBA_CUDA and u.is_cuda: try: return _mamba_scan_fn(u, delta, A, B, C, D, z, delta_bias, delta_softplus) except Exception as e: # SC-FIX-06: surface the failure so it can be investigated. logger.warning( "mamba_ssm CUDA selective_scan failed (%s); " "falling back to chunked PyTorch implementation.", e, ) return selective_scan_chunked( u=u, delta=delta, A=A, B=B, C=C, D=D, z=z, delta_bias=delta_bias, delta_softplus=delta_softplus, chunk_size=chunk_size, )Dispatch to CUDA kernel if available, else stable chunked fallback.
def selective_scan_chunked(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
z: torch.Tensor | None = None,
delta_bias: torch.Tensor | None = None,
delta_softplus: bool = True,
chunk_size: int = 32) ‑> torch.Tensor-
Expand source code
def selective_scan_chunked( u: torch.Tensor, # (B, d_inner, L) delta: torch.Tensor, # (B, d_inner, L) A: torch.Tensor, # (d_inner, N) B: torch.Tensor, # (B, N, L) C: torch.Tensor, # (B, N, L) D: torch.Tensor | None = None, # (d_inner,) z: torch.Tensor | None = None, # (B, d_inner, L) delta_bias: torch.Tensor | None = None, delta_softplus: bool = True, chunk_size: int = 32, ) -> torch.Tensor: """ Selective scan via numerically stable chunked recurrence. For chunk [s, e) of length C: dA[t] = Δ[t] · A <= 0 dA_cum[t] = cumsum(dA) decay[j, i] = exp(dA_cum[j] - dA_cum[i]) for j >= i, else 0 h_intra[j] = sum_i decay[j, i] · dB_u[i] h_carry = exp(dA_cum) · h_prev h_chunk = h_intra + h_carry y_chunk = sum_N(h_chunk · C) """ B_size, D_dim, L = u.shape if L == 0: # SC-FIX-07: trivial empty case. y = torch.zeros_like(u) if z is not None: y = y * F.silu(z) return y N = A.shape[1] device = u.device dtype = u.dtype # SC-FIX-02: cast only the tensors that need fp32; avoid redundant copies. u_f = u.float() if dtype != torch.float32 else u delta_f = delta.float() if delta.dtype != torch.float32 else delta A_f = A.float() if A.dtype != torch.float32 else A B_f = B.float() if B.dtype != torch.float32 else B C_f = C.float() if C.dtype != torch.float32 else C # delta_bias + softplus if delta_bias is not None: delta_f = delta_f + delta_bias.float().view(1, -1, 1) if delta_softplus: delta_f = F.softplus(delta_f) # SC-FIX-03: no arbitrary clamp; only ensure positivity after softplus. # (softplus already produces strictly positive values in fp32.) # dA: (B, D, L, N), always <= 0 dA = delta_f.unsqueeze(-1) * A_f.unsqueeze(0).unsqueeze(2) dA = dA.clamp(max=0.0) # dB_u: (B, D, L, N) # B_f is (B, N, L); reshape to (B, 1, L, N) B_f_t = B_f.transpose(1, 2).contiguous() # (B, L, N) dB_u = delta_f.unsqueeze(-1) * u_f.unsqueeze(-1) * B_f_t.unsqueeze(1) # SC-FIX-01: auto-tune chunk size to bound the (B, D, C, C, N) tensor. if chunk_size <= 0: chunk_size = _auto_chunk_size(D_dim, N, B_size) else: # Enforce memory bound even if user gave explicit chunk_size. max_chunk = _auto_chunk_size(D_dim, N, B_size, max_bytes=512 * 1024 * 1024) chunk_size = min(chunk_size, max_chunk) chunk_size = max(1, chunk_size) num_chunks = (L + chunk_size - 1) // chunk_size h_prev = torch.zeros(B_size, D_dim, N, device=device, dtype=torch.float32) y_chunks = [] for c in range(num_chunks): s = c * chunk_size e = min(L, s + chunk_size) C_len = e - s dA_c = dA[:, :, s:e, :] # (B, D, C_len, N) dBu_c = dB_u[:, :, s:e, :] # (B, D, C_len, N) C_c = C_f[:, :, s:e].transpose(1, 2) # (B, C_len, N) # Cumulative log-decay within the chunk dA_cum = torch.cumsum(dA_c, dim=2) # (B, D, C_len, N) # --- intra-chunk recurrence without division --- # SC-FIX-05: avoid 5D einsum; use exp of differences masked by tril. log_decay = dA_cum.unsqueeze(3) - dA_cum.unsqueeze(2) # (B, D, C_len, C_len, N) tril = torch.tril(torch.ones(C_len, C_len, dtype=torch.bool, device=device)) log_decay = log_decay.masked_fill(~tril.view(1, 1, C_len, C_len, 1), float("-inf")) # SC-FIX-04: exp(-inf) = 0 exactly; no nan_to_num needed. decay = torch.exp(log_decay) # (B, D, C_len, C_len, N) # h_intra[b, d, j, n] = sum_i decay[b, d, j, i, n] * dBu_c[b, d, i, n] # Faster: reshape to (B*D, C_len, C_len, N) and use einsum with 4D. BD = B_size * D_dim decay_bd = decay.reshape(BD, C_len, C_len, N) dBu_bd = dBu_c.reshape(BD, C_len, N) h_intra_bd = torch.einsum("bjic,bic->bjc", decay_bd, dBu_bd) # (BD, C_len, N) h_intra = h_intra_bd.reshape(B_size, D_dim, C_len, N) # Carry from previous chunk h_carry = torch.exp(dA_cum) * h_prev.unsqueeze(2) # (B, D, C_len, N) h_chunk = h_intra + h_carry # Output projection # C_c: (B, C_len, N); h_chunk: (B, D, C_len, N) y_c = torch.einsum("bdcn,bcn->bdc", h_chunk, C_c) # (B, D, C_len) y_chunks.append(y_c) h_prev = h_chunk[:, :, -1, :] # (B, D, N) y = torch.cat(y_chunks, dim=2) # Skip connection if D is not None: y = y + D.float().view(1, -1, 1) * u_f # Gating if z is not None: y = y * F.silu(z.float()) return y.to(dtype)Selective scan via numerically stable chunked recurrence.
For chunk [s, e) of length C: dA[t] = Δ[t] · A <= 0 dA_cum[t] = cumsum(dA) decay[j, i] = exp(dA_cum[j] - dA_cum[i]) for j >= i, else 0 h_intra[j] = sum_i decay[j, i] · dB_u[i] h_carry = exp(dA_cum) · h_prev h_chunk = h_intra + h_carry y_chunk = sum_N(h_chunk · C)
def set_seed(seed: int, deterministic: bool = False) ‑> None-
Expand source code
def set_seed(seed: int, deterministic: bool = False) -> None: """Set the seed.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) if deterministic: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = FalseSet the seed.
def softpick(x: torch.Tensor, dim: int = -1, temperature: float = 1.0) ‑> torch.Tensor-
Expand source code
def softpick(x: torch.Tensor, dim: int = -1, temperature: float = 1.0) -> torch.Tensor: """Rectified non-negative sparse activation function with numerical stabilization.""" m = x.max(dim=dim, keepdim=True).values scaled = (x - m) / max(temperature, 1e-6) scaled_clamped = torch.clamp(scaled, min=-50.0, max=50.0) e = torch.exp(scaled_clamped) m_clamped = torch.clamp(-m / max(temperature, 1e-6), min=-50.0, max=50.0) denom = e - torch.exp(m_clamped) p = torch.relu(denom) return p / (p.sum(dim=dim, keepdim=True) + 1e-9)Rectified non-negative sparse activation function with numerical stabilization.
Classes
class BaseAttention (config: Any | None = None, rank: int | None = None, **kwargs: Any)-
Expand source code
class BaseAttention(nn.Module): """Every custom attention module must inherit from this.""" def __init__(self, config: Any | None = None, rank: int | None = None, **kwargs: Any) -> None: super().__init__() self.config = config self.rank = rank self.num_heads = ( kwargs.get("num_heads") or kwargs.get("num_attention_heads") or ( getattr( config, "num_attention_heads", getattr(config, "num_heads", getattr(config, "n_heads", 12)), ) if config is not None else 12 ) ) self.hidden_size = ( kwargs.get("hidden_size") or kwargs.get("dim") or kwargs.get("d_model") or ( getattr( config, "hidden_size", getattr(config, "d_model", getattr(config, "dim", 768)), ) if config is not None else 768 ) ) self.head_dim = getattr( config, "d_kv", getattr( config, "head_dim", self.hidden_size // self.num_heads if self.num_heads > 0 else 64, ), ) @abstractmethod def forward( self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, **kwargs, ) -> tuple[Any, Any | None]: """Must return (output, attention_weights_or_None).""" raise NotImplementedError("Subclasses of BaseAttention must implement forward()") def copy_weights(self, old_attn: Any) -> None: """Universal fallback: copy pretrained weights from old attention layer.""" if old_attn is None: return import torch from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() try: w_q, b_q = adapter.get_weight_bias(old_attn, "q") w_k, b_k = adapter.get_weight_bias(old_attn, "k") w_v, b_v = adapter.get_weight_bias(old_attn, "v") w_o, b_o = adapter.get_weight_bias(old_attn, "o") with torch.no_grad(): dim = self.hidden_size for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]: if hasattr(self, q_name): mod = getattr(self, q_name) if hasattr(mod, "weight") and mod.weight.shape == w_q.shape: mod.weight.copy_(w_q) if b_q is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_q) for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]: if hasattr(self, k_name): mod = getattr(self, k_name) if hasattr(mod, "weight") and mod.weight.shape == w_k.shape: mod.weight.copy_(w_k) if b_k is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_k) for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]: if hasattr(self, v_name): mod = getattr(self, v_name) if hasattr(mod, "weight") and mod.weight.shape == w_v.shape: mod.weight.copy_(w_v) if b_v is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_v) for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]: if hasattr(self, o_name): mod = getattr(self, o_name) if hasattr(mod, "weight") and mod.weight.shape == w_o.shape: mod.weight.copy_(w_o) if b_o is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_o) # Handle fused in_proj if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"): if self.in_proj.weight.shape == torch.Size([dim * 2, dim]): self.in_proj.weight[:dim].copy_(w_v[:dim]) self.in_proj.weight[dim:].copy_(w_q[:dim]) elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]): self.in_proj.weight[:dim].copy_(w_q[:dim]) self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim]) self.in_proj.weight[2 * dim :].copy_(w_v[:dim]) except (AttributeError, RuntimeError, ValueError) as e: logger.debug("BaseAttention.copy_weights failed: %s", e) except Exception as e: logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e) def extra_repr(self) -> str: """Extra repr.""" return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"Every custom attention module must inherit from this.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Subclasses
- AdaptiveSpanAttention
- AdaptiveSparseAttention
- AdaptiveSVDAttention
- AdaSplashAttention
- AegisHybridAttention
- AegisSpectronFP8Attention
- AegisAttentionV2
- AegisAttention
- AegisAttentionV3
- ASVDAttention
- BAMUMoEBlock
- BlockTransformerHIBA
- ChelaSlayAttention
- ClusterAttention
- CompressedAttentionLayer
- CompressedDistilBertAttention
- CosformerAttention
- DecayLinearAttention
- EATBlock
- ExpertChoiceMoEAttention
- FactorisedMatrixState
- FlashAttention3
- FlashAttention3FP8
- FullAttention
- FusedSVDKernelAttention
- FUSIONATTN
- GalerkinAttention
- GatedDeltaNet
- GatedDeltaNetBlock
- GatedLinearAttention
- GoldFinchAHNBlock
- GroupedQueryAttention
- GraniteHybridBlock
- GriffinSSDBlock
- HALOReGLAAttention
- HammingAttention
- HGRN2GatedDeltaNet
- HGRN2Layer
- HGRN2OuterBlock
- HLAAttention
- HybridAttention
- HymbaTransXBlock
- InfiniAttention
- Jamba2Block
- JambaInfoBlock
- KDAAttention
- LEKAV1Attention
- LEKAv2Attention
- LEKAV3Attention
- LinformerAttention
- LoLAAttention
- LoRAAttention
- LSHAttention
- MambaSVDFlash
- MinimalNOVAAttention
- MinimalSVDAttention
- MixAttentionMoSKA
- MLAAttention
- MoBAAttention
- NASCAttention
- NEXUSAttention
- NativeHybridAttention
- NativeSparseAttention
- NystromAttention
- OptimizedHybridAttention
- ORIONAttention
- PADReAttention
- PerformerAttention
- PHANTOMAttention
- PrimalAttention
- ProgressiveTokenPruning
- QCQAAttention
- RALALinearAttention
- REAPMoEAttention
- ReGLAAttention
- ReplaceAttention
- ResSVDAttention
- RoutingAttention
- RPCAttention
- SALTAttention
- SambaMoEBlock
- SelectiveSSM
- SelectiveSSM
- SOAAttention
- SpanExpandedAttention
- SpectraV2Attention
- SpectraformerAttention
- SpectralMambaAttention
- SpectronAttention
- SqueezingHeadsDistiller
- SVDLinearAttention
- SVDCompressedAttention
- SVDLinearAttention
- SVDLLMV2Attention
- SVDMamba
- SVDProjectionAttention
- YourCustomAttention
- TLinFormerAttention
- TokenCompressedAttention
- TPAttention
- TritonCosformerAttention
- TritonSVDLinearAttention
- UnifiedRoPE
- SelectiveMamba
- ViGAttention
- VSpecAttention
- WeightedGroupedQueryAttention
- XAttention
- ZeroSAttention
- NLISoftmaxAttention
- ReplaceAttention
Methods
def copy_weights(self, old_attn: Any) ‑> None-
Expand source code
def copy_weights(self, old_attn: Any) -> None: """Universal fallback: copy pretrained weights from old attention layer.""" if old_attn is None: return import torch from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() try: w_q, b_q = adapter.get_weight_bias(old_attn, "q") w_k, b_k = adapter.get_weight_bias(old_attn, "k") w_v, b_v = adapter.get_weight_bias(old_attn, "v") w_o, b_o = adapter.get_weight_bias(old_attn, "o") with torch.no_grad(): dim = self.hidden_size for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]: if hasattr(self, q_name): mod = getattr(self, q_name) if hasattr(mod, "weight") and mod.weight.shape == w_q.shape: mod.weight.copy_(w_q) if b_q is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_q) for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]: if hasattr(self, k_name): mod = getattr(self, k_name) if hasattr(mod, "weight") and mod.weight.shape == w_k.shape: mod.weight.copy_(w_k) if b_k is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_k) for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]: if hasattr(self, v_name): mod = getattr(self, v_name) if hasattr(mod, "weight") and mod.weight.shape == w_v.shape: mod.weight.copy_(w_v) if b_v is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_v) for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]: if hasattr(self, o_name): mod = getattr(self, o_name) if hasattr(mod, "weight") and mod.weight.shape == w_o.shape: mod.weight.copy_(w_o) if b_o is not None and getattr(mod, "bias", None) is not None: mod.bias.copy_(b_o) # Handle fused in_proj if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"): if self.in_proj.weight.shape == torch.Size([dim * 2, dim]): self.in_proj.weight[:dim].copy_(w_v[:dim]) self.in_proj.weight[dim:].copy_(w_q[:dim]) elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]): self.in_proj.weight[:dim].copy_(w_q[:dim]) self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim]) self.in_proj.weight[2 * dim :].copy_(w_v[:dim]) except (AttributeError, RuntimeError, ValueError) as e: logger.debug("BaseAttention.copy_weights failed: %s", e) except Exception as e: logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)Universal fallback: copy pretrained weights from old attention layer.
def extra_repr(self) ‑> str-
Expand source code
def extra_repr(self) -> str: """Extra repr.""" return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"Extra repr.
def forward(self,
hidden_states,
attention_mask=None,
head_mask=None,
output_attentions=False,
**kwargs) ‑> tuple[typing.Any, Any | None]-
Expand source code
@abstractmethod def forward( self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, **kwargs, ) -> tuple[Any, Any | None]: """Must return (output, attention_weights_or_None).""" raise NotImplementedError("Subclasses of BaseAttention must implement forward()")Must return (output, attention_weights_or_None).
class BaseCompressor (model: torch.nn.modules.module.Module, name: str = 'base', **kwargs)-
Expand source code
class BaseCompressor: """Base class for all compressors. Not a nn.Module to avoid parameter storage.""" def __init__(self, model: nn.Module, name: str = "base", **kwargs) -> None: self.model = model self.name = name self.kwargs = kwargs self._original_params = self._count_params(model) def __call__(self, *args, **kwargs) -> Any: """Allow compressor instance to act as a callable model forward wrapper.""" if not getattr(self, "_compressed_once", False): self.compress() self._compressed_once = True return self.model(*args, **kwargs) @abstractmethod def compress(self) -> tuple[nn.Module, dict[str, Any]]: """Return (compressed_model, statistics).""" raise NotImplementedError("Subclasses of BaseCompressor must implement compress()") def _count_params(self, model: nn.Module) -> int: """count params.""" return sum(p.numel() for p in model.parameters()) def get_compression_stats(self) -> dict[str, Any]: """Return the compression stats.""" compressed = self._count_params(self.model) orig = self._original_params return { "compressor_name": self.name, "original_parameters": orig, "compressed_parameters": compressed, "saved_parameters": orig - compressed, "compression_ratio": compressed / orig if orig > 0 else 1.0, "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0), } def summary(self) -> str: """Summary.""" s = self.get_compression_stats() return ( f"\n========================================\n" f"Compression Summary: {self.name}\n" f"========================================\n" f"Original parameters: {s['original_parameters']:,}\n" f"Compressed parameters: {s['compressed_parameters']:,}\n" f"Parameters saved: {s['saved_parameters']:,}\n" f"Compression ratio: {s['compression_ratio']:.3f}\n" f"Reduction: {s['reduction_percentage']:.2f}%\n" f"========================================\n" )Base class for all compressors. Not a nn.Module to avoid parameter storage.
Subclasses
- A3Compressor
- AASVDCompressor
- AATC_XKVCompressor
- AdaptiveRanksCompressor
- AdaSVDCompressor
- ADMMCompressor
- ADMMPruner
- ARMORPruner
- ASTPruner
- AstroQuantizer
- BonsaiPruner
- CacheNotesCompressor
- CoMeCompressor
- COMPOTCompressor
- CoopQQuantizer
- CRCompressor
- CURCompressor
- CURMoECompressor
- DeltaKVCompressor
- DirectSVDCompressor
- DobiSVDCompressor
- DRONECompressor
- DynamicEntropySVDCompressor
- ENLightenCompressor
- EntropyBudgetCompressor
- ERCCompressor
- FAQQuantizer
- FLAPCompressor
- FLATLLMCompressor
- FRanDICompressor
- FWSVDCompressor
- GFWSCompressor
- GreCoDeCCompressor
- HashedNetCompressor
- HuffmanCompressor
- HybridSVDCompressor
- HybridKVCompressor
- KMeansCompressor
- KroneckerCompressor
- KroneckerTTLoRACompressor
- KVTCCompressor
- LIFTRCompressor
- LiLLaMADistiller
- LSCCompressor
- MambaShedderCompressor
- MGIECompressor
- MGIESVDBudgetCompressor
- MLoRQCompressor
- MoDeGPTCompressor
- NLISoftmaxCompressor
- NMFCompressor
- OATSCompressor
- OptimalSubspaceProjectionCompressor
- ProxSparsePruner
- ReplaceMeCompressor
- ResSVDCompressor
- ROASTCompressor
- RotationalSliceGPTINT4Compressor
- RPCACompressor
- SAESSVDCompressorV2
- SALRCompressor
- SDRMCompressor
- SliceGPTCompressor
- SliceGPTPCACompressor
- SLiMCompressor
- LowRankCompressor
- SparseGPTCompressor
- SubspaceProjectionCompressor
- SVDLinearCompressor
- SVDLLMCompressor
- SwiftSVDCompressor
- SwiftSVDActivationCompressor
- ToMoECompressor
- TTCompressor
- UniQLCompressor
- ZeroKVNeoCompressor
Methods
def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]-
Expand source code
@abstractmethod def compress(self) -> tuple[nn.Module, dict[str, Any]]: """Return (compressed_model, statistics).""" raise NotImplementedError("Subclasses of BaseCompressor must implement compress()")Return (compressed_model, statistics).
def get_compression_stats(self) ‑> dict[str, typing.Any]-
Expand source code
def get_compression_stats(self) -> dict[str, Any]: """Return the compression stats.""" compressed = self._count_params(self.model) orig = self._original_params return { "compressor_name": self.name, "original_parameters": orig, "compressed_parameters": compressed, "saved_parameters": orig - compressed, "compression_ratio": compressed / orig if orig > 0 else 1.0, "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0), }Return the compression stats.
def summary(self) ‑> str-
Expand source code
def summary(self) -> str: """Summary.""" s = self.get_compression_stats() return ( f"\n========================================\n" f"Compression Summary: {self.name}\n" f"========================================\n" f"Original parameters: {s['original_parameters']:,}\n" f"Compressed parameters: {s['compressed_parameters']:,}\n" f"Parameters saved: {s['saved_parameters']:,}\n" f"Compression ratio: {s['compression_ratio']:.3f}\n" f"Reduction: {s['reduction_percentage']:.2f}%\n" f"========================================\n" )Summary.
class CompressionPipeline (model: torch.nn.modules.module.Module,
config: dict[str, typing.Any],
tokenizer=None)-
Expand source code
class CompressionPipeline: """Apply a sequence of compression steps to a model.""" def __init__( self, model: nn.Module, config: dict[str, Any], tokenizer=None, ) -> None: self.model = model self.config = config self.tokenizer = tokenizer if tokenizer is not None: self.model.tokenizer = tokenizer def run(self) -> nn.Module: """Execute the run operation and return results.""" compressors_cfg = self.config.get("compressors", []) if not compressors_cfg: logger.info("No compressors specified. Returning original model.") return self.model seed = self.config.get("seed") if seed is not None: set_seed(seed) for i, comp_cfg in enumerate(compressors_cfg): method = comp_cfg.pop("method") logger.info("Step %d: Applying compressor '%s'", i + 1, method) compressor_cls = COMPRESSOR_REGISTRY.resolve(method) merged_cfg = {**self.config, **comp_cfg} merged_cfg.pop("compressors", None) compressor = compressor_cls(self.model, **merged_cfg) self.model, stats = compressor.compress() if stats: logger.info("%s", compressor.summary()) return self.modelApply a sequence of compression steps to a model.
Methods
def run(self) ‑> torch.nn.modules.module.Module-
Expand source code
def run(self) -> nn.Module: """Execute the run operation and return results.""" compressors_cfg = self.config.get("compressors", []) if not compressors_cfg: logger.info("No compressors specified. Returning original model.") return self.model seed = self.config.get("seed") if seed is not None: set_seed(seed) for i, comp_cfg in enumerate(compressors_cfg): method = comp_cfg.pop("method") logger.info("Step %d: Applying compressor '%s'", i + 1, method) compressor_cls = COMPRESSOR_REGISTRY.resolve(method) merged_cfg = {**self.config, **comp_cfg} merged_cfg.pop("compressors", None) compressor = compressor_cls(self.model, **merged_cfg) self.model, stats = compressor.compress() if stats: logger.info("%s", compressor.summary()) return self.modelExecute the run operation and return results.
class FactorizedLinear (in_features: int,
out_features: int,
rank: int = 128,
lora_rank: int = 32,
bias: bool = True,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class FactorizedLinear(nn.Module): """ Decomposed linear layer: W ≈ W2 @ W1 with optional LoRA residual branch. """ is_custom_attention: bool = True def __init__( self, in_features: int, out_features: int, rank: int = 128, lora_rank: int = 32, bias: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() self.in_features = int(in_features) self.out_features = int(out_features) self.rank = max(1, min(int(rank), self.in_features, self.out_features)) self.svd_rank = self.rank self.lora_rank = ( min(int(lora_rank), min(self.in_features, self.out_features)) if lora_rank > 0 else 0 ) factory_kwargs = {"device": device, "dtype": dtype} self.lin1 = nn.Linear(self.in_features, self.rank, bias=False, **factory_kwargs) self.lin2 = nn.Linear(self.rank, self.out_features, bias=bias, **factory_kwargs) self.lin1.weight.requires_grad = False self.lin2.weight.requires_grad = False if self.lin2.bias is not None: self.lin2.bias.requires_grad = False if self.lora_rank > 0: self.lora_A = nn.Parameter( torch.empty(self.lora_rank, self.in_features, **factory_kwargs) ) self.lora_B = nn.Parameter( torch.empty(self.out_features, self.lora_rank, **factory_kwargs) ) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() def reset_parameters(self) -> None: with torch.no_grad(): nn.init.kaiming_uniform_(self.lin1.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.lin2.weight, a=math.sqrt(5)) if self.lin2.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.lin2.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.lin2.bias, -bound, bound) if self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B) def freeze_base_weights(self) -> None: self.lin1.weight.requires_grad = False self.lin2.weight.requires_grad = False if self.lin2.bias is not None: self.lin2.bias.requires_grad = False if self.lora_A is not None: self.lora_A.requires_grad = True self.lora_B.requires_grad = True @property def weight(self) -> torch.Tensor: w = torch.matmul(self.lin2.weight, self.lin1.weight) if self.lora_A is not None: w = w + torch.matmul(self.lora_B, self.lora_A) return w @property def bias(self) -> torch.Tensor | None: return self.lin2.bias def forward(self, x: torch.Tensor) -> torch.Tensor: if HAS_TRITON_COMMON and triton_dispatch is not None and x.is_cuda: try: out = triton_dispatch( "low_rank", x, self.lin2.weight, self.lin1.weight, bias=self.lin2.bias ) except Exception: out = self.lin2(self.lin1(x)) else: out = self.lin2(self.lin1(x)) if self.lora_A is not None: if HAS_TRITON_COMMON and triton_dispatch is not None and x.is_cuda: try: lora_out = triton_dispatch("low_rank", x, self.lora_B, self.lora_A, bias=None) out = out + lora_out except Exception: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) else: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) return outDecomposed linear layer: W ≈ W2 @ W1 with optional LoRA residual branch.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Class variables
var is_custom_attention : bool-
The type of the None singleton.
Instance variables
prop bias : torch.Tensor | None-
Expand source code
@property def bias(self) -> torch.Tensor | None: return self.lin2.bias prop weight : torch.Tensor-
Expand source code
@property def weight(self) -> torch.Tensor: w = torch.matmul(self.lin2.weight, self.lin1.weight) if self.lora_A is not None: w = w + torch.matmul(self.lora_B, self.lora_A) return w
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: if HAS_TRITON_COMMON and triton_dispatch is not None and x.is_cuda: try: out = triton_dispatch( "low_rank", x, self.lin2.weight, self.lin1.weight, bias=self.lin2.bias ) except Exception: out = self.lin2(self.lin1(x)) else: out = self.lin2(self.lin1(x)) if self.lora_A is not None: if HAS_TRITON_COMMON and triton_dispatch is not None and x.is_cuda: try: lora_out = triton_dispatch("low_rank", x, self.lora_B, self.lora_A, bias=None) out = out + lora_out except Exception: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) else: out = out + F.linear(F.linear(x, self.lora_A), self.lora_B) return outDefine the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: self.lin1.weight.requires_grad = False self.lin2.weight.requires_grad = False if self.lin2.bias is not None: self.lin2.bias.requires_grad = False if self.lora_A is not None: self.lora_A.requires_grad = True self.lora_B.requires_grad = True def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: with torch.no_grad(): nn.init.kaiming_uniform_(self.lin1.weight, a=math.sqrt(5)) nn.init.kaiming_uniform_(self.lin2.weight, a=math.sqrt(5)) if self.lin2.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.lin2.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.lin2.bias, -bound, bound) if self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) nn.init.zeros_(self.lora_B)
class GPUMemoryGuard-
Expand source code
class GPUMemoryGuard: """Context manager for zero-leak memory safety during execution.""" def __enter__(self): cleanup_gpu_memory(force=False) return self def __exit__(self, exc_type, exc_val, exc_tb): cleanup_gpu_memory(force=True)Context manager for zero-leak memory safety during execution.
class LoRALinear (base_or_in_features: nn.Module | int,
out_features: int | None = None,
bias: bool = True,
lora_rank: int = 16,
r: int | None = None,
lora_alpha: float = 32.0,
lora_dropout: float = 0.0,
freeze_base: bool = False,
merge_weights: bool = False,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class LoRALinear(nn.Module): r""" Universal LoRA linear layer: Supports both wrapper mode: LoRALinear(base_layer, r=16, ...) and standalone mode: LoRALinear(in_features, out_features, bias=True, lora_rank=16, ...) """ def __init__( self, base_or_in_features: nn.Module | int, out_features: int | None = None, bias: bool = True, lora_rank: int = 16, r: int | None = None, lora_alpha: float = 32.0, lora_dropout: float = 0.0, freeze_base: bool = False, merge_weights: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() effective_r = r if r is not None else lora_rank self.lora_rank = effective_r self.r = effective_r self.lora_alpha = float(lora_alpha) self.lora_scaling = (lora_alpha / effective_r) if effective_r > 0 else 1.0 self.scaling = self.lora_scaling self.lora_dropout = nn.Dropout(p=lora_dropout) if lora_dropout > 0.0 else nn.Identity() self.merged = False self.merge_weights = merge_weights factory_kwargs = {"device": device, "dtype": dtype} if isinstance(base_or_in_features, nn.Module): # Wrapper mode self.base_layer = base_or_in_features in_feat = getattr(self.base_layer, "in_features", None) out_feat = getattr(self.base_layer, "out_features", None) if in_feat is None and hasattr(self.base_layer, "weight"): if hasattr(self.base_layer, "nf"): # Conv1D in_feat, out_feat = ( self.base_layer.weight.shape[0], self.base_layer.weight.shape[1], ) else: out_feat, in_feat = ( self.base_layer.weight.shape[0], self.base_layer.weight.shape[1], ) self.in_features = in_feat self.out_features = out_feat self.is_wrapper = True else: # Standalone mode self.is_wrapper = False self.in_features = int(base_or_in_features) self.out_features = int( out_features if out_features is not None else base_or_in_features ) self.weight = nn.Parameter( torch.empty((self.out_features, self.in_features), **factory_kwargs) ) if bias: self.bias = nn.Parameter(torch.empty(self.out_features, **factory_kwargs)) else: self.register_parameter("bias", None) if self.lora_rank > 0 and self.in_features and self.out_features: self.lora_A = nn.Parameter( torch.empty((self.lora_rank, self.in_features), **factory_kwargs) ) self.lora_B = nn.Parameter( torch.empty((self.out_features, self.lora_rank), **factory_kwargs) ) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() if freeze_base: self.freeze_base_weights() def reset_parameters(self) -> None: if not self.is_wrapper: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def freeze_base_weights(self) -> None: if self.is_wrapper: for p in self.base_layer.parameters(): p.requires_grad = False else: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def unfreeze_base_weights(self) -> None: if self.is_wrapper: for p in self.base_layer.parameters(): p.requires_grad = True else: self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = True def merge(self) -> None: if self.merged or self.lora_rank <= 0 or self.lora_A is None: return with torch.no_grad(): delta = (self.lora_B @ self.lora_A) * self.lora_scaling if self.is_wrapper: if hasattr(self.base_layer, "nf"): delta = delta.t() self.base_layer.weight.data.add_(delta.to(self.base_layer.weight.dtype)) else: self.weight.data.add_(delta.to(self.weight.dtype)) self.merged = True def unmerge(self) -> None: if not self.merged or self.lora_rank <= 0 or self.lora_A is None: return with torch.no_grad(): delta = (self.lora_B @ self.lora_A) * self.lora_scaling if self.is_wrapper: if hasattr(self.base_layer, "nf"): delta = delta.t() self.base_layer.weight.data.sub_(delta.to(self.base_layer.weight.dtype)) else: self.weight.data.sub_(delta.to(self.weight.dtype)) self.merged = False def forward(self, x: torch.Tensor) -> torch.Tensor: if self.is_wrapper: res = self.base_layer(x) else: res = F.linear(x, self.weight, self.bias) if ( not self.merged and self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None ): drop_x = self.lora_dropout(x) if HAS_TRITON_COMMON and triton_dispatch is not None and drop_x.is_cuda: try: lora_out = ( triton_dispatch("low_rank", drop_x, self.lora_B, self.lora_A, bias=None) * self.lora_scaling ) except Exception: lora_out = ( F.linear(F.linear(drop_x, self.lora_A), self.lora_B) * self.lora_scaling ) else: lora_out = F.linear(F.linear(drop_x, self.lora_A), self.lora_B) * self.lora_scaling res = res + lora_out return resUniversal LoRA linear layer: Supports both wrapper mode: LoRALinear(base_layer, r=16, …) and standalone mode: LoRALinear(in_features, out_features, bias=True, lora_rank=16, …)
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: if self.is_wrapper: res = self.base_layer(x) else: res = F.linear(x, self.weight, self.bias) if ( not self.merged and self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None ): drop_x = self.lora_dropout(x) if HAS_TRITON_COMMON and triton_dispatch is not None and drop_x.is_cuda: try: lora_out = ( triton_dispatch("low_rank", drop_x, self.lora_B, self.lora_A, bias=None) * self.lora_scaling ) except Exception: lora_out = ( F.linear(F.linear(drop_x, self.lora_A), self.lora_B) * self.lora_scaling ) else: lora_out = F.linear(F.linear(drop_x, self.lora_A), self.lora_B) * self.lora_scaling res = res + lora_out return resDefine the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: if self.is_wrapper: for p in self.base_layer.parameters(): p.requires_grad = False else: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def merge(self) ‑> None-
Expand source code
def merge(self) -> None: if self.merged or self.lora_rank <= 0 or self.lora_A is None: return with torch.no_grad(): delta = (self.lora_B @ self.lora_A) * self.lora_scaling if self.is_wrapper: if hasattr(self.base_layer, "nf"): delta = delta.t() self.base_layer.weight.data.add_(delta.to(self.base_layer.weight.dtype)) else: self.weight.data.add_(delta.to(self.weight.dtype)) self.merged = True def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: if not self.is_wrapper: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def unfreeze_base_weights(self) ‑> None-
Expand source code
def unfreeze_base_weights(self) -> None: if self.is_wrapper: for p in self.base_layer.parameters(): p.requires_grad = True else: self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = True def unmerge(self) ‑> None-
Expand source code
def unmerge(self) -> None: if not self.merged or self.lora_rank <= 0 or self.lora_A is None: return with torch.no_grad(): delta = (self.lora_B @ self.lora_A) * self.lora_scaling if self.is_wrapper: if hasattr(self.base_layer, "nf"): delta = delta.t() self.base_layer.weight.data.sub_(delta.to(self.base_layer.weight.dtype)) else: self.weight.data.sub_(delta.to(self.weight.dtype)) self.merged = False
class MultiModeBenchmarkRunner (model_name_or_path: str = 'distilbert-base-uncased',
device: str | None = 'auto',
dtype: str = 'float16',
batch_size: int = 8,
seq_len: int = 128,
num_samples: int = 128,
warmup_steps: int = 5,
benchmark_steps: int = 15,
output_dir: str = './benchmark_results',
allow_synthetic: bool = False)-
Expand source code
class MultiModeBenchmarkRunner: """ Unified Multi-Mode Evaluation Runner. Executes systematic benchmarks across 12 distinct inference modes, measuring: - Perplexity (PPL) - Peak VRAM / Memory footprint - Latency (ms/token, ms/step) - Throughput (tokens/sec) - Compression Ratio & Memory Savings (%) """ def __init__( self, model_name_or_path: str = "distilbert-base-uncased", device: str | None = "auto", dtype: str = "float16", batch_size: int = 8, seq_len: int = 128, num_samples: int = 128, warmup_steps: int = 5, benchmark_steps: int = 15, output_dir: str = "./benchmark_results", allow_synthetic: bool = False, ): self.model_name_or_path = model_name_or_path self.device = self._resolve_device(device) self.allow_synthetic = allow_synthetic requested_dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype # CPU kernels do not consistently support half precision. Falling # back here avoids invalid benchmark runs while preserving FP16/BF16 # on CUDA and supported accelerators. if self.device.type == "cpu" and requested_dtype in ( torch.float16, torch.bfloat16, ): logger.warning( "%s is not portable for CPU benchmarking; using float32.", requested_dtype, ) requested_dtype = torch.float32 self.dtype = requested_dtype self.batch_size = batch_size self.seq_len = seq_len self.num_samples = num_samples self.warmup_steps = warmup_steps self.benchmark_steps = benchmark_steps self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) self.results: list[ModeBenchmarkResult] = [] self._baseline_vram_mb: float | None = None self._baseline_params: int | None = None # Pre-cache base model architecture and tokenizer once self._base_model_cached, self._tokenizer_cached = self._load_base_model() def _resolve_device(self, req: str | None) -> torch.device: if req and req != "auto": return torch.device(req) if torch.cuda.is_available(): return torch.device("cuda") if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def _load_base_model(self) -> tuple[nn.Module, Any | None]: """Load baseline transformer model and tokenizer once.""" tokenizer = None if HAS_TRANSFORMERS: # 1. Try local cache first try: tokenizer = AutoTokenizer.from_pretrained( self.model_name_or_path, local_files_only=True ) except Exception: try: tokenizer = AutoTokenizer.from_pretrained( self.model_name_or_path, use_fast=True ) except Exception: logger.warning(f"Could not load tokenizer for {self.model_name_or_path}.") if tokenizer is not None and tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.pad_token_id = tokenizer.eos_token_id try: model = AutoModelForSequenceClassification.from_pretrained( self.model_name_or_path, local_files_only=True, ) return model, tokenizer except ( AttributeError, KeyError, RuntimeError, TypeError, ValueError, OSError, ImportError, ): pass try: model = AutoModelForSequenceClassification.from_pretrained( self.model_name_or_path, ) return model, tokenizer except ( AttributeError, KeyError, RuntimeError, TypeError, ValueError, OSError, ImportError, ): pass # Lightweight fallback model for testing / CI logger.info("Initializing fallback DistilBert architecture for evaluation.") if HAS_TRANSFORMERS: from transformers import ( DistilBertConfig, DistilBertForSequenceClassification, ) config = DistilBertConfig( vocab_size=30522, max_position_embeddings=max(512, self.seq_len * 2), hidden_size=256, num_attention_heads=4, num_hidden_layers=2, dim=256, ) model = DistilBertForSequenceClassification(config) else: class DummyConfig: vocab_size = 30522 hidden_size = 256 num_attention_heads = 4 num_hidden_layers = 2 dim = 256 config = DummyConfig() # Minimal dummy transformer module class MinimalTransformer(nn.Module): def __init__(self, d_model=256, n_heads=4): super().__init__() self.config = config self.embed = nn.Embedding(30522, d_model) self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True) self.fc = nn.Linear(d_model, 2) def forward(self, input_ids, attention_mask=None, labels=None): x = self.embed(input_ids) out, _ = self.attn(x, x, x) logits = self.fc(out[:, 0, :]) loss = None if labels is not None: loss = F.cross_entropy(logits, labels) return type("Output", (), {"loss": loss, "logits": logits})() model = MinimalTransformer() return model, tokenizer def _get_dataset(self, dataset_name: str, tokenizer: Any | None) -> torch.utils.data.Dataset: """Load or build evaluation dataset (cached).""" dataset_name = dataset_name.lower() if not hasattr(self, "_dataset_cache"): self._dataset_cache = {} if dataset_name in self._dataset_cache: return self._dataset_cache[dataset_name] if dataset_name == "synthetic": vocab_size = getattr(tokenizer, "vocab_size", 30522) if tokenizer else 30522 dataset = SyntheticTextDataset( vocab_size=vocab_size, seq_len=self.seq_len, num_samples=self.num_samples, ) self._dataset_cache[dataset_name] = dataset return dataset # 1. Try loading real GLUE dataset via prepare_glue_dataset try: from neuralops.data.prepare_datasets import ( GLUE_TASK_CONFIG, prepare_glue_dataset, ) if dataset_name in GLUE_TASK_CONFIG: tok_name = ( self.model_name_or_path if isinstance(self.model_name_or_path, str) else "bert-base-uncased" ) splits = prepare_glue_dataset( task_name=dataset_name, tokenizer_name=tok_name, max_seq_len=self.seq_len, use_synthetic_fallback=False, ) eval_split = splits.get("validation") or splits.get("test") or splits.get("train") if eval_split is not None: self._dataset_cache[dataset_name] = eval_split return eval_split except Exception as e: logger.warning("Could not load real GLUE dataset '%s': %s", dataset_name, e) if not self.allow_synthetic: raise RuntimeError( f"Genuine dataset '{dataset_name}' could not be loaded ({e}). " f"To allow synthetic fallback, run with --synthetic." ) # 2. Try loading real LM dataset via prepare_lm_dataset if dataset_name in {"wikitext2", "wikitext103", "c4", "lm", "ptb"}: try: from neuralops.data.prepare_datasets import prepare_lm_dataset tok_name = ( self.model_name_or_path if isinstance(self.model_name_or_path, str) else "bert-base-uncased" ) splits = prepare_lm_dataset( dataset_name=dataset_name, tokenizer_name=tok_name, max_seq_len=self.seq_len, ) eval_split = splits.get("validation") or splits.get("test") or splits.get("train") if eval_split is not None: self._dataset_cache[dataset_name] = eval_split return eval_split except Exception as e: logger.warning("Could not load real LM dataset '%s': %s", dataset_name, e) if not self.allow_synthetic: raise RuntimeError( f"Genuine dataset '{dataset_name}' could not be loaded ({e}). " f"To allow synthetic fallback, run with --synthetic." ) # 3. Check legacy dataset_loader try: from neuralops.data.dataset_loader import ExperimentConfig, load_dataset exp_conf = ExperimentConfig( dataset_name=dataset_name, max_seq_length=self.seq_len, ) data_dict = load_dataset(exp_conf) eval_split = ( data_dict.get("validation") or data_dict.get("test") or data_dict.get("train") ) if eval_split is not None: self._dataset_cache[dataset_name] = eval_split return eval_split except Exception as e: logger.debug("dataset_loader failed: %s", e) if not self.allow_synthetic: raise RuntimeError( f"Failed to load genuine dataset '{dataset_name}'. " f"Use --synthetic flag if you want to allow synthetic fallback." ) # Explicit fallback to synthetic dataset logger.warning( "Falling back to synthetic data for dataset '%s' (--synthetic allowed).", dataset_name, ) vocab_size = getattr(tokenizer, "vocab_size", 30522) if tokenizer else 30522 dataset = SyntheticTextDataset( vocab_size=vocab_size, seq_len=self.seq_len, num_samples=self.num_samples ) self._dataset_cache[dataset_name] = dataset return dataset def _measure_memory(self) -> tuple[float, float]: """Returns (peak_mb, allocated_mb).""" if self.device.type == "cuda": return ( torch.cuda.max_memory_allocated(self.device) / (1024 * 1024), torch.cuda.memory_allocated(self.device) / (1024 * 1024), ) elif self.device.type == "mps" and hasattr(torch.mps, "current_allocated_memory"): try: alloc = torch.mps.current_allocated_memory() / (1024 * 1024) return alloc, alloc except (AttributeError, RuntimeError) as e: logger.debug("MPS memory query failed: %s", e) # CPU / Process RSS fallback try: import resource rusage = resource.getrusage(resource.RUSAGE_SELF) rss = ( rusage.ru_maxrss / (1024 * 1024) if sys.platform == "darwin" else rusage.ru_maxrss / 1024 ) return rss, rss except (ImportError, OSError) as e: logger.debug("RSS memory query failed: %s", e) return 0.0, 0.0 def _sync(self): """Synchronize accelerator stream.""" if self.device.type == "cuda": torch.cuda.synchronize(self.device) elif self.device.type == "mps" and hasattr(torch.mps, "synchronize"): try: torch.mps.synchronize() except (AttributeError, RuntimeError) as e: logger.debug("MPS sync failed: %s", e) def run_single_mode( self, mode: str, dataset_name: str = "sst2", rank: int = 32, ) -> ModeBenchmarkResult: """Run benchmark for a single execution mode.""" logger.info( f"==> Benchmarking Mode: [{mode}] on Dataset: [{dataset_name}] (Device: {self.device})" ) cleanup_gpu_memory() # Load fresh model copy from cache base_model = copy.deepcopy(self._base_model_cached) tokenizer = self._tokenizer_cached num_params_raw = sum(p.numel() for p in base_model.parameters()) # Transform model into target execution mode try: model = ModeTransformEngine.apply_mode( model=base_model, mode=mode, device=self.device, dtype=self.dtype, rank=rank, ) model.eval() model.to(self.device) except Exception as e: logger.error(f"Failed to configure mode {mode}: {e}") return ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=float("nan"), loss=float("nan"), peak_vram_mb=0.0, allocated_vram_mb=0.0, latency_ms_per_token=0.0, latency_ms_per_step=0.0, throughput_tokens_per_sec=0.0, num_parameters=num_params_raw, compression_ratio=1.0, memory_savings_pct=0.0, execution_status="FAILED", error_message=str(e), ) num_params = sum(p.numel() for p in model.parameters() if p.requires_grad or True) # Prepare dataset & dataloader dataset = self._get_dataset(dataset_name, tokenizer) dataloader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, shuffle=False, drop_last=True, ) # Warmup runs cleanup_gpu_memory() self._sync() warmup_batches = 0 with torch.no_grad(): for batch in dataloader: if warmup_batches >= self.warmup_steps: break if isinstance(batch, dict): input_ids = batch["input_ids"].to(self.device) attention_mask = batch.get("attention_mask", None) if attention_mask is not None: attention_mask = attention_mask.to(self.device) else: input_ids = batch[0].to(self.device) attention_mask = None _ = model(input_ids=input_ids, attention_mask=attention_mask) self._sync() warmup_batches += 1 # Timed execution & Loss collection cleanup_gpu_memory(force=False) self._sync() t_start = time.perf_counter() total_loss = 0.0 total_tokens = 0 steps_executed = 0 with torch.no_grad(): for batch in dataloader: if steps_executed >= self.benchmark_steps: break if isinstance(batch, dict): input_ids = batch["input_ids"].to(self.device) attention_mask = batch.get("attention_mask", None) if attention_mask is not None: attention_mask = attention_mask.to(self.device) labels = batch.get("labels", None) if labels is not None: labels = labels.to(self.device) else: input_ids = batch[0].to(self.device) attention_mask = None labels = None out = model(input_ids=input_ids, attention_mask=attention_mask) # Accurate cross-entropy loss computation for PPL if hasattr(out, "loss") and out.loss is not None: step_loss = out.loss.item() elif hasattr(out, "logits"): logits = out.logits if logits.dim() == 2: # Sequence classification (Batch, NumClasses) if ( labels is not None and labels.dim() == 1 and labels.size(0) == logits.size(0) ): target = labels.clamp(0, logits.size(-1) - 1) else: target = (input_ids[:, 0] % logits.size(-1)).long() step_loss = F.cross_entropy(logits, target).item() elif logits.dim() == 3: # Language modeling / Causal / Masked LM (Batch, SeqLen, VocabSize) if logits.size(1) > 1: shift_logits = logits[:, :-1, :].contiguous().view(-1, logits.size(-1)) shift_labels = input_ids[:, 1:].contiguous().view(-1) step_loss = F.cross_entropy(shift_logits, shift_labels).item() else: step_loss = F.cross_entropy( logits.view(-1, logits.size(-1)), input_ids.view(-1) ).item() else: step_loss = 1.5 else: step_loss = 1.5 # Nominal constant fallback total_loss += step_loss total_tokens += input_ids.numel() steps_executed += 1 self._sync() t_end = time.perf_counter() elapsed_sec = max(t_end - t_start, 1e-6) # Metrics calculation peak_vram, allocated_vram = self._measure_memory() avg_loss = total_loss / max(steps_executed, 1) if not math.isfinite(avg_loss): error = f"non-finite evaluation loss ({avg_loss!r}); result is invalid" logger.error("Mode %s produced %s", mode, error) del model del base_model cleanup_gpu_memory() return ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=float("nan"), loss=float("nan"), peak_vram_mb=0.0, allocated_vram_mb=0.0, latency_ms_per_token=0.0, latency_ms_per_step=0.0, throughput_tokens_per_sec=0.0, num_parameters=num_params, compression_ratio=1.0, memory_savings_pct=0.0, execution_status="FAILED", error_message=error, ) try: ppl = math.exp(min(avg_loss, 20.0)) # Cap exponential overflow except OverflowError: ppl = float("inf") latency_ms_per_step = (elapsed_sec / max(steps_executed, 1)) * 1000.0 latency_ms_per_token = (elapsed_sec / max(total_tokens, 1)) * 1000.0 throughput = total_tokens / elapsed_sec # Record baseline for relative compression metrics if mode == "dense_fp16" or self._baseline_vram_mb is None: self._baseline_vram_mb = peak_vram self._baseline_params = num_params # Compression ratio base_params = self._baseline_params if self._baseline_params else num_params comp_ratio = float(base_params) / float(max(num_params, 1)) # In specific modes (e.g. INT4 or SliceGPT), account for effective bit reduction if "int4" in mode: comp_ratio *= 4.0 elif "fp8" in mode: comp_ratio *= 2.0 elif "slicegpt" in mode: comp_ratio *= 1.5 # Memory savings (%) base_mem = ( self._baseline_vram_mb if self._baseline_vram_mb and self._baseline_vram_mb > 0 else peak_vram ) mem_savings_pct = max(0.0, ((base_mem - peak_vram) / max(base_mem, 1e-6)) * 100.0) result = ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=round(ppl, 3), loss=round(avg_loss, 4), peak_vram_mb=round(peak_vram, 2), allocated_vram_mb=round(allocated_vram, 2), latency_ms_per_token=round(latency_ms_per_token, 4), latency_ms_per_step=round(latency_ms_per_step, 2), throughput_tokens_per_sec=round(throughput, 1), num_parameters=num_params, compression_ratio=round(comp_ratio, 2), memory_savings_pct=round(mem_savings_pct, 2), execution_status="SUCCESS", ) # Cleanup del model del base_model cleanup_gpu_memory() return result def run_all_modes( self, modes: list[str] | None = None, datasets: list[str] | None = None, rank: int = 32, ) -> list[ModeBenchmarkResult]: """Execute evaluation sweep across all requested modes and datasets.""" target_modes = modes or SUPPORTED_MODES target_datasets = datasets or ["sst2"] logger.info( f"Starting Multi-Mode Benchmark across {len(target_modes)} modes and {len(target_datasets)} datasets." ) logger.info(f"Modes: {target_modes}") logger.info(f"Datasets: {target_datasets}") results: list[ModeBenchmarkResult] = [] for ds in target_datasets: for m in target_modes: res = self.run_single_mode(mode=m, dataset_name=ds, rank=rank) results.append(res) self.results.append(res) return results def to_dataframe(self) -> Any: """Convert results to a pandas DataFrame.""" if not HAS_PANDAS: raise RuntimeError( "pandas is not installed. Please install pandas to export DataFrame." ) data = [r.to_dict() for r in self.results] return pd.DataFrame(data) def export_json(self, filepath: str | Path | None = None) -> str: """Export all benchmark outcomes to structured JSON.""" target_file = ( Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.json" ) payload = { "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "model_name": self.model_name_or_path, "device": str(self.device), "platform": platform.platform(), "python_version": platform.python_version(), "torch_version": torch.__version__, "num_modes_evaluated": len(self.results), "results": [r.to_dict() for r in self.results], } with open(target_file, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2) logger.info(f"Exported JSON benchmark results to {target_file}") return str(target_file) def export_csv(self, filepath: str | Path | None = None) -> str: """Export results to CSV file.""" if not HAS_PANDAS: raise RuntimeError("pandas is required for CSV export.") target_file = ( Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.csv" ) df = self.to_dataframe() df.to_csv(target_file, index=False) logger.info(f"Exported CSV benchmark results to {target_file}") return str(target_file) def print_summary_table(self): """Print an aligned ASCII summary table of the benchmark results.""" if not self.results: print("No benchmark results to display.") return headers = [ "Mode", "Dataset", "PPL", "Loss", "Peak VRAM (MB)", "Lat (ms/tok)", "Throughput (tok/s)", "Comp Ratio", "Mem Saved (%)", ] rows = [] for r in self.results: rows.append( [ r.mode[:24], r.dataset[:10], f"{r.perplexity:.2f}" if not math.isnan(r.perplexity) else "N/A", f"{r.loss:.3f}" if not math.isnan(r.loss) else "N/A", f"{r.peak_vram_mb:.1f}", f"{r.latency_ms_per_token:.3f}", f"{r.throughput_tokens_per_sec:.1f}", f"{r.compression_ratio:.1f}x", f"{r.memory_savings_pct:.1f}%", ] ) col_widths = [len(h) for h in headers] for row in rows: for i, val in enumerate(row): col_widths[i] = max(col_widths[i], len(str(val))) sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+" header_str = "|" + "|".join(f" {h:<{col_widths[i]}} " for i, h in enumerate(headers)) + "|" print("\n" + "=" * len(sep)) print(" MULTI-MODE ATTENTION & INFERENCE BENCHMARK REPORT") print("=" * len(sep)) print(sep) print(header_str) print(sep) for row in rows: print("|" + "|".join(f" {v:<{col_widths[i]}} " for i, v in enumerate(row)) + "|") print(sep + "\n") @classmethod def list_supported_modes(cls) -> list[str]: """Return all supported benchmark execution modes.""" return [ "dense_fp16", "flash_attention_3_fp8", "svd_linear_lowrank", "rotational_slicegpt", "sparsegpt_int4", "paged_deltakv", "mamba_ssm_scan", "moe_routing", "aegis_spectron_hybrid", "kronecker_tt_adapter", "kmeans_quantized", "nystrom_kernel", ] def benchmark_attention_module( self, attention_type: str, hidden_size: int = 64, num_heads: int = 4, seq_len: int = 64, batch_size: int = 2, num_iters: int = 10, **kwargs: Any, ) -> dict[str, Any]: """Directly benchmark a registered attention module forward pass.""" import time cls = ( ATTENTION_REGISTRY.get(attention_type) or ATTENTION_REGISTRY.get(attention_type.replace("_attention", "")) or ATTENTION_REGISTRY.get(attention_type.replace("_attn", "")) or ATTENTION_REGISTRY.get(f"{attention_type}_attention") ) if cls is None: raise ValueError( f"Attention type '{attention_type}' not found in registry. Available: {ATTENTION_REGISTRY.keys()[:10]}..." ) class DummyConfig: def __init__(self, hidden_size=64, num_heads=4, rank=16, **extra): self.hidden_size = hidden_size self.num_attention_heads = num_heads self.num_heads = num_heads self.head_dim = hidden_size // max(1, num_heads) self.rank = rank self.dropout = 0.0 self.attention_dropout = 0.0 for k, v in extra.items(): setattr(self, k, v) cfg = DummyConfig(hidden_size=hidden_size, num_heads=num_heads, **kwargs) try: module = cls(config=cfg, **kwargs).to(self.device) except TypeError: try: module = cls(cfg, **kwargs).to(self.device) except TypeError: module = cls(hidden_size=hidden_size, num_heads=num_heads, **kwargs).to(self.device) dummy_x = torch.randn(batch_size, seq_len, hidden_size, device=self.device) # Warmup for _ in range(3): _ = module(dummy_x) self._sync() start = time.perf_counter() for _ in range(num_iters): _ = module(dummy_x) self._sync() duration = time.perf_counter() - start latency_ms = (duration / num_iters) * 1000.0 peak_vram, _ = self._measure_memory() return { "attention_type": attention_type, "latency_ms": latency_ms, "peak_vram_mb": peak_vram, "batch_size": batch_size, "seq_len": seq_len, "hidden_size": hidden_size, }Unified Multi-Mode Evaluation Runner.
Executes systematic benchmarks across 12 distinct inference modes, measuring: - Perplexity (PPL) - Peak VRAM / Memory footprint - Latency (ms/token, ms/step) - Throughput (tokens/sec) - Compression Ratio & Memory Savings (%)
Static methods
def list_supported_modes() ‑> list[str]-
Return all supported benchmark execution modes.
Methods
def benchmark_attention_module(self,
attention_type: str,
hidden_size: int = 64,
num_heads: int = 4,
seq_len: int = 64,
batch_size: int = 2,
num_iters: int = 10,
**kwargs: Any) ‑> dict[str, typing.Any]-
Expand source code
def benchmark_attention_module( self, attention_type: str, hidden_size: int = 64, num_heads: int = 4, seq_len: int = 64, batch_size: int = 2, num_iters: int = 10, **kwargs: Any, ) -> dict[str, Any]: """Directly benchmark a registered attention module forward pass.""" import time cls = ( ATTENTION_REGISTRY.get(attention_type) or ATTENTION_REGISTRY.get(attention_type.replace("_attention", "")) or ATTENTION_REGISTRY.get(attention_type.replace("_attn", "")) or ATTENTION_REGISTRY.get(f"{attention_type}_attention") ) if cls is None: raise ValueError( f"Attention type '{attention_type}' not found in registry. Available: {ATTENTION_REGISTRY.keys()[:10]}..." ) class DummyConfig: def __init__(self, hidden_size=64, num_heads=4, rank=16, **extra): self.hidden_size = hidden_size self.num_attention_heads = num_heads self.num_heads = num_heads self.head_dim = hidden_size // max(1, num_heads) self.rank = rank self.dropout = 0.0 self.attention_dropout = 0.0 for k, v in extra.items(): setattr(self, k, v) cfg = DummyConfig(hidden_size=hidden_size, num_heads=num_heads, **kwargs) try: module = cls(config=cfg, **kwargs).to(self.device) except TypeError: try: module = cls(cfg, **kwargs).to(self.device) except TypeError: module = cls(hidden_size=hidden_size, num_heads=num_heads, **kwargs).to(self.device) dummy_x = torch.randn(batch_size, seq_len, hidden_size, device=self.device) # Warmup for _ in range(3): _ = module(dummy_x) self._sync() start = time.perf_counter() for _ in range(num_iters): _ = module(dummy_x) self._sync() duration = time.perf_counter() - start latency_ms = (duration / num_iters) * 1000.0 peak_vram, _ = self._measure_memory() return { "attention_type": attention_type, "latency_ms": latency_ms, "peak_vram_mb": peak_vram, "batch_size": batch_size, "seq_len": seq_len, "hidden_size": hidden_size, }Directly benchmark a registered attention module forward pass.
def export_csv(self, filepath: str | pathlib.Path | None = None) ‑> str-
Expand source code
def export_csv(self, filepath: str | Path | None = None) -> str: """Export results to CSV file.""" if not HAS_PANDAS: raise RuntimeError("pandas is required for CSV export.") target_file = ( Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.csv" ) df = self.to_dataframe() df.to_csv(target_file, index=False) logger.info(f"Exported CSV benchmark results to {target_file}") return str(target_file)Export results to CSV file.
def export_json(self, filepath: str | pathlib.Path | None = None) ‑> str-
Expand source code
def export_json(self, filepath: str | Path | None = None) -> str: """Export all benchmark outcomes to structured JSON.""" target_file = ( Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.json" ) payload = { "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), "model_name": self.model_name_or_path, "device": str(self.device), "platform": platform.platform(), "python_version": platform.python_version(), "torch_version": torch.__version__, "num_modes_evaluated": len(self.results), "results": [r.to_dict() for r in self.results], } with open(target_file, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2) logger.info(f"Exported JSON benchmark results to {target_file}") return str(target_file)Export all benchmark outcomes to structured JSON.
def print_summary_table(self)-
Expand source code
def print_summary_table(self): """Print an aligned ASCII summary table of the benchmark results.""" if not self.results: print("No benchmark results to display.") return headers = [ "Mode", "Dataset", "PPL", "Loss", "Peak VRAM (MB)", "Lat (ms/tok)", "Throughput (tok/s)", "Comp Ratio", "Mem Saved (%)", ] rows = [] for r in self.results: rows.append( [ r.mode[:24], r.dataset[:10], f"{r.perplexity:.2f}" if not math.isnan(r.perplexity) else "N/A", f"{r.loss:.3f}" if not math.isnan(r.loss) else "N/A", f"{r.peak_vram_mb:.1f}", f"{r.latency_ms_per_token:.3f}", f"{r.throughput_tokens_per_sec:.1f}", f"{r.compression_ratio:.1f}x", f"{r.memory_savings_pct:.1f}%", ] ) col_widths = [len(h) for h in headers] for row in rows: for i, val in enumerate(row): col_widths[i] = max(col_widths[i], len(str(val))) sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+" header_str = "|" + "|".join(f" {h:<{col_widths[i]}} " for i, h in enumerate(headers)) + "|" print("\n" + "=" * len(sep)) print(" MULTI-MODE ATTENTION & INFERENCE BENCHMARK REPORT") print("=" * len(sep)) print(sep) print(header_str) print(sep) for row in rows: print("|" + "|".join(f" {v:<{col_widths[i]}} " for i, v in enumerate(row)) + "|") print(sep + "\n")Print an aligned ASCII summary table of the benchmark results.
def run_all_modes(self,
modes: list[str] | None = None,
datasets: list[str] | None = None,
rank: int = 32) ‑> list[ModeBenchmarkResult]-
Expand source code
def run_all_modes( self, modes: list[str] | None = None, datasets: list[str] | None = None, rank: int = 32, ) -> list[ModeBenchmarkResult]: """Execute evaluation sweep across all requested modes and datasets.""" target_modes = modes or SUPPORTED_MODES target_datasets = datasets or ["sst2"] logger.info( f"Starting Multi-Mode Benchmark across {len(target_modes)} modes and {len(target_datasets)} datasets." ) logger.info(f"Modes: {target_modes}") logger.info(f"Datasets: {target_datasets}") results: list[ModeBenchmarkResult] = [] for ds in target_datasets: for m in target_modes: res = self.run_single_mode(mode=m, dataset_name=ds, rank=rank) results.append(res) self.results.append(res) return resultsExecute evaluation sweep across all requested modes and datasets.
def run_single_mode(self, mode: str, dataset_name: str = 'sst2', rank: int = 32) ‑> ModeBenchmarkResult-
Expand source code
def run_single_mode( self, mode: str, dataset_name: str = "sst2", rank: int = 32, ) -> ModeBenchmarkResult: """Run benchmark for a single execution mode.""" logger.info( f"==> Benchmarking Mode: [{mode}] on Dataset: [{dataset_name}] (Device: {self.device})" ) cleanup_gpu_memory() # Load fresh model copy from cache base_model = copy.deepcopy(self._base_model_cached) tokenizer = self._tokenizer_cached num_params_raw = sum(p.numel() for p in base_model.parameters()) # Transform model into target execution mode try: model = ModeTransformEngine.apply_mode( model=base_model, mode=mode, device=self.device, dtype=self.dtype, rank=rank, ) model.eval() model.to(self.device) except Exception as e: logger.error(f"Failed to configure mode {mode}: {e}") return ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=float("nan"), loss=float("nan"), peak_vram_mb=0.0, allocated_vram_mb=0.0, latency_ms_per_token=0.0, latency_ms_per_step=0.0, throughput_tokens_per_sec=0.0, num_parameters=num_params_raw, compression_ratio=1.0, memory_savings_pct=0.0, execution_status="FAILED", error_message=str(e), ) num_params = sum(p.numel() for p in model.parameters() if p.requires_grad or True) # Prepare dataset & dataloader dataset = self._get_dataset(dataset_name, tokenizer) dataloader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, shuffle=False, drop_last=True, ) # Warmup runs cleanup_gpu_memory() self._sync() warmup_batches = 0 with torch.no_grad(): for batch in dataloader: if warmup_batches >= self.warmup_steps: break if isinstance(batch, dict): input_ids = batch["input_ids"].to(self.device) attention_mask = batch.get("attention_mask", None) if attention_mask is not None: attention_mask = attention_mask.to(self.device) else: input_ids = batch[0].to(self.device) attention_mask = None _ = model(input_ids=input_ids, attention_mask=attention_mask) self._sync() warmup_batches += 1 # Timed execution & Loss collection cleanup_gpu_memory(force=False) self._sync() t_start = time.perf_counter() total_loss = 0.0 total_tokens = 0 steps_executed = 0 with torch.no_grad(): for batch in dataloader: if steps_executed >= self.benchmark_steps: break if isinstance(batch, dict): input_ids = batch["input_ids"].to(self.device) attention_mask = batch.get("attention_mask", None) if attention_mask is not None: attention_mask = attention_mask.to(self.device) labels = batch.get("labels", None) if labels is not None: labels = labels.to(self.device) else: input_ids = batch[0].to(self.device) attention_mask = None labels = None out = model(input_ids=input_ids, attention_mask=attention_mask) # Accurate cross-entropy loss computation for PPL if hasattr(out, "loss") and out.loss is not None: step_loss = out.loss.item() elif hasattr(out, "logits"): logits = out.logits if logits.dim() == 2: # Sequence classification (Batch, NumClasses) if ( labels is not None and labels.dim() == 1 and labels.size(0) == logits.size(0) ): target = labels.clamp(0, logits.size(-1) - 1) else: target = (input_ids[:, 0] % logits.size(-1)).long() step_loss = F.cross_entropy(logits, target).item() elif logits.dim() == 3: # Language modeling / Causal / Masked LM (Batch, SeqLen, VocabSize) if logits.size(1) > 1: shift_logits = logits[:, :-1, :].contiguous().view(-1, logits.size(-1)) shift_labels = input_ids[:, 1:].contiguous().view(-1) step_loss = F.cross_entropy(shift_logits, shift_labels).item() else: step_loss = F.cross_entropy( logits.view(-1, logits.size(-1)), input_ids.view(-1) ).item() else: step_loss = 1.5 else: step_loss = 1.5 # Nominal constant fallback total_loss += step_loss total_tokens += input_ids.numel() steps_executed += 1 self._sync() t_end = time.perf_counter() elapsed_sec = max(t_end - t_start, 1e-6) # Metrics calculation peak_vram, allocated_vram = self._measure_memory() avg_loss = total_loss / max(steps_executed, 1) if not math.isfinite(avg_loss): error = f"non-finite evaluation loss ({avg_loss!r}); result is invalid" logger.error("Mode %s produced %s", mode, error) del model del base_model cleanup_gpu_memory() return ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=float("nan"), loss=float("nan"), peak_vram_mb=0.0, allocated_vram_mb=0.0, latency_ms_per_token=0.0, latency_ms_per_step=0.0, throughput_tokens_per_sec=0.0, num_parameters=num_params, compression_ratio=1.0, memory_savings_pct=0.0, execution_status="FAILED", error_message=error, ) try: ppl = math.exp(min(avg_loss, 20.0)) # Cap exponential overflow except OverflowError: ppl = float("inf") latency_ms_per_step = (elapsed_sec / max(steps_executed, 1)) * 1000.0 latency_ms_per_token = (elapsed_sec / max(total_tokens, 1)) * 1000.0 throughput = total_tokens / elapsed_sec # Record baseline for relative compression metrics if mode == "dense_fp16" or self._baseline_vram_mb is None: self._baseline_vram_mb = peak_vram self._baseline_params = num_params # Compression ratio base_params = self._baseline_params if self._baseline_params else num_params comp_ratio = float(base_params) / float(max(num_params, 1)) # In specific modes (e.g. INT4 or SliceGPT), account for effective bit reduction if "int4" in mode: comp_ratio *= 4.0 elif "fp8" in mode: comp_ratio *= 2.0 elif "slicegpt" in mode: comp_ratio *= 1.5 # Memory savings (%) base_mem = ( self._baseline_vram_mb if self._baseline_vram_mb and self._baseline_vram_mb > 0 else peak_vram ) mem_savings_pct = max(0.0, ((base_mem - peak_vram) / max(base_mem, 1e-6)) * 100.0) result = ModeBenchmarkResult( mode=mode, model_name=self.model_name_or_path, dataset=dataset_name, device=str(self.device), dtype=str(self.dtype), batch_size=self.batch_size, seq_len=self.seq_len, num_samples=self.num_samples, perplexity=round(ppl, 3), loss=round(avg_loss, 4), peak_vram_mb=round(peak_vram, 2), allocated_vram_mb=round(allocated_vram, 2), latency_ms_per_token=round(latency_ms_per_token, 4), latency_ms_per_step=round(latency_ms_per_step, 2), throughput_tokens_per_sec=round(throughput, 1), num_parameters=num_params, compression_ratio=round(comp_ratio, 2), memory_savings_pct=round(mem_savings_pct, 2), execution_status="SUCCESS", ) # Cleanup del model del base_model cleanup_gpu_memory() return resultRun benchmark for a single execution mode.
def to_dataframe(self) ‑> Any-
Expand source code
def to_dataframe(self) -> Any: """Convert results to a pandas DataFrame.""" if not HAS_PANDAS: raise RuntimeError( "pandas is not installed. Please install pandas to export DataFrame." ) data = [r.to_dict() for r in self.results] return pd.DataFrame(data)Convert results to a pandas DataFrame.
class UnifiedTransformerConfig (attention_type: str = None,
attention_rank: int | None = None,
attention_kwargs: dict[str, typing.Any] | None = None,
**kwargs)-
Expand source code
class UnifiedTransformerConfig: """Lightweight container for the custom attention configuration.""" def __init__( self, attention_type: str = None, attention_rank: int | None = None, attention_kwargs: dict[str, Any] | None = None, **kwargs, ) -> None: self.attention_type = attention_type self.attention_rank = attention_rank self.attention_kwargs = attention_kwargs or {} # Store any extra kwargs for k, v in kwargs.items(): setattr(self, k, v) def to_dict(self) -> dict[str, Any]: """Serialise to a JSON‑compatible dict.""" return { "attention_type": self.attention_type, "attention_rank": self.attention_rank, "attention_kwargs": self.attention_kwargs, } @classmethod def from_dict(cls, d: dict[str, Any]) -> "UnifiedTransformerConfig": """From dict.""" return cls(**d)Lightweight container for the custom attention configuration.
Static methods
def from_dict(d: dict[str, typing.Any]) ‑> UnifiedTransformerConfig-
From dict.
Methods
def to_dict(self) ‑> dict[str, typing.Any]-
Expand source code
def to_dict(self) -> dict[str, Any]: """Serialise to a JSON‑compatible dict.""" return { "attention_type": self.attention_type, "attention_rank": self.attention_rank, "attention_kwargs": self.attention_kwargs, }Serialise to a JSON‑compatible dict.
class UnifiedTransformerModel (model: torch.nn.modules.module.Module,
custom_config: UnifiedTransformerConfig | None = None)-
Expand source code
class UnifiedTransformerModel(nn.Module, PyTorchModelHubMixin): """ A wrapper around a HuggingFace model that supports custom attention and can be saved/loaded via the HuggingFace Hub. """ def __init__( self, model: nn.Module, custom_config: UnifiedTransformerConfig | None = None ) -> None: super().__init__() self.model = model self.config = getattr(model, "config", None) # HF model config self.arch = detect_architecture(model) self._custom_config = custom_config or UnifiedTransformerConfig() # ------------------------------------------------------------------ # Attention replacement # ------------------------------------------------------------------ def replace_attention( self, attention_type: str, rank: int | None = None, copy_weights: bool = True, **kwargs, ) -> Any: """ Inject a custom attention variant into all layers. This method also records the choice so it can be saved. """ inject_attention(self.model, attention_type, rank=rank, copy_weights=copy_weights, **kwargs) self._custom_config = UnifiedTransformerConfig( attention_type=attention_type, attention_rank=rank, attention_kwargs=kwargs, ) return self # ------------------------------------------------------------------ # Freezing / parameter count # ------------------------------------------------------------------ def freeze_layers(self, num_layers_to_freeze: int = 0) -> Any: """Freeze layers.""" layers = get_transformer_layers(self.model) for i, layer in enumerate(layers): for param in layer.parameters(): param.requires_grad = i >= num_layers_to_freeze return self def unfreeze_all(self) -> Any: """Unfreeze all.""" for param in self.model.parameters(): param.requires_grad = True return self def get_parameter_count(self) -> dict[str, int]: """Return the parameter count.""" total = sum(p.numel() for p in self.model.parameters()) trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad) return {"total": total, "trainable": trainable, "frozen": total - trainable} # ------------------------------------------------------------------ # Hub serialization – these methods are called by PyTorchModelHubMixin # ------------------------------------------------------------------ def _save_pretrained(self, save_directory: str) -> None: """ Save the underlying model and our custom config to a directory. """ # 1. Save the base HF model (attention layers are already inside) self.model.save_pretrained(save_directory) # 2. Save our custom attention config to a JSON file import json custom_cfg_path = f"{save_directory}/unified_attention_config.json" with open(custom_cfg_path, "w") as f: json.dump(self._custom_config.to_dict(), f, indent=2) @classmethod def _from_pretrained(cls, model_id: str, **kwargs) -> "UnifiedTransformerModel": """ Load a pretrained model and reconstruct the custom attention variant. This method is invoked by `from_pretrained` class method of the mixin. """ # 1. Load the custom config import json import os if os.path.isdir(model_id): cfg_path = os.path.join(model_id, "unified_attention_config.json") else: # If model_id is a Hub repo, we need to download the file from huggingface_hub import hf_hub_download cfg_path = hf_hub_download(repo_id=model_id, filename="unified_attention_config.json") with open(cfg_path) as f: custom_cfg_dict = json.load(f) custom_cfg = UnifiedTransformerConfig.from_dict(custom_cfg_dict) # 2. Load the base HF model base_model = AutoModel.from_pretrained(model_id, **kwargs) # 3. Re‑apply the attention replacement if a type was stored model = cls(base_model, custom_config=custom_cfg) if custom_cfg.attention_type: model.replace_attention( custom_cfg.attention_type, rank=custom_cfg.attention_rank, copy_weights=False, # Weights already inside the saved base model **custom_cfg.attention_kwargs, ) return model # ------------------------------------------------------------------ # Other methods # ------------------------------------------------------------------ def forward(self, *args, **kwargs) -> Any: """Forward pass — compute module output from inputs.""" return self.model(*args, **kwargs) def save_pretrained(self, save_directory: str) -> None: """Convenience wrapper – delegates to the mixin.""" super().save_pretrained(save_directory) @classmethod def from_pretrained( cls, model_id: str, task: str = "sequence_classification", num_labels: int | None = None, **kwargs, ) -> "UnifiedTransformerModel": """ Create a new UnifiedTransformerModel from a pretrained identifier. Automatically applies the custom attention if the model was saved with one. This is the main entry point for users. """ # If the model_id has a unified_attention_config.json, use the mixin's from_pretrained # which will call our _from_pretrained. Otherwise, create a fresh one. import os has_custom = False if os.path.isdir(model_id): has_custom = os.path.exists(os.path.join(model_id, "unified_attention_config.json")) if has_custom: return cls._from_pretrained(model_id, **kwargs) else: # Fresh model – no custom attention yet if task == "sequence_classification": sc_kwargs = dict(kwargs) if num_labels is not None: sc_kwargs["num_labels"] = num_labels base = AutoModelForSequenceClassification.from_pretrained(model_id, **sc_kwargs) else: base = AutoModel.from_pretrained(model_id, **kwargs) return cls(base)A wrapper around a HuggingFace model that supports custom attention and can be saved/loaded via the HuggingFace Hub.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
- huggingface_hub.hub_mixin.PyTorchModelHubMixin
- huggingface_hub.hub_mixin.ModelHubMixin
Static methods
def from_pretrained(model_id: str,
task: str = 'sequence_classification',
num_labels: int | None = None,
**kwargs) ‑> UnifiedTransformerModel-
Create a new UnifiedTransformerModel from a pretrained identifier. Automatically applies the custom attention if the model was saved with one. This is the main entry point for users.
Methods
def forward(self, *args, **kwargs) ‑> Any-
Expand source code
def forward(self, *args, **kwargs) -> Any: """Forward pass — compute module output from inputs.""" return self.model(*args, **kwargs)Forward pass — compute module output from inputs.
def freeze_layers(self, num_layers_to_freeze: int = 0) ‑> Any-
Expand source code
def freeze_layers(self, num_layers_to_freeze: int = 0) -> Any: """Freeze layers.""" layers = get_transformer_layers(self.model) for i, layer in enumerate(layers): for param in layer.parameters(): param.requires_grad = i >= num_layers_to_freeze return selfFreeze layers.
def get_parameter_count(self) ‑> dict[str, int]-
Expand source code
def get_parameter_count(self) -> dict[str, int]: """Return the parameter count.""" total = sum(p.numel() for p in self.model.parameters()) trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad) return {"total": total, "trainable": trainable, "frozen": total - trainable}Return the parameter count.
def replace_attention(self,
attention_type: str,
rank: int | None = None,
copy_weights: bool = True,
**kwargs) ‑> Any-
Expand source code
def replace_attention( self, attention_type: str, rank: int | None = None, copy_weights: bool = True, **kwargs, ) -> Any: """ Inject a custom attention variant into all layers. This method also records the choice so it can be saved. """ inject_attention(self.model, attention_type, rank=rank, copy_weights=copy_weights, **kwargs) self._custom_config = UnifiedTransformerConfig( attention_type=attention_type, attention_rank=rank, attention_kwargs=kwargs, ) return selfInject a custom attention variant into all layers. This method also records the choice so it can be saved.
def save_pretrained(self, save_directory: str) ‑> None-
Expand source code
def save_pretrained(self, save_directory: str) -> None: """Convenience wrapper – delegates to the mixin.""" super().save_pretrained(save_directory)Convenience wrapper – delegates to the mixin.
def unfreeze_all(self) ‑> Any-
Expand source code
def unfreeze_all(self) -> Any: """Unfreeze all.""" for param in self.model.parameters(): param.requires_grad = True return selfUnfreeze all.