FillipandClaude Opus 4.5
5c495e65c4
Fixes #5: Add compatibility shim for find_pruneable_heads_and_indices which was removed in transformers 4.40.0. All 12 GPT2 model files now use a try/except fallback. Fixes #6: Add ultra low memory mode for 12GB VRAM systems. - New memory_mode dropdown: auto, normal, low, ultra - Auto mode detects available VRAM and selects best mode - Ultra mode uses ~6GB peak VRAM with aggressive cleanup - Added VRAM monitoring and logging during generation Bump version to 1.3.0 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
"""Compatibility utilities for different package versions."""
|
|
|
|
import torch
|
|
from typing import List, Set, Tuple
|
|
|
|
|
|
def get_find_pruneable_heads_and_indices():
|
|
"""
|
|
Get find_pruneable_heads_and_indices function with fallback for transformers>=4.40.
|
|
|
|
The function was removed from transformers.pytorch_utils in version 4.40.0.
|
|
This provides a compatible implementation for newer versions.
|
|
"""
|
|
try:
|
|
from transformers.pytorch_utils import find_pruneable_heads_and_indices
|
|
return find_pruneable_heads_and_indices
|
|
except ImportError:
|
|
def find_pruneable_heads_and_indices(
|
|
heads: List[int],
|
|
n_heads: int,
|
|
head_dim: int,
|
|
already_pruned_heads: Set[int],
|
|
) -> Tuple[Set[int], torch.LongTensor]:
|
|
"""
|
|
Finds the heads and their indices taking already_pruned_heads into account.
|
|
|
|
Args:
|
|
heads: List of the indices of heads to prune.
|
|
n_heads: The number of heads in the model.
|
|
head_dim: The dimension of each head.
|
|
already_pruned_heads: A set of already pruned heads.
|
|
|
|
Returns:
|
|
A tuple with the remaining heads to prune and the indices of heads to prune
|
|
taking already_pruned_heads into account.
|
|
"""
|
|
mask = torch.ones(n_heads, head_dim)
|
|
# Convert to set and remove already pruned heads
|
|
heads_to_prune = set(heads) - already_pruned_heads
|
|
for head in heads_to_prune:
|
|
mask[head] = 0
|
|
mask = mask.view(-1).contiguous().eq(1)
|
|
index = torch.arange(len(mask))[mask].long()
|
|
return heads_to_prune, index
|
|
|
|
return find_pruneable_heads_and_indices
|