rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/train/callbacks/freeze_unfreeze.py

@@ -1,10 +1,20 @@
 """FreezeUnfreeze callback."""
 
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+
 import torch
 from lightning.pytorch import LightningModule
 from lightning.pytorch.callbacks import BaseFinetuning
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 from torch.optim.optimizer import Optimizer
 
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
 
 class FreezeUnfreeze(BaseFinetuning):
     """Freezes a module and optionally unfreezes it after a number of epochs."""
@@ -14,7 +24,7 @@ class FreezeUnfreeze(BaseFinetuning):
         module_selector: list[str | int],
         unfreeze_at_epoch: int | None = None,
         unfreeze_lr_factor: float = 1,
-    ):
+    ) -> None:
         """Creates a new FreezeUnfreeze.
 
         Args:
@@ -30,6 +40,8 @@ class FreezeUnfreeze(BaseFinetuning):
         self.module_selector = module_selector
         self.unfreeze_at_epoch = unfreeze_at_epoch
         self.unfreeze_lr_factor = unfreeze_lr_factor
+        if unfreeze_at_epoch == 0:
+            raise ValueError("unfreeze_at_epoch cannot be 0")
 
     def _get_target_module(self, pl_module: LightningModule) -> torch.nn.Module:
         target_module = pl_module
@@ -40,18 +52,18 @@ class FreezeUnfreeze(BaseFinetuning):
                 target_module = getattr(target_module, k)
         return target_module
 
-    def freeze_before_training(self, pl_module: LightningModule):
+    def freeze_before_training(self, pl_module: LightningModule) -> None:
         """Freeze the model at the beginning of training.
 
         Args:
             pl_module: the LightningModule.
         """
-        print(f"freezing model at {self.module_selector}")
+        logger.info(f"freezing model at {self.module_selector}")
         self.freeze(self._get_target_module(pl_module))
 
     def finetune_function(
         self, pl_module: LightningModule, current_epoch: int, optimizer: Optimizer
-    ):
+    ) -> None:
         """Check whether we should unfreeze the model on each epoch.
 
         Args:
@@ -61,19 +73,338 @@ class FreezeUnfreeze(BaseFinetuning):
         """
         if self.unfreeze_at_epoch is None:
             return
-        if current_epoch != self.unfreeze_at_epoch:
-            return
-        print(
-            f"unfreezing model at {self.module_selector} since we are on epoch {current_epoch}"
+        elif current_epoch == self.unfreeze_at_epoch:
+            logger.info(
+                f"unfreezing model at {self.module_selector} since we are on epoch {current_epoch}"
+            )
+            self.unfreeze_and_add_param_group(
+                modules=self._get_target_module(pl_module),
+                optimizer=optimizer,
+                initial_denom_lr=self.unfreeze_lr_factor,
+            )
+            if "scheduler" in pl_module.schedulers:
+                scheduler = pl_module.schedulers["scheduler"]
+                if isinstance(scheduler, ReduceLROnPlateau):
+                    while len(scheduler.min_lrs) < len(optimizer.param_groups):
+                        logger.info(
+                            "appending to ReduceLROnPlateau scheduler min_lrs for unfreeze"
+                        )
+                        scheduler.min_lrs.append(scheduler.min_lrs[0])
+        elif current_epoch > self.unfreeze_at_epoch:
+            # always do this because overhead is minimal, and it allows restoring
+            # from a checkpoint (resuming a run) without messing up unfreezing
+            BaseFinetuning.make_trainable(self._get_target_module(pl_module))
+
+
+@dataclass
+class FTStage:
+    """Specification for a single fine-tuning stage.
+
+    Each stage is activated when the trainer reaches a specific epoch (`at_epoch`).
+    Within that stage, modules whose **qualified name** (from `named_modules()`)
+    matches any substring in `freeze_selectors` will be frozen, except those whose
+    name matches any substring in `unfreeze_selectors`, which are forced trainable.
+
+    freeze_selectors does not carry over to other stages. That is, if you freeze module
+    A for stage 1, it will not be frozen for stage 2 unless specified again in stage 2.
+    All stages indepedently update trainability of all modules specified or unspecified.
+
+    Args:
+        at_epoch: Epoch index at which to apply this stage (0-based).
+        freeze_selectors: Substrings; any module name containing any of these will
+            be frozen in this stage (unless also matched by `unfreeze_selectors`).
+        unfreeze_selectors: Substrings; any module name containing any of these
+            will be **unfrozen** (trainable) in this stage, overriding freezes.
+        unfreeze_lr_factor: When parameters become trainable and are **not yet**
+            part of the optimizer, a new param group is added with learning rate
+            `base_lr / unfreeze_lr_factor`. Use 1.0 to keep the base learning rate.
+        scale_existing_groups: If provided and not 1.0, multiply the learning rate
+            of **all existing optimizer param groups** by this factor at the moment
+            this stage is applied. Use this to calm down previously-trainable
+            parts (e.g., the head) when unfreezing deeper layers.
+            Set to ``None`` to leave existing groups unchanged.
+    """
+
+    at_epoch: int
+    freeze_selectors: Sequence[str]
+    unfreeze_selectors: Sequence[str]
+    unfreeze_lr_factor: float = 1.0
+    scale_existing_groups: float | None = None
+
+
+class MultiStageFineTuning(BaseFinetuning):
+    """Multi-stage fine-tuning with flexible name-based selection.
+
+    Behavior per stage:
+      1) Start from a **fully trainable** baseline.
+      2) Optionally **scale existing optimizer groups** via `scale_existing_groups`.
+      3) **Freeze** modules matching any `freeze_selectors`.
+      4) **Unfreeze** modules matching any `unfreeze_selectors` (overrides step 3).
+      5) For newly trainable parameters **not yet** in the optimizer, add a new
+         param group using `unfreeze_lr_factor` (lr = base_lr / factor).
+
+    Stages are applied exactly once at their `at_epoch`. The plan is recomputed
+    from scratch at each stage to keep behavior predictable on resume.
+    """
+
+    def __init__(self, stages: list[FTStage]) -> None:
+        """Multi-stage fine-tuning with flexible name-based selection.
+
+        Args:
+            stages: A sequence of stage specifications.
+
+        Raises:
+            ValueError: If two stages specify the same `at_epoch`.
+        """
+        super().__init__()
+        self.stages = stages
+
+        # Validate uniqueness of epochs and sort stages.
+        seen: set[int] = set()
+        for st in self.stages:
+            if st.at_epoch in seen:
+                raise ValueError(f"Duplicate at_epoch in stages: {st.at_epoch}")
+            if st.scale_existing_groups is not None and st.scale_existing_groups <= 0.0:
+                raise ValueError("scale_existing_groups, if set, must be > 0.")
+            seen.add(st.at_epoch)
+        self.stages.sort(key=lambda x: x.at_epoch)
+
+        self._applied_epochs: set[int] = set()
+
+    @staticmethod
+    def _freeze_unfreeze(mod: torch.nn.Module, freeze: bool) -> None:
+        """Freeze or unfreeze all parameters of a module without going through Lightning's flatten logic.
+
+        This is a workaround to avoid infinite recursion on ModuleDicts.
+
+        Args:
+            mod: The module to freeze.
+            freeze: Whether to freeze the module.
+        """
+        for p in mod.parameters(recurse=True):
+            p.requires_grad = not freeze
+
+    @staticmethod
+    def _names_matching(names: Iterable[str], selectors: Sequence[str]) -> set[str]:
+        """Return the subset of `names` that contains any of the given selectors.
+
+        Matching is done via simple substring checks (`sel in name`).
+
+        Args:
+            names: Iterable of qualified module names (e.g., from `named_modules()`).
+            selectors: Substrings to match against each name. Empty strings are ignored.
+
+        Returns:
+            A set of names from `names` that match at least one selector.
+        """
+        if not selectors:
+            return set()
+        sels: list[str] = [s for s in selectors if s]
+        out: set[str] = set()
+        for n in names:
+            if any(sel in n for sel in sels):
+                out.add(n)
+        return out
+
+    @staticmethod
+    def _modules_by_names(
+        root: torch.nn.Module, wanted: set[str]
+    ) -> list[torch.nn.Module]:
+        """Map qualified names to module objects.
+
+        Args:
+            root: The root module (e.g., your LightningModule).
+            wanted: Qualified names of submodules to retrieve.
+
+        Returns:
+            A list of modules corresponding to the given names that exist under `root`.
+        """
+        if not wanted:
+            return []
+        name_to_module: dict[str, torch.nn.Module] = dict(root.named_modules())
+        return [name_to_module[n] for n in wanted if n in name_to_module]
+
+    @staticmethod
+    def _existing_param_ids(optimizer: Optimizer) -> set[int]:
+        """Collect ids of all parameters already tracked by the optimizer.
+
+        Args:
+            optimizer: The optimizer to inspect.
+
+        Returns:
+            A set of parameter ids already tracked by the optimizer.
+        """
+        return {id(p) for g in optimizer.param_groups for p in g["params"]}
+
+    @staticmethod
+    def _iter_module_params(modules: list[torch.nn.Module]) -> list[torch.nn.Parameter]:
+        """Flatten parameters from a list of modules (no duplicates, trainable first).
+
+        Args:
+            modules: A list of modules to inspect.
+
+        Returns:
+            A list of parameters from the modules, in order of appearance.
+        """
+        seen: set[int] = set()
+        ordered: list[torch.nn.Parameter] = []
+        for m in modules:
+            for p in m.parameters():
+                if id(p) not in seen:
+                    seen.add(id(p))
+                    ordered.append(p)
+        return ordered
+
+    def _apply_stage(
+        self, pl_module: LightningModule, optimizer: Optimizer, stage: FTStage
+    ) -> None:
+        """Apply a single fine-tuning stage to `pl_module` and `optimizer`.
+
+        Order of operations:
+          1) Make everything trainable (baseline).
+          2) If `scale_existing_groups` is set, multiply LR of **existing** optimizer
+             groups by this factor (and update ReduceLROnPlateau `min_lrs` if present).
+          3) Freeze modules matched by `freeze_selectors` minus `unfreeze_selectors`.
+          4) Ensure modules matched by `unfreeze_selectors` are trainable.
+          5) Add new optimizer param groups for newly-trainable modules with LR
+             scaled by `unfreeze_lr_factor`.
+
+        Args:
+            pl_module: The LightningModule being trained.
+            optimizer: The optimizer currently used by the trainer.
+            stage: The stage specification to apply at the current epoch.
+
+        Returns:
+            None.
+        """
+        model: torch.nn.Module = pl_module
+        all_names: list[str] = [n for n, _ in model.named_modules()]
+
+        freeze_names: set[str] = self._names_matching(all_names, stage.freeze_selectors)
+        unfreeze_names: set[str] = self._names_matching(
+            all_names, stage.unfreeze_selectors
+        )
+
+        # 1) Baseline: everything trainable.
+        self._freeze_unfreeze(model, freeze=False)
+
+        # 2) Optionally scale existing optimizer groups (e.g., calm down the head).
+        if (
+            stage.scale_existing_groups is not None
+            and stage.scale_existing_groups != 1.0
+        ):
+            factor: float = stage.scale_existing_groups
+            for g in optimizer.param_groups:
+                old_lr = float(g.get("lr", 0.0))
+                g["lr"] = old_lr * factor
+            # Keep ReduceLROnPlateau bounds consistent if present.
+            if hasattr(pl_module, "schedulers") and "scheduler" in getattr(
+                pl_module, "schedulers", {}
+            ):
+                scheduler = pl_module.schedulers["scheduler"]
+                if isinstance(scheduler, ReduceLROnPlateau):
+                    scheduler.min_lrs = [float(m) * factor for m in scheduler.min_lrs]
+
+        # 3) Freeze matched, except those explicitly unfreezed.
+        to_freeze: set[str] = freeze_names - unfreeze_names
+        freeze_modules: list[torch.nn.Module] = self._modules_by_names(model, to_freeze)
+        if freeze_modules:
+            to_display = sorted(list(to_freeze))
+            logger.info(
+                f"[FT stage @ epoch {stage.at_epoch}] Freezing {len(freeze_modules)} modules "
+                f"(matched: {to_display[:2] + to_display[-2:]}{'...' if len(to_freeze) > 4 else ''})"
+            )
+            for m in freeze_modules:
+                self._freeze_unfreeze(m, freeze=True)
+
+        # 4) Ensure explicitly unfreezed modules are trainable.
+        unfreeze_modules: list[torch.nn.Module] = self._modules_by_names(
+            model, unfreeze_names
         )
-        self.unfreeze_and_add_param_group(
-            modules=self._get_target_module(pl_module),
-            optimizer=optimizer,
-            initial_denom_lr=self.unfreeze_lr_factor,
+        if unfreeze_modules:
+            to_display = sorted(list(unfreeze_names))
+            logger.info(
+                f"[FT stage @ epoch {stage.at_epoch}] Unfreezing {len(unfreeze_modules)} modules "
+                f"(matched: {to_display[:2] + to_display[-2:]}{'...' if len(unfreeze_names) > 4 else ''})"
+            )
+            for m in unfreeze_modules:
+                self._freeze_unfreeze(m, freeze=False)
+
+            # 5) Add *newly-trainable* params only (no duplicates)
+            denom: float = (
+                stage.unfreeze_lr_factor if stage.unfreeze_lr_factor != 1.0 else 1.0
+            )
+            all_params = self._iter_module_params(unfreeze_modules)
+            already = self._existing_param_ids(optimizer)
+            new_params = [
+                p for p in all_params if p.requires_grad and id(p) not in already
+            ]
+
+            if new_params:
+                # Use current "base" lr (after any scale_existing_groups) as the reference
+                base_lr = float(optimizer.param_groups[0].get("lr", 0.0))
+                group_lr = base_lr / denom if denom != 0 else base_lr
+                optimizer.add_param_group({"params": new_params, "lr": group_lr})
+
+                # Extend ReduceLROnPlateau.min_lrs to match param group count
+                if hasattr(pl_module, "schedulers") and "scheduler" in getattr(
+                    pl_module, "schedulers", {}
+                ):
+                    scheduler = pl_module.schedulers["scheduler"]
+                    if isinstance(scheduler, ReduceLROnPlateau):
+                        while len(scheduler.min_lrs) < len(optimizer.param_groups):
+                            logger.info(
+                                "Extending ReduceLROnPlateau.min_lrs for new param group"
+                            )
+                            scheduler.min_lrs.append(scheduler.min_lrs[0])
+
+        # Summary logging.
+        trainable, frozen = 0, 0
+        for p in model.parameters():
+            if p.requires_grad:
+                trainable += p.numel()
+            else:
+                frozen += p.numel()
+        logger.info(
+            f"[FT stage @ epoch {stage.at_epoch}] Trainable params: {trainable:,} | Frozen params: {frozen:,}"
         )
 
-        if "plateau" in pl_module.schedulers:
-            scheduler = pl_module.schedulers["plateau"]
-            while len(scheduler.min_lrs) < len(optimizer.param_groups):
-                print("appending to plateau scheduler min_lrs")
-                scheduler.min_lrs.append(scheduler.min_lrs[0])
+    def freeze_before_training(self, pl_module: LightningModule) -> None:
+        """Hook: Called by Lightning before the first training epoch.
+
+        If a stage is scheduled at epoch 0, we defer its application to the first
+        call of `finetune_function` (when the optimizer is available). Otherwise,
+        we simply log that training begins with a fully trainable model.
+
+        Args:
+            pl_module: The LightningModule being trained.
+        """
+        if any(st.at_epoch == 0 for st in self.stages):
+            logger.info(
+                "Stage scheduled for epoch 0 will be applied at the first finetune_function "
+                "call when the optimizer is available."
+            )
+        else:
+            logger.info("No stage at epoch 0; starting fully trainable by default.")
+
+    def finetune_function(
+        self, pl_module: LightningModule, current_epoch: int, optimizer: Optimizer
+    ) -> None:
+        """Hook: Called by Lightning at each epoch to adjust trainability.
+
+        Applies any stage whose `at_epoch` equals `current_epoch` and that has not
+        yet been applied in this run. Recomputes freeze/unfreeze decisions from
+        scratch for that stage.
+
+        Args:
+            pl_module: The LightningModule being trained.
+            current_epoch: The current epoch index (0-based).
+            optimizer: The optimizer currently used by the trainer.
+        """
+        for st in self.stages:
+            if st.at_epoch == current_epoch and st.at_epoch not in self._applied_epochs:
+                logger.info(
+                    f"Applying multi-stage fine-tuning plan at epoch {current_epoch}"
+                )
+                self._apply_stage(pl_module, optimizer, st)
+                self._applied_epochs.add(st.at_epoch)

rslearn/train/callbacks/gradients.py

@@ -0,0 +1,129 @@
+"""Gradient logging and surgery callbacks."""
+
+from typing import Any
+
+import torch
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.trainer import Trainer
+from torch.nn import Module
+from torch.optim import Optimizer
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class MiniPCGrad(Callback):
+    """PCGrad from https://arxiv.org/abs/2001.06782.
+
+    This is roughly equivalent to PCGrad but uses gradient accumulation to factorize
+    projections, so we can keep gradients orthogonal in O(1) memory instead of O(n).
+    This is still quite slow, requiring an extra copy of parameter gradients in memory.
+    """
+
+    def __init__(
+        self,
+        selectors: list[str],
+        deselectors: list[str] | None = None,
+        only_monitor: bool = False,
+    ) -> None:
+        """Initialize the callback.
+
+        Args:
+            selectors: Prefixes for selecting which parameters to operate on.
+            deselectors: Prefixes for deselecting which parameters to operate on. Applied after selectors.
+            only_monitor: If true, only log gradients, don't clip them.
+        """
+        self.selectors = selectors
+        self.deselectors = deselectors or []
+        self.only_monitor = only_monitor
+        self.prev_grads: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+    def on_train_batch_start(
+        self, trainer: Trainer, pl_module: Module, batch: Any, batch_idx: int
+    ) -> None:
+        """Save the dataset source each batch."""
+        self.dataset_source = batch[0][0]["dataset_source"]
+        self.batch_size = len(batch[0])
+
+    def on_before_optimizer_step(
+        self, trainer: Trainer, pl_module: Module, optimizer: Optimizer
+    ) -> None:
+        """Reset the previous gradients."""
+        self.prev_grads = {}
+
+    def on_after_backward(self, trainer: Trainer, pl_module: Module) -> None:
+        """Called after every loss.backward(), even under gradient accumulation.
+
+        Receives the accumulated gradients (i.e., accumulated + micro batch gradient).
+
+        Args:
+            trainer: The trainer object.
+            pl_module: The module object.
+        """
+        prev_grad_norms = []
+        micro_grad_norms = []
+        angles = []
+
+        eps = 1e-12  # numerical stability
+
+        for name, param in pl_module.named_parameters():
+            if param.grad is None:
+                continue
+            elif all(selector not in name for selector in self.selectors) or any(
+                deselector in name for deselector in self.deselectors
+            ):
+                continue
+
+            try:
+                prev_grad, prev_grad_norm = self.prev_grads[name]
+            except KeyError:
+                prev_grad = torch.zeros_like(param.grad, device=param.device)
+                prev_grad_norm = torch.tensor(0.0, device=param.device)
+
+            with torch.no_grad():
+                # current accumulated grad = prev_grad + micro_grad
+                micro_grad = param.grad - prev_grad
+                micro_grad_norm = micro_grad.norm()
+
+                micro_grad_norms.append(micro_grad_norm)
+                prev_grad_norms.append(prev_grad_norm)
+
+                # cosine of angle between micro and prev
+                denom = (micro_grad_norm * prev_grad_norm).clamp_min(eps)
+                if prev_grad_norm > 0 and micro_grad_norm > 0:
+                    dot = torch.dot(micro_grad.flatten(), prev_grad.flatten())
+                    cos_theta = dot / denom
+                    angles.append(cos_theta)
+
+                    if not self.only_monitor and dot < 0:
+                        # Remove the component of micro_grad along prev_grad
+                        proj_coeff = dot / (prev_grad_norm**2 + eps)
+                        micro_projection = micro_grad - proj_coeff * prev_grad
+                        # keep accumulated gradient as (prev + projected micro)
+                        param.grad = prev_grad + micro_projection
+                        logger.info(
+                            f"{name} (cos={cos_theta:.4f},dot={dot:.4f},prev_grad_norm={prev_grad_norm:.4f},micro_grad_norm={micro_grad_norm:.4f})"
+                        )
+
+                # store the latest accumulated gradient and its norm
+                self.prev_grads[name] = (param.grad.clone(), param.grad.norm())
+
+        log_prev_grad_norms = (
+            torch.stack(prev_grad_norms).norm()
+            if prev_grad_norms
+            else torch.tensor(0.0)
+        )
+        log_micro_grad_norms = (
+            torch.stack(micro_grad_norms).norm()
+            if micro_grad_norms
+            else torch.tensor(0.0)
+        )
+        log_angles = torch.stack(angles).mean() if angles else torch.tensor(0.0)
+
+        info = {
+            f"grads/{self.dataset_source}_prev_grad_norms": log_prev_grad_norms,
+            f"grads/{self.dataset_source}_micro_grad_norms": log_micro_grad_norms,
+            f"grads/{self.dataset_source}_angles": log_angles,
+        }
+        self.log_dict(info, on_step=True, on_epoch=False, batch_size=self.batch_size)

rslearn/train/callbacks/peft.py

@@ -0,0 +1,116 @@
+"""Parameter-efficient finetuning callbacks."""
+
+import torch
+import torch.nn.functional as F
+from lightning.pytorch import LightningModule
+from lightning.pytorch.callbacks import BaseFinetuning
+from torch.optim.optimizer import Optimizer
+
+
+class SplitProjection(torch.nn.Module):
+    """Split projection weights into trainable and frozen parts.
+
+    This module is used to split the projection weights into trainable and frozen parts.
+    The trainable part is used to compute the output, and the frozen part is used to
+    compute the output without gradients.
+    """
+
+    def __init__(self, dim: int, r: int = 8) -> None:
+        """Initialize the SplitProjection module.
+
+        Args:
+            dim: the dimension of the input and output
+            r: the number of trainable parameters
+        """
+        super().__init__()
+        self.dim = dim
+        self.r = r
+
+        # Register indices as buffers so they move to the correct device automatically
+        indices = torch.randperm(dim)
+        self.register_buffer("trainable_inds", indices[:r])
+        self.register_buffer("frozen_inds", indices[r:])
+
+        # Create parameter modules directly
+        self.trainable_w = torch.nn.Parameter(torch.empty(dim, r), requires_grad=True)
+        self.frozen_w = torch.nn.Parameter(
+            torch.empty(dim, dim - r), requires_grad=False
+        )
+        self.trainable_b = torch.nn.Parameter(torch.empty(r), requires_grad=True)
+        self.frozen_b = torch.nn.Parameter(torch.empty(dim - r), requires_grad=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of the SplitProjection module.
+
+        Args:
+            x: the input tensor
+
+        Returns:
+            the output tensor
+        """
+        trainable_out = F.linear(x, self.trainable_w, self.trainable_b)
+        frozen_out = F.linear(x, self.frozen_w, self.frozen_b)
+
+        output = torch.zeros(x.shape, device=x.device, dtype=trainable_out.dtype)
+        output[..., self.trainable_inds] = trainable_out  # type: ignore
+        output[..., self.frozen_inds] = frozen_out  # type: ignore
+
+        return output
+
+
+class APLA(BaseFinetuning):
+    """APLA (https://arxiv.org/pdf/2503.11335v2) finetuning callback."""
+
+    def __init__(self, r: int = 8) -> None:
+        """Initialize the APLA finetuning callback.
+
+        Args:
+            r: the number of trainable parameters
+        """
+        super().__init__()
+        self.r = r
+
+    def freeze_before_training(self, pl_module: LightningModule) -> None:
+        """Freeze the model before training.
+
+        Args:
+            pl_module: the LightningModule
+        """
+        print("splitting projection weights by monkeypatching")
+        model = pl_module.model
+        self.freeze(model.encoder[0])
+        n_trainable = 0
+        for layer in model.encoder[0].model.blocks:
+            if hasattr(layer, "attn"):
+                alpa_proj = SplitProjection(layer.attn.proj.weight.shape[0], r=self.r)
+                proj_weight = layer.attn.proj.weight.data.clone()
+                proj_bias = layer.attn.proj.bias.data.clone()
+
+                alpa_proj.trainable_w.data = proj_weight[alpa_proj.trainable_inds, :]
+                alpa_proj.frozen_w.data = proj_weight[alpa_proj.frozen_inds, :]
+
+                alpa_proj.trainable_b.data = proj_bias[alpa_proj.trainable_inds]
+                alpa_proj.frozen_b.data = proj_bias[alpa_proj.frozen_inds]
+
+                alpa_proj.trainable_w.requires_grad = True
+                alpa_proj.trainable_b.requires_grad = True
+                n_trainable += (
+                    alpa_proj.trainable_w.numel() + alpa_proj.trainable_b.numel()
+                )
+
+                layer.attn.proj = alpa_proj
+
+        print(f"n_trainable: {n_trainable / int(1e6)}M")
+
+    def finetune_function(
+        self, pl_module: LightningModule, current_epoch: int, optimizer: Optimizer
+    ) -> None:
+        """Do nothing here.
+
+        Args:
+            pl_module: the LightningModule
+            current_epoch: the current epoch
+            optimizer: the optimizer
+        """
+        # Maybe worth unfreezing down the line?
+        pass