pax_global_header00006660000000000000000000000064150125303170014507gustar00rootroot0000000000000052 comment=70b86af483841786e1cbdc1fcac9691f25fef3ce f4llm-main/000077500000000000000000000000001501253031700130135ustar00rootroot00000000000000f4llm-main/.gitignore000066400000000000000000000061031501253031700150030ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ cover/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder .pybuilder/ target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: # .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock # pdm # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. #pdm.lock # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it # in version control. # https://pdm.fming.dev/latest/usage/project/#working-with-version-control .pdm.toml .pdm-python .pdm-build/ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ f4llm-main/BaseEngine.py000066400000000000000000000166571501253031700154040ustar00rootroot00000000000000 import os import copy from visualization.visualizer import Visualizer, VisCallback from utils.register import registry from datas.base_data import FedBaseDataset from utils.general import is_petuning, metric_save from utils.serialization import SerializationTool class BaseEngine(object): def __init__(self, *args): """ Base Federated Trainer for all trainers. The trainer is run on the server or client side. It is responsible for building the dataset, model, metric, and communicators, and running the training process. Attributes: M (configs.ModelArguments): Model configuration D (configs.DataArguments): Data configuration T (configs.TrainingArguments): Training configuration F (configs.FLArguments): Federated Learning configuration is_fl (bool): Whether to use federated learning role (str): Role of the trainer, either "server" or "client" client_num (int): Number of clients param_list (list): Used for global update loss_list (list): Used for loss-aware aggregate round (int): Current training round phase (str): Phase of the training process logger (Logger): Logger object debug (bool): Whether to enable debug mode """ config = registry.get("config") self.M = config.M self.D = config.D self.T = config.T self.F = config.F self.is_fl = config.is_fl self.client_num = len(config.F.clients_id_list) self.param_list = [] # used for global update self.loss_list = [] # used for loss-aware aggregate self.round = 0 self.phase = registry.get("phase") self.logger = registry.get("logger") self.debug = registry.get("debug") def _build_data(self): self.logger.info(f"{self.role} building dataset ...") self.data = registry.get_data_class(self.D.data_name)() self.data.load_data() def _build_model(self): self.logger.info(f"{self.role} building model ...") self.model = registry.get_model_class(self.M.model_type)( task_name=self.D.task_name ).build_model() self.model_parameters = self.serialize_model_parameters() def _build_metric(self): self.logger.info(f"Metric name: {self.T.metric_name.upper()}, " f"is_decreased_valid_metric: " f"{self.T.is_decreased_valid_metric}") self.metric = registry.get_metric_class(self.T.metric_name)( self.data.tokenizer, self.T.is_decreased_valid_metric, self.T.save_outputs ) self.metric_name = self.metric.metric_name # build federated eval-function' args self._build_eval_mode() # build metric_log_and_line self._build_general_metric_logs() def _build_general_metric_logs(self): self.metric_log = { "model_type": self.M.model_type, "tuning_type": self.M.tuning_type, "clients_num": self.F.clients_num, "alpha": self.F.alpha, "task": self.D.task_name, "fl_algorithm": self.F.fl_algorithm, "data": self.D.data_name, "info": registry.get("metric_line")[0:-1], "train_logs": [], } # metric line self.metric_line = registry.get("metric_line") self._build_metric_logs() def _build_metric_logs(self): self.metric_log["eval_logs"] = {} self.global_test_best_metric = 0.0 self.global_valid_best_metric = \ float("inf") if self.T.is_decreased_valid_metric else -float("inf") def _build_eval_mode(self): self.eval_args = copy.deepcopy(self.T) self.eval_args.evaluation_strategy = "epoch" self.eval_args.predict_with_generate = True def _build_visualizer(self): if self.T.visualization is not None: self.visualizer = Visualizer( project=self.T.visualization.project, trial_name=self.T.visualization.trial, role=self.role, client_name=self.F.client_name, phase=self.phase, # name=registry.get('metric_line')[0:-1], key=self.T.visualization.key, config={ "seed": self.T.seed, "lr": self.T.learning_rate, "fl_algorithm": self.F.fl_algorithm, "clients_num": self.F.client_num_in_total, "total_round": self.F.rounds } ) self.logger.info(f"build visualizer: {self.T.visualization.project}({registry.get('metric_line')[0:-1]})") else: self.visualizer = None def run(self): raise NotImplementedError def get_dataset(self, client_id, mode="train"): """Get :class:`Dataset` for ``client_id``.""" train_dataset = self.data.train_dataset_dict[client_id] if mode == "train": valid_dataset_dict = self.data.valid_dataset_dict else: valid_dataset_dict = self.data.test_dataset_dict if client_id == -1 or self.F.pson: eval_dataset = valid_dataset_dict[client_id] else: eval_dataset = None if self.debug: if isinstance(train_dataset, FedBaseDataset): max_train_samples = min(len(train_dataset), self.D.max_train_samples) else: max_train_samples = range(min(len(train_dataset), self.D.max_train_samples)) train_dataset = train_dataset.select(max_train_samples) if eval_dataset is not None: max_eval_samples = min(len(eval_dataset), self.D.max_eval_samples) eval_dataset = eval_dataset.select(max_eval_samples) else: max_eval_samples = 0 self.logger.warning(f"Debug Mode Enable, Train Num: {max_train_samples}, " f"Eval Num: {max_eval_samples}") return train_dataset, eval_dataset def serialize_model_parameters(self): if is_petuning(self.M.tuning_type): model_parameters = SerializationTool.serialize_peft_model(self.model, self.M.tuning_type) else: model_parameters = SerializationTool.serialize_model(self.model) return model_parameters def deserialize_model(self, serialized_parameters): if is_petuning(self.M.tuning_type): SerializationTool.deserialize_peft_model(self.model, serialized_parameters, self.M.tuning_type) else: SerializationTool.deserialize_model(self.model, serialized_parameters) def metric_save(self): if self.debug: return metric_save(self, self.T, self.logger) def build_eval_cmd(self): base_opts = [ "main.py", "--do_eval", "--raw_dataset_path", self.D.raw_dataset_path, "--partition_dataset_path", self.D.partition_dataset_path, "--model_name_or_path", self.M.model_name_or_path, "--output_dir", os.path.dirname(self.T.output_dir), "--model_type", self.M.model_type, "--task_name", self.D.task_name, "--fl_algorithm", self.F.fl_algorithm, "--config_path", self.T.config_path, "--role", "client", "--times", self.T.times ] return base_opts @property def role(self): return self.F.role f4llm-main/FedBaseTrainer.py000066400000000000000000000342741501253031700162150ustar00rootroot00000000000000import os import copy import time import random import subprocess from visualization.visualizer import Visualizer, VisCallback from commus.message import Message from commus.communicator import gRPCCommunicationManager from contribs.centralized.miscs import CenEndEvalStepCallback from utils.register import registry from trainers.BaseEngine import BaseEngine from utils.general import metric_save, model_save from utils.general import setup_seed, cosine_learning_rate, LoadBalanceSampling, ClientSampling class BaseTrainer(BaseEngine): def __init__(self, *args): super().__init__(*args) def _build_selections(self): self.selections = ClientSampling( range(self.F.client_num_in_total), self.F.client_num_per_round, self.F.rounds, self.F.client_sample_type) def _build_communicators(self): if self.is_fl: self.logger.info(f"{self.role} building communicators ...") if self.role == "server": self.logger.debug(f"server build communicator") self.comm_manager = gRPCCommunicationManager( ip=self.F.server_ip, port=self.F.server_port, max_connection_num=self.F.num_sub ) self.logger.debug(f"server finish building communicator") else: time.sleep(random.randint(2, 10)) # wait for server self.logger.debug(f"subserver {self.F.client_name} build communicator") self.comm_manager = gRPCCommunicationManager( ip=self.F.client_ip, port=self.F.client_port, max_connection_num=1, ) self.comm_manager.add_communicator( communicator_id=self.F.server_ip, communicator_address='{}:{}'.format(self.F.server_ip, self.F.server_port) ) else: self.logger.info("local or central training") def _before_training(self): # set seed setup_seed(self.T.seed) # build dataset and dataloader self._build_data() # build federated model self._build_model() # build metric self._build_metric() # return computer metric # global model self.best_glo_params = self.serialize_model_parameters() # build client selection before building loc trainer self._build_selections() # build communicators self._build_communicators() # build visualizer self._build_visualizer() def run(self): """ Run the trainer according to the phase """ self.logger.critical(f" {self.role.upper()} {self.phase.upper()} START") if self.is_fl: self.server_run() if self.role == "server" else self.client_run() else: self.cen_train() self.on_cen_end() def server_join(self): client_num = 0 while client_num < self.F.num_sub: msg = self.comm_manager.receive() if msg.message_type == 100: client_num += 1 self.logger.info(f"Subserver {msg.content['client_ip']}:{msg.content['client_port']} join in.") self.comm_manager.add_communicator( communicator_id=msg.sender, communicator_address=f"{msg.content['client_ip']}:{msg.content['client_port']}") self.logger.info(f"Subserver {msg.sender} joined in ({client_num}/{self.F.num_sub}).") self.logger.info(list(self.comm_manager.communicators.keys())) self.logger.debug(f"all subserver ({len(self.comm_manager.communicators)}) connect") def server_run(self): self.server_join() self.server_valid() self.server_process() self.on_server_end() def server_process(self): while self.round < self.F.rounds: self.client_ids = self.selections[self.round] self.metric_log["train_logs"].append([0.0 for _ in range(self.F.client_num_in_total)]) self.logger.critical(f"Round {self.round + 1}/{self.F.rounds} start, Selected Clients: {self.client_ids}") balance_sampling = LoadBalanceSampling(self.client_ids, self.F.num_sub) client_ids = {} for i in range(self.F.num_sub): client_ids[i] = balance_sampling[i] self.comm_manager.send( Message( message_type=200, sender="0", receiver=list(self.comm_manager.communicators.keys()), content={ 'model': self.model_parameters, 'client_ids': client_ids, 'round': self.round } ) ) num_sub = 0 params_list, loss_list = [], [] while num_sub < self.F.num_sub: msg = self.comm_manager.receive() if msg.message_type == 200: num_sub += 1 for client_id, params in msg.content['model'].items(): params_list.append(params) loss_list.append(msg.content['loss'][client_id]) self.metric_log["train_logs"][self.round][client_id] = msg.content['loss'][client_id] # aggregation self.server_update(params_list, loss_list) def server_update(self, param_list, loss_list): assert len(param_list) <= self.F.client_num_per_round self.round += 1 should_eval, should_save = False, False if self.F.log_valid_len and self.round % self.F.log_valid_len == 0: should_eval = True if self.F.save_valid_len and self.round % self.F.save_valid_len == 0: should_save = True self.server_logging(loss_list) # Global Aggregation serialized_parameters = self.server_aggregator(param_list, loss_list) self.deserialize_model(serialized_parameters) if should_save and self.phase == "train" and not self.debug: checkpoint_file = os.path.join(self.T.checkpoint_dir, f"round-{self.round}") self.deserialize_model(serialized_parameters) model_save(self.model, self.eval_args, checkpoint_file) self.logger.debug(f"Model Saved in: {checkpoint_file}") registry.register("round", self.round) self.model_parameters = copy.deepcopy(serialized_parameters) def server_valid(self): if self.T.eval_during_train: eval_opts = self.build_eval_cmd() eval_opts.extend(["--eval_name", "ayn_local_eval", "--not_overwrite_args", "eval_name", "--checkpoint_file", f"{self.T.checkpoint_dir}"]) eval_opts = ["python"] + eval_opts setattr(self, "eval_opts", subprocess.Popen(eval_opts)) else: setattr(self, "eval_opts", None) def server_logging(self, loss_list): this_round_loss = sum(loss_list) / len(loss_list) self.logger.warning( f"FL={self.F.fl_algorithm}_Round={self.round}_ClientNum={len(loss_list)}_" f"Loss={this_round_loss:.3f}" ) if self.visualizer is not None: self.visualizer.log( global_round=self.round, data={ "train_loss": this_round_loss, # "loss_list": loss_list, } ) def server_aggregator(self, serialized_params_list, loss_list): """fl algorithm, default fedavg""" weights = self.get_agg_weight(loss_list) serialized_parameters = self.serialize_model_parameters() for key in serialized_parameters.keys(): serialized_parameters[key] = sum( [serialized_params_list[client][key] * weights[client] for client in range(len(serialized_params_list))]) return serialized_parameters def on_server_end(self): """Using best parameters for prediction""" if not self.debug: metric_save(self, self.T, self.logger) self.comm_manager.send( Message( message_type=101, sender="0", receiver=list(self.comm_manager.communicators.keys()), content={ '': '', } ) ) if self.eval_opts is not None: while self.eval_opts.returncode != 0: continue self.logger.critical(f"Train done, Please Eval and Test in {self.T.checkpoint_dir}") def client_join(self): self.comm_manager.send( Message( message_type=100, sender=self.F.client_name, receiver=[self.F.server_ip], content={ 'client_ip': self.F.client_ip, 'client_port': self.F.client_port } ) ) self.logger.debug(f"Subserver {self.F.client_ip}:{self.F.client_port} join in federated learning") def client_run(self): # client join in federated learning self.client_join() while True: msg = self.comm_manager.receive() if msg.message_type == 101: # quit federated learning self.on_client_end() break elif msg.message_type == 200: model_parameters = msg.content['model'] self.round = msg.content['round'] client_ids = msg.content['client_ids'][int(self.F.client_name)] self.client_process(client_ids, model_parameters) self.client_valid() self.client_logging() def client_process(self, client_ids, model_parameters): param_list, loss_list = {}, {} for idx in client_ids: train_loss = self.client_update( idx=idx, model_parameters=model_parameters ) updated_model_parameters = self.serialize_model_parameters() param_list[idx] = updated_model_parameters loss_list[idx] = train_loss self.comm_manager.send( Message( message_type=200, sender=self.F.client_name, receiver=[self.F.server_ip], content={ 'model': param_list, 'loss': loss_list } ) ) def client_update(self, idx, model_parameters, *args, **kwargs): self.logger.debug(f"\n{'=' * 37}\n>>> Subserver={self.F.client_name}_" f"Client={idx}_Round={self.round + 1} <<<\n{'=' * 37}") self.deserialize_model(model_parameters) train_dataset, eval_dataset = self.get_dataset(idx) # manually schedule the learning rate self.T.learning_rate = cosine_learning_rate( self.round, self.F.rounds, self.eval_args.learning_rate, 1e-6) # Initialize local Trainer vis_cb = VisCallback(self.visualizer) train_op = registry.get_loctrainer(self.T.local_trainer_name)( model=self.model, args=self.T, train_dataset=train_dataset, tokenizer=self.data.tokenizer, data_collator=self.data.coll_fn(self.model), compute_metrics=self.metric.calculate_metric, callbacks=[vis_cb] # optimizers ) train_result = train_op.train() del train_op train_loss = round(train_result.training_loss, 3) self.logger.info(f">>> Subserver={self.F.client_name}_Client={idx}_lr=" f"{self.T.learning_rate * 10000:.2f}e-4_Loss={train_loss}") return train_loss def client_valid(self): """Valid on Client""" def client_logging(self): """Logging on Client""" def on_client_end(self): self.logger.critical(f"Subserver {self.F.client_name} Train done") def get_agg_weight(self, loss_list, epsilon=1e-8): if self.F.weight_type == "num": weights = [self.data.train_examples_num_dict[client_id] for client_id in self.client_ids] elif self.F.weight_type == "loss": adjusted_losses = [max(loss, epsilon) for loss in loss_list] inverse_weights = [1 / loss for loss in adjusted_losses] weights = inverse_weights else: weights = [1.0 for _ in range(len(loss_list))] total = sum(weights) normalized_weights = [weight / total for weight in weights] self.logger.info(f"This round clients' weights: {[round(weight, 3) for weight in normalized_weights]}") return normalized_weights def cen_train(self, client_id=-1): # get local train and eval dataset train_dataset, eval_dataset = self.get_dataset(client_id) # set some parameters total_steps = len(train_dataset) / registry.get("total_bs") * self.T.num_train_epochs if self.F.log_valid_len: self.T.greater_is_better = False if self.T.is_decreased_valid_metric else True self.T.eval_steps = max(int(total_steps / self.F.log_valid_len), 1) if self.F.save_valid_len: self.T.save_strategy = "steps" self.T.save_steps = max(int(total_steps / self.T.num_train_epochs), 1) self.T.save_total_limit = int(self.T.num_train_epochs) self.T.evaluation_strategy = 'no' # Initialize Centralized Trainer train_op = registry.get_loctrainer(self.T.local_trainer_name)( model=self.model, args=self.T, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=self.data.tokenizer, data_collator=self.data.coll_fn(self.model), compute_metrics=self.metric.calculate_metric ) train_op.add_callback(CenEndEvalStepCallback( train_op, metric_name=self.metric_name) ) train_op.train() def on_cen_end(self): """Using best parameters for prediction""" if not self.debug: metric_save(self, self.T, self.logger) self.logger.critical(f"Train done, Please Eval and Test in {self.T.checkpoint_dir}") f4llm-main/LICENSE000066400000000000000000000020541501253031700140210ustar00rootroot00000000000000MIT License Copyright (c) 2024 ZhuoZhuoBye Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. f4llm-main/LocBaseDPO.py000066400000000000000000000165531501253031700152520ustar00rootroot00000000000000import torch import inspect from torch import nn from peft import PeftModel import torch.nn.functional as F from transformers import Trainer from typing import Any, Dict, List, Optional, Tuple, Union IGNORE_INDEX = -100 class LocalDPOTrainer(Trainer): def __init__(self, **kwargs): super().__init__(**kwargs) def get_log_probs(self, logits, labels): log_probs = F.log_softmax(logits, dim=-1) log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)) return log_probs_labels.squeeze(-1) def masked_mean(self, data, mask, dim=None, eps=1e-8): data = data * mask if dim is not None: return data.sum(dim=dim) / (mask.sum(dim=dim) + eps) else: return data.sum() / (mask.sum() + eps) def get_entropy(self, logits, mask): probs = torch.nn.functional.softmax(logits, dim=-1) log_probs = F.log_softmax(logits, dim=-1) entropy = self.masked_mean(-torch.sum(probs * log_probs, dim=-1), mask) return entropy def _set_signature_columns_if_needed(self): # important self.label_names/default_label_names != "labels" if self._signature_columns is None: # Inspect model forward signature to keep only the arguments it accepts. signature = inspect.signature(self.model.forward) self._signature_columns = list(signature.parameters.keys()) # self.logger.info(f"The following columns {self._signature_columns} are accepted.") if "input_ids" not in self._signature_columns: self._signature_columns += ["input_ids"] # Labels may be named label or label_ids, the default data collator handles that. self._signature_columns += list(set(["label", "label_ids", "labels"] + self.label_names)) # self.logger.info(f"The following columns {self._signature_columns} are accepted.") def get_model_output(self, model, inputs, is_ref_model=False): if is_ref_model: if isinstance(model, nn.parallel.DistributedDataParallel): with model.module.disable_adapter(): chosen_logits = model(input_ids=inputs["chosen_input_ids"], attention_mask=inputs["chosen_attention_mask"], return_dict=True).logits rejected_logits = model(input_ids=inputs["rejected_input_ids"], attention_mask=inputs["rejected_attention_mask"], return_dict=True).logits elif isinstance(model, PeftModel): with model.disable_adapter(): chosen_logits = model(input_ids=inputs["chosen_input_ids"], attention_mask=inputs["chosen_attention_mask"], return_dict=True).logits rejected_logits = model(input_ids=inputs["rejected_input_ids"], attention_mask=inputs["rejected_attention_mask"], return_dict=True).logits else: raise AttributeError( f" model object [{model.__class__.__name__}] has no attribute [disable_adapter] " ) else: chosen_logits = model(input_ids=inputs["chosen_input_ids"], attention_mask=inputs["chosen_attention_mask"], return_dict=True).logits rejected_logits = model(input_ids=inputs["rejected_input_ids"], attention_mask=inputs["rejected_attention_mask"], return_dict=True).logits chosen_labels, rejected_labels = inputs["chosen_labels"], inputs["rejected_labels"] chosen_action_masks = chosen_labels.ne(IGNORE_INDEX).long() rejected_action_masks = rejected_labels.ne(IGNORE_INDEX).long() chosen_log_probs = self.get_log_probs(chosen_logits[:, :-1, :], inputs["chosen_input_ids"][:, 1:]) rejected_log_probs = self.get_log_probs(rejected_logits[:, :-1, :], inputs["rejected_input_ids"][:, 1:]) if self.args.average_log_prob: chosen_logps = self.masked_mean(chosen_log_probs, chosen_action_masks[:, 1:], dim=-1) rejected_logps = self.masked_mean(rejected_log_probs, rejected_action_masks[:, 1:], dim=-1) else: chosen_logps = (chosen_log_probs * chosen_action_masks[:, 1:]).sum(dim=-1) rejected_logps = (rejected_log_probs * rejected_action_masks[:, 1:]).sum(dim=-1) chosen_entropy = self.get_entropy(chosen_logits[:, :-1, :], chosen_action_masks[:, 1:]) rejected_entropy = self.get_entropy(rejected_logits[:, :-1, :], rejected_action_masks[:, 1:]) return chosen_entropy, rejected_entropy, chosen_logps, rejected_logps def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): chosen_entropy, rejected_entropy, chosen_logps, rejected_logps = self.get_model_output(model, inputs) with torch.no_grad(): ref_chosen_entropy, ref_rejected_entropy, ref_chosen_logps, ref_rejected_logps = self.get_model_output \ (model, inputs, is_ref_model=True) chosen_ratio = self.args.dpo_beta * (chosen_logps - ref_chosen_logps) rejected_ratio = self.args.dpo_beta * (rejected_logps - ref_rejected_logps) pi_ratio = self.args.dpo_beta * (chosen_logps - rejected_logps) ref_ratio = self.args.dpo_beta * (ref_chosen_logps - ref_rejected_logps) if self.args.reference_free: ref_ratio = 0 loss = -torch.nn.functional.logsigmoid(pi_ratio - ref_ratio).mean() outputs = dict( chosen_reward=chosen_ratio.detach(), # shape: (batch_size,) rejected_reward=rejected_ratio.detach(), pi_ratio=pi_ratio.detach(), ref_ratio=ref_ratio.detach(), chosen_entropy=chosen_entropy.detach(), rejected_entropy=rejected_entropy.detach(), ref_chosen_entropy=ref_chosen_entropy.detach(), ref_rejected_entropy=ref_rejected_entropy.detach(), chosen_ce_loss=-chosen_logps.detach(), rejected_ce_loss=-rejected_logps.detach(), ) return (loss, outputs) if return_outputs else loss 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[float], Optional[torch.Tensor], Optional[torch.Tensor]]: 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 = [] with torch.no_grad(): loss, outputs = self.compute_loss(model, inputs, return_outputs=True) loss = loss.detach() logits = tuple(v for k, v in outputs.items() if k in ["chosen_reward", "rejected_reward"]) if prediction_loss_only: return loss, None, None logits = torch.stack(logits, dim=1) # chosen first, therefore, we can use label_ids as the label for rewardben if "labels" in inputs: labels = inputs["labels"] # set_name info else: labels = torch.zeros(logits.shape[0]).to(logits.device) return loss, logits, labels f4llm-main/LocBaseRM.py000066400000000000000000000146341501253031700151440ustar00rootroot00000000000000import os import inspect from typing import Optional import torch from peft import get_peft_model_state_dict from transformers import Trainer, PreTrainedModel from transformers.modeling_utils import unwrap_model from utils.register import registry IGNORE_INDEX = -100 WEIGHTS_NAME = "adapter_model.bin" TRAINING_ARGS_NAME = "training_args.bin" class LocalRMTrainer(Trainer): def __init__(self, **kwargs): super().__init__(**kwargs) self.args.use_last_reward = getattr(self.args, "use_last_reward", False) self.args.clm_loss_weight = getattr(self.args, "clm_loss_weight", 0.2) # self.logger = registry.get("logger") def get_state_dict(self, model): pretrained_model_state_dict = model.pretrained_model.state_dict() v_head_state_dict = model.v_head.state_dict() for k, v in v_head_state_dict.items(): pretrained_model_state_dict[f"v_head.{k}"] = v return pretrained_model_state_dict def _save(self, output_dir: Optional[str] = None, state_dict=None): # If we are executing this function, we are the process zero, so we don't check for that. output_dir = output_dir if output_dir is not None else self.args.output_dir os.makedirs(output_dir, exist_ok=True) # logger.info(f"Saving model checkpoint to {output_dir}") if not isinstance(self.model, PreTrainedModel): if state_dict is None: state_dict = self.get_state_dict(self.model) if isinstance(unwrap_model(self.model), PreTrainedModel): unwrap_model(self.model).save_pretrained(output_dir, state_dict=state_dict) else: # logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") adapter_state_dict = get_peft_model_state_dict(unwrap_model(self.model).pretrained_model) # add v_head (v_head not in modules_to_save) v_head_state_dict = self.model.v_head.state_dict() for k, v in v_head_state_dict.items(): adapter_state_dict[f"v_head.{k}"] = v torch.save(adapter_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) try: unwrap_model(self.model).pretrained_model.peft_config.save_pretrained(output_dir) except AttributeError: unwrap_model(self.model).pretrained_model.peft_config['default'].save_pretrained(output_dir) else: self.model.save_pretrained(output_dir, state_dict=state_dict) if self.tokenizer is not None: self.tokenizer.save_pretrained(output_dir) # Good practice: save your training arguments together with the trained model torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) def _set_signature_columns_if_needed(self): # important self.label_names/default_label_names != "labels" if self._signature_columns is None: # Inspect model forward signature to keep only the arguments it accepts. signature = inspect.signature(self.model.forward) self._signature_columns = list(signature.parameters.keys()) # self.logger.info(f"The following columns {self._signature_columns} are accepted.") if "input_ids" not in self._signature_columns: self._signature_columns += ["input_ids"] # Labels may be named label or label_ids, the default data collator handles that. self._signature_columns += list(set(["label", "label_ids", "labels"] + self.label_names)) # self.logger.info(f"The following columns {self._signature_columns} are accepted.") def compute_loss(self, model, inputs, return_outputs=False): _, chosen_clm_loss, chosen_value = model( input_ids=inputs["chosen_input_ids"], attention_mask=inputs["chosen_attention_mask"], labels=inputs["chosen_labels"], return_dict=True) _, _, rejected_value = model( input_ids=inputs["rejected_input_ids"], attention_mask=inputs["rejected_attention_mask"], return_dict=True) chosen_labels, rejected_labels = inputs["chosen_labels"], inputs["rejected_labels"] chosen_action_masks = chosen_labels.ne(IGNORE_INDEX).long() rejected_action_masks = rejected_labels.ne(IGNORE_INDEX).long() chosen_value = chosen_value * chosen_action_masks rejected_value = rejected_value * rejected_action_masks batch_size = chosen_value.shape[0] chosen_seq_lengths = (torch.ne(inputs["chosen_input_ids"], self.tokenizer.pad_token_id).sum(-1) - 1).to \ (chosen_value.device) rejected_seq_lengths = (torch.ne(inputs["rejected_input_ids"], self.tokenizer.pad_token_id).sum(-1) - 1).to \ (rejected_value.device) chosen_end_token_value = chosen_value[ torch.arange(batch_size, device=chosen_value.device), chosen_seq_lengths] rejected_end_token_value = rejected_value[ torch.arange(batch_size, device=rejected_value.device), rejected_seq_lengths] if self.args.use_last_reward: loss1 = -torch.nn.functional.logsigmoid(chosen_end_token_value - rejected_end_token_value).mean() else: loss1 = -torch.nn.functional.logsigmoid(chosen_value - rejected_value).mean() loss2 = self.args.clm_loss_weight * chosen_clm_loss loss = loss1 + loss2 outputs = dict( chosen_end_token_value=chosen_end_token_value, # shape: (batch_size,) rejected_end_token_value=rejected_end_token_value, # shape: (batch_size,) ) return (loss, outputs) if return_outputs else loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys): 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 = [] with torch.no_grad(): loss, outputs = self.compute_loss(model, inputs, return_outputs=True) loss = loss.detach() logits = tuple(v for k, v in outputs.items() if k in ["chosen_end_token_value", "rejected_end_token_value"]) if prediction_loss_only: return (loss, None, None) logits = torch.stack(logits, dim=1) labels = torch.zeros(logits.shape[0]).to(logits.device) return loss, logits, labels f4llm-main/LocBaseSFT.py000066400000000000000000000143461501253031700152620ustar00rootroot00000000000000 import torch from torch import nn from torch.utils.data import Dataset from typing import Any, Dict, List, Optional, Tuple, Union from transformers import Trainer from transformers.trainer_utils import PredictionOutput # from transformers.deepspeed import is_deepspeed_zero3_enabled def build_generator_config(args, gen_kwargs): gen_kwargs = gen_kwargs.copy() if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: gen_kwargs["max_new_tokens"] = args.generation_max_length # gen_kwargs["num_beams"] = ( gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else args.generation_num_beams ) gen_kwargs["do_sample"] = ( gen_kwargs["do_sample"] if gen_kwargs.get("do_sample") is not None else args.do_sample ) gen_kwargs["temperature"] = ( gen_kwargs["temperature"] if gen_kwargs.get("temperature") is not None else args.temperature ) gen_kwargs["top_p"] = ( gen_kwargs["top_p"] if gen_kwargs.get("top_p") is not None else args.top_p ) return gen_kwargs class LocalSFTTrainer(Trainer): def evaluate( self, eval_dataset: Optional[Dataset] = None, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "eval", **gen_kwargs ) -> Dict[str, float]: gen_kwargs = build_generator_config(self.args, gen_kwargs) self._gen_kwargs = gen_kwargs return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) def predict( self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test", **gen_kwargs ) -> PredictionOutput: gen_kwargs = build_generator_config(self.args, gen_kwargs) self._gen_kwargs = gen_kwargs return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) 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[float], Optional[torch.Tensor], Optional[torch.Tensor]]: if not self.args.predict_with_generate or prediction_loss_only: return super().prediction_step( model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys ) has_labels = "labels" in inputs inputs = self._prepare_inputs(inputs) # XXX: adapt synced_gpus for fairscale as well gen_kwargs = self._gen_kwargs.copy() if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: gen_kwargs["max_length"] = self.model.config.max_length gen_kwargs["num_beams"] = ( gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams ) # default_synced_gpus = True if is_deepspeed_zero3_enabled() else False default_synced_gpus = False gen_kwargs["synced_gpus"] = ( gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus ) if "attention_mask" in inputs: gen_kwargs["attention_mask"] = inputs.get("attention_mask", None) if "position_ids" in inputs: gen_kwargs["position_ids"] = inputs.get("position_ids", None) if "global_attention_mask" in inputs: gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None) # prepare generation inputs # some encoder-decoder models can have varying encoder's and thus # varying model input names if hasattr(self.model, "encoder") and self.model.encoder.main_input_name != self.model.main_input_name: generation_inputs = inputs[self.model.encoder.main_input_name] else: generation_inputs = inputs[self.model.main_input_name] gen_kwargs["input_ids"] = generation_inputs generated_tokens = self.model.generate(**gen_kwargs) generated_tokens = generated_tokens[:, generation_inputs.size()[-1]:] # in case the batch is shorter than max length, the output should be padded if gen_kwargs.get("max_length") is not None and generated_tokens.shape[-1] < gen_kwargs["max_length"]: generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"]) elif gen_kwargs.get("max_new_tokens") is not None and generated_tokens.shape[-1] < ( gen_kwargs["max_new_tokens"] + 1 ): generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_new_tokens"] + 1) loss = None if self.args.prediction_loss_only: return (loss, None, None) if has_labels: labels = inputs["labels"] if gen_kwargs.get("max_length") is not None and labels.shape[-1] < gen_kwargs["max_length"]: labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"]) elif gen_kwargs.get("max_new_tokens") is not None and labels.shape[-1] < ( gen_kwargs["max_new_tokens"] + 1 ): labels = self._pad_tensors_to_max_len(labels, (gen_kwargs["max_new_tokens"] + 1)) else: labels = None return (loss, generated_tokens, labels) def _pad_tensors_to_max_len(self, tensor, max_length): if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"): # If PAD token is not defined at least EOS token has to be defined pad_token_id = ( self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id ) else: if self.model.config.pad_token_id is not None: pad_token_id = self.model.config.pad_token_id else: raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors") padded_tensor = pad_token_id * torch.ones( (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device ) padded_tensor[:, : tensor.shape[-1]] = tensor return padded_tensor f4llm-main/README.md000066400000000000000000000046651501253031700143050ustar00rootroot00000000000000

F4LLM

SMILELAB Licence Build

[//]: # (

) [//]: # (

A Comprehensive Toolkit for Large Language Models in Federated Learning

) [//]: # (

)

## Abstract The rapid success of Large Language Models (LLMs) has unlocked vast potential for AI applications in privacy-sensitive domains. However, the traditional centralized training of LLMs poses significant challenges due to privacy concerns regarding collecting sensitive data from diverse sources. This paper offers a promising and privacy-enhancing solution for LLMs: collaboratively training LLMs via Federated Learning (FL) across multiple clients, eliminating the need for raw data transmission. To this end, we present F4LLM, a new and comprehensive toolbox that supports the entire Federated LLM pipeline, from Continuous pre-training to alignment and LLM evaluation. F4LLM employs gRPC as the communication protocol to support various widely-used FL algorithms, ensuring efficient development and benchmarking in geo-distributed FL environments. Moreover, F4LLM offers both open-form and closed-form evaluation options via the efficient inference tool vLLM. ## Folder Structure ```grapha ├── userhome │ └── data | | └── safe_rlhf │ ├── pretrained │ │ └── nlp │ ├── output │ └── code │ └── f4llm ``` This repo should be cloned into f4llm: ```bash mkdir userhome cd userhome mkdir data mkdir code mkdir pretrain cd pretrain mkdir nlp cd .. cd code git clone https://github.com/iezhuozhuo/f4llm.git ``` ## Usage ### Setup We recommend that the `python` version of the runtime environment be `3.10+`, and `pytorch` be `2.4+`. ```bash pip install -r requirements ``` ### Run F4LLM with FL algorithms ```bash bash ./scripts/run.sh {your_file_path}/userhome {code_name} {model_name} fedavg {task_name} {port} {gpu_id} # example: bash ./scripts/run_all.sh /userhome f4llm tinyllama fedavg safe_rlhf 10001 001 ``` ## License [MIT license](https://github.com/iezhuozhuo/f4llm?tab=MIT-1-ov-file#)f4llm-main/main.py000066400000000000000000000011371501253031700143130ustar00rootroot00000000000000""" Fine-tuning the library models for sequence to sequence. """ from configs import build_config from utils.general import setup_imports from utils.register import registry def run_experiment(config): phase = registry.get("phase") if phase == "train": # federated llm training engine = registry.get_fedtrainer(config.F.fl_algorithm)() else: # llm inference engine = registry.get_eval_class(config.T.eval_name)() engine.run() def main(): setup_imports() config = build_config() run_experiment(config) if __name__ == "__main__": main() f4llm-main/requirement.txt000066400000000000000000000004301501253031700161110ustar00rootroot00000000000000# torch2.4.0 cuda 12.2 transformers==4.45.2 deepspeed==0.11.1 loguru omegaconf jieba rouge_chinese gpustat deepspeed xformers sentencepiece peft==0.5.0 bert-score rouge pandas grpcio grpcio-tools trl watchdog tabulate bitsandbytes datasets pympler watchdog vllm wandb scikit-learn