Knowledge Distillation: How to Train a 1.5B Model That Matches Your 7B

# Example: Text classification
label = [0, 0, 0, 1, 0]  # Class 4
# Model learns: "This is class 4, others are wrong"

# Teacher's output (temperature=2)
soft_targets = [0.05, 0.12, 0.18, 0.48, 0.17]
# Model learns: "Mostly class 4, but 3 is plausible, 2 somewhat, 1 unlikely, 0 rare"

import torch
import torch.nn.functional as F
 
logits = torch.tensor([2.0, 1.0, 0.5, 0.2])
 
print("T=1 (standard):", F.softmax(logits / 1, dim=0))
print("T=2 (soft):", F.softmax(logits / 2, dim=0))
print("T=5 (very soft):", F.softmax(logits / 5, dim=0))
 
# T=1 (standard): tensor([0.5580, 0.2047, 0.1225, 0.0918])
# T=2 (soft): tensor([0.4018, 0.2451, 0.1866, 0.1665])
# T=5 (very soft): tensor([0.2886, 0.2524, 0.2331, 0.2259])

# Teacher probabilities (T=2)
{
  "Paris": 0.82,
  "paris": 0.08,  # Dark knowledge: capitalization variant
  "Lyon": 0.03,   # Dark knowledge: other French cities
  "France": 0.02, # Dark knowledge: common error
  "the": 0.01,
  ...
}

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
 
# Load teacher model (Llama-7B)
teacher_name = "meta-llama/Llama-2-7b-hf"
 
print("Loading teacher model...")
teacher = AutoModelForCausalLM.from_pretrained(
    teacher_name,
    torch_dtype=torch.float16,
    device_map="auto",  # Automatic device placement
)
tokenizer = AutoTokenizer.from_pretrained(teacher_name)
tokenizer.pad_token = tokenizer.eos_token
 
# Freeze teacher (never update during distillation)
teacher.eval()
for param in teacher.parameters():
    param.requires_grad = False
 
print(f"Teacher parameters: {teacher.num_parameters() / 1e9:.2f}B")
print(f"Teacher memory: {teacher.get_memory_footprint() / 1e9:.2f}GB")
# Teacher parameters: 6.74B
# Teacher memory: 13.48GB (FP16)

from datasets import load_dataset
 
def evaluate_model(model, tokenizer, num_samples=100):
    """Quick eval on validation set."""
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation")
    
    total_loss = 0
    total_tokens = 0
    
    model.eval()
    with torch.no_grad():
        for i, example in enumerate(dataset):
            if i >= num_samples:
                break
            
            text = example["text"]
            if len(text.strip()) == 0:
                continue
            
            inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
            
            outputs = model(**inputs, labels=inputs["input_ids"])
            total_loss += outputs.loss.item() * inputs["input_ids"].numel()
            total_tokens += inputs["input_ids"].numel()
    
    perplexity = torch.exp(torch.tensor(total_loss / total_tokens))
    return perplexity.item()
 
teacher_ppl = evaluate_model(teacher, tokenizer)
print(f"Teacher perplexity: {teacher_ppl:.2f}")
# Teacher perplexity: 8.79

from transformers import LlamaConfig
 
# Teacher config
teacher_config = teacher.config
print("Teacher architecture:")
print(f"  hidden_size: {teacher_config.hidden_size}")
print(f"  num_layers: {teacher_config.num_hidden_layers}")
print(f"  num_heads: {teacher_config.num_attention_heads}")
print(f"  intermediate_size: {teacher_config.intermediate_size}")
 
# Output:
#   hidden_size: 4096
#   num_layers: 32
#   num_heads: 32
#   intermediate_size: 11008

# Student config (1.5B target)
student_config = LlamaConfig(
    vocab_size=teacher_config.vocab_size,  # Same vocab
    hidden_size=1536,         # 37.5% of teacher (4096)
    intermediate_size=4096,   # Maintain ~2.7× ratio
    num_hidden_layers=24,     # 75% of teacher depth (32)
    num_attention_heads=12,   # Proportional reduction
    num_key_value_heads=4,    # GQA for efficiency (3× reduction)
    max_position_embeddings=teacher_config.max_position_embeddings,
    rope_theta=teacher_config.rope_theta,
    rms_norm_eps=teacher_config.rms_norm_eps,
    tie_word_embeddings=False,
)
 
from transformers import LlamaForCausalLM
 
student = LlamaForCausalLM(student_config)
print(f"Student parameters: {student.num_parameters() / 1e9:.2f}B")
print(f"Student memory: {student.get_memory_footprint() / 1e9:.2f}GB")
# Student parameters: 1.52B
# Student memory: 3.04GB (FP16)

tiny_config = LlamaConfig(
    hidden_size=1024,
    intermediate_size=2816,
    num_hidden_layers=20,
    num_attention_heads=8,
    num_key_value_heads=2,
)
# ~500M parameters

balanced_config = LlamaConfig(
    hidden_size=2048,
    intermediate_size=5632,
    num_hidden_layers=28,
    num_attention_heads=16,
    num_key_value_heads=8,
)
# ~2.5B parameters

import torch.nn.functional as F
 
class DistillationLoss(torch.nn.Module):
    """
    Combined distillation + task loss for language modeling.
    
    Args:
        temperature: Softening parameter for distillation
        alpha: Weight for distillation loss (0=task only, 1=distill only)
        reduction: How to reduce loss ('mean' or 'batchmean')
    """
    def __init__(self, temperature=2.0, alpha=0.7, reduction='batchmean'):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.reduction = reduction
        
    def forward(self, student_logits, teacher_logits, labels):
        """
        Args:
            student_logits: [batch, seq_len, vocab_size]
            teacher_logits: [batch, seq_len, vocab_size]
            labels: [batch, seq_len] (with -100 for padding)
        
        Returns:
            loss: Scalar
            metrics: Dict with breakdown
        """
        # Flatten for loss computation
        batch_size, seq_len, vocab_size = student_logits.shape
        student_logits_flat = student_logits.view(-1, vocab_size)
        teacher_logits_flat = teacher_logits.view(-1, vocab_size)
        labels_flat = labels.view(-1)
        
        # Create mask for non-padding tokens
        mask = (labels_flat != -100).float()
        num_tokens = mask.sum()
        
        # Task loss: Standard cross-entropy with true labels
        loss_ce = F.cross_entropy(
            student_logits_flat,
            labels_flat,
            ignore_index=-100,
            reduction='sum'
        ) / num_tokens
        
        # Distillation loss: KL divergence with temperature scaling
        # Apply temperature scaling
        student_soft = F.log_softmax(student_logits_flat / self.temperature, dim=-1)
        teacher_soft = F.softmax(teacher_logits_flat / self.temperature, dim=-1)
        
        # Compute KL divergence
        loss_kd = F.kl_div(
            student_soft,
            teacher_soft,
            reduction='none'
        ).sum(dim=-1)  # Sum over vocabulary
        
        # Apply mask and reduce
        loss_kd = (loss_kd * mask).sum() / num_tokens
        
        # Scale by T² to compensate for temperature
        loss_kd = loss_kd * (self.temperature ** 2)
        
        # Combined loss
        loss = self.alpha * loss_kd + (1 - self.alpha) * loss_ce
        
        metrics = {
            "loss": loss.item(),
            "loss_kd": loss_kd.item(),
            "loss_ce": loss_ce.item(),
            "alpha": self.alpha,
        }
        
        return loss, metrics
 
# Instantiate
distill_criterion = DistillationLoss(temperature=2.0, alpha=0.7)

# Test distillation loss
batch_size, seq_len, vocab_size = 2, 10, 32000
 
student_logits = torch.randn(batch_size, seq_len, vocab_size)
teacher_logits = torch.randn(batch_size, seq_len, vocab_size)
labels = torch.randint(0, vocab_size, (batch_size, seq_len))
labels[:, -2:] = -100  # Padding
 
loss, metrics = distill_criterion(student_logits, teacher_logits, labels)
print(f"Total loss: {metrics['loss']:.4f}")
print(f"  KD loss: {metrics['loss_kd']:.4f} (weight: {metrics['alpha']})")
print(f"  CE loss: {metrics['loss_ce']:.4f} (weight: {1 - metrics['alpha']})")

from datasets import load_dataset
from torch.utils.data import DataLoader
 
def prepare_dataset(tokenizer, max_length=512, num_proc=8):
    """Load and tokenize training data."""
    
    # Load dataset (using SlimPajama subset for demo)
    dataset = load_dataset(
        "cerebras/SlimPajama-627B",
        split="train",
        streaming=True  # Stream for large datasets
    )
    
    def tokenize_function(examples):
        """Tokenize and create labels."""
        # Tokenize text
        tokenized = tokenizer(
            examples["text"],
            truncation=True,
            max_length=max_length,
            padding="max_length",
            return_tensors="pt"
        )
        
        # Labels are input_ids (for causal LM)
        tokenized["labels"] = tokenized["input_ids"].clone()
        
        # Mask padding in labels
        tokenized["labels"][tokenized["input_ids"] == tokenizer.pad_token_id] = -100
        
        return tokenized
    
    # Tokenize dataset
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset.column_names,
    )
    
    return tokenized_dataset
 
# Prepare data
train_dataset = prepare_dataset(tokenizer)
 
# Create dataloader
train_dataloader = DataLoader(
    train_dataset,
    batch_size=4,  # Adjust based on GPU memory
    shuffle=False,  # Streaming dataset
    num_workers=4,
)

from transformers import get_cosine_schedule_with_warmup
from torch.cuda.amp import autocast, GradScaler
from tqdm import tqdm
 
def train_distillation(
    student,
    teacher,
    train_dataloader,
    criterion,
    num_steps=100000,
    learning_rate=2e-4,
    warmup_steps=2000,
    gradient_accumulation_steps=8,
    max_grad_norm=1.0,
    log_interval=100,
    save_interval=5000,
    output_dir="./student_checkpoints"
):
    """
    Complete distillation training loop.
    
    Args:
        student: Student model (trainable)
        teacher: Teacher model (frozen)
        train_dataloader: DataLoader for training data
        criterion: DistillationLoss instance
        num_steps: Total training steps
        learning_rate: Peak learning rate
        warmup_steps: Learning rate warmup
        gradient_accumulation_steps: Accumulate gradients
        max_grad_norm: Gradient clipping threshold
        log_interval: Steps between logging
        save_interval: Steps between checkpoints
        output_dir: Directory for checkpoints
    """
    import os
    os.makedirs(output_dir, exist_ok=True)
    
    # Setup optimizer
    optimizer = torch.optim.AdamW(
        student.parameters(),
        lr=learning_rate,
        betas=(0.9, 0.95),
        weight_decay=0.1,
        eps=1e-8
    )
    
    # Learning rate scheduler
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=num_steps
    )
    
    # Mixed precision training
    scaler = GradScaler()
    
    # Move models to device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    student = student.to(device)
    teacher = teacher.to(device)
    
    # Training state
    student.train()
    teacher.eval()
    
    global_step = 0
    total_loss = 0
    
    progress_bar = tqdm(total=num_steps, desc="Training")
    
    dataloader_iter = iter(train_dataloader)
    
    while global_step < num_steps:
        optimizer.zero_grad()
        
        # Gradient accumulation loop
        accum_metrics = {"loss": 0, "loss_kd": 0, "loss_ce": 0}
        
        for accum_step in range(gradient_accumulation_steps):
            # Get batch
            try:
                batch = next(dataloader_iter)
            except StopIteration:
                dataloader_iter = iter(train_dataloader)
                batch = next(dataloader_iter)
            
            # Move to device
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)
            
            # Mixed precision forward pass
            try:
                with autocast():
                    # Teacher forward (no gradients)
                    with torch.no_grad():
                        teacher_outputs = teacher(
                            input_ids=input_ids,
                            attention_mask=attention_mask,
                        )
                        teacher_logits = teacher_outputs.logits
                    
                    # Student forward
                    student_outputs = student(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                    )
                    student_logits = student_outputs.logits
                    
                    # Check for shape mismatch (common distillation error)
                    if teacher_logits.shape != student_logits.shape:
                        print(f"Shape mismatch: teacher {teacher_logits.shape} vs student {student_logits.shape}")
                        optimizer.zero_grad()
                        continue
                    
                    # Compute distillation loss
                    loss, metrics = criterion(student_logits, teacher_logits, labels)
                    
                    # Check for NaN loss (indicates numerical instability)
                    if torch.isnan(loss):
                        print(f"NaN loss detected at step {global_step}. Skipping batch.")
                        optimizer.zero_grad()
                        continue
                    
                    # Scale loss for gradient accumulation
                    loss = loss / gradient_accumulation_steps
                
                # Backward pass with scaling
                scaler.scale(loss).backward()
                
            except RuntimeError as e:
                if "out of memory" in str(e):
                    print(f"OOM at step {global_step}. Clearing cache and skipping batch.")
                    torch.cuda.empty_cache()
                    optimizer.zero_grad()
                    continue
                else:
                    raise e
            
            # Accumulate metrics
            for key in accum_metrics:
                accum_metrics[key] += metrics[key] / gradient_accumulation_steps
        
        # Gradient clipping
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(student.parameters(), max_grad_norm)
        
        # Optimizer step
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
        
        # Update tracking
        global_step += 1
        total_loss += accum_metrics["loss"]
        
        # Logging
        if global_step % log_interval == 0:
            avg_loss = total_loss / log_interval
            lr = scheduler.get_last_lr()[0]
            
            progress_bar.set_postfix({
                "loss": f"{avg_loss:.4f}",
                "kd": f"{accum_metrics['loss_kd']:.4f}",
                "ce": f"{accum_metrics['loss_ce']:.4f}",
                "lr": f"{lr:.2e}"
            })
            
            total_loss = 0
        
        # Checkpointing
        if global_step % save_interval == 0:
            checkpoint_path = os.path.join(output_dir, f"checkpoint-{global_step}")
            student.save_pretrained(checkpoint_path)
            tokenizer.save_pretrained(checkpoint_path)
            print(f"\nSaved checkpoint to {checkpoint_path}")
        
        progress_bar.update(1)
    
    progress_bar.close()
    
    # Final save
    final_path = os.path.join(output_dir, "final")
    student.save_pretrained(final_path)
    tokenizer.save_pretrained(final_path)
    print(f"Training complete! Final model saved to {final_path}")
 
# Run training
train_distillation(
    student=student,
    teacher=teacher,
    train_dataloader=train_dataloader,
    criterion=distill_criterion,
    num_steps=100000,
    learning_rate=2e-4,
    gradient_accumulation_steps=8,
)

def progressive_distillation(
    teacher_name="meta-llama/Llama-2-7b-hf",
    final_size=1.5e9,
    num_stages=3
):
    """
    Progressive distillation: teacher → intermediate → student.
    
    Example: 7B → 3.5B → 1.5B
    """
    # Load teacher
    teacher = AutoModelForCausalLM.from_pretrained(teacher_name)
    teacher_size = teacher.num_parameters()
    
    # Calculate intermediate sizes
    sizes = [teacher_size]
    ratio = (final_size / teacher_size) ** (1 / num_stages)
    for i in range(num_stages):
        sizes.append(sizes[-1] * ratio)
    
    print("Progressive distillation stages:")
    for i, size in enumerate(sizes):
        print(f"  Stage {i}: {size/1e9:.2f}B parameters")
    
    # Distill stage by stage
    current_teacher = teacher
    for stage in range(1, num_stages + 1):
        print(f"\n=== Stage {stage}: {sizes[stage-1]/1e9:.2f}B → {sizes[stage]/1e9:.2f}B ===")
        
        # Create student config
        student_config = create_config_for_size(sizes[stage])
        student = LlamaForCausalLM(student_config)
        
        # Train distillation
        train_distillation(
            student=student,
            teacher=current_teacher,
            train_dataloader=train_dataloader,
            criterion=DistillationLoss(temperature=2.0, alpha=0.7),
            num_steps=30000,  # Fewer steps per stage
        )
        
        # This student becomes next teacher
        current_teacher = student
        current_teacher.eval()
        for param in current_teacher.parameters():
            param.requires_grad = False
    
    return current_teacher  # Final student

class MultiTeacherDistillationLoss(torch.nn.Module):
    """Distill from multiple teachers."""
    
    def __init__(self, num_teachers=2, temperature=2.0, alpha=0.7):
        super().__init__()
        self.num_teachers = num_teachers
        self.temperature = temperature
        self.alpha = alpha
    
    def forward(self, student_logits, teacher_logits_list, labels):
        """
        Args:
            teacher_logits_list: List of [batch, seq, vocab] from each teacher
        """
        # Average teacher distributions
        teacher_soft_avg = torch.zeros_like(teacher_logits_list[0])
        for teacher_logits in teacher_logits_list:
            teacher_soft = F.softmax(teacher_logits / self.temperature, dim=-1)
            teacher_soft_avg += teacher_soft / self.num_teachers
        
        # Student distribution
        student_soft = F.log_softmax(student_logits / self.temperature, dim=-1)
        
        # KL divergence
        loss_kd = F.kl_div(student_soft, teacher_soft_avg, reduction='batchmean')
        loss_kd *= (self.temperature ** 2)
        
        # Task loss
        loss_ce = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            labels.view(-1),
            ignore_index=-100
        )
        
        loss = self.alpha * loss_kd + (1 - self.alpha) * loss_ce
        return loss, {"loss_kd": loss_kd.item(), "loss_ce": loss_ce.item()}
 
# Usage: Combine Llama-7B and Mistral-7B
teachers = [
    AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf"),
    AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1"),
]
 
criterion = MultiTeacherDistillationLoss(num_teachers=2)

class TaskSpecificDistillation:
    """Distillation with task-specific data and loss."""
    
    def __init__(self, task="summarization"):
        self.task = task
        
    def get_dataset(self):
        """Load task-specific dataset."""
        if self.task == "summarization":
            return load_dataset("cnn_dailymail", "3.0.0")
        elif self.task == "qa":
            return load_dataset("squad_v2")
        elif self.task == "code":
            return load_dataset("codeparrot/github-code")
        # ... more tasks
    
    def compute_task_loss(self, student_output, teacher_output, labels):
        """Task-specific loss beyond standard distillation."""
        # Standard distillation
        base_loss, metrics = self.distill_criterion(student_output, teacher_output, labels)
        
        # Add task-specific objectives
        if self.task == "summarization":
            # Encourage conciseness
            length_penalty = torch.mean(student_output.sum(dim=-1))
            base_loss += 0.1 * length_penalty
        
        return base_loss, metrics
 
# Fine-tune for code generation
code_distiller = TaskSpecificDistillation(task="code")
code_dataset = code_distiller.get_dataset()
# Train with task-specific loss

def evaluate_student(student, teacher, tokenizer, benchmarks=["wikitext", "lambada"]):
    """
    Compare student to teacher across multiple benchmarks.
    """
    from lm_eval import evaluator
    
    results = {}
    
    for benchmark in benchmarks:
        print(f"\nEvaluating on {benchmark}...")
        
        # Evaluate teacher
        teacher_results = evaluator.simple_evaluate(
            model=teacher,
            tasks=[benchmark],
            num_fewshot=0,
        )
        
        # Evaluate student
        student_results = evaluator.simple_evaluate(
            model=student,
            tasks=[benchmark],
            num_fewshot=0,
        )
        
        # Compare
        teacher_score = teacher_results["results"][benchmark]["acc"]
        student_score = student_results["results"][benchmark]["acc"]
        retention = (student_score / teacher_score) * 100
        
        results[benchmark] = {
            "teacher": teacher_score,
            "student": student_score,
            "retention": retention
        }
        
        print(f"  Teacher: {teacher_score:.2%}")
        print(f"  Student: {student_score:.2%}")
        print(f"  Retention: {retention:.1f}%")
    
    return results
 
# Run evaluation
eval_results = evaluate_student(student, teacher, tokenizer)

def grid_search_distillation(temperatures=[1.5, 2.0, 3.0], alphas=[0.5, 0.7, 0.9]):
    """
    Find optimal temperature and alpha for your dataset.
    """
    best_score = 0
    best_params = {}
    
    for temp in temperatures:
        for alpha in alphas:
            print(f"\nTrying T={temp}, α={alpha}")
            
            # Create fresh student
            student = LlamaForCausalLM(student_config)
            
            # Train with these hyperparameters
            criterion = DistillationLoss(temperature=temp, alpha=alpha)
            train_distillation(
                student=student,
                teacher=teacher,
                train_dataloader=train_dataloader,
                criterion=criterion,
                num_steps=10000,  # Short run for search
            )
            
            # Evaluate
            score = evaluate_student(student, teacher, tokenizer)["wikitext"]["retention"]
            
            if score > best_score:
                best_score = score
                best_params = {"temperature": temp, "alpha": alpha}
            
            print(f"  Retention: {score:.1f}%")
    
    print(f"\nBest params: {best_params} (retention: {best_score:.1f}%)")
    return best_params

# Save final model
student.save_pretrained("./student-1.5B-final")
tokenizer.save_pretrained("./student-1.5B-final")
 
# Quantize to INT8 for deployment
from transformers import AutoModelForCausalLM
import torch
 
# Load and quantize
student_int8 = AutoModelForCausalLM.from_pretrained(
    "./student-1.5B-final",
    device_map="auto",
    load_in_8bit=True,  # Automatic INT8 quantization
)
 
print(f"FP16 size: {student.get_memory_footprint() / 1e9:.2f}GB")
print(f"INT8 size: {student_int8.get_memory_footprint() / 1e9:.2f}GB")
# FP16 size: 3.04GB
# INT8 size: 1.52GB (2× reduction)

import time
 
def benchmark_inference(model, tokenizer, num_iterations=100):
    """Measure tokens/second."""
    prompt = "The quick brown fox jumps over the lazy dog. " * 10
    
    # Warm up
    for _ in range(5):
        _ = model.generate(**tokenizer(prompt, return_tensors="pt").to(model.device), max_new_tokens=50)
    
    # Benchmark
    start = time.time()
    total_tokens = 0
    for _ in range(num_iterations):
        output = model.generate(
            **tokenizer(prompt, return_tensors="pt").to(model.device),
            max_new_tokens=50,
            do_sample=False
        )
        total_tokens += 50
    elapsed = time.time() - start
    
    tokens_per_second = total_tokens / elapsed
    return tokens_per_second
 
print("Teacher (7B FP16):", benchmark_inference(teacher, tokenizer), "tok/s")
print("Student (1.5B FP16):", benchmark_inference(student, tokenizer), "tok/s")
print("Student (1.5B INT8):", benchmark_inference(student_int8, tokenizer), "tok/s")
# Teacher (7B FP16): 18 tok/s
# Student (1.5B FP16): 58 tok/s (3.2× faster)
# Student (1.5B INT8): 82 tok/s (4.6× faster)

# Configuration for 90%+ retention
config = {
    "student_size": 0.25,  # 25% of teacher
    "temperature": 2.0,
    "alpha": 0.8,  # Heavy distillation
    "num_steps": 200000,
    "batch_size": 2,
    "gradient_accumulation": 16,
    "learning_rate": 1e-4,  # Lower LR for stability
}

# Fast training, acceptable quality
config = {
    "student_size": 0.15,  # Smaller student
    "temperature": 3.0,  # More aggressive
    "alpha": 0.6,
    "num_steps": 50000,
    "batch_size": 8,
    "gradient_accumulation": 4,
    "learning_rate": 3e-4,
}

# Train on single GPU
config = {
    "student_size": 0.2,
    "temperature": 2.0,
    "alpha": 0.7,
    "num_steps": 100000,
    "batch_size": 1,  # Small batch
    "gradient_accumulation": 32,  # High accumulation
    "learning_rate": 2e-4,
    "mixed_precision": True,  # Essential
    "gradient_checkpointing": True,  # Save memory
}