| from __future__ import annotations
|
|
|
| import typing as T
|
| from dataclasses import dataclass, fields
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
|
|
| @dataclass
|
| class TTTConfig:
|
| lr: float = 4e-4
|
| steps: int = 30
|
| ags: int = 16
|
| batch_size: int = 2
|
| mask_ratio: float = 0.15
|
| crop_size: int = 1024
|
| bert_leave_prob: float = 0.1
|
| bert_replace_prob: float = 0.1
|
| optimizer: str = "sgd"
|
| momentum: float = 0.0
|
| weight_decay: float = 0.0
|
| seed: int | None = 0
|
| lora_rank: int = 8
|
| lora_alpha: float = 32.0
|
| lora_target_replace_module: str | None = None
|
| lora_target_modules: tuple[str, ...] | None = None
|
| initial_state_reset: bool = True
|
| automatic_best_state_reset: bool = False
|
| eval_each_step: bool = False
|
| gradient_clip: bool = False
|
| gradient_clip_max_norm: float = 1.0
|
|
|
| @classmethod
|
| def from_kwargs(cls, **kwargs: T.Any) -> "TTTConfig":
|
| valid_names = {field.name for field in fields(cls)}
|
| unknown_names = set(kwargs) - valid_names
|
| assert len(unknown_names) == 0, f"Unknown TTTConfig fields: {sorted(unknown_names)}"
|
| return cls(**kwargs)
|
|
|
| def merged(self, overrides: T.Mapping[str, T.Any] | "TTTConfig" | None) -> "TTTConfig":
|
| if overrides is None:
|
| return self
|
| if isinstance(overrides, TTTConfig):
|
| return overrides
|
| values = {field.name: self.__dict__[field.name] for field in fields(self)}
|
| for name, value in overrides.items():
|
| assert name in values, f"Unknown TTTConfig field: {name}"
|
| values[name] = value
|
| return TTTConfig(**values)
|
|
|
| def verify(self) -> None:
|
| assert self.lr > 0.0, "TTT learning rate must be positive."
|
| assert self.steps >= 1, "TTT steps must be >= 1."
|
| assert self.ags >= 1, "TTT gradient accumulation steps must be >= 1."
|
| assert self.batch_size >= 1, "TTT batch_size must be >= 1."
|
| assert 0.0 < self.mask_ratio <= 1.0, "TTT mask_ratio must be in (0, 1]."
|
| assert self.crop_size >= 1, "TTT crop_size must be >= 1."
|
| assert self.lora_rank >= 1, "TTT v1 is LoRA-only, so lora_rank must be >= 1."
|
| assert self.lora_alpha > 0.0, "TTT lora_alpha must be positive."
|
| assert self.optimizer in {"adamw", "sgd"}, "TTT optimizer must be 'adamw' or 'sgd'."
|
| assert 0.0 <= self.bert_leave_prob <= 1.0, "bert_leave_prob must be in [0, 1]."
|
| assert 0.0 <= self.bert_replace_prob <= 1.0, "bert_replace_prob must be in [0, 1]."
|
| assert self.bert_leave_prob + self.bert_replace_prob <= 1.0, (
|
| "bert_leave_prob + bert_replace_prob must be <= 1."
|
| )
|
| if self.gradient_clip:
|
| assert self.gradient_clip_max_norm > 0.0, "gradient_clip_max_norm must be positive."
|
|
|
|
|
| class LoraInjectedLinear(nn.Module):
|
| def __init__(self, linear: nn.Module, rank: int, alpha: float) -> None:
|
| super().__init__()
|
| weight = linear._parameters["weight"]
|
| assert weight.ndim == 2, "LoRA can only wrap 2D linear weights."
|
| self.linear = linear
|
| self.linear.requires_grad_(False)
|
| self.rank = rank
|
| self.scale = alpha
|
| in_features = weight.shape[1]
|
| out_features = weight.shape[0]
|
| self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32)
|
| self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32)
|
| self.lora_down.to(device=weight.device)
|
| self.lora_up.to(device=weight.device)
|
| nn.init.normal_(self.lora_down.weight, std=1.0 / rank)
|
| nn.init.zeros_(self.lora_up.weight)
|
|
|
| @property
|
| def weight(self) -> torch.Tensor:
|
| return self.linear._parameters["weight"]
|
|
|
| @property
|
| def bias(self) -> torch.Tensor | None:
|
| return self.linear._parameters["bias"]
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| base = self.linear(x)
|
| delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale
|
| return base + delta.to(dtype=base.dtype)
|
|
|
|
|
| class FastPLMTestTimeTrainingMixin:
|
| def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None:
|
| base_config = TTTConfig()
|
| self._ttt_cfg = base_config.merged(ttt_config)
|
| self._ttt_cfg.verify()
|
| self._ttt_initialized = False
|
| self._ttt_initial_state: list[dict[str, torch.Tensor]] | None = None
|
|
|
| @property
|
| def ttt_config(self) -> TTTConfig:
|
| if "_ttt_cfg" not in self.__dict__:
|
| self.init_ttt()
|
| return self._ttt_cfg
|
|
|
| def _ttt_get_trainable_modules(self) -> list[nn.Module]:
|
| return [self]
|
|
|
| def _ttt_get_frozen_modules(self) -> list[nn.Module]:
|
| return []
|
|
|
| def _ttt_tokenize(
|
| self,
|
| seq: str | list[str] | None = None,
|
| input_ids: torch.Tensor | None = None,
|
| **kwargs: T.Any,
|
| ) -> torch.Tensor | dict[str, torch.Tensor]:
|
| del kwargs
|
| if input_ids is not None:
|
| return input_ids
|
| assert seq is not None, "Pass either seq or input_ids for TTT."
|
| tokenized = self.tokenizer(seq, return_tensors="pt", padding=True)
|
| return tokenized["input_ids"]
|
|
|
| def _ttt_mask_token(self) -> int:
|
| return int(self.tokenizer.mask_token_id)
|
|
|
| def _ttt_padding_token(self) -> int:
|
| return int(self.tokenizer.pad_token_id)
|
|
|
| def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
| tokenizer = self.tokenizer
|
| special_ids = set(tokenizer.all_special_ids)
|
| vocab_size = int(self.config.vocab_size)
|
| ids = [idx for idx in range(vocab_size) if idx not in special_ids]
|
| assert len(ids) > 0, "TTT replacement token set is empty."
|
| return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
|
|
|
| def _ttt_predict_logits(
|
| self,
|
| batch: torch.Tensor | dict[str, torch.Tensor],
|
| **kwargs: T.Any,
|
| ) -> torch.Tensor:
|
| del kwargs
|
| if isinstance(batch, dict):
|
| output = self(**batch)
|
| return output.logits
|
| attention_mask = batch.ne(self._ttt_padding_token())
|
| output = self(input_ids=batch, attention_mask=attention_mask)
|
| return output.logits
|
|
|
| def _ttt_eval_step(
|
| self,
|
| step: int,
|
| loss: float,
|
| seq: str | list[str] | None = None,
|
| input_ids: torch.Tensor | None = None,
|
| **kwargs: T.Any,
|
| ) -> tuple[dict[str, T.Any], float | None]:
|
| del step, loss, seq, input_ids, kwargs
|
| return {}, None
|
|
|
| def _ttt_is_lora_target(
|
| self,
|
| name: str,
|
| full_name: str,
|
| module: nn.Module,
|
| active: bool,
|
| target_modules: tuple[str, ...] | None,
|
| ) -> bool:
|
| if not active:
|
| return False
|
| if isinstance(module, LoraInjectedLinear):
|
| return False
|
| if (
|
| target_modules is not None
|
| and name not in target_modules
|
| and full_name not in target_modules
|
| ):
|
| return False
|
| if isinstance(module, nn.Linear):
|
| return True
|
| if "weight" not in module._parameters:
|
| return False
|
| weight = module._parameters["weight"]
|
| if weight is None or weight.ndim != 2:
|
| return False
|
| return "Linear" in module.__class__.__name__
|
|
|
| def _ttt_inject_lora(self) -> int:
|
| cfg = self.ttt_config
|
| cfg.verify()
|
| target_class = cfg.lora_target_replace_module
|
| target_modules = cfg.lora_target_modules
|
| wrapped = 0
|
|
|
| def inject(module: nn.Module, prefix: str, active: bool) -> None:
|
| nonlocal wrapped
|
| for name, child in list(module.named_children()):
|
| full_name = f"{prefix}.{name}" if prefix else name
|
| child_active = active
|
| if target_class is not None:
|
| child_active = active or child.__class__.__name__ == target_class
|
| if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules):
|
| setattr(
|
| module,
|
| name,
|
| LoraInjectedLinear(child, rank=cfg.lora_rank, alpha=cfg.lora_alpha),
|
| )
|
| wrapped += 1
|
| continue
|
| inject(child, full_name, child_active)
|
|
|
| for trainable_module in self._ttt_get_trainable_modules():
|
| inject(trainable_module, "", target_class is None)
|
| assert wrapped > 0, "TTT LoRA injection did not find any target modules."
|
| return wrapped
|
|
|
| def _ttt_lora_modules(self) -> list[LoraInjectedLinear]:
|
| return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)]
|
|
|
| def _ttt_lora_parameters(self) -> list[nn.Parameter]:
|
| params: list[nn.Parameter] = []
|
| for module in self._ttt_lora_modules():
|
| params.extend(module.lora_down.parameters())
|
| params.extend(module.lora_up.parameters())
|
| assert len(params) > 0, "TTT has no LoRA parameters."
|
| return params
|
|
|
| def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]:
|
| snapshot = []
|
| for module in self._ttt_lora_modules():
|
| snapshot.append(
|
| {
|
| "lora_down.weight": module.lora_down.weight.detach().clone(),
|
| "lora_up.weight": module.lora_up.weight.detach().clone(),
|
| }
|
| )
|
| assert len(snapshot) > 0, "TTT has no LoRA state to snapshot."
|
| return snapshot
|
|
|
| def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None:
|
| modules = self._ttt_lora_modules()
|
| assert len(modules) == len(state), "TTT LoRA state/module count mismatch."
|
| with torch.no_grad():
|
| for module, module_state in zip(modules, state):
|
| module.lora_down.weight.copy_(module_state["lora_down.weight"])
|
| module.lora_up.weight.copy_(module_state["lora_up.weight"])
|
|
|
| def _ttt_ensure_initialized(self) -> None:
|
| if "_ttt_cfg" not in self.__dict__:
|
| self.init_ttt()
|
| if self._ttt_initialized:
|
| return
|
| self._ttt_inject_lora()
|
| self._ttt_initial_state = self._ttt_snapshot_lora_state()
|
| self._ttt_initialized = True
|
|
|
| def ttt_reset(self) -> None:
|
| self._ttt_ensure_initialized()
|
| assert self._ttt_initial_state is not None, "TTT initial state is not available."
|
| self._ttt_restore_lora_state(self._ttt_initial_state)
|
|
|
| def _ttt_make_optimizer(self) -> torch.optim.Optimizer:
|
| cfg = self.ttt_config
|
| params = self._ttt_lora_parameters()
|
| if cfg.optimizer == "sgd":
|
| return torch.optim.SGD(
|
| params,
|
| lr=cfg.lr,
|
| momentum=cfg.momentum,
|
| weight_decay=cfg.weight_decay,
|
| )
|
| return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay)
|
|
|
| def _ttt_to_device(
|
| self,
|
| batch: torch.Tensor | dict[str, torch.Tensor],
|
| device: torch.device,
|
| ) -> torch.Tensor | dict[str, torch.Tensor]:
|
| if isinstance(batch, dict):
|
| return {name: tensor.to(device) for name, tensor in batch.items()}
|
| return batch.to(device)
|
|
|
| def _ttt_input_ids_from_batch(
|
| self,
|
| batch: torch.Tensor | dict[str, torch.Tensor],
|
| ) -> torch.Tensor:
|
| if isinstance(batch, dict):
|
| return batch["input_ids"]
|
| return batch
|
|
|
| def _ttt_set_input_ids(
|
| self,
|
| batch: torch.Tensor | dict[str, torch.Tensor],
|
| input_ids: torch.Tensor,
|
| ) -> torch.Tensor | dict[str, torch.Tensor]:
|
| if isinstance(batch, dict):
|
| updated = dict(batch)
|
| updated["input_ids"] = input_ids
|
| return updated
|
| return input_ids
|
|
|
| def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
|
| pad_token = self._ttt_padding_token()
|
| mask = input_ids.ne(pad_token)
|
| special_ids = set(self.tokenizer.all_special_ids)
|
| for special_id in special_ids:
|
| mask = mask & input_ids.ne(int(special_id))
|
| return mask
|
|
|
| def _ttt_sample_crop(
|
| self,
|
| batch: torch.Tensor | dict[str, torch.Tensor],
|
| generator: torch.Generator,
|
| ) -> torch.Tensor | dict[str, torch.Tensor]:
|
| input_ids = self._ttt_input_ids_from_batch(batch)
|
| cfg = self.ttt_config
|
| if input_ids.shape[1] <= cfg.crop_size:
|
| return batch
|
| high = input_ids.shape[1] - cfg.crop_size + 1
|
| start = int(
|
| torch.randint(
|
| high,
|
| (1,),
|
| generator=generator,
|
| device=input_ids.device,
|
| ).item()
|
| )
|
| end = start + cfg.crop_size
|
| if isinstance(batch, dict):
|
| cropped = {}
|
| for name, tensor in batch.items():
|
| if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]:
|
| cropped[name] = tensor[:, start:end]
|
| else:
|
| cropped[name] = tensor
|
| return cropped
|
| return input_ids[:, start:end]
|
|
|
| def _ttt_sample_batch(
|
| self,
|
| tokenized: torch.Tensor | dict[str, torch.Tensor],
|
| generator: torch.Generator,
|
| ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]:
|
| cfg = self.ttt_config
|
| batch = self._ttt_sample_crop(tokenized, generator)
|
| input_ids = self._ttt_input_ids_from_batch(batch)
|
| rows = torch.randint(
|
| input_ids.shape[0],
|
| (cfg.batch_size,),
|
| generator=generator,
|
| device=input_ids.device,
|
| )
|
| if isinstance(batch, dict):
|
| sampled: torch.Tensor | dict[str, torch.Tensor] = {}
|
| for name, tensor in batch.items():
|
| if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]:
|
| sampled[name] = tensor.index_select(0, rows)
|
| else:
|
| sampled[name] = tensor
|
| else:
|
| sampled = input_ids.index_select(0, rows)
|
|
|
| sampled_ids = self._ttt_input_ids_from_batch(sampled)
|
| labels = sampled_ids.clone()
|
| non_special = self._ttt_non_special_mask(sampled_ids)
|
| label_mask = torch.zeros_like(non_special)
|
| for row_idx in range(sampled_ids.shape[0]):
|
| candidate_positions = torch.where(non_special[row_idx])[0]
|
| if candidate_positions.numel() == 0:
|
| continue
|
| num_mask = max(1, int(round(candidate_positions.numel() * cfg.mask_ratio)))
|
| order = torch.randperm(
|
| candidate_positions.numel(),
|
| generator=generator,
|
| device=sampled_ids.device,
|
| )
|
| chosen = candidate_positions[order[:num_mask]]
|
| label_mask[row_idx, chosen] = True
|
| labels = labels.masked_fill(~label_mask, -100)
|
|
|
| masked_ids = sampled_ids.clone()
|
| chosen_positions = torch.where(label_mask)
|
| if chosen_positions[0].numel() > 0:
|
| random_values = torch.rand(
|
| chosen_positions[0].shape,
|
| generator=generator,
|
| device=sampled_ids.device,
|
| )
|
| leave = random_values < cfg.bert_leave_prob
|
| replace = (random_values >= cfg.bert_leave_prob) & (
|
| random_values < cfg.bert_leave_prob + cfg.bert_replace_prob
|
| )
|
| mask = ~(leave | replace)
|
| if mask.any():
|
| masked_ids[
|
| chosen_positions[0][mask],
|
| chosen_positions[1][mask],
|
| ] = self._ttt_mask_token()
|
| if replace.any():
|
| replacement_tokens = self._ttt_replacement_tokens(sampled_ids)
|
| replacement_idx = torch.randint(
|
| replacement_tokens.shape[0],
|
| (int(replace.sum().item()),),
|
| generator=generator,
|
| device=sampled_ids.device,
|
| )
|
| masked_ids[
|
| chosen_positions[0][replace],
|
| chosen_positions[1][replace],
|
| ] = replacement_tokens[replacement_idx]
|
|
|
| return self._ttt_set_input_ids(sampled, masked_ids), labels
|
|
|
| def ttt(
|
| self,
|
| seq: str | list[str] | None = None,
|
| input_ids: torch.Tensor | None = None,
|
| ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None,
|
| **kwargs: T.Any,
|
| ) -> dict[str, T.Any]:
|
| if ttt_config is not None:
|
| if "_ttt_initialized" in self.__dict__ and self._ttt_initialized:
|
| next_cfg = self.ttt_config.merged(ttt_config)
|
| assert next_cfg.lora_rank == self.ttt_config.lora_rank, (
|
| "Changing lora_rank after TTT initialization is not supported."
|
| )
|
| assert next_cfg.lora_alpha == self.ttt_config.lora_alpha, (
|
| "Changing lora_alpha after TTT initialization is not supported."
|
| )
|
| assert (
|
| next_cfg.lora_target_replace_module
|
| == self.ttt_config.lora_target_replace_module
|
| ), "Changing LoRA target class after TTT initialization is not supported."
|
| assert next_cfg.lora_target_modules == self.ttt_config.lora_target_modules, (
|
| "Changing LoRA target modules after TTT initialization is not supported."
|
| )
|
| self._ttt_cfg = next_cfg
|
| else:
|
| self.init_ttt(ttt_config)
|
|
|
| self._ttt_ensure_initialized()
|
| cfg = self.ttt_config
|
| if cfg.initial_state_reset:
|
| self.ttt_reset()
|
|
|
| device = next(self.parameters()).device
|
| tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs)
|
| tokenized = self._ttt_to_device(tokenized, device)
|
| generator_device = device if device.type == "cuda" else torch.device("cpu")
|
| generator = torch.Generator(device=generator_device)
|
| if cfg.seed is not None:
|
| generator.manual_seed(cfg.seed)
|
|
|
| module_modes = {module: module.training for module in self.modules()}
|
| requires_grad = {param: param.requires_grad for param in self.parameters()}
|
| losses: list[float] = []
|
| step_metrics: list[dict[str, T.Any]] = []
|
| best_state: list[dict[str, torch.Tensor]] | None = None
|
| best_metric: float | None = None
|
| best_step = 0
|
|
|
| try:
|
| self.train()
|
| for param in self.parameters():
|
| param.requires_grad_(False)
|
| for param in self._ttt_lora_parameters():
|
| param.requires_grad_(True)
|
|
|
| optimizer = self._ttt_make_optimizer()
|
| optimizer.zero_grad(set_to_none=True)
|
| total_micro_steps = cfg.steps * cfg.ags
|
| for micro_step in range(total_micro_steps):
|
| batch, labels = self._ttt_sample_batch(tokenized, generator)
|
| logits = self._ttt_predict_logits(batch, **kwargs)
|
| labels = labels.to(device=logits.device)
|
| loss = F.cross_entropy(
|
| logits.reshape(-1, logits.shape[-1]),
|
| labels.reshape(-1),
|
| ignore_index=-100,
|
| )
|
| (loss / cfg.ags).backward()
|
| if (micro_step + 1) % cfg.ags != 0:
|
| continue
|
|
|
| if cfg.gradient_clip:
|
| torch.nn.utils.clip_grad_norm_(
|
| self._ttt_lora_parameters(),
|
| cfg.gradient_clip_max_norm,
|
| )
|
| optimizer.step()
|
| optimizer.zero_grad(set_to_none=True)
|
| step = (micro_step + 1) // cfg.ags
|
| loss_value = float(loss.detach().item())
|
| losses.append(loss_value)
|
| if cfg.eval_each_step:
|
| metrics, metric = self._ttt_eval_step(
|
| step=step,
|
| loss=loss_value,
|
| seq=seq,
|
| input_ids=input_ids,
|
| **kwargs,
|
| )
|
| if len(metrics) > 0:
|
| step_metrics.append(metrics)
|
| if metric is not None and (
|
| best_metric is None or metric > best_metric
|
| ):
|
| best_metric = metric
|
| best_step = step
|
| best_state = self._ttt_snapshot_lora_state()
|
|
|
| if cfg.automatic_best_state_reset and best_state is not None:
|
| self._ttt_restore_lora_state(best_state)
|
| finally:
|
| for param, value in requires_grad.items():
|
| param.requires_grad_(value)
|
| for module, training in module_modes.items():
|
| module.train(training)
|
|
|
| return {
|
| "losses": losses,
|
| "step_metrics": step_metrics,
|
| "best_step": best_step,
|
| "best_metric": best_metric,
|
| }
|
|
|
