盘点目前最常用的四种语言模型压缩技术

pip install torch transformers bitsandbytes
1.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Specify the model you want to use
model_name = "gpt2"  # You can replace this with any other LLM model
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Load the model with 8-bit quantization using bitsandbytes
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,  # Enable 8-bit quantization
    device_map="auto"   # Automatically allocate to available device (CPU/GPU)
)
# Example text for inference
input_text = "Weight Quantization is an efficient technique for compressing language models."
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
# Generate text
with torch.no_grad():
    output_ids = model.generate(input_ids, max_length=50)
# Decode and print the generated text
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.

import torch
import torch.nn.utils.prune as prune
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the pretrained model and tokenizer
model_name = "gpt2"  # You can replace this with any other LLM model
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Define a pruning method (here we use L1 unstructured pruning)
def prune_model_layer(layer, amount=0.3):
    # Prune 30% of the weights with the lowest L1 norm in the linear layers
    for name, module in layer.named_modules():
        if isinstance(module, torch.nn.Linear):
            prune.l1_unstructured(module, name="weight", amount=amount)
            print(f"Pruned layer {name} with amount {amount}")
# Apply pruning to all transformer layers in the model
for layer in model.transformer.h:
    prune_model_layer(layer, amount=0.3)  # Prune 30% of the weights
# Check the sparsity of the model
total_params = 0
pruned_params = 0
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        total_params += module.weight.nelement()
        pruned_params += torch.sum(module.weight == 0).item()
print(f"Total parameters: {total_params}")
print(f"Pruned parameters: {pruned_params}")
print(f"Sparsity: {pruned_params / total_params:.2%}")
# Test the pruned model on a sample input
input_text = "Pruning is an effective way to compress language models."
input_ids = tokenizer(input_text, return_tensors="pt").input_ids
# Generate text using the pruned model
with torch.no_grad():
    output_ids = model.generate(input_ids, max_length=50)
# Decode and print the generated text
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import torch.nn.functional as F

# Load the teacher (large) and student (smaller) models
teacher_model_name = "gpt2"  # You can replace this with any large LLM
student_model_name = "tiny-gpt2"  # A smaller variant to act as the student
# Load the teacher model and tokenizer
teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model_name).to("cuda")
teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
# Load the student model and tokenizer
student_model = AutoModelForCausalLM.from_pretrained(student_model_name).to("cuda")
student_tokenizer = AutoTokenizer.from_pretrained(student_model_name)
# Load a dataset for training (e.g., Wikitext for language modeling)
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
# Set training parameters
learning_rate = 5e-5
epochs = 3
optimizer = torch.optim.AdamW(student_model.parameters(), lr=learning_rate)
# Set temperature for softening probabilities
temperature = 2.0
alpha = 0.5  # Weighting factor for combining loss functions
# Training loop for knowledge distillation
for epoch in range(epochs):
    for i, example in enumerate(dataset):
        # Get the input text
        input_text = example["text"]
        
        # Skip empty lines
        if not input_text.strip():
            continue
        
        # Tokenize the input text for the teacher and student models
        teacher_inputs = teacher_tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=32).to("cuda")
        student_inputs = student_tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=32).to("cuda")
        
        # Get teacher predictions (soft labels)
        with torch.no_grad():
            teacher_outputs = teacher_model(**teacher_inputs)
            teacher_logits = teacher_outputs.logits / temperature
            teacher_probs = F.softmax(teacher_logits, dim=-1)
        
        # Get student predictions
        student_outputs = student_model(**student_inputs)
        student_logits = student_outputs.logits
        
        # Calculate distillation loss (Kullback-Leibler divergence)
        distillation_loss = F.kl_div(
            input=F.log_softmax(student_logits / temperature, dim=-1),
            target=teacher_probs,
            reduction="batchmean",
            log_target=False
        ) * (temperature ** 2)
        
        # Calculate student task loss (Cross-Entropy with true labels)
        target_labels = student_inputs["input_ids"]
        task_loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), target_labels.view(-1), ignore_index=student_tokenizer.pad_token_id)
        
        # Combined loss
        loss = alpha * distillation_loss + (1 - alpha) * task_loss
        
        # Backpropagation and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Print training progress
        if i % 100 == 0:
            print(f"Epoch [{epoch + 1}/{epochs}], Step [{i}], Loss: {loss.item():.4f}")
print("Knowledge distillation completed!")
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.

import torch
import numpy as np
from sklearn.cluster import KMeans

def apply_weight_sharing(model, num_clusters=16):
    # Iterate through each parameter in the model
    for name, param in model.named_parameters():
        if param.requires_grad:  # Only consider trainable parameters
            # Flatten the weights into a 1D array for clustering
            weights = param.data.cpu().numpy().flatten().reshape(-1, 1)
            # Apply k-means clustering
            kmeans = KMeans(n_clusters=num_clusters)
            kmeans.fit(weights)
            # Replace weights with their corresponding cluster centroids
            cluster_centroids = kmeans.cluster_centers_
            labels = kmeans.labels_
            # Map the original weights to their shared values
            shared_weights = np.array([cluster_centroids[label] for label in labels]).reshape(param.data.shape)
            # Update the model's parameters with the shared weights
            param.data = torch.tensor(shared_weights, dtype=param.data.dtype).to(param.device)
    return model
# Example usage with a pre-trained model
from transformers import GPT2LMHeadModel
model = GPT2LMHeadModel.from_pretrained("gpt2")
model = apply_weight_sharing(model, num_clusters=16)  # Apply weight sharing with 16 clusters
print("Weight sharing applied to the model!")
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.