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 = ['' 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 = ['' for i in range(n_mem_tokens)] decoder_tokenizer.add_special_tokens({'additional_special_tokens': mem_tokens + ['', '', '']}) decoder_tokenizer.mem_tokens = mem_tokens else: decoder_tokenizer.add_special_tokens({'additional_special_tokens': ['', '', '', '']}) decoder_tokenizer.mem_tokens = [''] * 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 = '' # token for autoencoding on decoder side decoder_tokenizer.ae_token_id = decoder_tokenizer.convert_tokens_to_ids('') decoder_tokenizer.enc_token = '' # token for autoencoding on compressor side decoder_tokenizer.sep_token = '' # sep token between document decoder_tokenizer.sep_token_id = decoder_tokenizer.convert_tokens_to_ids('') # If kbtc training, we add another one yet if cfg.kbtc_training: decoder_tokenizer.add_special_tokens({'additional_special_tokens': ['']}) decoder_tokenizer.kbtc_token = '' decoder_tokenizer.kbtc_token_id = decoder_tokenizer.convert_tokens_to_ids('') # 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')