optimize beacon speed (new)

Author: namespace-Pt

Long_LLM/activation_beacon/new/README.md

@@ -20,7 +20,7 @@ conda activate beacon
 
 # You may need to adjust the cuda version
 conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
-pip install transformers==4.39.3 deepspeed accelerate datasets peft pandas seaborn rouge fuzzywuzzy jieba python-Levenshtein
+pip install transformers deepspeed accelerate datasets peft pandas seaborn rouge fuzzywuzzy jieba python-Levenshtein
 pip install flash-attn --no-build-isolation
 ```
 

Long_LLM/activation_beacon/new/src/__init__.py

@@ -2,7 +2,7 @@
 from .chat import apply_chat_template
 from .args import ModelArgs
 from .data import Data
-from .modeling_utils import evaluate_perplexity, evaluate_generation, evaluate_nll, move_to_device
+from .modeling_utils import evaluate_perplexity, evaluate_generation, evaluate_nll, move_to_device, get_shifted_labels
 
 import logging
 logging.basicConfig(
@@ -12,7 +12,7 @@
 )
 
 
-def get_model_and_tokenizer(model_args, device="cpu", evaluation_mode=True, return_tokenizer_only=False, **kwargs):
+def get_model_and_tokenizer(model_args, device="cpu", evaluation_mode=True, return_tokenizer_only=False, **kwargs):    
     import torch
     import transformers
     from dataclasses import asdict
@@ -97,8 +97,6 @@ def get_model_and_tokenizer(model_args, device="cpu", evaluation_mode=True, retu
     for k, v in model_args_dict.items():
         if k.startswith("beacon") and v is not None:
             beacon_kwargs[k] = v
-        elif k.startswith("retrieval") and v is not None:
-            beacon_kwargs[k] = v
 
     # use architecture attribute to distinguish different models
     probe_config = AutoConfig.from_pretrained(

Long_LLM/activation_beacon/new/src/args.py

@@ -29,7 +29,7 @@ class ModelArgs:
     )
 
     model_name_or_path: str = field(
-        default='meta-llama/Llama-2-7b-chat-hf',
+        default='Qwen/Qwen2-7B-Instruct',
         metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}
     )
     padding_side: str = field(
@@ -161,7 +161,7 @@ class ModelArgs:
     )
     beacon_param: Optional[List[str]] = field(
         default=None,
-        metadata={'help': 'The introduced parameters for beacon.'}
+        metadata={'help': 'The introduced parameters for beacon. {q, k, v, o}'}
     )
     beacon_embed_init: str = field(
         default="eos",
@@ -183,18 +183,6 @@ class ModelArgs:
         default=None,
         metadata={'help': 'How many windows to run in parallel?'}
     )
-    retrieval_method: Optional[str] = field(
-        default=None,
-        metadata={'help': 'How to retrieve? {bm25}'}
-    )
-    retrieval_topk: Optional[int] = field(
-        default=None,
-        metadata={'help': 'How many windows to retrieve?'}
-    )
-    retrieval_key_length: Optional[int] = field(
-        default=None,
-        metadata={'help': 'The key sequence length in retrieval.'}
-    )
 
     max_new_tokens: Optional[int] = field(
         default=None,

Long_LLM/activation_beacon/new/src/data.py

@@ -14,13 +14,19 @@
 logger = logging.get_logger(__name__)
 
 
-# RETRIEVAL_CAND = [(1024,1), (512,2), (256,4), (128,8), (512,1), (256,2), (128,4)]
-RETRIEVAL_CAND = [(1024,1)]
-
 
 class Data:
+    def _process_pretrain_data(data, indices):
+        outputs = {"labels": [], "index": [], "length": []}
+        for input_ids, index in zip(data['input_ids'], indices):
+            outputs["index"].append(index)
+            outputs["length"].append(len(input_ids))
+            # NOTE: the labels will be automatically generated in Trainer._prepare_inputs
+            outputs["labels"].append(None)
+        return outputs
+
     def _process_language_modeling(data, indices, tokenizer, min_length, max_length):
-        outputs = {'input_ids': [], 'attention_mask': [], "labels": [], "length": [], "index": []}
+        outputs = {'input_ids': [], "labels": [], "length": [], "index": []}
 
         for i, text in enumerate(data['text']):
             # truncate text for faster processing
@@ -33,18 +39,20 @@ def _process_language_modeling(data, indices, tokenizer, min_length, max_length)
                 for k, v in encoded.items():
                     encoded[k] = v[:max_length]
 
-            encoded["labels"] = encoded["input_ids"].copy()
+            # NOTE: the labels will be automatically generated in Trainer._prepare_inputs
+            encoded["labels"] = None
 
             for k, v in encoded.items():
-                outputs[k].append(v)
+                if k in outputs:
+                    outputs[k].append(v)
             # length is required for grouping
             outputs["length"].append(len(encoded['input_ids']))
             outputs["index"].append(indices[i])
 
         return outputs
 
     def _process_instruction_tuning(data, indices, tokenizer, chat_template, min_length, max_length, eval_mode=False):
-        outputs = {'input_ids': [], 'attention_mask': [], "labels": [], "length": [], "index": []}
+        outputs = {'input_ids': [], "labels": [], "length": [], "index": []}
 
         for i, source in enumerate(data['conversations']):
             if source[0]["role"] != 'user':
@@ -69,6 +77,11 @@ def _process_instruction_tuning(data, indices, tokenizer, chat_template, min_len
                 add_generation_prompt=eval_mode, 
             ).encoded
 
+            # NOTE: shift the labels in advance
+            # labels = encoded["labels"][1:]
+            # labels.append(-100)
+            # encoded["labels"] = labels
+
             # skip data that not fall in between min_length and max_length
             if min_length is not None and len(encoded["input_ids"]) < min_length:
                 continue
@@ -79,13 +92,14 @@ def _process_instruction_tuning(data, indices, tokenizer, chat_template, min_len
                 encoded["labels"] = labels
 
             for k, v in encoded.items():
-                outputs[k].append(v)
+                if k in outputs:
+                    outputs[k].append(v)
             outputs['length'].append(len(encoded['input_ids']))
             outputs['index'].append(indices[i])
 
         return outputs
 
-    def prepare_train_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", seed=42, cache_dir=None, load_from_cache_file=None):
+    def prepare_train_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", seed=42, cache_dir=None, load_from_cache_file=None, ignore_index=False, ignore_length=False):
         if data_files is None:
             return None
 
@@ -115,6 +129,7 @@ def prepare_train_data(data_files=None, tokenizer=None, max_length=4096, min_len
             if os.path.isdir(data_file) and os.path.exists(os.path.join(data_file, "dataset_info.json")):
                 # the dataset may be save_to_disk in advance
                 dataset = datasets.load_from_disk(data_file)
+                dataset = dataset.map(Data._process_pretrain_data, batched=True, num_proc=32, batch_size=32, with_indices=True)
 
             else:
                 # the dataset is a json file
@@ -145,16 +160,18 @@ def prepare_train_data(data_files=None, tokenizer=None, max_length=4096, min_len
                 dataset = dataset.train_test_split(max_sample_num, seed=seed)["test"]
 
             # index column is useless in training
-            if "index" in dataset.column_names:
+            if "index" in dataset.column_names and ignore_index:
                 dataset = dataset.remove_columns(["index"])
+            if "length" in dataset.column_names and ignore_length:
+                dataset = dataset.remove_columns(["length"])
 
             train_datasets.append(dataset)
 
         dataset = datasets.concatenate_datasets(train_datasets)
 
         return dataset
 
-    def prepare_eval_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", max_eval_num=None, cache_dir=None, seed=42, load_from_cache_file=None):
+    def prepare_eval_data(data_files=None, tokenizer=None, max_length=4096, min_length=512, chat_template="vicuna", max_eval_num=None, cache_dir=None, seed=42, load_from_cache_file=None, ignore_index=False, ignore_length=False):
         if data_files is None:
             return None
 
@@ -186,4 +203,9 @@ def prepare_eval_data(data_files=None, tokenizer=None, max_length=4096, min_leng
             raise ValueError(f"Found neither 'text' nor 'conversations' in the training data!")
 
         dataset = dataset.map(process_fn, batched=True, num_proc=32, remove_columns=dataset.column_names, with_indices=True, load_from_cache_file=load_from_cache_file)
+        if "index" in dataset.column_names and ignore_index:
+            dataset = dataset.remove_columns(["index"])
+        if "length" in dataset.column_names and ignore_length:
+            dataset = dataset.remove_columns(["length"])
+
         return dataset