AlignX / models /task_feature_matrix.py

Gautam Kashyap

Upload AlignX llama2_7b checkpoints and model card

00dd625 verified 4 months ago

5.62 kB

"""Task-Feature Matrix computation for AlignX Stage 1.
T_a = W1 @ compress(Δθ_a) + W2 @ F_a
Δθ_a: task vector (compressed LoRA delta to latent_dim)
F_a:  alignment feature vector (mean hidden state, size hidden_dim)
T_a:  k-dim latent alignment representation
"""

import torch
import torch.nn as nn
from typing import Dict, Optional


def compress_task_vector(
    delta_state_dict: Dict[str, torch.Tensor],
    target_dim: int = 4096,
) -> torch.Tensor:
    """Flatten and mean-pool all parameter deltas to a fixed-size vector.
    LoRA deltas are small (~4M params for rank-8) so we can flatten them
    and chunk-average down to `target_dim` to make W1 tractable.
    """
    all_params = []
    for v in delta_state_dict.values():
        if v.dtype in (torch.float32, torch.float16, torch.bfloat16):
            all_params.append(v.detach().cpu().float().flatten())

    if not all_params:
        return torch.zeros(target_dim)

    flat = torch.cat(all_params)
    total = flat.shape[0]

    if total >= target_dim:
        chunk = total // target_dim
        remainder = total - chunk * target_dim
        trimmed = flat[:chunk * target_dim]
        compressed = trimmed.view(target_dim, chunk).mean(dim=1)
    else:
        # Pad with zeros
        compressed = torch.zeros(target_dim)
        compressed[:total] = flat

    return compressed.float()


def extract_alignment_features(
    model,
    dataset,
    layer_idx: int = -1,
    device: str = "cuda",
    max_samples: int = 200,
    batch_size: int = 4,
) -> torch.Tensor:
    """Extract mean hidden-state vector F_a at transformer layer `layer_idx`.
    Runs a forward pass over `dataset` samples, collecting the hidden state
    from the chosen layer, then returns the mean across all samples and tokens.
    """
    import torch
    from torch.utils.data import DataLoader

    model.eval()
    # Only call .to(device) if the model isn't already dispatched across multiple devices
    # (dispatch_model sets model.hf_device_map; calling .to() on a dispatched model
    # would un-dispatch it and may OOM a single GPU).
    if not getattr(model, "hf_device_map", None):
        model.to(device)

    hidden_dim = model.config.hidden_size
    feature_accum = torch.zeros(hidden_dim, device=device)
    count = 0

    # Identify the target layer
    layers = _get_transformer_layers(model)
    target_layer = layers[layer_idx]

    captured = []

    def _hook(module, input, output):
        # output can be tuple (hidden_state, ...) or just tensor
        hs = output[0] if isinstance(output, tuple) else output
        # hs: (batch, seq_len, hidden_dim)
        captured.append(hs.detach().mean(dim=(0, 1)).cpu())  # (hidden_dim,)

    handle = target_layer.register_forward_hook(_hook)

    # When the model is dispatched across GPUs (hf_device_map set), inputs must go
    # to cuda:0 (where the embedding layer lives), not to the `device` argument.
    input_device = "cuda:0" if getattr(model, "hf_device_map", None) else device

    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
    with torch.no_grad():
        for i, batch in enumerate(loader):
            if count >= max_samples:
                break
            input_ids = batch["input_ids"].to(input_device)
            attention_mask = batch["attention_mask"].to(input_device)
            try:
                model(input_ids=input_ids, attention_mask=attention_mask)
            except Exception:
                pass
            count += input_ids.shape[0]

    handle.remove()

    if captured:
        F_a = torch.stack(captured).mean(0)  # (hidden_dim,)
    else:
        F_a = torch.zeros(hidden_dim)

    print(f"[FeatureExtraction] Extracted F_a of shape {F_a.shape} from {count} samples")
    return F_a.float()


def _get_transformer_layers(model):
    """Return the list of transformer layers for common model families."""
    # Handles LLaMA, Mistral, Gemma, DeepSeek
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    raise ValueError(f"Cannot locate transformer layers for {type(model)}")


class TaskFeatureMatrix(nn.Module):
    """T_a = W1 @ Δθ_a + W2 @ F_a
    Maps the compressed task vector and alignment feature vector
    into a shared k-dimensional latent alignment space.
    """

    def __init__(self, k: int = 256, compress_dim: int = 4096, hidden_dim: int = 4096):
        super().__init__()
        self.k = k
        self.compress_dim = compress_dim
        self.hidden_dim = hidden_dim
        # W1: (k, compress_dim)  maps compressed delta-theta to latent space
        self.W1 = nn.Linear(compress_dim, k, bias=False)
        # W2: (k, hidden_dim)    maps feature vector to latent space
        self.W2 = nn.Linear(hidden_dim, k, bias=False)
        nn.init.xavier_uniform_(self.W1.weight)
        nn.init.xavier_uniform_(self.W2.weight)

    def forward(
        self,
        delta_theta_compressed: torch.Tensor,
        F_a: torch.Tensor,
    ) -> torch.Tensor:
        """
        delta_theta_compressed: (compress_dim,) — compressed task vector
        F_a: (hidden_dim,) — alignment feature vector
        Returns T_a: (k,)
        """
        T_a = self.W1(delta_theta_compressed) + self.W2(F_a)
        return T_a

    def compute_and_cache(
        self,
        delta_theta_compressed: torch.Tensor,
        F_a: torch.Tensor,
    ) -> torch.Tensor:
        with torch.no_grad():
            return self.forward(delta_theta_compressed, F_a)

	"""Task-Feature Matrix computation for AlignX Stage 1.

	T_a = W1 @ compress(Δθ_a) + W2 @ F_a

	Δθ_a: task vector (compressed LoRA delta to latent_dim)
	F_a: alignment feature vector (mean hidden state, size hidden_dim)
	T_a: k-dim latent alignment representation
	"""

	import torch
	import torch.nn as nn
	from typing import Dict, Optional


	def compress_task_vector(
	delta_state_dict: Dict[str, torch.Tensor],
	target_dim: int = 4096,
	) -> torch.Tensor:
	"""Flatten and mean-pool all parameter deltas to a fixed-size vector.

	LoRA deltas are small (~4M params for rank-8) so we can flatten them
	and chunk-average down to `target_dim` to make W1 tractable.
	"""
	all_params = []
	for v in delta_state_dict.values():
	if v.dtype in (torch.float32, torch.float16, torch.bfloat16):
	all_params.append(v.detach().cpu().float().flatten())

	if not all_params:
	return torch.zeros(target_dim)

	flat = torch.cat(all_params)
	total = flat.shape[0]

	if total >= target_dim:
	chunk = total // target_dim
	remainder = total - chunk * target_dim
	trimmed = flat[:chunk * target_dim]
	compressed = trimmed.view(target_dim, chunk).mean(dim=1)
	else:
	# Pad with zeros
	compressed = torch.zeros(target_dim)
	compressed[:total] = flat

	return compressed.float()


	def extract_alignment_features(
	model,
	dataset,
	layer_idx: int = -1,
	device: str = "cuda",
	max_samples: int = 200,
	batch_size: int = 4,
	) -> torch.Tensor:
	"""Extract mean hidden-state vector F_a at transformer layer `layer_idx`.

	Runs a forward pass over `dataset` samples, collecting the hidden state
	from the chosen layer, then returns the mean across all samples and tokens.
	"""
	import torch
	from torch.utils.data import DataLoader

	model.eval()
	# Only call .to(device) if the model isn't already dispatched across multiple devices
	# (dispatch_model sets model.hf_device_map; calling .to() on a dispatched model
	# would un-dispatch it and may OOM a single GPU).
	if not getattr(model, "hf_device_map", None):
	model.to(device)

	hidden_dim = model.config.hidden_size
	feature_accum = torch.zeros(hidden_dim, device=device)
	count = 0

	# Identify the target layer
	layers = _get_transformer_layers(model)
	target_layer = layers[layer_idx]

	captured = []

	def _hook(module, input, output):
	# output can be tuple (hidden_state, ...) or just tensor
	hs = output[0] if isinstance(output, tuple) else output
	# hs: (batch, seq_len, hidden_dim)
	captured.append(hs.detach().mean(dim=(0, 1)).cpu()) # (hidden_dim,)

	handle = target_layer.register_forward_hook(_hook)

	# When the model is dispatched across GPUs (hf_device_map set), inputs must go
	# to cuda:0 (where the embedding layer lives), not to the `device` argument.
	input_device = "cuda:0" if getattr(model, "hf_device_map", None) else device

	loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
	with torch.no_grad():
	for i, batch in enumerate(loader):
	if count >= max_samples:
	break
	input_ids = batch["input_ids"].to(input_device)
	attention_mask = batch["attention_mask"].to(input_device)
	try:
	model(input_ids=input_ids, attention_mask=attention_mask)
	except Exception:
	pass
	count += input_ids.shape[0]

	handle.remove()

	if captured:
	F_a = torch.stack(captured).mean(0) # (hidden_dim,)
	else:
	F_a = torch.zeros(hidden_dim)

	print(f"[FeatureExtraction] Extracted F_a of shape {F_a.shape} from {count} samples")
	return F_a.float()


	def _get_transformer_layers(model):
	"""Return the list of transformer layers for common model families."""
	# Handles LLaMA, Mistral, Gemma, DeepSeek
	if hasattr(model, "model") and hasattr(model.model, "layers"):
	return model.model.layers
	if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
	return model.transformer.h
	raise ValueError(f"Cannot locate transformer layers for {type(model)}")


	class TaskFeatureMatrix(nn.Module):
	"""T_a = W1 @ Δθ_a + W2 @ F_a

	Maps the compressed task vector and alignment feature vector
	into a shared k-dimensional latent alignment space.
	"""

	def __init__(self, k: int = 256, compress_dim: int = 4096, hidden_dim: int = 4096):
	super().__init__()
	self.k = k
	self.compress_dim = compress_dim
	self.hidden_dim = hidden_dim
	# W1: (k, compress_dim) maps compressed delta-theta to latent space
	self.W1 = nn.Linear(compress_dim, k, bias=False)
	# W2: (k, hidden_dim) maps feature vector to latent space
	self.W2 = nn.Linear(hidden_dim, k, bias=False)
	nn.init.xavier_uniform_(self.W1.weight)
	nn.init.xavier_uniform_(self.W2.weight)

	def forward(
	self,
	delta_theta_compressed: torch.Tensor,
	F_a: torch.Tensor,
	) -> torch.Tensor:
	"""
	delta_theta_compressed: (compress_dim,) — compressed task vector
	F_a: (hidden_dim,) — alignment feature vector
	Returns T_a: (k,)
	"""
	T_a = self.W1(delta_theta_compressed) + self.W2(F_a)
	return T_a

	def compute_and_cache(
	self,
	delta_theta_compressed: torch.Tensor,
	F_a: torch.Tensor,
	) -> torch.Tensor:
	with torch.no_grad():
	return self.forward(delta_theta_compressed, F_a)