11. Prompt Tuning and Prefix Tuning — Lightweight Alternatives to LoRA

graph TB
    subgraph "Prompt Tuning (Input Only)"
        PT_SOFT["Soft prompt P<br>[p₁, p₂, ..., p₂₀]<br>Learnable: 80KB"]
        PT_INPUT["Input tokens<br>[x₁, x₂, ..., xₙ]"]
        PT_EMB["Combined embeddings<br>[P; Embed(X)]"]

        PT_SOFT --> PT_EMB
        PT_INPUT --> PT_EMB

        PT_L1["Layer 1: Self-Attention<br>Soft tokens attend ↔ input tokens<br>(frozen weights)"]
        PT_L2["Layer 2-31: Same<br>Soft prompt influence propagates<br>through residual stream<br>(all frozen)"]
        PT_L32["Layer 32: Final<br>Soft prompt effect diluted<br>through 32 layers"]

        PT_EMB --> PT_L1 --> PT_L2 --> PT_L32
    end

    subgraph "Prefix Tuning (Every Layer)"
        PF_INPUT2["Input tokens<br>[x₁, x₂, ..., xₙ]"]

        PF_L1["Layer 1:<br>K = [K¹_prefix; K_input]<br>V = [V¹_prefix; V_input]<br>Direct prefix influence ✅"]
        PF_L2["Layer 2-31:<br>K = [Kˡ_prefix; K_input]<br>V = [Vˡ_prefix; V_input]<br>Fresh prefix at every layer ✅"]
        PF_L32["Layer 32:<br>K = [K³²_prefix; K_input]<br>V = [V³²_prefix; V_input]<br>Full control at output ✅"]

        PF_INPUT2 --> PF_L1 --> PF_L2 --> PF_L32
    end

    style PT_SOFT fill:#fff9c4
    style PF_L1 fill:#c8e6c9
    style PF_L2 fill:#c8e6c9
    style PF_L32 fill:#c8e6c9

graph LR
    subgraph "How Input Tokens See Soft Prompt"
        Q["Query: 'recommend'<br>q = W_Q · h_recommend"]
        K1["Key: p₁<br>k = W_K · p₁<br>Learned to attract queries<br>about recommendations"]
        K2["Key: p₂<br>k = W_K · p₂<br>Learned to encode<br>'manga expert' persona"]
        K3["Key: 'recommend'<br>k = W_K · h_recommend"]
        K4["Key: 'manga'<br>k = W_K · h_manga"]

        Q -->|"α₁ = 0.25"| K1
        Q -->|"α₂ = 0.20"| K2
        Q -->|"α₃ = 0.30"| K3
        Q -->|"α₄ = 0.25"| K4
    end

    subgraph "After Training"
        NOTE["Soft prompt tokens capture<br>~45% of attention in early layers<br>~15% in later layers<br><br>They 'inject' task context that<br>the frozen model processes<br>as if it were real input"]
    end

    style K1 fill:#fff9c4
    style K2 fill:#fff9c4

sequenceDiagram
    participant SP as Soft Prompt P (80KB, trainable)
    participant FZ as Frozen Model (8B params)
    participant LOSS as Loss Function

    Note over SP,LOSS: Forward Pass
    SP->>FZ: Prepend soft tokens to input embeddings
    FZ->>FZ: Layer 1: Q,K,V computed (frozen weights)<br>Soft tokens participate in attention
    FZ->>FZ: Layer 2-31: Same. Soft prompt<br>influence propagates via residual stream
    FZ->>FZ: Layer 32: Final representation
    FZ->>LOSS: Logits → Cross-entropy loss

    Note over SP,LOSS: Backward Pass
    LOSS->>FZ: ∂L/∂logits
    FZ->>FZ: Gradients flow backward through<br>all 32 layers (no weight updates!)
    FZ->>SP: ∂L/∂P — only these gradients<br>are used for optimization
    SP->>SP: P ← P - η · ∂L/∂P<br>(AdamW, lr=3e-1, 0.001% of model)

graph TB
    subgraph "MangaAssist Multi-Persona Routing"
        ROUTER["Intent Router<br>(from doc 01)"]

        P1["Persona: Manga Expert<br>Soft prompt: 20 tokens<br>80KB, top-k genre rec"]
        P2["Persona: Returns Agent<br>Soft prompt: 15 tokens<br>60KB, policy-focused"]
        P3["Persona: JP Culture Guide<br>Soft prompt: 30 tokens<br>120KB, cultural context"]
        P4["Persona: General Assistant<br>No soft prompt<br>(base SFT behavior)"]

        MODEL["Llama 3 8B (frozen)<br>Loaded ONCE in memory<br>~16GB (quantized)"]

        ROUTER -->|"manga_recommendation"| P1
        ROUTER -->|"return_refund"| P2
        ROUTER -->|"cultural_question"| P3
        ROUTER -->|"other"| P4

        P1 --> MODEL
        P2 --> MODEL
        P3 --> MODEL
        P4 --> MODEL
    end

    subgraph "Memory Comparison"
        LORA_MEM["LoRA approach:<br>4 adapters × 50MB = 200MB<br>+ model reload per switch"]
        SOFT_MEM["Soft prompt approach:<br>4 prompts × 80KB = 320KB<br>No model reload, just<br>prepend different prefix"]
    end

    style SOFT_MEM fill:#c8e6c9
    style LORA_MEM fill:#ffcdd2

graph TB
    subgraph "Reparameterization trick (Li & Liang 2021)"
        Z["Learnable seed matrix<br>z ∈ ℝ^(m × d')<br>d' = 512 (small)"]
        MLP["2-layer MLP<br>512 → 2048 → 8192<br>Generates prefix KV pairs"]
        KV["Per-layer KV:<br>K_prefix ∈ ℝ^(20 × 4096)<br>V_prefix ∈ ℝ^(20 × 4096)"]

        Z --> MLP --> KV

        TRAIN["During Training:<br>Optimize z and MLP params<br>More stable than direct<br>per-layer optimization"]
        INFER["At Inference:<br>Run MLP once, cache KV<br>Discard MLP, store only<br>generated prefix matrices"]

        KV --> TRAIN
        KV --> INFER
    end

    style Z fill:#fff9c4
    style INFER fill:#c8e6c9

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, PromptTuningConfig, PromptTuningInit, TaskType
from torch.utils.data import Dataset, DataLoader


class PromptTuningTrainer:
    """
    Train soft prompts for task-specific adaptation.
    Entire model remains frozen; only soft prompt tokens are learned.
    """

    def __init__(
        self,
        model_name: str = "meta-llama/Llama-3-8b-hf",
        num_virtual_tokens: int = 20,
        task_name: str = "manga_expert",
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        # Load frozen model
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            load_in_4bit=True,
        )

        # Configure prompt tuning
        peft_config = PromptTuningConfig(
            task_type=TaskType.CAUSAL_LM,
            num_virtual_tokens=num_virtual_tokens,
            # Initialize from text — gives meaningful starting point
            prompt_tuning_init=PromptTuningInit.TEXT,
            prompt_tuning_init_text=(
                "You are a manga expert chatbot for an online bookstore. "
                "Be conversational, knowledgeable, and helpful."
            ),
            tokenizer_name_or_path=model_name,
        )

        self.model = get_peft_model(self.model, peft_config)
        self.task_name = task_name

        # Verify: only soft prompt is trainable
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        total = sum(p.numel() for p in self.model.parameters())
        print(f"Trainable: {trainable:,} / {total:,} = {100*trainable/total:.4f}%")

    def train(self, data: list[dict], epochs: int = 10, lr: float = 3e-1):
        """
        Train soft prompt on task-specific data.
        High learning rate (0.3) is standard for prompt tuning —
        the soft prompt embeddings are randomly initialized and need
        large updates to become meaningful.
        """
        dataset = SoftPromptDataset(data, self.tokenizer)
        loader = DataLoader(dataset, batch_size=8, shuffle=True)

        # Only soft prompt params
        optimizer = torch.optim.AdamW(
            [p for p in self.model.parameters() if p.requires_grad],
            lr=lr,
            weight_decay=0.01,
        )

        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=epochs * len(loader),
        )

        self.model.train()
        for epoch in range(epochs):
            total_loss = 0
            for batch in loader:
                batch = {k: v.to(self.model.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                loss = outputs.loss

                loss.backward()
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                total_loss += loss.item()

            avg_loss = total_loss / len(loader)
            print(f"Epoch {epoch+1}: loss={avg_loss:.4f}")

    def save_adapter(self, path: str):
        """Save only the soft prompt (~80KB)."""
        self.model.save_pretrained(path)
        print(f"Saved soft prompt adapter to {path}")

    def load_adapter(self, path: str):
        """Load a soft prompt adapter (instant swap, no model reload)."""
        from peft import PeftModel
        self.model = PeftModel.from_pretrained(self.model, path)


class SoftPromptDataset(Dataset):
    def __init__(self, data: list[dict], tokenizer, max_length: int = 512):
        self.examples = []
        for item in data:
            text = f"User: {item['prompt']}\nAssistant: {item['response']}"
            encoded = tokenizer(
                text, truncation=True, max_length=max_length,
                padding="max_length", return_tensors="pt",
            )
            encoded["labels"] = encoded["input_ids"].clone()
            self.examples.append({k: v.squeeze() for k, v in encoded.items()})

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        return self.examples[idx]

from peft import PrefixTuningConfig


class PrefixTuningTrainer:
    """
    Prefix-tuning: learnable KV pairs at every transformer layer.
    More expressive than prompt tuning but still lightweight.
    """

    def __init__(
        self,
        model_name: str = "meta-llama/Llama-3-8b-hf",
        num_virtual_tokens: int = 20,
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            load_in_4bit=True,
        )

        peft_config = PrefixTuningConfig(
            task_type=TaskType.CAUSAL_LM,
            num_virtual_tokens=num_virtual_tokens,
            # Reparameterization for stable training
            prefix_projection=True,
            encoder_hidden_size=512,  # MLP hidden size (d' in the math)
        )

        self.model = get_peft_model(self.model, peft_config)

        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"Trainable: {trainable:,} ({trainable * 4 / 1024 / 1024:.1f} MB)")

    def train(self, data: list[dict], epochs: int = 5, lr: float = 1e-2):
        """
        Train prefix-tuning.
        Lower LR than prompt tuning (0.01 vs 0.3) because more parameters.
        """
        dataset = SoftPromptDataset(data, self.tokenizer)
        loader = DataLoader(dataset, batch_size=4, shuffle=True)

        optimizer = torch.optim.AdamW(
            [p for p in self.model.parameters() if p.requires_grad],
            lr=lr,
        )

        self.model.train()
        for epoch in range(epochs):
            total_loss = 0
            for batch in loader:
                batch = {k: v.to(self.model.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                loss = outputs.loss

                loss.backward()
                torch.nn.utils.clip_grad_norm_(
                    [p for p in self.model.parameters() if p.requires_grad],
                    max_norm=1.0,
                )
                optimizer.step()
                optimizer.zero_grad()
                total_loss += loss.item()

            print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}")

from peft import PromptEncoderConfig, PromptEncoderReparameterizationType


class PTuningV2Trainer:
    """
    P-Tuning v2: deep continuous prompts at every layer.
    Best balance of quality and parameter efficiency for medium-sized models.
    """

    def __init__(
        self,
        model_name: str = "meta-llama/Llama-3-8b-hf",
        num_virtual_tokens: int = 20,
        encoder_hidden_size: int = 256,
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            load_in_4bit=True,
        )

        peft_config = PromptEncoderConfig(
            task_type=TaskType.CAUSAL_LM,
            num_virtual_tokens=num_virtual_tokens,
            encoder_reparameterization_type=PromptEncoderReparameterizationType.MLP,
            encoder_hidden_size=encoder_hidden_size,
        )

        self.model = get_peft_model(self.model, peft_config)

        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        print(f"Trainable: {trainable:,}")


class MultiAdapterInference:
    """
    Serve multiple personas with different soft prompts.
    The base model loads once; only the soft prompt prefix changes per request.
    """

    def __init__(self, model_name: str = "meta-llama/Llama-3-8b-hf"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.base_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            load_in_4bit=True,
        )

        # Load all adapters (each ~80KB-5MB)
        self.adapters: dict[str, torch.Tensor] = {}

    def load_adapter(self, name: str, path: str):
        """Load a soft prompt adapter from disk."""
        from peft import PeftModel
        model = PeftModel.from_pretrained(self.base_model, path)
        # Extract the soft prompt embeddings
        for param_name, param in model.named_parameters():
            if "prompt" in param_name and param.requires_grad:
                self.adapters[name] = param.data.clone()
                break
        print(f"Loaded adapter '{name}' ({self.adapters[name].numel() * 2 / 1024:.1f}KB)")

    def generate(self, prompt: str, adapter_name: str, max_new_tokens: int = 256):
        """Generate with a specific soft prompt adapter."""
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.base_model.device)

        if adapter_name in self.adapters:
            # Prepend soft prompt embeddings
            soft_prompt = self.adapters[adapter_name].unsqueeze(0)
            input_embeds = self.base_model.get_input_embeddings()(inputs["input_ids"])
            combined_embeds = torch.cat([soft_prompt, input_embeds], dim=1)

            # Extend attention mask
            prefix_mask = torch.ones(
                1, soft_prompt.size(1), device=inputs["attention_mask"].device,
            )
            combined_mask = torch.cat([prefix_mask, inputs["attention_mask"]], dim=1)

            outputs = self.base_model.generate(
                inputs_embeds=combined_embeds,
                attention_mask=combined_mask,
                max_new_tokens=max_new_tokens,
            )
        else:
            outputs = self.base_model.generate(**inputs, max_new_tokens=max_new_tokens)

        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

import sagemaker
from sagemaker.huggingface import HuggingFace


def launch_prompt_tuning_job(
    adapter_type: str = "prompt_tuning",
    num_virtual_tokens: int = 20,
    persona: str = "manga_expert",
):
    """Launch prompt tuning on SageMaker."""
    session = sagemaker.Session()
    role = sagemaker.get_execution_role()

    hyperparameters = {
        "adapter_type": adapter_type,
        "num_virtual_tokens": num_virtual_tokens,
        "persona": persona,
        "epochs": 10 if adapter_type == "prompt_tuning" else 5,
        "learning_rate": 3e-1 if adapter_type == "prompt_tuning" else 1e-2,
        "batch_size": 8,
    }

    estimator = HuggingFace(
        entry_point="train_soft_prompt.py",
        source_dir="./src",
        instance_type="ml.g5.xlarge",  # Single GPU sufficient
        instance_count=1,
        role=role,
        transformers_version="4.45.0",
        pytorch_version="2.4.0",
        py_version="py311",
        hyperparameters=hyperparameters,
    )

    estimator.fit({
        "train": f"s3://mangaassist-data/preference/{persona}/train/",
        "eval": f"s3://mangaassist-data/preference/{persona}/eval/",
    })

    return estimator