change: update skipbert mechanism

template_FL/src/fedllm/client_app.py

@@ -2,6 +2,7 @@
 
 import os
 import warnings
+import logging
 from typing import Dict, Tuple
 
 import torch
@@ -13,8 +14,16 @@ from flwr.common.config import unflatten_dict
 from flwr.common.typing import NDArrays, Scalar
 from omegaconf import DictConfig
 
-from transformers import TrainingArguments, DataCollatorForSeq2Seq, Trainer, EarlyStoppingCallback, BertForSequenceClassification, GenerationConfig
+from transformers import (
+    TrainingArguments, 
+    DataCollatorForSeq2Seq, 
+    Trainer, 
+    EarlyStoppingCallback, 
+    # BertForSequenceClassification,
+    GenerationConfig
+)
 
+                          
 from trl import SFTTrainer, SFTConfig
 from deepspeed.profiling.flops_profiler import get_model_profile
 from deepspeed.accelerator import get_accelerator
@@ -40,6 +49,11 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true"
 os.environ["RAY_DISABLE_DOCKER_CPU_WARNING"] = "1"
 warnings.filterwarnings("ignore", category=UserWarning)
 
+logging.getLogger("flwr").setLevel(logging.INFO)
+logging.getLogger("ClientAppActor").setLevel(logging.INFO)
+
+logger = logging.getLogger(__name__)
+
 
 def input_constructor(batch_size, seq_len, tokenizer):
     fake_seq = ""
@@ -80,13 +94,15 @@ class FlowerClient(NumPyClient):
         self,
         model_cfg: DictConfig,
         train_cfg: DictConfig,
-        mates_args: DictConfig,
+        mates_cfg: DictConfig,
+        skipbert_cfg: DictConfig,
         trainset,
         valset,
         num_rounds,
     ):  # pylint: disable=too-many-arguments
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.train_cfg = train_cfg
+        self.skipbert_cfg = skipbert_cfg
         
         self.training_arguments = TrainingArguments(**train_cfg.training_arguments)
         # self.training_arguments = SFTConfig(**train_cfg.training_arguments, max_seq_length=train_cfg.seq_length) 
@@ -94,17 +110,18 @@ class FlowerClient(NumPyClient):
         self.num_rounds = num_rounds
         self.trainset = trainset
         self.valset = valset
-        self.mates_args = mates_args
+        self.mates_cfg = mates_cfg
         self.holdoutset = None
         self.refset = None
-        self.data_influence_model = None
+        self.teacher_data_influence_model = None
+        self.student_data_influence_model = None
         self.data_influence_tokenizer = None
 
         # instantiate model
         self.model, self.tokenizer = get_model(model_cfg)
         
-        if self.mates_args.state:
-            self.data_influence_model, self.data_influence_tokenizer = get_data_influence_model(model_cfg)      
+        if self.mates_cfg.state:
+            self.teacher_data_influence_model, self.student_data_influence_model ,self.data_influence_tokenizer = get_data_influence_model(model_cfg, skipbert_cfg)      
         
         # (
         #     self.data_collator, 
@@ -129,8 +146,8 @@ class FlowerClient(NumPyClient):
         # Replace -100 with pad token id in labels
         labels_ids[labels_ids == -100] = self.tokenizer.pad_token_id
         
-        print(f"Shape of predictions: {np.shape(pred_ids)}")
-        print(f"Shape of labels: {np.shape(labels_ids)}")
+        # print(f"Shape of predictions: {np.shape(pred_ids)}")
+        # print(f"Shape of labels: {np.shape(labels_ids)}")
         
         # Decode predictions and labels
         pred_str = self.tokenizer.batch_decode(
@@ -165,6 +182,7 @@ class FlowerClient(NumPyClient):
             .map(
                 lambda x: generate_and_tokenize_prompt(x, **tmp_dict),
                 num_proc=8,
+                remove_columns=['instruction', 'input', 'output']
             )
         )
 
@@ -175,16 +193,17 @@ class FlowerClient(NumPyClient):
             .map(
                 lambda x: generate_and_tokenize_prompt(x, **tmp_dict),
                 num_proc=8,
+                remove_columns=['instruction', 'input', 'output']
             )
         )
 
         # Create holdoutset and refset if state is True
-        if self.mates_args.state:
+        if self.mates_cfg.state:
             trainset_size = len(self.trainset)
 
             # Calculate sizes for holdout and reference sets
-            holdout_size = int(trainset_size * self.mates_args.holdout_ratio)
-            ref_size = int(trainset_size * self.mates_args.reference_ratio)
+            holdout_size = int(trainset_size * self.mates_cfg.holdout_ratio)
+            ref_size = int(trainset_size * self.mates_cfg.reference_ratio)
 
             # Shuffle the trainset to ensure randomness
             shuffled_indices = list(range(trainset_size))
@@ -199,16 +218,17 @@ class FlowerClient(NumPyClient):
             self.refset = self.trainset.select(ref_indices)
 
             print(f"Holdoutset size: {len(self.holdoutset)}, Refset size: {len(self.refset)}")
+            # logger.info(f"Holdoutset size: {len(self.holdoutset)}, Refset size: {len(self.refset)}")
             
 
     def fit(
         self, parameters: NDArrays, config: Dict[str, Scalar]
     ) -> Tuple[NDArrays, int, Dict]:
         """Implement distributed fit function for a given client."""
-        if self.mates_args.state and int(config["current_round"]) != 1:
+        if self.mates_cfg.state and int(config["current_round"]) != 1:
             main_model_params, data_influence_model_params = split_models(parameters)
             set_parameters(self.model, main_model_params)
-            set_parameters_bert(self.data_influence_model, data_influence_model_params)
+            set_parameters_bert(self.teacher_data_influence_model, data_influence_model_params)
         else:
             set_parameters(self.model, parameters)
 
@@ -259,18 +279,20 @@ class FlowerClient(NumPyClient):
             args=self.training_arguments,
             data_collator=self.data_collator,
             compute_metrics=self.compute_metrics, 
-            mates_args=self.mates_args,
-            data_influence_model=self.data_influence_model,
+            mates_cfg=self.mates_cfg,
+            skipbert_cfg=self.skipbert_cfg,
+            teacher_data_influence_model=self.teacher_data_influence_model,
+            student_data_influence_model=self.student_data_influence_model,
             data_influence_tokenizer=self.data_influence_tokenizer,
         )
 
         # Train the model
         results = trainer.train()
         
-        if self.mates_args.state:
+        if self.mates_cfg.state:
             # After training
             main_model_params = get_parameters(self.model)
-            data_influence_model_params = model_parameters_to_ndarrays(self.data_influence_model)
+            data_influence_model_params = model_parameters_to_ndarrays(self.teacher_data_influence_model)
             final_model_params = concatenate_models_with_marker(main_model_params, data_influence_model_params)
         else:
             final_model_params = get_parameters(self.model)
@@ -286,7 +308,7 @@ class FlowerClient(NumPyClient):
               detailed=False,
             )
             flops2, macs2, params2 = get_model_profile(
-              self.data_influence_model,
+              self.teacher_data_influence_model,
               kwargs=input_constructor(batch_size, seq_len, self.data_influence_tokenizer),
               print_profile=True,
               detailed=False,
@@ -315,11 +337,16 @@ def client_fn(context: Context) -> FlowerClient:
         client_set = load_data_homo(partition_id, num_partitions, cfg.dataset.name)
     else:
         client_set = load_data_hete(partition_id)
+        
+        
+    cfg.skipbert.att_layer_maps = [int(s) for s in cfg.skipbert.att_layer_maps.split(', ')]
+    cfg.skipbert.hid_layer_maps = [int(k) for k in cfg.skipbert.hid_layer_maps.split(', ')]
 
     return FlowerClient(
         cfg.model,
         cfg.train,
         cfg.mates,
+        cfg.skipbert,
         client_set['train'],
         client_set['test'],
         num_rounds,

template_FL/src/fedllm/dataset.py

@@ -9,7 +9,7 @@ import pandas as pd
 
 FDS = None  # Cache FederatedDataset
 client_id_ds = None
-global_test_set_homo = None
+# global_test_set_homo = None
 
 def split_train_test(dataset, test_size):
     # Split the dataset into train and test sets

template_FL/src/fedllm/models.py

@@ -11,7 +11,16 @@ from peft import (
     set_peft_model_state_dict,
 )
 from peft.utils import prepare_model_for_kbit_training
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainerCallback, BertForSequenceClassification
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer, 
+    BitsAndBytesConfig, 
+    TrainerCallback, 
+    # BertForSequenceClassification,
+    BertConfig,
+)
+
+from .skipbert.modeling import BertForSequenceClassification, SkipBertForSequenceClassification
 
 from flwr.common.typing import NDArrays
 from transformers.trainer_callback import TrainerControl, TrainerState
@@ -90,24 +99,71 @@ def get_model(model_cfg: DictConfig):
 
     return get_peft_model(model, peft_config), tokenizer
 
-def get_data_influence_model(model_cfg: DictConfig):
+def get_custom_config(teacher_name, skipbert_cfg: DictConfig):
+    num_labels = 1 # Set number of labels to 1 for regression or single-class tasks
+
+
+    teacher_config = BertConfig.from_pretrained(teacher_name)
+    teacher_config.num_labels = num_labels
+    teacher_config.fit_size = teacher_config.hidden_size
+
+    student_config = BertConfig.from_pretrained(skipbert_cfg.student_model)
+    student_config.num_labels = num_labels
+    student_config.fit_size = teacher_config.hidden_size
+
+    if skipbert_cfg.num_layers_student > 0:
+        student_config.num_hidden_layers = skipbert_cfg.num_layers_student
+    if skipbert_cfg.num_full_hidden_layers_student > 0:
+        student_config.num_full_hidden_layers = skipbert_cfg.num_full_hidden_layers_student
+    else:
+        student_config.num_full_hidden_layers = student_config.num_hidden_layers
+
+    student_config.task_type = skipbert_cfg.output_mode
+    student_config.n_gram_left = skipbert_cfg.n_gram_left
+    student_config.n_gram_right = skipbert_cfg.n_gram_right
+    #     student_config.plot_mode = 'plot_passive'
+    student_config.plot_mode = 'force_compute'
+    student_config.ngram_masking = 0.
+    if not hasattr(student_config, 'enter_hidden_size'):
+        student_config.enter_hidden_size = student_config.hidden_size
+    if not hasattr(student_config, 'max_num_entries'):
+        student_config.max_num_entries = 100000
+
+    return teacher_config, student_config
+    
+
+def get_data_influence_model(model_cfg: DictConfig, skipbert_cfg: DictConfig):
     use_cuda = torch.cuda.is_available()
     device_map = torch.device("cuda:0" if use_cuda else "cpu")
-
+    teacher_name = "bert-base-uncased"
+    
+    teacher_config, student_config = get_custom_config(teacher_name=teacher_name, skipbert_cfg=skipbert_cfg)
+    
     # Load model with num_labels=1
-    model = BertForSequenceClassification.from_pretrained(
-        "bert-base-uncased",
-        num_labels=1,  # Set number of labels to 1 for regression or single-class tasks
+    teacher_model = BertForSequenceClassification.from_pretrained(
+        teacher_name, config=teacher_config
     ).to(device_map)
     
-    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
+    student_model = SkipBertForSequenceClassification.from_pretrained(
+        skipbert_cfg.student_model, config=student_config, 
+        do_fit=skipbert_cfg.do_fit, share_param=skipbert_cfg.share_param).to(device_map)
+
+    if skipbert_cfg.freeze_lower_layers:
+        student_model.freeze_shallow_layers()
+    
+    tokenizer = AutoTokenizer.from_pretrained(teacher_name, do_lower_case=skipbert_cfg.do_lower_case, use_fast=True)
     
     if use_cuda:
-        model = prepare_model_for_kbit_training(
-            model, use_gradient_checkpointing=model_cfg.gradient_checkpointing
+        teacher_model = prepare_model_for_kbit_training(
+            teacher_model, use_gradient_checkpointing=model_cfg.gradient_checkpointing
+        )
+        
+        student_model = prepare_model_for_kbit_training(
+            student_model, use_gradient_checkpointing=model_cfg.gradient_checkpointing
         )
+        
 
-    return model, tokenizer
+    return teacher_model, student_model, tokenizer
 
 
 def set_parameters(model, parameters: NDArrays) -> None:

template_FL/src/fedllm/server_app.py

@@ -1,9 +1,11 @@
 """flowertune-llm: A Flower / FlowerTune app."""
 
 import os
+import sys
 import torch
 import wandb
 import numpy as np
+import pandas as pd
 from dotenv import load_dotenv
 from datetime import datetime
 from tqdm import tqdm
@@ -12,9 +14,11 @@ from transformers import DataCollatorForSeq2Seq, DataCollatorWithPadding, Traini
 from .trainer import ManualTrainer
 from transformers.integrations import WandbCallback
 from torch.utils.data import DataLoader
+import flwr
 from flwr.common import Context, ndarrays_to_parameters
 from flwr.common.config import unflatten_dict
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.common.logger import FLOWER_LOGGER
 # from flwr.server.strategy import FedAvg
 from omegaconf import DictConfig
 
@@ -27,6 +31,12 @@ from .metrics import exact_match, f1, get_rouge_score
 
 from datasets import load_dataset, Dataset
 from sklearn.model_selection import train_test_split
+import logging
+import uuid
+
+logging.getLogger("flwr").setLevel(logging.INFO)
+logging.getLogger("ClientAppActor").setLevel(logging.INFO)
+logging.getLogger("Trainer").setLevel(logging.INFO)
 
 
 load_dotenv(".env")
@@ -36,6 +46,79 @@ os.environ["WANDB_NAME"] = os.getenv("WANDB_NAME")
 os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
 # os.environ["WANDB_LOG_MODEL"] = "checkpoint"
 
+
+class SessionIDFilter(logging.Filter):
+    """Adds a session_id to log records."""
+    def __init__(self, session_id):
+        super().__init__()
+        self.session_id = session_id
+
+    def filter(self, record):
+        record.session_id = self.session_id
+        return True
+
+
+def configure_logging():
+    # Generate a unique session ID for this run
+    session_id = str(uuid.uuid4())
+
+    # Define log format with session ID and process ID
+    log_format = (
+        "%(asctime)s - %(session_id)s - %(process)d - %(name)s - "
+        "%(levelname)s - %(message)s"
+    )
+
+    # Create a FileHandler and attach the SessionIDFilter to it
+    file_handler = logging.FileHandler("main.log", mode='a')  # Append mode
+    formatter = logging.Formatter(log_format)
+    file_handler.setFormatter(formatter)
+    file_handler.addFilter(SessionIDFilter(session_id))  # Add filter to the handler
+    
+    # Console handler: logs to stdout (you can also log to stderr)
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(logging.DEBUG)
+    console_handler.setFormatter(formatter)
+    
+
+    # Configure root logger to use this handler
+    logging.basicConfig(
+        level=logging.INFO,
+        handlers=[
+            file_handler,
+            console_handler,
+        ],  # Use the filtered handler
+    )
+    
+    # if not any(
+    #     isinstance(handler, logging.FileHandler) and handler.baseFilename == file_handler.baseFilename
+    #     for handler in FLOWER_LOGGER.handlers
+    # ):
+    #     FLOWER_LOGGER.addHandler(file_handler)
+        
+    for handler in FLOWER_LOGGER.handlers:
+        FLOWER_LOGGER.addHandler(file_handler)
+        
+    # Get the logger for the ClientAppActor module and attach the same file handler
+    client_actor_logger = logging.getLogger("flwr.simulation.ray_transport.ray_actor")
+    # if not any(
+    #     isinstance(handler, logging.FileHandler) and handler.baseFilename == file_handler.baseFilename
+    #     for handler in client_actor_logger.handlers
+    # ):
+    #     client_actor_logger.addHandler(file_handler)
+    
+    for handler in client_actor_logger.handlers:
+        client_actor_logger.addHandler(file_handler)
+        
+    # Explicitly configure Ray's logger to propagate
+    ray_logger = logging.getLogger("ray")  # Ray's parent logger
+    for handler in ray_logger.handlers:
+        ray_logger.addHandler(file_handler)
+
+
+    # Log the start of the session
+    logger = logging.getLogger(__name__)
+    logger.info("===== Application Started =====")
+
 class LLMSampleCB(WandbCallback):
     def __init__(self, trainer, test_dataset, task, num_samples=10, max_new_tokens=256, log_model="checkpoint"):
         "A CallBack to log samples a wandb.Table during training"
@@ -84,7 +167,7 @@ class LLMSampleCB(WandbCallback):
 
 
 
-def test_model(dataset, model, tokenizer, train_cfg, tmp_dict, sround, mates_args, task):
+def test_model(dataset, model, tokenizer, train_cfg, tmp_dict, sround, mates_args, skipbert_cfg, task):
     
     wandb.init(
         project='FL@CSS25',
@@ -157,8 +240,10 @@ def test_model(dataset, model, tokenizer, train_cfg, tmp_dict, sround, mates_arg
         args=training_arguments,
         data_collator=data_collator,
         compute_metrics=compute_metrics, 
-        mates_args=mates_args,
-        data_influence_model=None,
+        mates_cfg=mates_args,
+        skipbert_cfg=skipbert_cfg,
+        teacher_data_influence_model=None,
+        student_data_influence_model=None,
         data_influence_tokenizer=None,
     )
     
@@ -185,7 +270,7 @@ def test_model(dataset, model, tokenizer, train_cfg, tmp_dict, sround, mates_arg
 # Get function that will be executed by the strategy's evaluate() method
 # Here we use it to save global model checkpoints
 
-def get_evaluate_fn(train_cfg, model_cfg, dataset_cfg, save_every_round, total_round, total_nodes, save_path, mates_args):
+def get_evaluate_fn(train_cfg, model_cfg, dataset_cfg, save_every_round, total_round, total_nodes, save_path, mates_args, skipbert_cfg):
     """Return an evaluation function for saving global model."""
 
     def evaluate(server_round: int, parameters, config):
@@ -208,12 +293,14 @@ def get_evaluate_fn(train_cfg, model_cfg, dataset_cfg, save_every_round, total_r
             }
             if dataset_cfg.type == 'homo':
                 ds = load_dataset(dataset_cfg.name)
+                option = 'test' if 'test' in ds else 'train'
+                df = pd.DataFrame(ds[option])
                 _, test = train_test_split(
-                    ds, test_size=0.09, shuffle=True, random_state=42
+                    df, test_size=0.09, shuffle=True, random_state=42
                 )
                 global_test_set_homo = Dataset.from_pandas(test).remove_columns(['__index_level_0__'])
                 
-                loss, metrics = test_model(global_test_set_homo, model, tokenizer, train_cfg, tmp_dict, server_round, mates_args, 'homo')
+                loss, metrics = test_model(global_test_set_homo, model, tokenizer, train_cfg, tmp_dict, server_round, mates_args, skipbert_cfg, 'homo')
                 total_loss = loss
                 result_metric = {'homo_f1': metrics['homo_f1']}
             else:
@@ -225,7 +312,7 @@ def get_evaluate_fn(train_cfg, model_cfg, dataset_cfg, save_every_round, total_r
                 
                 for task in ['general', 'finance', 'math', 'medical', 'code']:
                     ds = global_test_set_hete[task]
-                    loss, metrics = test_model(ds, model, tokenizer, train_cfg, tmp_dict, server_round, mates_args, task)
+                    loss, metrics = test_model(ds, model, tokenizer, train_cfg, tmp_dict, server_round, mates_args, skipbert_cfg, task)
                     list_loss.append(loss)
                     
                     list_f1[f'{task}_f1'] = metrics[f'{task}_f1']
@@ -273,6 +360,8 @@ def fit_weighted_average(metrics):
 def server_fn(context: Context):
     """Construct components that set the ServerApp behaviour."""
     # Create output directory given current timestamp
+    configure_logging()
+    logger = logging.getLogger(__name__)
     current_time = datetime.now()
     folder_name = current_time.strftime("%Y-%m-%d_%H-%M-%S")
     save_path = os.path.join(os.getcwd(), f"results/{folder_name}")
@@ -296,7 +385,7 @@ def server_fn(context: Context):
         fit_metrics_aggregation_fn=fit_weighted_average,
         initial_parameters=init_model_parameters,
         evaluate_fn=get_evaluate_fn(
-            cfg.train, cfg.model, cfg.dataset, cfg.train.save_every_round, num_rounds, num_nodes, save_path, cfg.mates
+            cfg.train, cfg.model, cfg.dataset, cfg.train.save_every_round, num_rounds, num_nodes, save_path, cfg.mates, cfg.skipbert
         ),
         use_mates=cfg.mates.state
     )

template_FL/src/fedllm/skipbert/modeling.py

@@ -474,6 +474,7 @@ class ShallowSkipping(nn.Module):
     
     @torch.jit.script
     def merge_ngrams(input_ids, ngram_hidden_states, aux_embeddings):
+        # batch_size, seq_length = input_ids.shape
         batch_size, seq_length = input_ids.shape
         lens = (input_ids!=0).sum(1)
         hidden_state = torch.zeros([batch_size, seq_length, ngram_hidden_states.size(-1)], dtype=ngram_hidden_states.dtype, device=ngram_hidden_states.device)
@@ -562,7 +563,6 @@ class ShallowSkipping(nn.Module):
     ):
         
         device = model.device
-        
         batch_size, seq_length = input_ids.shape
         aux_embeddings = model.embeddings.position_embeddings2.weight[:seq_length].unsqueeze(0)
         aux_embeddings = aux_embeddings + model.embeddings.token_type_embeddings2(token_type_ids)

template_FL/src/fedllm/skipbert/plot.py

@@ -117,8 +117,8 @@ class Plot:
         self.max_num_entries = max_num_entries
         self.hidden_size = hidden_size
         
-        self.trigram_to_id, self.id_to_trigram = self.build_hash_table('input_ids_tri_gram.memmap', max_num_entries)
-        self.orig_trigram_hidden_states =  _read_or_create_memmap("plot_hidden_states_tri_gram.memmap", dtype='float16', shape=(max_num_entries, 3, hidden_size))
+        self.trigram_to_id, self.id_to_trigram = self.build_hash_table('./input_ids_tri_gram.memmap', max_num_entries)
+        self.orig_trigram_hidden_states =  _read_or_create_memmap("./plot_hidden_states_tri_gram.memmap", dtype='float16', shape=(max_num_entries, 3, hidden_size))
 
     def build_hash_table(self, path, max_num_entries):
         n_gram = 3

template_FL/src/fedllm/skipbert/trainer.py

@@ -0,0 +1,691 @@
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
+import torch
+from torch import nn
+from torch.nn import MSELoss, CrossEntropyLoss
+import copy
+import numpy as np
+from transformers import (
+    # BertForSequenceClassification, 
+    GenerationConfig, 
+    AutoTokenizer,
+    Trainer,
+    get_scheduler, 
+    EarlyStoppingCallback,
+    TrainingArguments
+)
+from transformers.trainer_utils import (
+    EvaluationStrategy,
+    IntervalStrategy,
+)
+from transformers.trainer_pt_utils import nested_detach
+from transformers.utils import is_sagemaker_mp_enabled
+from transformers.training_args import OptimizerNames
+from typing import Dict, List, Optional, Any, Union, Tuple, Callable
+import numpy as np
+from torch.utils.data import Dataset, DataLoader
+from accelerate.utils import (
+        AutocastKwargs,
+        DistributedDataParallelKwargs,
+        DistributedType,
+    )
+from datasets import Dataset
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+import wandb
+import logging
+
+logging.getLogger("Trainer").setLevel(logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def compute_metrics_skipbert(pred):
+    """
+    Compute metrics for model evaluation
+    """
+    labels = pred.label_ids
+    
+    preds = pred.predictions
+    
+    if len(preds[0]) >= 2:
+        preds = torch.tensor(preds.argmax(-1))
+        labels = torch.tensor(labels)
+        
+        acc = accuracy_score(labels, preds)
+        precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
+        return {
+            'accuracy': acc,
+            'f1': f1,
+            'precision': precision,
+            'recall': recall
+        }
+    else:
+        labels = torch.tensor(pred.label_ids[:, np.newaxis])
+        preds = torch.tensor(pred.predictions)
+        
+        # MSE
+        mse = nn.MSELoss()
+        mse_loss = mse(labels, preds)
+
+        #RMSE
+        rmse = torch.sqrt(mse_loss)
+
+        # MAE
+        mae = nn.L1Loss()
+        mae_loss = mae(labels, preds)
+
+
+        return {
+            'mse': mse_loss,
+            'rmse': rmse,
+            'mae': mae_loss,
+        }
+
+
+# Create custom Trainer for training SkipBERT
+class SkipBertTrainer(Trainer):
+    def __init__(
+        self,
+        student_model: nn.Module,
+        teacher_model: Optional[nn.Module] = None,
+        train_dataset: Optional[Dataset] = None,
+        eval_dataset: Optional[Dataset] = None,
+        args: Optional[TrainingArguments] = None,
+        data_collator: Optional[Callable] = None,
+        compute_metrics: Optional[Callable] = None,
+        alpha: float = 0.5,
+        temperature: float = 2.0,
+        beta: float = 1.0,
+        use_logits: bool = True,
+        use_att: bool = True,
+        use_rep: bool = True,
+        use_embedding: bool = True,
+        att_layer_maps: Optional[List[int]] = None,
+        hid_layer_maps: Optional[List[int]] = None,
+        epochs_no_cls: int = 0,
+        reduce_T: int = 1,
+        output_mode: str = 'classification',
+        num_masked_layers_teacher: int = 0,
+        num_masked_last_layers_teacher: int = 0,
+        fp16: bool = False,
+        num_full_hidden_layers_student: int = 0,
+        **kwargs,
+    ):
+        """
+        Initialize SkipBERT Trainer with knowledge distillation capabilities.
+        
+        Args:
+            student_model: The student model to be trained
+            teacher_model: The teacher model for knowledge distillation
+            train_dataset: Training dataset
+            eval_dataset: Evaluation dataset
+            args: Training arguments
+            alpha: Balance between distillation loss and cross-entropy loss
+            temperature: Temperature for softening probability distributions
+            beta: Weighting factor for different loss components
+            use_logits: Whether to use logits-based distillation
+            use_att: Whether to use attention-based distillation
+            use_rep: Whether to use representation-based distillation
+            use_embedding: Whether to use embedding-based distillation
+        """
+        # Set default training arguments if not provided
+        if args is None:
+            args = TrainingArguments(
+                output_dir="./results",
+                num_train_epochs=3,
+                per_device_train_batch_size=2,
+                per_device_eval_batch_size=2,
+                logging_dir='./logs',
+                evaluation_strategy=EvaluationStrategy.EPOCH,
+                save_strategy=IntervalStrategy.EPOCH,
+            )
+        
+        
+        # Call parent constructor
+        super().__init__(
+            model=student_model,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            data_collator=data_collator,
+            compute_metrics=compute_metrics,
+            **kwargs
+        )
+        
+        # Store additional knowledge distillation parameters
+        self.teacher_model = teacher_model
+        self.alpha = alpha
+        self.temperature = temperature
+        self.beta = beta
+        self.use_logits = use_logits
+        self.use_att = use_att
+        self.use_rep = use_rep
+        self.use_embedding = use_embedding
+        self.att_layer_maps = att_layer_maps or []
+        self.hid_layer_maps = hid_layer_maps or []
+        self.epochs_no_cls = epochs_no_cls
+        self.reduce_T = reduce_T
+        self.output_mode = output_mode
+        self.num_masked_layers_teacher = num_masked_layers_teacher
+        self.num_masked_last_layers_teacher = num_masked_last_layers_teacher
+        self.num_full_hidden_layers_student = num_full_hidden_layers_student
+        self.tr_att_loss = 0
+        self.tr_rep_loss = 0
+        self.tr_cls_loss = 0
+        self.list_att_loss = []
+        self.list_rep_loss = []
+        self.list_embed_loss = []
+      
+        # Prepare FP16 if enabled
+        self.fp16 = fp16
+        if fp16:
+            try:
+                from apex import amp
+            except ImportError:
+                raise ImportError(
+                    "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
+                )
+            
+            # Initialize amp
+            self.model, self.optimizer = amp.initialize(
+                self.model, 
+                self.optimizer, 
+                opt_level='01'
+            )
+            
+            # Half precision for teacher model if exists
+            if self.teacher_model is not None:
+                self.teacher_model = self.teacher_model.half()
+        
+        # Loss functions
+        self.loss_mse = MSELoss()
+    
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        """
+        How the loss is computed by Trainer. By default, all models return the loss in the first element.
+
+        Subclass and override for custom behavior.
+        """
+        
+        # Separate labels from inputs
+        labels = inputs.pop("labels")
+        
+        if self.model_accepts_loss_kwargs:
+            loss_kwargs = {}
+            if num_items_in_batch is not None:
+                loss_kwargs["num_items_in_batch"] = num_items_in_batch
+            inputs = {**inputs, **loss_kwargs}
+        
+        # Forward pass through student model
+        student_logits, student_atts, student_reps = model(**inputs)
+        student_reps = student_reps[-self.num_full_hidden_layers_student-1:]
+
+        # Forward pass through teacher model
+        self.teacher_model.eval()
+        with torch.no_grad():
+            # teacher_logits, teacher_atts, teacher_reps = self.teacher_model(**inputs)
+            teacher_outputs = self.teacher_model(**inputs, output_hidden_states=True, output_attentions=True)
+            teacher_logits, teacher_atts, teacher_reps = teacher_outputs.logits, teacher_outputs.attentions, teacher_outputs.hidden_states
+            start, end = self.num_masked_layers_teacher, -1 * self.num_masked_layers_teacher if self.num_masked_layers_teacher != 0 else None
+            teacher_reps = teacher_reps[start:end]
+            
+        # Save past state if it exists
+        # TODO: this needs to be fixed and made cleaner later.
+        if self.args.past_index >= 0:
+            self._past = student_outputs[self.args.past_index]
+            
+        # Compute losses
+        att_loss, rep_loss = 0., 0.
+        
+        # ---------------------------
+        if labels is not None:
+            
+            
+            # --------------------------- 
+            if self.att_layer_maps is None:
+                teacher_layer_num = len(teacher_atts)
+                student_layer_num = len(student_atts)
+                assert teacher_layer_num % student_layer_num == 0
+                layers_per_block = int(teacher_layer_num / student_layer_num)
+                new_teacher_atts = [
+                    teacher_atts[(i * 1) * layers_per_block - 1]
+                    for i in range(student_layer_num)
+                ]
+                assert len(student_atts) == len(new_teacher_atts)
+
+            else:
+                new_teacher_atts = []
+                for t2s in self.att_layer_maps:
+                    if t2s >= 0:
+                        new_teacher_atts.append(teacher_atts[t2s])
+                    else:
+                        new_teacher_atts.append(None)
+
+            # ----------------------------
+
+            for student_att, teacher_att in zip(student_atts, new_teacher_atts):
+                if teacher_att is None:
+                    continue
+                student_att = torch.where(
+                    student_att <= 1e-2,
+                    torch.zeros_like(student_att),
+                    student_att
+                )
+
+                teacher_att = torch.where(
+                    teacher_att <= 1e-2,
+                    torch.zeros_like(teacher_att),
+                    teacher_att
+                )
+
+                att_loss += self.loss_mse(student_att, teacher_att)
+
+            # ---------------------------
+
+            if self.hid_layer_maps is None:
+                teacher_layer_num = len(teacher_atts) - 1
+                student_layer_num = len(student_atts) - 1
+                assert teacher_layer_num % student_layer_num == 0
+                layers_per_block = int(teacher_layer_num / student_layer_num)
+                new_teacher_reps = [
+                    teacher_reps[i * layers_per_block]
+                    for i in range(student_layer_num + 1)
+                ]
+                assert len(new_student_reps) == len(new_teacher_reps)
+            else:
+                new_student_reps = student_reps
+                new_teacher_reps = []
+                for t2s in self.hid_layer_maps:
+                    if t2s >= 0:
+                        new_teacher_reps.append(teacher_reps[t2s])
+                    else:
+                        new_teacher_reps.append(None)
+
+            # ---------------------------
+
+            for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps):
+                if teacher_rep is None:
+                    continue
+                tmp_loss = self.loss_mse(student_rep, teacher_rep)
+                rep_loss += tmp_loss
+
+            self.tr_att_loss += att_loss.item()
+            self.tr_rep_loss += rep_loss.item()
+
+            # ---------------------------
+            embedding_loss = 0
+            if self.use_embedding:
+                embedding_loss = self.loss_mse(
+                    student_reps[0], teacher_reps[0]
+                )
+
+            # ---------------------------
+
+            # ---------------------------
+
+            if self.use_logits and self.state.epoch >= self.epochs_no_cls:
+                if isinstance(student_logits, tuple) or \
+                    isinstance(student_logits, list):
+                    cls_loss = None
+                    _scale = 0.
+                    for il, logits in enumerate(student_logits):
+                        _loss, _, _ = self._compute_distillation_loss(
+                            student_logits, student_atts, student_reps,
+                            teacher_logits, teacher_atts, teacher_reps,
+                            labels
+                        )
+                        if cls_loss is None:
+                            cls_loss = _loss
+                        else:
+                            cls_loss = _loss * (il + 1.) + cls_loss
+                        _scale += il + 1
+
+                    cls_loss = cls_loss * (1. / _scale)
+
+                else:
+                    cls_loss, kd_loss, ce_loss = self._compute_distillation_loss(
+                            student_logits, student_atts, student_reps,
+                            teacher_logits, teacher_atts, teacher_reps,
+                            labels
+                        )
+                self.tr_cls_loss += cls_loss.item()
+
+            else:
+                cls_loss = 0
+
+            # ---------------------------
+
+
+            check = self.state.epoch >= self.epochs_no_cls
+            self.beta = self.beta * check + (1 - check) * 1.
+
+            # ---------------------------
+
+            if self.use_embedding and \
+                 self.use_att and \
+                 self.use_rep:
+                loss = self.beta * (rep_loss + att_loss + embedding_loss) + cls_loss
+
+            elif self.use_att and self.use_rep:
+                loss = self.beta * (rep_loss + att_loss) + cls_loss
+
+            elif self.use_embedding and self.use_att:
+                loss = self.beta * (att_loss + embedding_loss) + cls_loss
+
+            elif self.use_embedding and self.use_rep:
+                loss = self.beta * (rep_loss + embedding_loss) + cls_loss
+
+            elif self.use_att and \
+                not self.use_embedding and \
+                not self.use_rep:
+                loss = self.beta * att_loss + cls_loss
+
+            elif self.use_rep and \
+                not self.use_embedding and \
+                not self.use_att:
+                loss = self.beta * rep_loss + cls_loss
+
+            else:
+                loss = cls_loss
+                
+        
+            # --------------------------- 
+        
+        else:
+            if isinstance(outputs, dict) and "loss" not in outputs:
+                raise ValueError(
+                    "The model did not return a loss from the inputs, only the following keys: "
+                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
+                )
+            # We don't use .loss here since the model may return tuples instead of ModelOutput.
+            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
+        
+        # ---------------------------
+        
+        # ---------------------------
+        if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs:
+            loss *= self.accelerator.num_processes
+            rep_loss *= self.accelerator.num_processes
+            att_loss *= self.accelerator.num_processes
+            embedding_loss *= self.accelerator.num_processes
+            self.list_att_loss.append(att_loss.item())
+            self.list_rep_loss.append(rep_loss.item())
+            self.list_embed_loss.append(embedding_loss.item())
+        # ---------------------------
+        
+        
+        # ---------------------------
+        
+        # Ensure logits are properly formatted for evaluation metrics
+        logits = student_logits
+        if return_outputs:
+            # Ensure student_logits has the correct shape [batch_size, num_classes]
+            if isinstance(student_logits, (tuple, list)):
+                logits = student_logits[-1]
+            else:
+                logits = student_logits
+
+            # If logits is 1D, reshape it to 2D
+            if len(logits.shape) == 1:
+                logits = logits.unsqueeze(0)
+
+            # Ensure we have [batch_size, num_classes] shape
+            if len(logits.shape) != 2:
+                raise ValueError(f"Unexpected logits shape: {logits.shape}. Expected [batch_size, num_classes]")
+            
+            if self.output_mode == "classification": # Classification
+                loss = nn.functional.cross_entropy(labels.view(-1), logits.view(-1, len(logits[0])), reduction="mean")
+            
+            elif self.output_mode == "regression":  # Regression
+                # print(f"Return output -  student: {nn.functional.softmax(student_logits, dim=0).view(-1)}, labels: {labels.view(-1)}")
+                loss = self.loss_mse(labels.view(-1), logits.view(-1))
+
+            
+        # ---------------------------
+        # print(f"loss: {loss}, att_loss: {att_loss}, rep_loss: {rep_loss}, embed_loss: {embedding_loss}, Train {return_outputs}")
+        return (loss, logits) if return_outputs else loss
+
+    def _compute_distillation_loss(
+        self, 
+        student_logits, student_atts, student_reps,
+        teacher_logits, teacher_atts, teacher_reps,
+        labels
+    ):
+        """
+        Compute comprehensive knowledge distillation loss.
+        
+        Args:
+            student_*: Student model's outputs
+            teacher_*: Teacher model's outputs
+            labels: Ground truth labels
+        
+        Returns:
+            Computed loss
+        """
+
+        # Classification/distillation loss
+        if self.output_mode == "classification": # Classification
+            # Similar to previous implementation's distillation loss
+            if teacher_logits is not None:
+                student_likelihood = nn.functional.log_softmax(student_logits / self.temperature, dim=-1)
+                targets_prob = nn.functional.softmax(teacher_logits / self.temperature, dim=-1)
+                d_loss = (-targets_prob * student_likelihood).mean() * (self.temperature ** 2) / self.reduce_T
+            else:
+                d_loss = 0
+        # Standard cross-entropy/MSE loss
+            nll_loss = nn.functional.cross_entropy(student_logits, labels, reduction="mean")
+            
+        elif self.output_mode == "regression":  # Regression
+            # student_likelihood = nn.functional.softmax(student_logits, dim=0)
+            # teacher_likelihood = nn.functional.softmax(teacher_logits, dim=0)
+            student_likelihood = torch.tensor(student_logits)
+            teacher_likelihood = torch.tensor(teacher_logits)
+            d_loss = self.loss_mse(student_likelihood.view(-1), teacher_likelihood.view(-1))
+            nll_loss = self.loss_mse(teacher_likelihood.view(-1), labels.view(-1))
+        else:
+            assert output_mode in ["classification", "regression"]
+            d_loss = 0.
+            nll_loss = 0.
+        tol_loss = self.alpha * d_loss + (1 - self.alpha) * nll_loss
+        return tol_loss, d_loss, nll_loss
+
+    def train(
+        self,
+        resume_from_checkpoint: Optional[str] = None,
+        trial: Optional[Dict[str, Any]] = None,
+        ignore_keys_for_eval: Optional[List[str]] = None,
+        **kwargs
+    ):
+        """
+        Train method with explicit configuration for knowledge distillation training.
+        
+        Args:
+            resume_from_checkpoint: Optional checkpoint to resume training
+            trial: Optional hyperparameter trial configuration
+            ignore_keys_for_eval: Keys to ignore during evaluation
+        """
+        # Prepare teacher model if exists
+        if self.teacher_model is not None:
+            self.teacher_model.to(self.args.device)
+            self.teacher_model.eval()  # Ensure teacher is in eval mode
+
+        # Call parent train method
+        return super().train(
+            resume_from_checkpoint=resume_from_checkpoint,
+            trial=trial,
+            ignore_keys_for_eval=ignore_keys_for_eval,
+            **kwargs
+        )
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None
+    ) -> torch.Tensor:
+        """
+        Perform a training step on a batch of inputs.
+
+        Subclass and override to inject custom behavior.
+
+        Args:
+            model (`nn.Module`):
+                The model to train.
+            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
+                The inputs and targets of the model.
+
+                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
+                argument `labels`. Check your model's documentation for all accepted arguments.
+
+        Return:
+            `torch.Tensor`: The tensor with training loss on this batch.
+        """
+        model.train()
+        if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
+            self.optimizer.train()
+
+        inputs = self._prepare_inputs(inputs)
+        if is_sagemaker_mp_enabled():
+            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
+            return loss_mb.reduce_mean().detach().to(self.args.device)
+
+        with self.compute_loss_context_manager():
+            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
+            
+        del inputs
+        if (
+            self.args.torch_empty_cache_steps is not None
+            and self.state.global_step % self.args.torch_empty_cache_steps == 0
+        ):
+            if is_torch_xpu_available():
+                torch.xpu.empty_cache()
+            elif is_torch_mlu_available():
+                torch.mlu.empty_cache()
+            elif is_torch_musa_available():
+                torch.musa.empty_cache()
+            elif is_torch_npu_available():
+                torch.npu.empty_cache()
+            elif is_torch_mps_available(min_version="2.0"):
+                torch.mps.empty_cache()
+            else:
+                torch.cuda.empty_cache()
+
+        kwargs = {}
+
+        # For LOMO optimizers you need to explicitly use the learnign rate
+        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
+            kwargs["learning_rate"] = self._get_learning_rate()
+
+        if self.args.n_gpu > 1:
+            loss = loss.mean()  # mean() to average on multi-gpu parallel training
+
+        if self.use_apex:
+            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
+                scaled_loss.requires_grad = True
+                scaled_loss.backward()
+            
+            if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
+                nn.utils.clip_grad_norm_(amp.master_params(self.optimizer[0]), 1.0)
+            
+            
+        else:
+            # Finally we need to normalize the loss for reporting
+            loss.requires_grad = True
+            if not self.model_accepts_loss_kwargs and self.compute_loss_func is None:
+                loss = loss / self.args.gradient_accumulation_steps
+
+            # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled
+            # https://github.com/huggingface/transformers/pull/35808
+            if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
+                kwargs["scale_wrt_gas"] = False
+
+            self.accelerator.backward(loss, **kwargs)
+            
+            if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
+                # nn.utils.clip_grad_norm_(student_model.parameters(), 1.0)
+                nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+
+            return loss.detach()
+
+    def evaluate(
+        self,
+        eval_dataset: Optional[Dataset] = None,
+        ignore_keys: Optional[List[str]] = None,
+        metric_key_prefix: str = "eval",
+        **kwargs
+    ) -> Dict[str, float]:
+        """
+        Evaluation method with custom metrics computation.
+        
+        Args:
+            eval_dataset: Optional evaluation dataset
+            ignore_keys: Keys to ignore during evaluation
+            metric_key_prefix: Prefix for metrics
+        
+        Returns:
+            Dictionary of evaluation metrics
+        """
+        # Use parent's evaluate method with optional customizations
+        return super().evaluate(
+            eval_dataset=eval_dataset,
+            ignore_keys=ignore_keys,
+            metric_key_prefix=metric_key_prefix,
+            **kwargs
+        )
+    
+    def prediction_step(
+        self,
+        model: nn.Module,
+        inputs: Dict[str, Union[torch.Tensor, Any]],
+        prediction_loss_only: bool,
+        ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Override prediction step to handle the model's output format correctly.
+        """
+        has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
+        
+        return_loss = inputs.get("return_loss", None)
+        if return_loss is None:
+            return_loss = self.can_return_loss
+        loss_without_labels = True if len(self.label_names) == 0 and return_loss else False
+        
+
+        inputs = self._prepare_inputs(inputs)
+        if ignore_keys is None:
+            if hasattr(self.model, "config"):
+                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
+            else:
+                ignore_keys = []
+                
+        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
+        if has_labels or loss_without_labels:
+            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
+            if len(labels) == 1:
+                labels = labels[0]
+        else:
+            labels = None
+        
+        with torch.no_grad():
+            loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+            loss = loss.mean().detach()
+
+            # Get logits from outputs
+            if isinstance(outputs, dict):
+                logits = outputs["logits"]
+            else:
+                # logits = outputs[0]
+                logits = outputs
+            
+            
+            # Ensure logits has correct shape [batch_size, num_classes]
+            if len(logits.shape) == 1:
+                logits = logits.unsqueeze(0)
+
+        if prediction_loss_only:
+            return (loss, None, None)
+        
+        if labels is not None:
+            labels = labels.detach()
+        
+        logits = nested_detach(logits)
+        if len(logits.shape) == 1:
+            logits = logits[0]
+
+        # print(f"Validation loss: {loss}")
+        return (loss, logits, labels)

template_FL/src/fedllm/trainer.py

@@ -3,14 +3,48 @@ from torch.utils.data import DataLoader
 import torch
 import copy
 import numpy as np
-from transformers import BertForSequenceClassification, GenerationConfig, AutoTokenizer
+from transformers import (
+    # BertForSequenceClassification, 
+    GenerationConfig, 
+    AutoTokenizer,
+    Trainer,
+    get_scheduler, 
+    EarlyStoppingCallback,
+    TrainingArguments,
+    DataCollatorWithPadding
+)
+from datasets import Dataset
+from .skipbert.trainer import compute_metrics_skipbert, SkipBertTrainer
+
 import inspect
 import logging
 import wandb
 from tqdm import tqdm
+import time
+from functools import partial
+
+logging.getLogger("Trainer").setLevel(logging.INFO)
 
 logger = logging.getLogger(__name__)
 
+
+def time_format(runtime, logger):
+    if runtime < 60:
+        logger.info(f'Runtime: {runtime:.2f} seconds')
+    elif runtime < 3600:  # Less than one hour
+        minutes = runtime / 60
+        logger.info(f'Runtime: {minutes:.2f} minutes')
+    else:
+        hours = runtime / 3600
+        logger.info(f'Runtime: {hours:.2f} hours')
+
+
+
+def convert_to_tokens_reg(data, tokenizer, max_seq_length, device):
+    input_tokenzied = tokenizer(data['text'], truncation=True, padding=True, max_length=max_seq_length, return_tensors="pt")
+    input_tokenzied['labels'] = torch.tensor(data['label'], dtype=torch.float32).reshape(-1, 1)
+    return input_tokenzied
+
 class ManualLLMSampleCB:
     def __init__(self, model, tokenizer, task, num_samples=10, max_new_tokens=256):
         self.model = model
@@ -61,12 +95,14 @@ class ManualLLMSampleCB:
     def log_samples_to_wandb(self, dataset):
         samples_table = self.create_samples_table(dataset)
         wandb.log({"sample_predictions": samples_table})
-
+        
+        
+        
 
 class ManualTrainer:
     def __init__(
         self, model, tokenizer, train_dataset, val_dataset, holdout_dataset, reference_dataset,
-        args, data_collator, compute_metrics, mates_args, data_influence_model, data_influence_tokenizer
+        args, data_collator, compute_metrics, mates_cfg, skipbert_cfg, teacher_data_influence_model, student_data_influence_model, data_influence_tokenizer
     ):
         self.accelerator = Accelerator()
         self.model = model
@@ -74,10 +110,13 @@ class ManualTrainer:
         self.args = args
         self.data_collator = data_collator
         self.compute_metrics = compute_metrics
-        self.mates_args = mates_args
-        self.data_influence_model = data_influence_model
+        self.mates_cfg = mates_cfg
+        self.skipbert_cfg = skipbert_cfg
+        self.teacher_data_influence_model = teacher_data_influence_model
+        self.student_data_influence_model = student_data_influence_model
         self.data_influence_tokenizer = data_influence_tokenizer
-
+        
+        
         # Remove unused columns from datasets
         if train_dataset:
             self.train_dataset = self._remove_unused_columns(train_dataset, "training")     
@@ -105,13 +144,13 @@ class ManualTrainer:
         else:
             self.val_loader = None
 
-        if self.mates_args.state:
+        if self.mates_cfg.state:
             self.holdout_dataset = self._remove_unused_columns(holdout_dataset, "holdout")
             self.reference_dataset = self._remove_unused_columns(reference_dataset, "reference")
 
             self.holdout_loader = DataLoader(
                 self.holdout_dataset,
-                batch_size=self.mates_args.holdout_batch_size,
+                batch_size=self.mates_cfg.holdout_batch_size,
                 shuffle=True,
                 collate_fn=self.data_collator,
                 drop_last=self.args.dataloader_drop_last
@@ -119,7 +158,7 @@ class ManualTrainer:
 
             self.reference_loader = DataLoader(
                 self.reference_dataset,
-                batch_size=self.mates_args.reference_batch_size,
+                batch_size=self.mates_cfg.reference_batch_size,
                 shuffle=False,
                 collate_fn=self.data_collator,
                 drop_last=self.args.dataloader_drop_last
@@ -136,11 +175,56 @@ class ManualTrainer:
             self.model, self.optimizer, self.full_train_loader, self.val_loader
         )
 
-        if self.mates_args.state:
+        ### Define for MATEs ###
+        if self.mates_cfg.state:
             # Prepare holdout and reference loaders for Accelerator
-            self.data_influence_model, self.holdout_loader, self.reference_loader = self.accelerator.prepare(
-                self.data_influence_model, self.holdout_loader, self.reference_loader
+            self.teacher_data_influence_model, self.holdout_loader, self.reference_loader = self.accelerator.prepare(
+                self.teacher_data_influence_model, self.holdout_loader, self.reference_loader
             )
+            self.student_data_influence_model = self.accelerator.prepare(self.student_data_influence_model)
+        ######
+            
+        ### Define for SkipBERT ###
+        self.skipbert_train_args = TrainingArguments(
+            output_dir=self.skipbert_cfg.output_dir,
+            learning_rate=self.skipbert_cfg.learning_rate,
+            num_train_epochs=self.skipbert_cfg.num_train_epochs,
+            per_device_train_batch_size=self.skipbert_cfg.train_batch_size,
+            gradient_accumulation_steps=self.skipbert_cfg.gradient_accumulation_steps,
+            per_device_eval_batch_size=self.skipbert_cfg.eval_batch_size,
+            eval_accumulation_steps=self.skipbert_cfg.eval_accumulation_steps,
+            max_steps=self.skipbert_cfg.max_steps,
+            logging_steps = 10,
+            evaluation_strategy=self.skipbert_cfg.evaluation_strategy,
+            save_strategy=self.skipbert_cfg.save_strategy,
+            lr_scheduler_type=self.skipbert_cfg.lr_scheduler_type,
+            warmup_steps=self.skipbert_cfg.warmup_steps,
+            weight_decay=self.skipbert_cfg.weight_decay,
+            logging_dir=self.skipbert_cfg.logging_dir,
+            report_to='wandb',
+            run_name='skipbert',
+            do_train=self.skipbert_cfg.do_train,
+            do_eval=self.skipbert_cfg.do_eval,
+            dataloader_drop_last=False,
+            ddp_find_unused_parameters=False,
+            group_by_length=True,
+            load_best_model_at_end = True
+        )
+        
+        # Prepare custom optimizer student model's parameters
+        if self.student_data_influence_model is not None:
+            no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+            self.student_optimizer_grouped_parameters = [
+                {
+                    'params': [p for n, p in self.student_data_influence_model.named_parameters() if not any(nd in n for nd in no_decay)], 
+                    'weight_decay': 0.01
+                },
+                {
+                    'params': [p for n, p in self.student_data_influence_model.named_parameters() if any(nd in n for nd in no_decay)], 
+                    'weight_decay': 0.0
+                }
+            ]
+        ######
 
     def _remove_unused_columns(self, dataset, description=None):
         """
@@ -188,15 +272,15 @@ class ManualTrainer:
 
         for epoch in range(self.args.num_train_epochs):
             # Check if it's time to update the data influence model and state is True
-            if self.mates_args.state and epoch % self.mates_args.update_data_influence_model_step == 0:
+            if self.mates_cfg.state and epoch % self.mates_cfg.update_data_influence_model_step == 0:
                 print("Updating the data influence model and selecting high-quality data...")
-                logger.info("Updating the data influence model and selecting high-quality data...")
+                # logger.info("Updating the data influence model and selecting high-quality data...")
                 self.update_data_influence_model()
 
             # Filter high-quality data using the data influence model
             high_quality_indices = self.select_high_quality_data(
                 dataset_size=len(self.train_dataset),
-                selection_fraction=self.mates_args.selection_fraction,
+                selection_fraction=self.mates_cfg.selection_fraction,
             )
             self.train_loader = self.accelerator.prepare(
                 self.create_filtered_dataloader(high_quality_indices)
@@ -224,16 +308,16 @@ class ManualTrainer:
                 epoch_loss += loss.item()
 
                 if (step + 1) % self.args.logging_steps == 0:
-                    # print(f"Step {step + 1}: Train Loss = {epoch_loss / (step + 1):.4f}")
-                    logger.info(f"Step {step + 1}: Train Loss = {epoch_loss / (step + 1):.4f}")
+                    print(f"Step {step + 1}: Train Loss = {epoch_loss / (step + 1):.4f}")
+                    # logger.info(f"Step {step + 1}: Train Loss = {epoch_loss / (step + 1):.4f}")
 
             avg_epoch_loss = epoch_loss / len(self.train_loader)
             training_loss.append(avg_epoch_loss)
 
             val_results = self.evaluate()
 
-            # print(f"Epoch {epoch + 1}: Train Loss = {avg_epoch_loss:.4f}, Val Loss = {val_results['eval_loss']:.4f}")
-            logger,info(f"Epoch {epoch + 1}: Train Loss = {avg_epoch_loss:.4f}, Val Loss = {val_results['eval_loss']:.4f}")
+            print(f"Epoch {epoch + 1}: Train Loss = {avg_epoch_loss:.4f}, Val Loss = {val_results['eval_loss']:.4f}")
+            # logger.info(f"Epoch {epoch + 1}: Train Loss = {avg_epoch_loss:.4f}, Val Loss = {val_results['eval_loss']:.4f}")
 
             # Early stopping logic
             if val_results["eval_loss"] < best_val_loss:
@@ -243,6 +327,7 @@ class ManualTrainer:
                 early_stopping_counter += 1
                 if early_stopping_counter >= early_stopping_patience:
                     print("Early stopping triggered")
+                    # logger.info("Early stopping triggered")
                     break
             
         return {"training_loss": sum(training_loss) / len(training_loss), "best_val_loss": best_val_loss}
@@ -252,14 +337,20 @@ class ManualTrainer:
         Use the data influence model to predict quality scores and select high-quality data indices.
         """
         print("Selecting high-quality data using the data influence model...")
+        # logger.info("Selecting high-quality data using the data influence model...")
 
         # Predict influence scores for the entire dataset
         influence_scores = []
-        self.data_influence_model.eval()
+        self.student_data_influence_model.eval()
         influence_optimizer = self.accelerator.prepare(
-            torch.optim.AdamW(self.data_influence_model.parameters(), lr=self.args.learning_rate)
+            torch.optim.AdamW(
+                self.student_optimizer_grouped_parameters,
+                lr=self.skipbert_train_args.learning_rate,
+            )
         )
         i = 0
+        
+        start_time = time.perf_counter()
         with torch.no_grad():
             for batch in self.full_train_loader:  # Full dataset loader        
                 text = self.tokenizer.batch_decode(
@@ -278,12 +369,13 @@ class ManualTrainer:
 
                 # Train the data influence model
                 influence_optimizer.zero_grad()
-                outputs = self.data_influence_model(
+                logits, attn_outputs, hidn_output = self.student_data_influence_model(
                     input_ids=bert_inputs['input_ids'],
                     attention_mask=bert_inputs['attention_mask'],
                 )
                 
-                influence_scores.extend(outputs.logits.squeeze(-1).cpu().numpy())
+                
+                influence_scores.extend(logits.squeeze(-1).cpu().numpy())
                 
                 i += 1
                 
@@ -293,6 +385,7 @@ class ManualTrainer:
         # Normalize influence scores and apply Gumbel-Top-$k$ selection
         influence_scores = np.array(influence_scores)
         print(">> Influence scores shape:", influence_scores.shape)
+        # logger.info(">> Influence scores shape:", influence_scores.shape)
 
         # Add Gumbel noise for diversity
         rng = np.random.default_rng()
@@ -303,7 +396,12 @@ class ManualTrainer:
         selection_size = int(len(influence_scores)*selection_fraction)
         high_quality_indices = np.argpartition(-influence_scores, selection_size)[:selection_size]
         print(f"Selected {len(high_quality_indices)} high-quality samples.")
-
+        # logger.info(f"Selected {len(high_quality_indices)} high-quality samples.")
+        
+        end_time = time.perf_counter()
+        runtime = round((end_time - start_time), 2)
+        time_format(runtime, logger)
+        
         return high_quality_indices
 
     def create_filtered_dataloader(self, indices):
@@ -311,6 +409,7 @@ class ManualTrainer:
         Create a new dataloader with only the selected high-quality data.
         """
         print("Creating a filtered dataloader with selected high-quality data...")
+        # logger.info("Creating a filtered dataloader with selected high-quality data...")
         subset_dataset = torch.utils.data.Subset(self.train_dataset, indices)
         return torch.utils.data.DataLoader(
             subset_dataset,
@@ -325,16 +424,18 @@ class ManualTrainer:
         # Train a copy of the model on holdout data and validate on reference data
         copied_model = copy.deepcopy(self.model)
         copied_model.train()
+        self.accelerator.state._reset_state()
+        self.accelerator = Accelerator()
         optimizer = self.accelerator.prepare(
             torch.optim.Adam(copied_model.parameters(), lr=self.args.learning_rate)
         )
         holdout_reference_pairs = []
 
-        # print("Starting to collect holdout-reference pairs...")
-        logger.info("Starting to collect holdout-reference pairs...")
+        print("Starting to collect holdout-reference pairs...")
+        # logger.info("Starting to collect holdout-reference pairs...")
         for step, holdout_batch in enumerate(self.holdout_loader):
-            # print(f"Processing holdout batch {step+1}/{len(self.holdout_loader)}...")
-            logger.info(f"Processing holdout batch {step+1}/{len(self.holdout_loader)}...")
+            print(f"Processing holdout batch {step+1}/{len(self.holdout_loader)}...")
+            # logger.info(f"Processing holdout batch {step+1}/{len(self.holdout_loader)}...")
             
             optimizer.zero_grad()
             outputs = copied_model(
@@ -352,7 +453,7 @@ class ManualTrainer:
             optimizer.step()
 
             print(f"Evaluating reference losses at step {step}...")
-            logger.info(f"Evaluating reference losses at step {step}...")
+            # logger.info(f"Evaluating reference losses at step {step}...")
             
             copied_model.eval()
             reference_losses = []
@@ -373,42 +474,138 @@ class ManualTrainer:
 
         # Train the data influence model using the generated pairs
         print("Starting to train the data influence model...")
-        logger.info("Starting to train the data influence model...")
+        # logger.info("Starting to train the data influence model...")
         
-        self.data_influence_model.train()
-        influence_optimizer = torch.optim.AdamW(self.data_influence_model.parameters(), lr=self.args.learning_rate)
-
-        for step, (text, score) in enumerate(holdout_reference_pairs):
+        self.teacher_data_influence_model.train()
+        influence_optimizer = torch.optim.AdamW(self.teacher_data_influence_model.parameters(), lr=self.args.learning_rate)
+        
+        list_texts, list_score = [], []
+        batch_size = 0
+        # Convert to Dataset objective
+        for texts, score in holdout_reference_pairs:
+            if batch_size == 0:
+                batch_size = len(texts)
+            list_texts.extend(texts)
+            list_score.extend([score] * len(texts))
+        
+        
+        holdout_reference_pairs = {'text': list_texts, 'label': list_score}
+        holdout_reference_pairs = Dataset.from_dict(holdout_reference_pairs)
+        
+        # Wrap the function with partial
+        convert_func = partial(
+            convert_to_tokens_reg,
+            tokenizer=self.data_influence_tokenizer,
+            max_seq_length=self.skipbert_cfg.max_seq_length,
+            device=self.accelerator.device
+        )
+        holdout_reference_pairs_loader = DataLoader(
+            holdout_reference_pairs.map(
+                convert_func,
+                batched=True,
+                num_proc=8,
+                remove_columns=holdout_reference_pairs.column_names
+            ), 
+            batch_size=batch_size,
+            collate_fn=DataCollatorWithPadding(tokenizer=self.data_influence_tokenizer, padding=True, max_length=self.skipbert_cfg.max_seq_length),  # Use the same collate function
+            drop_last=self.args.dataloader_drop_last
+        )
+        loss_mse = torch.nn.MSELoss()
+        
+        for step, batch_input in enumerate(holdout_reference_pairs_loader):
             # Tokenize the text using the BERT tokenizer
-            bert_inputs = self.data_influence_tokenizer(
-                text,
-                truncation=True,
-                padding='max_length',
-                max_length=256,
-                return_tensors='pt'
-            ).to(self.accelerator.device)
-
-            # Convert score to tensor and enable gradients
-            score_tensor = torch.tensor([score], device=self.accelerator.device, dtype=torch.float32, requires_grad=True)
+            batch_input = {k: v.to('cuda:0') for k, v in batch_input.items()} # cuda:0
+#             text, score = row['text'], row['label']
+#             bert_inputs = self.data_influence_tokenizer(
+#                 text,
+#                 truncation=True,
+#                 padding='max_length',
+#                 max_length=256,
+#                 return_tensors='pt'
+#             ).to(self.accelerator.device)
+
+#             # Convert score to tensor and enable gradients
+#             score_tensor = torch.tensor([score] * len(text), device=self.accelerator.device, dtype=torch.float32, requires_grad=True)
             
             # Train the data influence model
             influence_optimizer.zero_grad()
-            outputs = self.data_influence_model(
-                input_ids=bert_inputs['input_ids'],
-                attention_mask=bert_inputs['attention_mask'],
-                labels=score_tensor
+#             outputs = self.teacher_data_influence_model(
+#                 input_ids=bert_inputs['input_ids'],
+#                 attention_mask=bert_inputs['attention_mask'],
+#                 labels=score_tensor
+#             )
+
+            outputs = self.teacher_data_influence_model(
+                **batch_input
             )
-            influence_loss = outputs.loss
-
+            
+            influence_loss = loss_mse(batch_input['labels'].view(-1), outputs.logits.view(-1))
+            # print(f"Loss: {influence_loss} - require_grad: {influence_loss.grad_fn}")
+            
+            influence_loss.requires_grad = True
             influence_loss.backward()
             influence_optimizer.step()
-
+            
             if step % 50 == 0:
                 print(f"[Influence Training] Step {step}: Loss = {influence_loss.item():.4f}")
-                logger.info(f"[Influence Training] Step {step}: Loss = {influence_loss.item():.4f}")
+                # logger.info(f"[Influence Training] Step {step}: Loss = {influence_loss.item():.4f}")
+        
         
+        ### Distillation for SkipBERT ###
+        train_converted = holdout_reference_pairs.map(
+            convert_func,
+            batched=True,
+            num_proc=8,
+            remove_columns=holdout_reference_pairs.column_names
+        )
         
-        # Distillation for SkipBERT
+
+        # Call parent constructor with custom optimizer
+        optimizer = torch.optim.AdamW(
+            self.student_optimizer_grouped_parameters, 
+            lr=self.skipbert_train_args.learning_rate,
+        )
+
+        scheduler = get_scheduler(
+            name=self.skipbert_train_args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=self.skipbert_train_args.warmup_steps,
+            # num_training_steps=training_args.max_steps
+            num_training_steps=100/(self.skipbert_train_args.per_device_train_batch_size * self.skipbert_train_args.gradient_accumulation_steps)
+        )
+        
+        # Initialize the trainer
+        trainer = SkipBertTrainer(
+            student_model=self.student_data_influence_model,
+            teacher_model=self.teacher_data_influence_model,
+            args=self.skipbert_train_args,
+            train_dataset=train_converted,
+            eval_dataset=train_converted.shuffle().select(range(min(len(train_converted),10))),
+            compute_metrics=compute_metrics_skipbert,
+            # SkipBERT specific arguments
+            alpha=0.5,
+            temperature=2.0,
+            beta=1.0,
+            use_logits=self.skipbert_cfg.use_logits,
+            use_att=self.skipbert_cfg.use_att,
+            use_rep=self.skipbert_cfg.use_rep,
+            use_embedding=self.skipbert_cfg.use_embedding,
+            att_layer_maps=self.skipbert_cfg.att_layer_maps,
+            hid_layer_maps=self.skipbert_cfg.hid_layer_maps,
+            epochs_no_cls=self.skipbert_cfg.epochs_no_cls,
+            reduce_T=self.skipbert_cfg.reduce_T,
+            output_mode=self.skipbert_cfg.output_mode, # 'classification' or 'regression'
+            num_masked_layers_teacher=self.skipbert_cfg.num_masked_layers_teacher,
+            num_masked_last_layers_teacher=self.skipbert_cfg.num_masked_last_layers_teacher,
+            fp16=self.skipbert_cfg.fp16,
+            num_full_hidden_layers_student=self.skipbert_cfg.num_full_hidden_layers_student,
+            tokenizer=self.data_influence_tokenizer,
+            optimizers=(optimizer,scheduler),
+            callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
+        )
+
+        # Train the model
+        trainer.train()
         
 
 
@@ -459,7 +656,8 @@ class ManualTrainer:
         metrics = self.compute_metrics({"predictions": padded_preds, "label_ids": padded_labels})
 
         metrics.update({"eval_loss": val_loss / len(self.val_loader)})
-        print("Validation Metrics:", metrics)
+        print(f"Validation Metrics: {metrics}")
+        # logger.info(f"Validation Metrics: {metrics}")
 
         if wandb_sample:
             # Sample Logging

template_FL/src/pyproject.toml

@@ -37,14 +37,14 @@ num-server-rounds = 2
 num-supernodes = 10
 
 # Define dataset
-dataset.type = 'hete' # type = ['homo','hete']
+dataset.type = 'homo' # type = ['homo','hete']
 dataset.name = "vicgalle/alpaca-gpt4"
 
 # Define model settings
 model.name = "Qwen/Qwen2.5-1.5B-Instruct" 
 model.quantization = 4
 model.gradient-checkpointing = true
-model.flash_attention = false
+model.flash-attention = false
 
 ### Use MATES ###
 mates.state = true
@@ -60,67 +60,66 @@ mates.selection-fraction = 0.4
 
 # Model setting
 skipbert.student-model = "bert-base-uncased"
-skipbert.num_layers_student = 12
-skipbert.num_full_hidden_layers_student = 6
-skipbert.num_masked_layers_teacher = 0
-skipbert.num_masked_last_layers_teacher = 0
+skipbert.output-mode = "regression"
+skipbert.num-layers-student = 12
+skipbert.num-full-hidden-layers-student = 6
+skipbert.num-masked-layers-teacher = 0
+skipbert.num-masked-last-layers-teacher = 0
 
 # Training hyperparameters
-skipbert.train_batch_size = 8
-skipbert.gradient_accumulation_steps = 2
-skipbert.eval_batch_size = 8
-skipbert.eval_accumulation_steps = 2
-skipbert.learning_rate = 2.0e-5
-skipbert.num_train_epochs = 10
-skipbert.eval_step = 10
-skipbert.max_seq_length = 128
-skipbert.weight_decay = 1.0e-4
-skipbert.warmup_steps = 100 # 500
-skipbert.do_train = true
-skipbert.do_eval = true
-skipbert.max_steps = -1
-skipbert.evaluation_strategy = "epoch"
-skipbert.save_strategy = "epoch"
-skipbert.lr_scheduler_type = "cosine" # or 'warmup_linear'
-skipbert.logging_dir = './skipbert_logs'
-skipbert.output_dir = "./skipbert_results"
-skipbert.report_to = 'wandb'
+skipbert.train-batch-size = 4
+skipbert.gradient-accumulation-steps = 1
+skipbert.eval-batch-size = 4
+skipbert.eval-accumulation-steps = 1
+skipbert.learning-rate = 2.0e-5
+skipbert.num-train-epochs = 10
+skipbert.eval-step = 10
+skipbert.max-seq-length = 256
+skipbert.weight-decay = 1.0e-4
+skipbert.warmup-steps = 100 # 500
+skipbert.do-train = true
+skipbert.do-eval = true
+skipbert.do-predict = false
+skipbert.max-steps = -1
+skipbert.evaluation-strategy = "epoch"
+skipbert.save-strategy = "epoch"
+skipbert.lr-scheduler-type = "cosine" # or 'warmup-linear'
+skipbert.logging-dir = './skipbert-logs'
+skipbert.output-dir = "./skipbert-results"
+skipbert.report-to = 'wandb'
 
 # Knowledge distillation parameters
 skipbert.beta = 0.01
 skipbert.T = 1.0
 skipbert.alpha = 1.0
-skipbert.reduce_T = 1.0
-skipbert.epochs_no_cls = 5
+skipbert.reduce-T = 1.0
+skipbert.epochs-no-cls = 5
 
 # Training schedule and features
-skipbert.freeze_lower_layers = true
+skipbert.freeze-lower-layers = true
 
 # Feature usage flags
-skipbert.use_logits = true
-skipbert.use_att = true
-skipbert.use_rep = true
-skipbert.use_embedding = false
+skipbert.use-logits = true
+skipbert.use-att = true
+skipbert.use-rep = true
+skipbert.use-embedding = false
 
 # Training modes
-skipbert.do_train = true
-skipbert.do_eval = true
-skipbert.do_predict = false
-skipbert.do_fit = false
+skipbert.do-fit = false
 skipbert.fp16 = false
-skipbert.no_pretrain = false
-skipbert.use_init_weight = false
-skipbert.share_param = true
-skipbert.do_lower_case = true
-skipbert.no_cuda = false
+skipbert.no-pretrain = false
+skipbert.use-init-weight = false
+skipbert.share-param = true
+skipbert.do-lower-case = true
+skipbert.no-cuda = false
 
 # N-gram settings
-skipbert.n_gram_left = 1
-skipbert.n_gram_right = 1
+skipbert.n-gram-left = 1
+skipbert.n-gram-right = 1
 
 # Layer mappings
-skipbert.att_layer_maps: [1, 3, 5, 7, 9, 11]
-skipbert.hid_layer_maps: [6, 7, 8, 9, 10, 11, 12]
+skipbert.att-layer-maps = "1, 3, 5, 7, 9, 11"
+skipbert.hid-layer-maps = "6, 7, 8, 9, 10, 11, 12"
 
 ### END ###
 
@@ -138,8 +137,8 @@ train.save-every-round = 5
 train.learning-rate-max = 5e-5
 train.learning-rate-min = 1e-6
 train.seq-length = 256
-train.prompt_template_name = "alpaca"
-train.train_on_inputs = true
+train.prompt-template-name = "alpaca"
+train.train-on-inputs = true
 train.verbose = false
 
 # Define training agruments for HF Trainer
@@ -164,7 +163,7 @@ train.training-arguments.eval-strategy = "epoch"
 train.training-arguments.save-strategy = "epoch"
 train.training-arguments.ddp-find-unused-parameters = false
 train.training-arguments.group-by-length = true
-train.training-arguments.load_best_model_at_end = true
+train.training-arguments.load-best-model-at-end = true
 train.training-arguments.report-to = "wandb"
 
 # Define local training settings