pisco-mistral / modelling_pisco.py
maxoul's picture
Upload COCOM
e423d32 verified
import warnings
import os
import torch
import gc
from torch import nn
from jinja2.exceptions import TemplateError
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, PreTrainedModel, PretrainedConfig, AutoModel, AutoConfig
from huggingface_hub import hf_hub_download
def get_first_layers_model(base_model_name: str, n_layers: int, attn_implementation: str = 'flash_attention_2'):
"""
Builds a model comprising only the n_layers first layer of the base_model_name
(it keeps the embedding and head layers)
"""
full_model = AutoModelForCausalLM.from_pretrained(base_model_name)
# Create a new config for a model with fewer layers (e.g., 3 layers)
custom_config = AutoConfig.from_pretrained(base_model_name)
custom_config.num_hidden_layers = n_layers
first_layers_model = AutoModelForCausalLM.from_config(config=custom_config, attn_implementation=attn_implementation, torch_dtype=torch.bfloat16)
# Load the state dict of the full model
full_state_dict = full_model.state_dict()
custom_state_dict = first_layers_model.state_dict()
kept_state_dict = {k:v for k,v in full_state_dict.items() if k in custom_state_dict}
first_layers_model.load_state_dict(kept_state_dict, strict=True)
del full_model
torch.cuda.empty_cache()
gc.collect()
return first_layers_model
def get_every_n_layer_model(base_model_name: str, every_n_layer: int, attn_implementation: str = 'flash_attention_2'):
"""
Builds a model comprising 1 every every_n_layer layer of the base_model_name
(it keeps the embedding and head layers)
"""
full_model = AutoModelForCausalLM.from_pretrained(base_model_name)
n_kept_layers = full_model.config.num_hidden_layers // every_n_layer
print(f'New model with 1/{every_n_layer} from {base_model_name} will have {n_kept_layers} layers')
custom_config = AutoConfig.from_pretrained(base_model_name)
custom_config.num_hidden_layers = n_kept_layers
custom_model = AutoModelForCausalLM.from_config(config=custom_config,
attn_implementation=attn_implementation,
torch_dtype=torch.bfloat16)
full_state_dict = full_model.state_dict()
custom_state_dict = custom_model.state_dict()
# Filter out every Nth layer and rename to form a new state dict
kept_state_dict = {}
for key, value in full_state_dict.items():
if ".layers." in key:
# Extract layer index
layer_idx = int(key.split(".layers.")[1].split(".")[0])
# Check if it's an Nth layer
if layer_idx % every_n_layer == 0:
# Adjust layer index for the smaller model
new_layer_idx = layer_idx // every_n_layer
# print('replacing', f".layers.{layer_idx}.", f".layers.{new_layer_idx}.")
new_key = key.replace(f".layers.{layer_idx}.", f".layers.{new_layer_idx}.")
if new_key in custom_state_dict:
kept_state_dict[new_key] = value
else:
# Keep non-layer-specific parameters
if key in custom_state_dict:
kept_state_dict[key] = value
# Load the filtered state dict into the custom model
custom_model.load_state_dict(kept_state_dict, strict=True)
del full_model
torch.cuda.empty_cache()
gc.collect()
return custom_model
class MistralTrimmed(torch.nn.Module):
"""
Trimmed version of base models for faster compression
NB: the name 'MistralTrimmed' suggests it just works with mistral but NO in fact most LLMs are supported !
"""
def __init__(self,
n_layers: int = 15,
every_n_layer: int = None,
rms_norm: bool = False,
base_model_name: str = 'mistralai/Mistral-7B-Instruct-v0.2',
attn_implementation: str = 'flash_attention_2'):
"""
you can either specify
- n_layers to some number: we take the n_layers first layers of the base model.
- every_n_layer to some number: in that case we take 1/N layer of the base model
The base_model_name is the name of the model from which this model is built.
"""
assert (n_layers is None) ^ (every_n_layer is None), 'Cannot specify both n_layers and every_n_layer for MistralTrimmed'
super().__init__()
self.n_layers = n_layers
self.every_n_layer = every_n_layer
self.base_model_name = base_model_name
if n_layers is not None:
self.custom_model = get_first_layers_model(self.base_model_name,
n_layers,
attn_implementation=attn_implementation)
else:
self.custom_model = get_every_n_layer_model(self.base_model_name,
every_n_layer,
attn_implementation=attn_implementation)
self.custom_model = self.custom_model.bfloat16()
self.custom_model.cuda()
if rms_norm:
print('Compressor keeps its original rms norm')
else:
print('De-activating RMS norm in compressor')
# We deactivate the norm: we don't need it here since we want to manipulate stuff within embed space
# see https://github.com/huggingface/transformers/blob/v4.45.0/src/transformers/models/mistral/modeling_mistral.py#L699
self.custom_model.model.norm = nn.Identity()
# Piping useful methods:
self.add_adapter = self.custom_model.add_adapter
self.set_adapter = self.custom_model.set_adapter
self.load_adapter = self.custom_model.load_adapter
self.num_parameters = self.custom_model.num_parameters
self.resize_token_embeddings = self.custom_model.resize_token_embeddings
self.get_input_embeddings = self.custom_model.get_input_embeddings
self.get_adapter_state_dict = self.custom_model.get_adapter_state_dict
# self.custom_model.gradient_checkpointing_enable()
# del self.custom_model.lm_head # THIS FAILS since some models have tie_embeddings=True !
# gc.collect()
# torch.cuda.empty_cache()
def forward(self, input_ids, attention_mask=None):
return self.custom_model.model(input_ids, attention_mask, output_hidden_states=True) # we call the .model attribute of the causal LM to avoid the cost of the LM head ! nice huh ?
def __call__(self, input_ids, attention_mask=None, output_hidden_states=True):
return self.forward(input_ids, attention_mask)
class AbstractCompressor(nn.Module):
def __init__(self, compr_model_name: str, compr_rate: int, decoder_hidden_size: int):
super().__init__()
self.compr_model_name = compr_model_name
self.compr_rate = compr_rate
self.decoder_hidden_size = decoder_hidden_size
def forward(self, input_ids, attention_mask, generation_top_k):
"""
input_ids of shape (batch_size, top_k, seq_length)
attention_mask of shape (batch_size, top_k, seq_length)
generation_top_k: the number of docs
"""
raise NotImplementedError
def save_pretrained(self, save_directory):
raise NotImplementedError
def load_pretrained(self, load_directory):
raise NotImplementedError
class BertCompressor(AbstractCompressor):
def __init__(self,
compr_model_name: str,
compr_rate: int,
decoder_hidden_size: int,
mlp_hidden_dim: int = 8192,
use_mlp: bool = True,
doc_max_length : int = 128,
**kwargs):
# TODO use the device_map
super().__init__(compr_model_name=compr_model_name, compr_rate=compr_rate, decoder_hidden_size=decoder_hidden_size)
if compr_model_name == 'mistral_trimmed':
assert 'compr_n_layers' in kwargs
self.model = MistralTrimmed(n_layers=kwargs['compr_n_layers'],
every_n_layer=kwargs['compr_every_n_layer'],
rms_norm=kwargs['compr_rms_norm'],
base_model_name=kwargs['compr_base_model_name'],
attn_implementation=kwargs['attn_implementation'])
self.tokenizer = AutoTokenizer.from_pretrained(self.model.base_model_name)
self.hidden_size = self.model.custom_model.config.hidden_size
else:
self.model = AutoModel.from_pretrained(compr_model_name, torch_dtype=torch.bfloat16, device_map='auto')
self.tokenizer = AutoTokenizer.from_pretrained(compr_model_name, use_fast=True)
self.tokenizer.padding_side = "left"
self.hidden_size = self.model.config.hidden_size
print('Base compressor nb parameters', self.model.num_parameters())
self.mlp_hidden_dim = mlp_hidden_dim
self.use_mlp = use_mlp
self.doc_max_length = doc_max_length
if self.use_mlp:
self.mlp = nn.Sequential(
nn.Linear(self.hidden_size, self.mlp_hidden_dim),
nn.ReLU(),
nn.Linear(self.mlp_hidden_dim, decoder_hidden_size)
).bfloat16()
self.mlp.cuda()
self.n_emb = self.doc_max_length // self.compr_rate
mem_tokens = ['<MEM' + str(i) + '>' for i in range(self.n_emb)]
self.tokenizer.add_special_tokens({'additional_special_tokens': mem_tokens})
self.tokenizer.mem_tokens = mem_tokens
self.tokenizer.mem_token_ids = [self.tokenizer.convert_tokens_to_ids(elt) for elt in self.tokenizer.mem_tokens]
self.tokenizer.mem_token_ids_pt = torch.LongTensor(self.tokenizer.mem_token_ids)
self.model.resize_token_embeddings(len(self.tokenizer))
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token_id = self.tokenizer.bos_token_id
if not use_mlp:
assert decoder_hidden_size == self.hidden_size, f'Mlp mandatory is hidden sizes not equal: {decoder_hidden_size} vs {self.hidden_size}'
self.lora = False
self.lora_name = 'compr_adapter'
def prepare_mem_tokens_optimization(self):
assert self.lora, 'should only be called with lora.'
self.model.get_input_embeddings().weight.requires_grad = True
# Applying a hook zero-ing the gradients except for the mem token:
def hook(grad):
mask = torch.zeros_like(grad)
mask[self.tokenizer.mem_token_ids] = 1.0
return grad * mask
self.model.get_input_embeddings().weight.register_hook(hook)
def set_lora(self, peft_config):
self.model.add_adapter(peft_config, self.lora_name)
self.model.set_adapter(self.lora_name)
self.lora = True
self.prepare_mem_tokens_optimization()
def forward(self, input_ids, attention_mask):
assert input_ids.size() == attention_mask.size()
assert len(input_ids.size()) == 2
batch_size_times_top_k = input_ids.size(0)
last_hidden_states = self.model(input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True).hidden_states[-1]
# Getting the hidden states at the mem token positions, as for regular cocom:
mask = torch.isin(input_ids, self.tokenizer.mem_token_ids_pt.to(input_ids.device))
selected_n_tokens = last_hidden_states[mask].reshape(last_hidden_states.size(0), -1, last_hidden_states.size(-1))
assert selected_n_tokens.size() == (batch_size_times_top_k, self.n_emb, self.hidden_size), f"{selected_n_tokens.size()} vs {(batch_size_times_top_k, self.n_emb, self.hidden_size)}"
if self.use_mlp:
selected_n_tokens = self.mlp(selected_n_tokens) # now of shape (batch_size, top_k, decoder_hidden_size)
assert selected_n_tokens.size() == (batch_size_times_top_k, self.n_emb, self.decoder_hidden_size), f"{selected_n_tokens.size()} vs {(batch_size_times_top_k, self.n_emb, self.decoder_hidden_size)}"
return selected_n_tokens
def get_lora_path_from_directory(self, directory):
return os.path.join(directory, 'compressor_adapters.pth')
def get_compressor_path_from_directory(self, directory):
return os.path.join(directory, 'compressor.pth')
def get_mlp_path_from_directory(self, directory):
return os.path.join(directory, 'mlp.pth')
def get_first_layer_path_from_directory(self, directory):
return os.path.join(directory, 'first_layer.pth')
def get_first_layer_state_dict(self) -> dict:
out = {}
for k, v in self.model.named_parameters():
if 'embed_tokens.weight' in k:
out[k] = v.cpu()
assert len(out) == 1, len(out) # We should get exactly one layer here
return out
def save_pretrained(self, save_directory):
"""
Here we just save mlp state_dict and model state_dict
Config is handled in cocom model.
"""
if not os.path.exists(save_directory):
os.makedirs(save_directory)
# Save MLP weights
if self.use_mlp:
mlp_path = self.get_mlp_path_from_directory(directory=save_directory)
torch.save(self.mlp.state_dict(), mlp_path)
# Saving the model
if not self.lora: # full training: save the full dict:
model_path = self.get_compressor_path_from_directory(directory=save_directory)
torch.save(self.model.state_dict(), model_path)
else: # lora training of the compressor
# We save the first layer:
first_layer_state_dict = self.get_first_layer_state_dict()
torch.save(first_layer_state_dict, self.get_first_layer_path_from_directory(directory=save_directory))
# We save the adapters:
adapter_state_dict = {k: v.cpu() for k, v in self.model.get_adapter_state_dict(self.lora_name).items()}
torch.save(adapter_state_dict, self.get_lora_path_from_directory(directory=save_directory))
def load_adapter(self, load_directory, peft_config):
assert peft_config is not None
map_location = torch.device("cpu") if not torch.cuda.is_available else None
adapter_state_dict = torch.load(self.get_lora_path_from_directory(directory=load_directory), map_location=map_location, weights_only=True)
print('loading compr adapter onto compressor model from', self.get_lora_path_from_directory(directory=load_directory))
self.model.load_adapter(peft_config=peft_config, adapter_name=self.lora_name, adapter_state_dict=adapter_state_dict)
self.lora = True
self.prepare_mem_tokens_optimization()
def load_first_layer(self, load_directory):
map_location = torch.device("cpu") if not torch.cuda.is_available else None
first_layer_state_dict = torch.load(self.get_first_layer_path_from_directory(load_directory), map_location=map_location, weights_only=True)
assert len(first_layer_state_dict.keys()) == 1
self.model.load_state_dict(first_layer_state_dict, strict=False)
def load_pretrained(self, load_directory, lora: bool = False, peft_config=None):
"""
Loading the state dicts.
:lora: if True then the compressor was trained using lora: we just need to load the adapters
if False, the compressor was fully trained: we load it fully.
"""
if self.use_mlp:
mlp_path = self.get_mlp_path_from_directory(directory=load_directory)
self.mlp.load_state_dict(torch.load(mlp_path, weights_only=True))
if lora:
self.load_first_layer(load_directory)
self.load_adapter(load_directory, peft_config)
else:
model_path = self.get_compressor_path_from_directory(directory=load_directory)
self.model.load_state_dict(torch.load(model_path, weights_only=True))
def prepare_inputs(self, texts, max_length, q_texts=None):
if q_texts is not None: # Query-dependent here:
assert len(texts) == len(q_texts), f"{len(texts)} == {len(q_texts)}"
if self.compr_model_name == 'mistral_trimmed':
# No special token, just formulating:
texts_to_encode = [ '\nQuery:\n' + query + 'Document:\n' + text for text, query in zip(texts, q_texts)]
inp_enc = self.tokenizer(texts_to_encode,
return_tensors='pt',
padding='max_length',
max_length=max_length + 8, # some margin for query/doc stuff + bos / eos
truncation=True,
add_special_tokens=True)
else:
inp_enc = self.tokenizer(q_texts, # we put the query in first position
texts,
return_tensors='pt',
padding='max_length',
max_length=max_length + 3,
truncation='only_second',
add_special_tokens=True)
else:
inp_enc = self.tokenizer(texts, return_tensors='pt', padding='max_length', max_length=max_length + 2, truncation=True, add_special_tokens=True)
inp_enc['input_ids'], inp_enc['attention_mask'] = add_memory_tokens_to_inputs(inp_enc['input_ids'],
inp_enc['attention_mask'],
self.n_emb,
tokenizer=self.tokenizer)
return inp_enc
def add_memory_tokens_to_inputs(input_ids: torch.Tensor, attention_mask: torch.Tensor, n_mem_tokens: int, tokenizer):
"""
Concatenate the input ids with n_mem_tokens mem_tokens and update the corresponding attention mask
"""
assert len(tokenizer.mem_tokens) == n_mem_tokens, f"{len(tokenizer.mem_tokens)} VS {n_mem_tokens}"
mem_tokens = torch.stack([tokenizer.mem_token_ids_pt] * input_ids.size(0), 0)
assert len(mem_tokens.size()) == 2
assert len(mem_tokens) == input_ids.size(0)
assert len(mem_tokens[0]) == n_mem_tokens
#mem_tokens = torch.full((input_ids.size(0), n_mem_tokens), tokenizer.mem_token_id, dtype=torch.long)
input_ids = torch.cat([input_ids, mem_tokens], dim=1)
attention_mask = torch.cat([attention_mask, torch.ones(input_ids.size(0), n_mem_tokens)], dim=1)
return input_ids, attention_mask
class COCOMConfig(PretrainedConfig):
model_type = "COCOM"
def __init__(self,
decoder_model_name: str = "meta-llama/Llama-2-7b-chat-hf",
doc_max_length: int = 128,
quantization: str = 'no',
sep: bool = False,
compr_model_name: str = "google-bert/bert-base-uncased",
compr_rate: int = 64,
compr_n_layers: int = None, # only for surgical mistral compressor
compr_every_n_layer: int = None,
compr_base_model_name: str = 'mistralai/Mistral-7B-Instruct-v0.2',
compr_rms_norm: bool = False, # only for surgical mistral compressor: if true, rms norm applied on h-s
compr_mlp_hidden_dim: int = 8096,
compr_use_mlp: bool = True,
lora: bool = False, # lora on decoder (and decoder as compr)
lora_compressor: bool = False, # lora only on the compressor if it exists
training_form: str = "both",
lora_r: int = 16,
lora_r_compressor: int = None,
load_adapters: bool = True,
kbtc_training: bool = False,
optimize_mem_tokens: bool = False,
different_mem_tokens: bool = False,
attn_implementation: str = 'flash_attention_2',
device_map = None,
**kwargs):
super().__init__(**kwargs)
self.decoder_model_name = decoder_model_name # model name of decoder
self.doc_max_length = doc_max_length # the maximum length of document that can be used by this model (it is used to compute number of mem tokens !)
self.quantization = quantization # quantization, could be no, int4, int8
self.sep = sep # boolean type, whether to use sep token
self.compr_model_name = compr_model_name # model name of compressor
self.compr_rate = compr_rate # compression rate
self.compr_use_mlp = compr_use_mlp
self.compr_mlp_hidden_dim = compr_mlp_hidden_dim
self.compr_n_layers = compr_n_layers
self.compr_every_n_layer = compr_every_n_layer
self.compr_base_model_name = compr_base_model_name
self.compr_rms_norm = compr_rms_norm
self.lora = lora # boolean type, whether to use lora trsining
self.lora_compressor = lora_compressor
self.training_form = training_form # training form, could be compressor: training only comprssor; both: training both
# Or both_separately: training both with separate adapters
self.lora_r = lora_r # lora_r for lora training, we use 16 throughout the experiment.
self.lora_r_compressor = lora_r_compressor or lora_r # defaulting to same lora as decoder.
self.load_adapters = load_adapters # used to load pretrained model: we first load without adapters, and then load them from file.
self.optimize_mem_tokens = optimize_mem_tokens
self.different_mem_tokens = different_mem_tokens
self.kbtc_training = kbtc_training
self.device_map = device_map
self.attn_implementation = attn_implementation
if training_form == 'compressor':
assert compr_model_name is not None and not self.lora
class COCOM(PreTrainedModel):
config_class = COCOMConfig
def __init__(self, cfg):
super().__init__(cfg)
self.decoder_model_name = cfg.decoder_model_name
self.decoder = self.create_decoder(cfg)
self.doc_max_length = cfg.doc_max_length
print('Base decoder nb parameters', self.decoder.num_parameters())
self.compr_model_name = cfg.compr_model_name
self.training_form = cfg.training_form
self.lora = cfg.lora
self.adapter_keys = []
self.compr = None
# when compr_model_name is not set, then means using a decoder-based compressor, otherwise a bert based compressor
if cfg.compr_model_name is not None:
# case bert based compressor
print('Instantiating compressor ', cfg.compr_model_name)
self.compr = BertCompressor(cfg.compr_model_name,
cfg.compr_rate,
doc_max_length=self.doc_max_length,
decoder_hidden_size=self.decoder.config.hidden_size,
mlp_hidden_dim=cfg.compr_mlp_hidden_dim,
compr_n_layers=cfg.compr_n_layers,
compr_every_n_layer=cfg.compr_every_n_layer,
compr_base_model_name=cfg.compr_base_model_name,
compr_rms_norm=cfg.compr_rms_norm,
use_mlp=cfg.compr_use_mlp,
attn_implementation=cfg.attn_implementation)
# set lora adaptors on decoder model
if cfg.lora:
peft_config = self.get_peft_config(lora_r=cfg.lora_r)
if cfg.load_adapters:
self.decoder.add_adapter(peft_config, 'decoder_adapter')
self.decoder.set_adapter('decoder_adapter') # active adapter by default
self.adapter_keys.append('decoder_adapter')
# Create separate adapters (if not BERT compressor and training_form == 'both_separately')
if self.training_form == 'both_separately' and self.compr is None:
if cfg.load_adapters:
self.decoder.add_adapter(peft_config, 'encoder_adapter')
self.adapter_keys.append('encoder_adapter')
# set lora adapters on compressor model:
if cfg.lora_compressor and self.compr is not None and cfg.load_adapters:
peft_config = self.get_peft_config(lora_r=cfg.lora_r_compressor)
self.compr.set_lora(peft_config)
self.decoder_tokenizer = COCOM.create_decoder_tokenizer(cfg)
# resize the tokenizer embedding
self.decoder.resize_token_embeddings(len(self.decoder_tokenizer))
self.decoder.generation_config.top_p = None
self.decoder.generation_config.temperature = None
self.decoder.generation_config.pad_token_id = self.decoder_tokenizer.pad_token_id
# self.decoder.gradient_checkpointing_enable()
# if self.compr is not None:
# self.compr.gradient_checkpointing_enable()
# other settings
self.generation_top_k = 1
self.sep = cfg.sep
self.compr_rate = cfg.compr_rate
self.local_rank = os.getenv('LOCAL_RANK', '0')
self.n_mem_tokens = self.doc_max_length // self.compr_rate # crucial!
if self.lora:
for adapter_key in self.adapter_keys:
self.decoder.set_adapter(adapter_key)
print(f'Adapter {adapter_key} trainable parameters: {self.num_parameters(only_trainable=True)}')
# We need to activate all adapters so that they are both trained...
self.set_all_adapters()
else:
print(f'Total trainable parameters: {self.num_parameters(only_trainable=True)}')
if self.compr is not None:
print(f'Compressor number of parameters: {self.compr.model.num_parameters(only_trainable=True)}')
self.prepare_mem_tokens_optimization()
def prepare_mem_tokens_optimization(self):
if self.config.optimize_mem_tokens:
if self.compr is None:
# Enforcing gradients for input embeddings (even if lora)
self.decoder.get_input_embeddings().weight.requires_grad = True
# Applying a hook zero-ing the gradients except for the mem token:
def hook(grad):
mask = torch.zeros_like(grad)
mask[self.decoder_tokenizer.mem_token_ids] = 1.0
return grad * mask
self.decoder.get_input_embeddings().weight.register_hook(hook)
def set_all_adapters(self):
if len(self.adapter_keys) > 0:
self.decoder.set_adapter(self.adapter_keys)
@staticmethod
def create_decoder_tokenizer(cfg: COCOMConfig):
decoder_tokenizer = AutoTokenizer.from_pretrained(cfg.decoder_model_name, use_fast=True, padding_side='left')
# define special tokens
n_mem_tokens = cfg.doc_max_length // cfg.compr_rate
if cfg.different_mem_tokens:
# estimation fo the number of memory tokens needed:
mem_tokens = ['<MEM' + str(i) + '>' for i in range(n_mem_tokens)]
decoder_tokenizer.add_special_tokens({'additional_special_tokens': mem_tokens + ['<AE>', '<ENC>', '<SEP>']})
decoder_tokenizer.mem_tokens = mem_tokens
else:
decoder_tokenizer.add_special_tokens({'additional_special_tokens': ['<MEM>', '<AE>', '<ENC>', '<SEP>']})
decoder_tokenizer.mem_tokens = ['<MEM>'] * n_mem_tokens
decoder_tokenizer.mem_token_ids = [decoder_tokenizer.convert_tokens_to_ids(elt) for elt in decoder_tokenizer.mem_tokens]
decoder_tokenizer.mem_token_ids_pt = torch.LongTensor(decoder_tokenizer.mem_token_ids) # required later on for operations on tensors
decoder_tokenizer.ae_token = '<AE>' # token for autoencoding on decoder side
decoder_tokenizer.ae_token_id = decoder_tokenizer.convert_tokens_to_ids('<AE>')
decoder_tokenizer.enc_token = '<ENC>' # token for autoencoding on compressor side
decoder_tokenizer.sep_token = '<SEP>' # sep token between document
decoder_tokenizer.sep_token_id = decoder_tokenizer.convert_tokens_to_ids('<SEP>')
# If kbtc training, we add another one yet
if cfg.kbtc_training:
decoder_tokenizer.add_special_tokens({'additional_special_tokens': ['<KBTC>']})
decoder_tokenizer.kbtc_token = '<KBTC>'
decoder_tokenizer.kbtc_token_id = decoder_tokenizer.convert_tokens_to_ids('<KBTC>')
# if pad token exists then use pad token, othrwise bos token
if decoder_tokenizer.pad_token_id is None:
decoder_tokenizer.pad_token_id = decoder_tokenizer.bos_token_id
return decoder_tokenizer
def get_peft_config(self, lora_r: int) -> LoraConfig:
"""
Builds the peft config
"""
return LoraConfig(task_type="CAUSAL_LM", r=lora_r, lora_alpha=2*lora_r, target_modules='all-linear', lora_dropout=0.1)
def create_decoder(self, cfg):
"""
Loads the base decoder.
"""
if torch.cuda.is_available():
if cfg.quantization == "no":
return AutoModelForCausalLM.from_pretrained(
cfg.decoder_model_name,
torch_dtype=torch.bfloat16,
attn_implementation=self.config.attn_implementation,
# low_cpu_mem_usage = True,
device_map=cfg.device_map
)
elif cfg.quantization == "int4":
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type='nf4',
bnb_4bit_compute_dtype='bfloat16',
# low_cpu_mem_usage = True,
)
return AutoModelForCausalLM.from_pretrained(
cfg.decoder_model_name,
quantization_config=quant_config,
attn_implementation=self.config.attn_implementation,
torch_dtype=torch.bfloat16,
resume_download=True,
# low_cpu_mem_usage = True,
trust_remote_code=True,
device_map=cfg.device_map
)
elif cfg.quantization == "int8":
quant_config = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_enable_fp32_cpu_offload=True,
bnb_4bit_compute_dtype='bfloat16',
# low_cpu_mem_usage = True,
)
return AutoModelForCausalLM.from_pretrained(
cfg.decoder_model_name,
quantization_config=quant_config,
attn_implementation=self.config.attn_implementation,
torch_dtype=torch.bfloat16,
resume_download=True,
# low_cpu_mem_usage = True,
trust_remote_code=True,
device_map=cfg.device_map
)
else:
raise NotImplementedError()
else:
return AutoModelForCausalLM.from_pretrained(
cfg.decoder_model_name,
torch_dtype=torch.bfloat16,
resume_download=True,
# low_cpu_mem_usage = True,
trust_remote_code=True,
device_map=cfg.device_map
)
def compress(self, enc_input_ids, enc_attention_mask):
if self.compr:
return self.compr(enc_input_ids, enc_attention_mask)
else:
return self.compr_decoder(enc_input_ids, enc_attention_mask)
def replace_emb(self, compressed_embs, dec_input_ids):
"""
Compression logic (either with decoder or with dedicated compressor)
"""
indices = range(0, compressed_embs.size(0) + 1, self.generation_top_k)
input_embeds = self.replace_embeddings(compressed_embs, dec_input_ids, indices)
return input_embeds
def compr_decoder(self, input_ids, attention_mask):
"""
Compression using the decoder
"""
assert input_ids.size() == attention_mask.size(), f"{input_ids.size()} vs {attention_mask.size()}"
# Switch adapter if we are training two different ones:
if 'encoder_adapter' in self.adapter_keys:
self.decoder.set_adapter('encoder_adapter')
emb = self.decoder(input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True).hidden_states[-1]
mask = torch.isin(input_ids, self.decoder_tokenizer.mem_token_ids_pt.to(input_ids.device))
return emb[mask].reshape(emb.size(0), -1, emb.size(-1))
def prepare_encoder_inputs_to_decoder(self, texts, max_length, q_texts=None):
if q_texts is not None:
texts_to_encode = [self.decoder_tokenizer.enc_token + self.decoder_tokenizer.bos_token + '\nQuery:\n' + query + 'Document:\n' + text + self.decoder_tokenizer.eos_token
for text, query in zip(texts, q_texts)]
inp_enc = self.decoder_tokenizer(texts_to_encode, return_tensors='pt', padding='max_length', max_length=max_length + 8, truncation=True, add_special_tokens=False)
else:
inp_enc = [self.decoder_tokenizer.enc_token + self.decoder_tokenizer.bos_token + text + self.decoder_tokenizer.eos_token for text in texts]
inp_enc = self.decoder_tokenizer(inp_enc, return_tensors='pt', padding="max_length", max_length=max_length+3, truncation=True, add_special_tokens=False)
num_mem_tokens = self.doc_max_length // self.compr_rate
assert num_mem_tokens == len(self.decoder_tokenizer.mem_tokens)
inp_enc['input_ids'], inp_enc['attention_mask'] = add_memory_tokens_to_inputs(inp_enc['input_ids'],
inp_enc['attention_mask'],
num_mem_tokens,
tokenizer=self.decoder_tokenizer)
return inp_enc
def prepare_encoder_inputs(self, texts: list[str], max_length: int, q_texts: list[str] = None):
"""
Create the inputs to the encoder, for compression.
"""
if q_texts is not None:
assert len(texts) == len(q_texts), f"{len(texts)} == {len(q_texts)}"
# Case where the encoder is the decoder with adapter:
if self.compr is None:
return self.prepare_encoder_inputs_to_decoder(texts, max_length, q_texts)
# Case where the encoder is a separate network:
else:
return self.compr.prepare_inputs(texts, max_length, q_texts)
def replace_embeddings(self, compressed_embs, dec_input_ids, indices):
"""
Replace memory tokens in the decoder input to with the compressed embeddings
"""
inputs_embeds = self.decoder.get_input_embeddings()(dec_input_ids)
num_embs = compressed_embs.size(1)
if self.sep:
slot_len = num_embs + 1
else:
slot_len = num_embs
# get first mem_token indices
first_mem_token_indices = torch.argmax((dec_input_ids == self.decoder_tokenizer.mem_token_ids[0]).int(), dim=1)
batch_size = inputs_embeds.size(0)
# for each example in batch, replace them with compressed embeddings
for i in range(batch_size):
for j in range(indices[i], indices[i + 1]):
start_idx = first_mem_token_indices[i].item() + (j-indices[i]) * slot_len
assert inputs_embeds[i, start_idx:start_idx + num_embs, :].size() == compressed_embs[j].size(), \
f"{inputs_embeds[i, start_idx:start_idx + num_embs, :].size()} VS {compressed_embs[j].size()}"
inputs_embeds[i, start_idx:start_idx + num_embs, :] = compressed_embs[j]
return inputs_embeds
def forward(self,
enc_input_ids: torch.LongTensor = None,
enc_attention_mask: torch.LongTensor = None,
dec_input_ids: torch.LongTensor = None,
dec_attention_mask: torch.LongTensor = None,
labels: torch.LongTensor = None):
"""
enc_input_ids: stores the contexts, should be flattened from all queries before input, can be of shape:
- (batch_size*generation_top_k, enc_token_length)
- (batch_size, generation_top_k, enc_token_length)
enc_attention_mask: attention mask of enc_input_ids, same shape as enc_input_ids
dec_input_ids: stores the prompts (including mem tokens), dimention (batch_size, dec_token_length)
dec_attention_mask: attention mask of dec_input_ids
"""
assert enc_input_ids.size() == enc_attention_mask.size(), f"{enc_input_ids.size()} vs {enc_attention_mask.size()}"
if len(enc_input_ids.size()) == 3: # likely from bergen: we just flatten all of this to perform encoding in one batch
batch_size, top_k, seq_length = enc_input_ids.size()
enc_input_ids = enc_input_ids.view(batch_size * top_k, seq_length)
enc_attention_mask = enc_attention_mask.view(batch_size * top_k, seq_length)
# Here, we should have top_k times more elements in enc_input_ids than in dec_input_ids
assert enc_input_ids.size(0) == dec_input_ids.size(0) * self.generation_top_k, \
f"{enc_input_ids.size(0)} VS {dec_input_ids.size(0)} with generation_top_k={self.generation_top_k}"
# Perform compression with gradient tracking
compressed_embs = self.compress(enc_input_ids, enc_attention_mask)
inputs_embeds = self.replace_emb(compressed_embs, dec_input_ids)
# if training_form is compressor, then detach the inputs_embeds, to make gradient not count in decoder
if (self.training_form == "compressor") and (self.compr is None):
inputs_embeds = inputs_embeds.detach()
# decoding
if 'decoder_adapter' in self.adapter_keys:
self.decoder.set_adapter('decoder_adapter')
decoder_outputs = self.decoder(inputs_embeds=inputs_embeds, attention_mask=dec_attention_mask, labels=labels)
# At end of forward, we need to activate all adapters so that they are both trained...
self.set_all_adapters()
return {"loss": decoder_outputs.loss, "logits": decoder_outputs.logits}
def generate(self, model_input, max_new_tokens=128, return_doc_embeddings: bool = False):
enc_input_ids, enc_attention_mask, dec_input_ids, dec_attention_mask = model_input['enc_input_ids'], model_input['enc_attention_mask'], model_input['dec_input_ids'], model_input['dec_attention_mask']
assert enc_input_ids.size() == enc_attention_mask.size()
if len(enc_input_ids.size()) == 3: # likely from bergen: we just flatten all of this to perform encoding in one batch
batch_size, top_k, seq_length = enc_input_ids.size()
enc_input_ids = enc_input_ids.view(batch_size * top_k, seq_length)
enc_attention_mask = enc_attention_mask.view(batch_size * top_k, seq_length)
# Here, we should have top_k times more elements in enc_input_ids than in dec_input_ids
assert enc_input_ids.size(0) == dec_input_ids.size(0) * self.generation_top_k, \
f"{enc_input_ids.size(0)} VS {dec_input_ids.size(0)} with generation_top_k={self.generation_top_k}"
compressed_embs = self.compress(enc_input_ids.to('cuda'), enc_attention_mask.to('cuda'))
inputs_embeds = self.replace_emb(compressed_embs, dec_input_ids.to('cuda'))
# Switch adapter if we are training two different ones:
if 'decoder_adapter' in self.adapter_keys:
self.decoder.set_adapter('decoder_adapter')
output_ids = self.decoder.generate(
inputs_embeds=inputs_embeds.to("cuda"),
attention_mask=dec_attention_mask.to("cuda"),
do_sample=False,
top_p=None,
max_new_tokens=max_new_tokens
)
decoded = self.decoder_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
if return_doc_embeddings:
# Compressed_embds is of shape (batch_size*top_k, n_mem_tokens, hidden_dim)
# We reshape to batch_size, top_k, n_mem_tokens, hidden_dim
assert batch_size is not None
assert top_k is not None
compressed_embs = compressed_embs.view(batch_size, top_k, compressed_embs.size(1), compressed_embs.size(2))
return decoded, compressed_embs
else:
return decoded
def get_all_adapters_state_dict(self):
"""
Return the state dicts of the adapters
Used for saving so we go to cpu automatically
"""
return {key: {k:v.cpu() for k, v in self.decoder.get_adapter_state_dict(key).items()} for key in self.adapter_keys}
def load_adapter_from_state_dict(self, peft_config: LoraConfig, adapter_name: str, adapter_state_dict: dict) -> None:
"""
Creates an adapter from the state dict (used to load from pretrained)
"""
# assert adapter_name not in self.adapter_keys, f'Adapter {adapter_name} already exists'
print(f'loading adapter {adapter_name}')
self.decoder.load_adapter(peft_config=peft_config, adapter_name=adapter_name, adapter_state_dict=adapter_state_dict)
self.adapter_keys.append(adapter_name)
def get_decoder_first_and_last_layer_state_dict(self) -> dict:
"""
Just getting the first and last layers: the only ones which change when adding tokens
Used to save the model so we automatically move to cpu.
"""
out = {}
for k, v in self.decoder.named_parameters():
if 'lm_head.weight' in k or 'embed_tokens.weight' in k:
out[k] = v.cpu()
# assert len(out) == 2, len(out) # We should get both the embedding layer and the head layer.
return out
def save_pretrained(self, save_directory: str, **kwargs):
"""
Save only the LoRA adapters and their configurations.
"""
if self.lora:
if not os.path.exists(save_directory):
os.makedirs(save_directory)
# Save the LoRA adapter weights
torch.save(self.get_all_adapters_state_dict(), os.path.join(save_directory, "adapters.pth"))
# Save the first and last layers of decoder (because of diffs with tokens !)
torch.save(self.get_decoder_first_and_last_layer_state_dict(), os.path.join(save_directory, "decoder_first_last_layers.pth"))
# Save the bert compressor if it exists
if self.compr_model_name is not None:
self.compr.save_pretrained(os.path.join(save_directory, 'compressor'))
# Save the configuration
self.config.save_pretrained(save_directory)
else:
super().save_pretrained(save_directory, **kwargs)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
"""
Loading: to take care of checkpoints containing only lora and not base model.
"""
# Load the configuration
config = COCOMConfig.from_pretrained(pretrained_model_name_or_path)
config.attn_implementation = kwargs.get('attn_implementation', config.attn_implementation)
map_location = torch.device("cpu") if not torch.cuda.is_available() else None
if config.lora:
# We need to delay the construction of the adapters (otherwise peft complains)
config.load_adapters = False
if 'device_map' in kwargs:
config.device_map = kwargs['device_map']
# Initialize the model
model = cls(config)
# Loading first and last layers (they might have changed due to extra tokens)
try:
# If loading from Hugging Face Hub
first_and_last_layers_path = hf_hub_download(
repo_id=pretrained_model_name_or_path,
filename="decoder_first_last_layers.pth"
)
except Exception as e:
# If loading from a local directory
first_and_last_layers_path = os.path.join(pretrained_model_name_or_path, "decoder_first_last_layers.pth")
if os.path.exists(first_and_last_layers_path):
first_and_last_decoder_state_dict = torch.load(first_and_last_layers_path, map_location=map_location, weights_only=True)
for key in first_and_last_decoder_state_dict:
assert key in model.decoder.state_dict()
model.decoder.load_state_dict(first_and_last_decoder_state_dict, strict=False)
else:
print('FIRST AND LAST LAYER NOT FOUND (ok for some old models):', first_and_last_layers_path)
peft_config = model.get_peft_config(lora_r=config.lora_r)
# Load the LoRA adapters (if the file exists)
try:
# If loading from Hugging Face Hub
adapters_path = hf_hub_download(
repo_id=pretrained_model_name_or_path,
filename="adapters.pth"
)
except Exception as e:
# If loading from a local directory
adapters_path = os.path.join(pretrained_model_name_or_path, "adapters.pth")
if os.path.exists(adapters_path):
adapters_state_dict = torch.load(adapters_path, map_location=map_location, weights_only=True)
for key, val in adapters_state_dict.items():
model.load_adapter_from_state_dict(peft_config=peft_config, adapter_name=key, adapter_state_dict=val)
else:
warnings.warn(f'I see lora on that PISCO model, but {adapters_path} does not exist, it may be normal \
for recent versions of transformers, be aware.')
# If there is a compressor, it's been built: we just need to load the state dict or the adapters:
if config.compr_model_name is not None:
model.compr.load_pretrained(os.path.join(pretrained_model_name_or_path, 'compressor'),
lora=config.lora_compressor,
peft_config=model.get_peft_config(lora_r=config.lora_r_compressor))
model.set_all_adapters()
model.config.load_adapters = True
return model
else:
return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
def generate_from_text(self, questions: list[str], documents: list[list[str]], max_new_tokens: int = 128) -> list[str]:
"""
Generates answers from documents (via compression then decoding)
questions: list of string
documents: list of list of strings (they should all be of equal length: the nb of doc for each question)
"""
self.generation_top_k = len(documents[0])
assert len(documents) == len(questions)
assert all([len(context) == len(documents[0]) for context in documents])
flat_documents = sum(documents, [])
model_input = {}
# Creating encoder inputs:
input_encoder = self.prepare_encoder_inputs(flat_documents, max_length=128)
device = self.decoder.device
model_input['enc_input_ids'], model_input['enc_attention_mask'] = input_encoder['input_ids'].to(device), input_encoder['attention_mask'].to(device)
# Creating decoder inputs
instr = [self.blend_prompt_and_memory_tokens(query=q) for q in questions]
inp_dec = self.decoder_tokenizer(instr, return_tensors='pt', padding="longest", add_special_tokens=False, truncation=True, max_length=2048)
model_input['dec_input_ids'], model_input['dec_attention_mask'] = inp_dec['input_ids'].to(device), inp_dec['attention_mask'].to(device)
# Generation
return self.generate(model_input, max_new_tokens=max_new_tokens)
def generate_from_compressed_documents_and_questions(self, questions: list[str], compressed_documents: torch.Tensor, max_new_tokens: int = 128) -> list[str]:
"""
Generates answers from compressed documents
questions: list of string
compressed_documents: torch tensor, its first dimension should be a multiple of len(questions)
"""
self.generation_top_k = compressed_documents.size(0) // len(questions)
assert compressed_documents.size(0) % self.generation_top_k == 0, f"{compressed_documents.size(0)} {self.generation_top_k}"
# Creating decoder inputs
instr = [self.blend_prompt_and_memory_tokens(query=q) for q in questions]
inp_dec = self.decoder_tokenizer(instr, return_tensors='pt', padding="longest", add_special_tokens=False, truncation=True, max_length=2048)
device = self.decoder.device
dec_input_ids, dec_attention_mask = inp_dec['input_ids'].to(device), inp_dec['attention_mask'].to(device)
# Creating input decoder embeddings from prompt + compressed documents
inputs_embeds = self.replace_emb(compressed_documents, dec_input_ids)
# Activating decoder generator:
if 'decoder_adapter' in self.adapter_keys:
self.decoder.set_adapter('decoder_adapter')
output_ids = self.decoder.generate(
inputs_embeds=inputs_embeds,
attention_mask=dec_attention_mask,
generation_config=self.generation_config,
max_new_tokens=max_new_tokens
)
# de-tokenizing
return self.decoder_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
def compress_documents(self, documents: list[str]) -> torch.Tensor:
"""
Compress a list of documents
"""
input_encoder = self.prepare_encoder_inputs(documents, max_length=128)
enc_input_ids = input_encoder['input_ids'].to(self.decoder.device)
attention_mask = input_encoder['attention_mask'].to(self.decoder.device)
return self.compress(enc_input_ids=enc_input_ids, enc_attention_mask=attention_mask)
def blend_prompt_and_memory_tokens(self, query: str):
"""
Takes care of blending the prompt with the memory tokens:
Also returns, if a label is provided, the position of the first token index of the label (for loss comp later on)
(Used for the HUB version)
"""
mem_tokens_str = ''.join(self.decoder_tokenizer.mem_tokens) + self.decoder_tokenizer.sep_token
# proper names for "eval" call, don't remove these lines
docs = mem_tokens_str * self.generation_top_k
question = query
prompt_system = 'You are a helpful assistant. Your task is to extract relevant information from provided documents and to answer to questions as briefly as possible.'
prompt_user = f"Background:\n{docs}\n\nQuestion:{question}"
# Prepare the messages with system and user roles
messages = [
{"role": "system", "content": prompt_system},
{"role": "user", "content": prompt_user.replace(':\ ', ': ')}
]
# Attempt to apply the system role and catch if it's not supported
try:
prompt = self.decoder_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
except TemplateError as e:
# Catch the error related to system role and handle it (e.g. gemma)
if "System role not supported" in str(e):
# Remove system role and proceed with only the user role
messages = [{"role": "user", "content": messages[0]['content'] + '\n' + messages[1]['content']}]
# Apply template again without system role
prompt = self.decoder_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
else:
# Re-raise the exception if it's unrelated to system role
raise e
return prompt
if __name__ == '__main__':
cfg = COCOMConfig(decoder_model_name='mistralai/Mistral-7B-Instruct-v0.2',
compr_model_name = "mistral_trimmed",
compr_rate = 64,
compr_n_layers = 5,
compr_mlp_hidden_dim = 8096,
compr_use_mlp = False,
lora = True, # lora on decoder (and decoder as compr)
lora_compressor = True, # lora only on the compressor if it exists
training_form = "both",
load_adapters = True,
kbtc_training = False,
optimize_mem_tokens = True,
different_mem_tokens = True,
attn_implementation = 'flash_attention_2')
cocom = COCOM(cfg)
cocom.save_pretrained('test_ckpt')
del cocom
torch.cuda.empty_cache()
import gc
gc.collect()
cocom = COCOM.from_pretrained('test_ckpt')