update the project to work with torch==2.10

models/__init__.py

@@ -8,7 +8,7 @@
 from utils.utils import get_label
 from utils.minus_utils import model_layer_switch, lora_to_linear, lora_to_distill, lora_to_prunelora, linear_to_lora
 from trainer.model_arch import get_ffn1, get_mha_proj
-from utils.alpaca_utils import smart_tokenizer_and_embedding_resize
+# from utils.alpaca_utils import smart_tokenizer_and_embedding_resize
 from .modeling_bert import CoFiBertForSequenceClassification, AdaPBertForQuestionAnswering
 # from .modeling_roberta_backup import CoFiRobertaForSequenceClassification, NewRobertaForQuestionAnswering
 from .modeling_roberta import CoFiRobertaForSequenceClassification, NewRobertaForQuestionAnswering

models/modeling_mt5.py

@@ -7,7 +7,10 @@
 from dataclasses import dataclass
 from typing import Union
 
-from transformers.file_utils import ModelOutput
+try:
+    from transformers.utils import ModelOutput
+except ImportError:
+    from transformers.file_utils import ModelOutput
 
 from transformers.models.mt5.modeling_mt5 import (
     MT5ForConditionalGeneration,
@@ -1763,4 +1766,4 @@ def custom_forward(*inputs):
             hidden_states=all_hidden_states,
             attentions=all_attentions,
             cross_attentions=all_cross_attentions,
-        )
+        )

models/modeling_outputs.py

@@ -3,7 +3,10 @@
 from dataclasses import dataclass
 from typing import Optional, Tuple
 
-from transformers.file_utils import ModelOutput
+try:
+    from transformers.utils import ModelOutput
+except ImportError:
+    from transformers.file_utils import ModelOutput
 from transformers.modeling_outputs import SequenceClassifierOutput, QuestionAnsweringModelOutput, BaseModelOutputWithPastAndCrossAttentions
 
 
@@ -40,4 +43,4 @@ class NewSequenceClassifierOutput(SequenceClassifierOutput):
     masked_loss: Optional[torch.FloatTensor] = None
 
 class AdaPBaseModelOutputWithPastAndCrossAttentions(BaseModelOutputWithPastAndCrossAttentions):
-    masked_hidden_states: torch.FloatTensor = None
+    masked_hidden_states: torch.FloatTensor = None

requirements.txt

@@ -1,15 +1,17 @@
-datasets==2.10.0
-deepspeed==0.8.0
-matplotlib==3.7.1
-numpy==1.24.3
-ortools==9.6.2534
-pandas==1.5.2
-scikit_learn==1.1.3
-scipy==1.10.1
-seaborn==0.12.2
-tqdm==4.65.0
-transformers==4.28.1
-nltk==3.8.1
-rouge-score==0.1.2
-torch==1.10.2+cu113
---extra-index-url https://download.pytorch.org/whl/cu113
+torch>=2.2
+transformers==4.45.0
+accelerate>=1.0
+datasets>=3.0
+evaluate>=0.4.3
+deepspeed>=0.16
+tokenizers>=0.20
+numpy>=1.26
+pandas>=2.2
+scikit-learn>=1.5
+scipy>=1.13
+matplotlib>=3.9
+seaborn>=0.13
+tqdm>=4.66
+nltk>=3.9
+rouge-score>=0.1.2
+ortools>=9.10

run_minus_seq2seq_training.py

@@ -15,7 +15,7 @@
 from transformers import (HfArgumentParser, EvalPrediction, DataCollatorForSeq2Seq, set_seed)
 from torch.nn.utils.rnn import pad_sequence
 from deepspeed.profiling.flops_profiler import get_model_profile
-from datasets import load_metric
+import evaluate
 from models.model_args import ModelArguments
 from utils.utils import *
 from utils.minus_utils import efficiency_testing, input_constructor, compare_parameters
@@ -109,7 +109,7 @@ def main():
         model.hidden_mask = model.hidden_mask.to(training_args.device)
 
     if 'wmt' in task_name:
-        metric = load_metric("sacrebleu")
+        metric = evaluate.load("sacrebleu")
         gen_prefix = "eval"
 
         def postprocess_text(preds, labels):
@@ -140,7 +140,7 @@ def compute_metrics(eval_preds):
             result = {k: round(v, 4) for k, v in result.items()}
             return result
     else:
-        metric = load_metric("rouge")
+        metric = evaluate.load("rouge")
         def compute_metrics(eval_pred):
             predictions, labels = eval_pred
             decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
@@ -269,4 +269,4 @@ def compute_metrics(eval_pred):
 
 
 if __name__ == '__main__':
-    main()
+    main()

run_minus_squad_training.py

@@ -12,7 +12,7 @@
 
 from transformers import (HfArgumentParser, EvalPrediction, default_data_collator, set_seed)
 from deepspeed.profiling.flops_profiler import get_model_profile
-from datasets import load_metric
+import evaluate
 from transformers import SquadDataTrainingArguments
 from models.model_args import ModelArguments
 from utils.utils import *
@@ -153,9 +153,9 @@ def post_processing_function(examples, features, predictions):
 
     # Get the metric function
     if IS_SQUAD_V2:
-        metric = load_metric("squad_v2")
+        metric = evaluate.load("squad_v2")
     else:
-        metric = load_metric("squad")
+        metric = evaluate.load("squad")
 
     def compute_metrics(p: EvalPrediction):
         return metric.compute(predictions=p.predictions, references=p.label_ids)
@@ -248,4 +248,4 @@ def compute_metrics(p: EvalPrediction):
         logger.info(f"Parameter variables not changed after pruning: {same_vars}")
 
 if __name__ == '__main__':
-    main()
+    main()

run_minus_training.py

@@ -1,3 +1,4 @@
+import torch
 import os
 import json
 os.environ["WANDB_DISABLED"] = "true"
@@ -9,11 +10,15 @@
 torch.backends.cudnn.allow_tf32 = True
 import time
 transformers.logging.set_verbosity_error()
-
-
+import evaluate
+from sklearn.metrics import f1_score, matthews_corrcoef
+from scipy.stats import pearsonr, spearmanr
 from transformers import (HfArgumentParser, EvalPrediction, default_data_collator, DataCollatorWithPadding, set_seed)
-from deepspeed.profiling.flops_profiler import get_model_profile
-from datasets import load_metric
+try:
+    from deepspeed.profiling.flops_profiler import get_model_profile
+except ImportError:
+    get_model_profile = None
+# from datasets import load_metric
 from args import DataTrainingArguments
 from models.model_args import ModelArguments
 from utils.utils import *
@@ -95,9 +100,9 @@ def main():
 
     # Get the metric function
     if data_args.task_name is not None:
-        metric = load_metric("glue", data_args.task_name, experiment_id='elastictuning' + data_args.task_name + str(time.time()))
+        metric = evaluate.load("glue", data_args.task_name, experiment_id='elastictuning' + data_args.task_name + str(time.time()))
     else:
-        metric = load_metric("accuracy", experiment_id='elastictuning' + data_args.task_name + str(time.time()))
+        metric = evaluate.load("accuracy", experiment_id='elastictuning' + data_args.task_name + str(time.time()))
     # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
     # predictions and label_ids field) and has to return a dictionary string to float.
     def compute_metrics(p: EvalPrediction):
@@ -115,7 +120,31 @@ def compute_metrics(p: EvalPrediction):
             preds = list(map(lambda x: label2id[tuple(x)] if tuple(x) in label2id else -1, preds.tolist()))
             labels = list(map(lambda x: label2id[tuple(x)], labels.tolist()))
         if data_args.task_name is not None:
-            result = metric.compute(predictions=preds, references=labels)
+            try:
+                result = metric.compute(predictions=preds, references=labels)
+            except ValueError as e:
+                if "Unable to avoid copy while creating an array as requested" not in str(e):
+                    raise
+                preds_np = np.asarray(preds)
+                labels_np = np.asarray(labels)
+                if data_args.task_name in {"sst2", "mnli", "qnli", "rte"}:
+                    result = {"accuracy": float((preds_np == labels_np).mean())}
+                elif data_args.task_name in {"mrpc", "qqp"}:
+                    result = {
+                        "accuracy": float((preds_np == labels_np).mean()),
+                        "f1": float(f1_score(labels_np, preds_np)),
+                    }
+                elif data_args.task_name == "cola":
+                    result = {"matthews_correlation": float(matthews_corrcoef(labels_np, preds_np))}
+                elif data_args.task_name == "stsb":
+                    p = np.asarray(preds_np, dtype=np.float64)
+                    y = np.asarray(labels_np, dtype=np.float64)
+                    result = {
+                        "pearson": float(pearsonr(p, y)[0]),
+                        "spearmanr": float(spearmanr(p, y)[0]),
+                    }
+                else:
+                    raise
             if len(result) > 1:
                 result["combined_score"] = np.mean(list(result.values())).item()
             return result
@@ -162,14 +191,17 @@ def compute_metrics(p: EvalPrediction):
     if getattr(model, 'hidden_mask', None) is not None:
         model.hidden_mask = model.hidden_mask.to(training_args.device)
 
-    flops, macs, params = get_model_profile(
-        model,
-        kwargs={k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer, output_seq_len=2).items()} if MODEL_GENERATIVE else {k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer).items()},
-        print_profile=True,
-        detailed=True,
-        output_file=os.path.join(training_args.output_dir, 'pretrain_deepspeed_profile.txt'),
-    )
-    torch.cuda.reset_peak_memory_stats()
+    if get_model_profile is not None:
+        flops, macs, params = get_model_profile(
+            model,
+            kwargs={k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer, output_seq_len=2).items()} if MODEL_GENERATIVE else {k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer).items()},
+            print_profile=True,
+            detailed=True,
+            output_file=os.path.join(training_args.output_dir, 'pretrain_deepspeed_profile.txt'),
+        )
+        torch.cuda.reset_peak_memory_stats()
+    else:
+        logger.warning("deepspeed is not installed; skipping pre-training FLOPs/MACs profiling.")
 
     seq_len = 170 if IS_SQUAD else avg_seq_length(data_args.task_name)
     training_args.task_name = data_args.task_name
@@ -226,14 +258,15 @@ def compute_metrics(p: EvalPrediction):
         if isinstance(module, LoRALayer):
             module.eval()
 
-    flops, macs, params = get_model_profile(
-        model,
-        kwargs={k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer, output_seq_len=2).items()} if MODEL_GENERATIVE else {k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer).items()},        print_profile=True,
-        detailed=True,
-        output_file=os.path.join(training_args.output_dir, 'deepspeed_profile.txt'),
-    )
-    efficiency_results['model_flops'] = flops
-    efficiency_results['model_macs'] = macs
+    if get_model_profile is not None:
+        flops, macs, params = get_model_profile(
+            model,
+            kwargs={k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer, output_seq_len=2).items()} if MODEL_GENERATIVE else {k: v.to(model.device) for k, v in input_constructor(training_args.per_device_eval_batch_size, data_args.max_seq_length, tokenizer).items()},        print_profile=True,
+            detailed=True,
+            output_file=os.path.join(training_args.output_dir, 'deepspeed_profile.txt'),
+        )
+        efficiency_results['model_flops'] = flops
+        efficiency_results['model_macs'] = macs
 
     json.dump(efficiency_results, open(os.path.join(training_args.output_dir, 'efficiency_results.json'), 'w'), indent=4, sort_keys=True)
     run_report = gen_run_report(training_args.output_dir)
@@ -253,4 +286,4 @@ def compute_metrics(p: EvalPrediction):
         logger.info(f"Parameter variables not changed after pruning: {same_vars}")
 
 if __name__ == '__main__':
-    main()
+    main()

run_pruning.py

@@ -4,7 +4,7 @@
 import torch
 
 from tqdm import tqdm
-from datasets import load_metric
+import evaluate
 from utils import avg_seq_length
 from transformers import HfArgumentParser, TrainingArguments, DataCollatorWithPadding
 from torch.utils.data import DataLoader, RandomSampler
@@ -81,9 +81,9 @@ def test(model, eval_dataloader, head_mask, intermediate_mask, metric, data_args
 
     # Get the metric function
     if data_args.task_name is not None:
-        metric = load_metric("glue", data_args.task_name)
+        metric = evaluate.load("glue", data_args.task_name)
     else:
-        metric = load_metric("accuracy")
+        metric = evaluate.load("accuracy")
 
 
     if prune_mode == 'random':
@@ -170,4 +170,4 @@ def test(model, eval_dataloader, head_mask, intermediate_mask, metric, data_args
     # pre-pruning eval_results: {'accuracy': 0.84625}
     # cofi eval_results: {'accuracy': 0.8052}
 
-    json.dump(accuracy_by_ratios, open(os.path.join(training_args.output_dir, '%s_mask_%s_%s_nolora.json' % (prune_mode, mask_mode, test_mode)), 'w'))
+    json.dump(accuracy_by_ratios, open(os.path.join(training_args.output_dir, '%s_mask_%s_%s_nolora.json' % (prune_mode, mask_mode, test_mode)), 'w'))

scripts/adaptpruning/roberta_base_mrpc_momentum.sh

@@ -55,7 +55,7 @@ pre_pruning_tuning_steps=200
 sparsity_warmup_epochs=4
 
 learning_rate=2e-4
-training_batch_size=32
+training_batch_size=16
 num_train_epochs=120
 warmup_param_tuning_config=q:0-11,v:0-11,i:0-11
 teacher_param_tuning_config=q:0-11,v:0-11,i:0-11