2️⃣ LoRA Fine-Tuning

class Lora(nn.Module):
    """
    Module that implements the basic LoRA block.
    - Input: tensor of shape (..., [inst], d_in) and returns a tensor of shape (..., inst, d_out).
    - Calculated intermediate activations of shape (..., inst, rank)
    - Output: tensor of shape (..., inst, d_out)
    """

    A: nn.Parameter  # (n_inst, d_in, rank)
    B: nn.Parameter  # (n_inst, rank, d_out)

    def __init__(
        self,
        d_in: int = 768,
        d_out: int = 768,
        rank: int = 4,
        lora_alpha: float = 32,
        n_inst: int | None = None,
        dtype: t.dtype | None = None,
    ):
        """
        Initialize the weights of the LoRA block.
        - The A block should be initialized with kaiming uniform with a=sqrt(5)
        - The B block should be initialized with zeros.
        """
        super().__init__()
        self.rank = rank
        self.d_in = d_in
        self.d_out = d_out
        self.n_inst = 1 if n_inst is None else n_inst
        self.lora_alpha = lora_alpha
        self.dtype = dtype

        # Define the model parameters here
        self.A = nn.Parameter(t.empty(self.n_inst, d_in, rank, dtype=dtype))
        self.B = nn.Parameter(t.zeros(self.n_inst, rank, d_out, dtype=dtype))

        nn.init.kaiming_uniform_(self.A, a=5**0.5)

    def forward(self, x: Float[Tensor, "... inst d_in"]) -> Float[Tensor, "... inst d_out"]:
        """
        Computes the forward pass of the LoRA block f(x) = (x @ A) @ B * lora_alpha / rank
        Args:
            x: Tensor of shape (..., inst, d_in)
        Returns:
            out (..., inst, d_out) such that out[..., i, :] = (x[..., i] @ A[i]) @ B[i] * lora_alpha / rank
        """
        if x.dtype != self.dtype:
            x = x.to(self.dtype)
        assert x.shape[-2] == self.n_inst or x.shape[-2] == 1, (
            f"Expected inst dim {self.n_inst} or 1, got {x.shape[-2]}. (input shape was {x.shape=})"
        )


        # force order of operations (x A) B
        tmp = einops.einsum(x, self.A, "... inst d_in, inst d_in rank -> ... inst rank")
        out = einops.einsum(tmp, self.B, "... inst rank, inst rank d_out -> ... inst d_out")

        return out * self.lora_alpha / self.rank


model = HookedTransformer.from_pretrained("pythia-14m")
tests_lora.testing_lora(Lora)

model = transformer_lens.HookedTransformer.from_pretrained("meta-llama/Llama-3.2-1B")
print(f"{model.cfg.n_heads=}")
print(f"{model.cfg.n_key_value_heads=}")
print(f"Group size: {model.cfg.n_heads // model.cfg.n_key_value_heads}")
print(f"{(model.W_K[0][:4] == model.W_K[0][0]).all()=}") 

print(f"{model.blocks[0].attn._W_K.shape=}")
print(f"{model.blocks[0].attn.W_K.shape=}")

class LoraHooks(nn.Module):
    """
    Defines the LoRA hooks needed for the Attention Layers of the transformer.
    (Could be modified to add LoRA to the MLP layers)
    """

    lora_q: Lora
    lora_k: Lora
    lora_v: Lora
    lora_o: Lora
    cache_qkv_in: Float[Tensor, "batch pos d_model"] = None
    cache_z: Float[Tensor, "batch pos n_heads d_head"] = None

    def __init__(
        self,
        layer_idx: int,
        cfg: HookedTransformerConfig,
        lora_alpha: float = 32,
        rank: int = 4,
        dtype: t.dtype = None,
    ):
        super().__init__()
        self.layer_idx = layer_idx
        self.rank = rank
        self.lora_alpha = lora_alpha
        self.dtype = dtype

        self.n_qo_heads = n_qo_heads = cfg.n_heads
        self.n_kv_heads = n_kv_heads = cfg.n_key_value_heads if cfg.n_key_value_heads is not None else cfg.n_heads
        d_model, d_head = cfg.d_model, cfg.d_head

        self.lora_q = Lora(d_model, d_head, n_inst=n_qo_heads, rank=rank, lora_alpha=lora_alpha, dtype=dtype)
        self.lora_k = Lora(d_model, d_head, n_inst=n_kv_heads, rank=rank, lora_alpha=lora_alpha, dtype=dtype)
        self.lora_v = Lora(d_model, d_head, n_inst=n_kv_heads, rank=rank, lora_alpha=lora_alpha, dtype=dtype)
        self.lora_o = Lora(d_head, d_model, n_inst=n_qo_heads, rank=rank, lora_alpha=lora_alpha, dtype=dtype)

    def store_hook_attn_normalized(self, normalized: Float[Tensor, "batch pos d_model"], hook: HookPoint) -> None:
        """
        Cache the input to query/key/value.
        """
        self.cache_qkv_in = normalized

    def store_hook_z(self, z: Float[Tensor, "batch pos n_heads d_head"], hook: HookPoint) -> None:
        """
        Cache the input to $W_O$.
        """
        self.cache_z = z

    def list_fwd_hooks(self) -> list[tuple[str, Callable]]:
        """
        Returns a list of hook_point names and functions to call for the forward pass of
        the model using LoRA.
        """
        fwd_hooks = []
        # Attention Hooks qkv
        fwd_hooks.append((f"blocks.{self.layer_idx}.ln1.hook_normalized", self.store_hook_attn_normalized))
        fwd_hooks.append((f"blocks.{self.layer_idx}.attn.hook_q", self.lora_hook_qkv))
        fwd_hooks.append((f"blocks.{self.layer_idx}.attn.hook_k", self.lora_hook_qkv))
        fwd_hooks.append((f"blocks.{self.layer_idx}.attn.hook_v", self.lora_hook_qkv))
        # Attention Hooks z/out
        fwd_hooks.append((f"blocks.{self.layer_idx}.attn.hook_z", self.store_hook_z))
        fwd_hooks.append((f"blocks.{self.layer_idx}.hook_attn_out", self.lora_hook_out))

        return fwd_hooks

    def lora_hook_qkv(
        self, qkv_hook_out: Float[Tensor, "batch pos n_heads d_head"], hook: HookPoint
    ) -> Float[Tensor, "batch pos n_heads d_head"]:
        """
        Applies the LoRA modules to query/key/value, based on the hook location.
        Args:
            hook_qkv_out: Float[Tensor, "batch pos n_heads d_head"]
                The original output from query/key/value.
            hook: HookPoint
        Returns:
            The original output from query/key/value, plus the output from the corresponding LoRA module.
        """

        hook_location = hook.name.split(".")[-1]

        qkv_in = self.cache_qkv_in
        qkv_in_repeated = einops.repeat(qkv_in, "batch pos d_model -> batch pos n_inst d_model", n_inst=1)

        if hook_location == "hook_q":
            return qkv_hook_out + self.lora_q(qkv_in_repeated)
        elif hook_location == "hook_k":
            return qkv_hook_out + self.lora_k(qkv_in_repeated)
        elif hook_location == "hook_v":
            return qkv_hook_out + self.lora_v(qkv_in_repeated)
        else:
            raise ValueError(f"Invalid hook location: {hook_location}")

    def lora_hook_out(
        self, attn_out: Float[Tensor, "batch pos n_heads d_head"], hook: HookPoint
    ) -> Float[Tensor, "batch pos n_heads d_head"]:
        """
        Applies the LoRA modules to the output projection matrix W_O in the attention layer.
        The output of the LoRA module is computed per head, so we sum over heads before adding
        to the activation `attn_out`.

        Args:
            attn_out: Float[Tensor, "batch pos n_heads d_head"]
                The output from the attention layer.
            hook: HookPoint
        Returns:
            The original output from the attention layer, plus the output from the LoRA module.
        """

        lora_result = self.lora_o(self.cache_z)
        lora_attn_out = einops.einsum(lora_result, "... n_heads d_model -> ... d_model")
        return attn_out + lora_attn_out

class TransformerWithValueHeadLora(HookedTransformerWithValueHead):
    lora: nn.ModuleList
    lora_fwd_hooks: list[tuple[str, Callable]]
    dtype: t.dtype
    device: t.device
    use_value_head: bool

    def base_model_params(self):
        return (p for name, p in self.named_parameters() if "value_head" not in name and "lora" not in name)

    def lora_params(self):
        return self.lora.parameters()

    # we use these for compatibility with get_optimizer_and_scheduler
    def get_base_model_trainable_params(self):
        return self.lora_params()

    def get_value_head_params(self):
        return (p for name, p in self.named_parameters() if "value_head" in name)

    @classmethod
    def from_pretrained(cls, *args, lora_alpha: float = 32, rank: int = 4, **kwargs):
        model = super(TransformerWithValueHeadLora, cls).from_pretrained(*args, **kwargs)
        model.setup_lora(lora_alpha=lora_alpha, rank=rank, **kwargs)

        for param in model.base_model_params():
            param.requires_grad = False

        return model

    def setup_lora(self, lora_alpha: float = 32, rank: int = 4, **kwargs):
        """
        Initializes LoRA (Low-Rank Adaptation) for all attention layers in the transformer.

        Steps of this function are:
           - Creates a LoraHooks module for each transformer layer
           - Creates the list of forward hooks for all layers
        """

        self.lora = nn.ModuleList(
            [LoraHooks(layer_idx, self.cfg, lora_alpha, rank) for layer_idx in range(len(self.blocks))]
        ).to(device)

        # create list of all hooks for all layers
        self.lora_fwd_hooks = []
        for layer_idx in range(len(self.blocks)):
            self.lora_fwd_hooks.extend(self.lora[layer_idx].list_fwd_hooks())

    @property
    def fwd_hooks(self):
        return self.lora_fwd_hooks + [self.value_head_hook]

    def forward_with_value_head(
        self, tokens: Int[Tensor, "batch seq"]
    ) -> tuple[Float[Tensor, "batch seq d_vocab"], Float[Tensor, "batch seq"]]:
        """
        Forward pass with LoRA enabled, including the value head outputs.

        Args:
            tokens: Int[Tensor, "batch seq"]
                The input tokens to the transformer.
        Returns:
            logits: Float[Tensor, "batch seq d_vocab"]
                The logits of the transformer.
            value: Float[Tensor, "batch seq"]
                The value head outputs for each token.
        """

        with self.hooks(fwd_hooks=self.fwd_hooks):
            logits = self.forward(tokens)
        value = self.value_head_output
        return logits, value

    @t.no_grad()
    def generate(self, tokens: Int[Tensor, "batch seq"], **kwargs) -> Int[Tensor, "batch seq"]:
        """
        We override the generate method to use the LoRA hooks applied so that we don't need to update the previous training code.
        This function should call generate on the parent class (HookedTransformer), but with the LoRA hooks applied.
        We don't need to return the value head outputs during generation.

        Args:
            tokens: Int[Tensor, "batch seq"]
                The input tokens to the transformer.
            **kwargs:
                Additional keyword arguments to pass to the base class generate method.
        Returns:
            gen_tokens: Int[Tensor, "batch gen_len"]
                The generated tokens.
        """

        with self.hooks(fwd_hooks=self.lora_fwd_hooks):
            gen_tokens = super().generate(tokens, **kwargs)
        return gen_tokens


model = TransformerWithValueHeadLora.from_pretrained("pythia-14m").to(device)
tests_lora.test_lora_fwd_hooks_list(model)
tests_lora.test_lora_model_forward_methods(model)
print("All tests for TransformerWithValueHeadLora passed!")