diff --git a/llm-inference.py b/llm-inference.py new file mode 100644 index 0000000000..3c3a6cabcf --- /dev/null +++ b/llm-inference.py @@ -0,0 +1,77 @@ +#%% +from transformers.models.llama import LlamaForCausalLM, LlamaTokenizer +from transformers import GenerationConfig + +#path = '/home/oamontoy/workspace/weights-llama-2-7B' +path = '/home/oamontoy/workspace/weights-llama-2-7B-chat' +tokenizer = LlamaTokenizer.from_pretrained(path) +model = LlamaForCausalLM.from_pretrained(path) +# %% +from peft import PeftModel +model = PeftModel.from_pretrained(model, "dominguesm/alpaca-lora-ptbr-7b") +# %% + +def generate_prompt(instruction, input=None): + if input: + return f"""Below is a statement that describes a task, paired with an input that provides more context. Write a response that appropriately completes the request. + +### instruction: +{instruction} + +### input: +{input} + +### response:""" + else: + return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request. + +### instruction: +{instruction} + +### response:""" + +# %% +from pprint import pprint +# %% +generation_config = GenerationConfig( + temperature=0.1, + top_p=0.75, + num_beams=4, +) + +def evaluate(instruction, input=None): + prompt = generate_prompt(instruction, input) + inputs = tokenizer(prompt, return_tensors="pt") + input_ids = inputs["input_ids"] + generation_output = model.generate( + input_ids=input_ids, + generation_config=generation_config, + return_dict_in_generate=True, + output_scores=True, + max_new_tokens=256 + ) + for s in generation_output.sequences: + output = tokenizer.decode(s) + pprint("response: " + output.split("### response:")[1].strip()) +# %% +#evaluate(input("instruction: ")) +# %% +instruction = 'print a long paragraph of giberish' +prompt = generate_prompt(instruction, None) +inputs = tokenizer(prompt, return_tensors="pt") +inputs +#%% +input_ids = inputs["input_ids"] +# %% +generation_output = model.generate( + input_ids=input_ids, + generation_config=generation_config, + return_dict_in_generate=True, + output_scores=True, + max_new_tokens=256 + ) +# %% +for s in generation_output.sequences: + output = tokenizer.decode(s) + pprint("response: " + output.split("### response:")[1].strip()) +# %% diff --git a/llm-lora_classification copy.py b/llm-lora_classification copy.py new file mode 100644 index 0000000000..19c73e0c5a --- /dev/null +++ b/llm-lora_classification copy.py @@ -0,0 +1,178 @@ +#%% +import numpy as np +import torch +from datasets import load_dataset, load_metric +from peft import LoraConfig, TaskType, get_peft_model +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm + +from transformers import (AutoModelForSequenceClassification, AutoTokenizer, + DataCollatorWithPadding, Trainer, TrainingArguments, + get_scheduler) + +# %% +task = "mrpc" +dataset = load_dataset("glue", task) +padding_side = "right" + +metric = load_metric('glue', task) +def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + +#%% +model_checkpoint = "roberta-base" +tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + +def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + +tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], +) +tokenized_datasets = tokenized_datasets.rename_column("label", "labels") +tokenized_datasets.set_format("torch") +data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + +# %% +trainer_model = AutoModelForSequenceClassification.from_pretrained( + model_checkpoint, return_dict=True) + +trainer_peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all" +) + + +model = AutoModelForSequenceClassification.from_pretrained( + model_checkpoint, return_dict=True) +peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all" +) + +trainer_model = get_peft_model(trainer_model, trainer_peft_config) +trainer_model.print_trainable_parameters() + +model = get_peft_model(model, peft_config) +model.print_trainable_parameters() + +#%% + +BATCH_SIZE = 32 +NUM_EPOCHS = 5 +LR = 1e-3 +WEIGHT_DECAY = 0.01 +#%% + +training_args = TrainingArguments( + output_dir="roberta-large-lora-seq1_tests", + learning_rate=LR, + per_device_train_batch_size=BATCH_SIZE, + per_device_eval_batch_size=BATCH_SIZE, + num_train_epochs=NUM_EPOCHS, + weight_decay=WEIGHT_DECAY, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, +) +trainer = Trainer( + model=trainer_model, + args=training_args, + #train_dataset=tokenized_datasets["train"], + train_dataset=tokenized_datasets["train"].select(range(1000)), + #eval_dataset=tokenized_datasets["test"], + eval_dataset=tokenized_datasets["test"].select(range(1000)), + tokenizer=tokenizer, + compute_metrics=compute_metrics, + data_collator=data_collator +) + +trainer.train() +#%% +#model.save_pretrained('saved_model') + +#%% + +#accelerator = Accelerator() +#accelerator.wait_for_everyone() + +def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + +from torch import nn + +from transformers.trainer_pt_utils import get_parameter_names + +ALL_LAYERNORM_LAYERS = [nn.LayerNorm] +decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS) +decay_parameters = [name for name in decay_parameters if "bias" not in name] + +optimizer_grouped_parameters = [ + { + "params": [ + p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": WEIGHT_DECAY, + }, + { + "params": [ + p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + }, + ] +optimizer = AdamW(optimizer_grouped_parameters, lr=LR) + +#%% +#train_dataloader = DataLoader(tokenized_datasets["train"].select(range(1500)), shuffle=True, collate_fn=data_collator, batch_size=BATCH_SIZE) +train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=data_collator, batch_size=BATCH_SIZE) +#valid_dataloader = DataLoader(tokenized_datasets["test"].select(range(500)), collate_fn=data_collator, batch_size=BATCH_SIZE) +valid_dataloader = DataLoader(tokenized_datasets["test"], collate_fn=data_collator, batch_size=BATCH_SIZE) + +lr_scheduler = get_scheduler( + name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=len(train_dataloader) * NUM_EPOCHS +) +def eval_metrics(model, dataloader, device='cpu'): + model.eval() + for batch in dataloader: + batch = {k: v.to(device) for k, v in batch.items()} + + with torch.no_grad(): + outputs = model(**batch) + + logits = outputs.logits + predictions = torch.argmax(logits, dim=-1) + metric.add_batch(predictions=predictions, references=batch["labels"]) + model.train() + return metric.compute() + +#model, optimizer, train_dataloader, valid_dataloader, lr_scheduler = accelerator.prepare( +# model, optimizer, train_dataloader, valid_dataloader, lr_scheduler +# ) +#%% +for epoch in range(NUM_EPOCHS): + model.train() + losses = [] + for batch in tqdm(train_dataloader): + outputs = model(**batch) + loss = outputs.loss + loss.backward() + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + #model.zero_grad() + losses.append(loss.detach()) + + accuracy = eval_metrics(model, valid_dataloader) + print(accuracy, 'loss: ', np.mean(losses)) +# %% diff --git a/llm-lora_classification.py b/llm-lora_classification.py new file mode 100644 index 0000000000..461a1bfcbf --- /dev/null +++ b/llm-lora_classification.py @@ -0,0 +1,179 @@ +#%% +import numpy as np +import torch +from datasets import load_dataset, load_metric +from peft import LoraConfig, TaskType, get_peft_model +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm + +from transformers import (AutoModelForSequenceClassification, AutoTokenizer, + DataCollatorWithPadding, Trainer, TrainingArguments, + get_scheduler) + +# %% +task = "mrpc" +dataset = load_dataset("glue", task) +padding_side = "right" + +metric = load_metric('glue', task) +def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + +#%% +model_checkpoint = "roberta-base" +tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + +def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + +tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], +) +tokenized_datasets = tokenized_datasets.rename_column("label", "labels") +tokenized_datasets.set_format("torch") +data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + +# %% +trainer_model = AutoModelForSequenceClassification.from_pretrained( + model_checkpoint, return_dict=True) + +trainer_peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all" +) + + +model = AutoModelForSequenceClassification.from_pretrained( + model_checkpoint, return_dict=True) +peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all" +) + +trainer_model = get_peft_model(trainer_model, trainer_peft_config) +trainer_model.print_trainable_parameters() + +model = get_peft_model(model, peft_config) +model.print_trainable_parameters() + +#%% + +BATCH_SIZE = 32 +NUM_EPOCHS = 5 +LR = 1e-3 +WEIGHT_DECAY = 0.01 +#%% + +training_args = TrainingArguments( + output_dir="roberta-large-lora-seq1_tests", + learning_rate=LR, + per_device_train_batch_size=BATCH_SIZE, + per_device_eval_batch_size=BATCH_SIZE, + num_train_epochs=NUM_EPOCHS, + weight_decay=WEIGHT_DECAY, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, +) +trainer = Trainer( + model=trainer_model, + args=training_args, + #train_dataset=tokenized_datasets["train"], + train_dataset=tokenized_datasets["train"].select(range(1000)), + #eval_dataset=tokenized_datasets["test"], + eval_dataset=tokenized_datasets["test"].select(range(1000)), + tokenizer=tokenizer, + compute_metrics=compute_metrics, + data_collator=data_collator +) + +trainer.train() +#%% +#model.save_pretrained('saved_model') + +#%% + +#accelerator = Accelerator() +#accelerator.wait_for_everyone() + +def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + +from torch import nn + +from transformers.trainer_pt_utils import get_parameter_names + +ALL_LAYERNORM_LAYERS = [nn.LayerNorm] +decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS) +decay_parameters = [name for name in decay_parameters if "bias" not in name] + +optimizer_grouped_parameters = [ + { + "params": [ + p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": WEIGHT_DECAY, + }, + { + "params": [ + p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + }, + ] +optimizer = AdamW(optimizer_grouped_parameters, lr=LR) + +#%% +train_dataloader = DataLoader(tokenized_datasets["train"].select(range(1000)), shuffle=True, collate_fn=data_collator, batch_size=BATCH_SIZE) +#train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=data_collator, batch_size=BATCH_SIZE) +valid_dataloader = DataLoader(tokenized_datasets["test"].select(range(1000)), collate_fn=data_collator, batch_size=BATCH_SIZE) +#valid_dataloader = DataLoader(tokenized_datasets["test"], collate_fn=data_collator, batch_size=BATCH_SIZE) + +lr_scheduler = get_scheduler( + name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=len(train_dataloader) * NUM_EPOCHS +) +def eval_metrics(model, dataloader, device='cpu'): + model.eval() + for batch in dataloader: + batch = {k: v.to(device) for k, v in batch.items()} + + with torch.no_grad(): + outputs = model(**batch) + + logits = outputs.logits + predictions = torch.argmax(logits, dim=-1) + metric.add_batch(predictions=predictions, references=batch["labels"]) + model.train() + return metric.compute() + +#model, optimizer, train_dataloader, valid_dataloader, lr_scheduler = accelerator.prepare( +# model, optimizer, train_dataloader, valid_dataloader, lr_scheduler +# ) +#%% +for epoch in range(NUM_EPOCHS): + model.train() + losses = [] + for batch in tqdm(train_dataloader): + + outputs = model(**batch) + loss = outputs.loss + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(),1.0) + optimizer.step() + lr_scheduler.step() + #optimizer.zero_grad() + model.zero_grad() + losses.append(loss.detach()) + accuracy = eval_metrics(model, valid_dataloader) + print(accuracy, 'loss: ', np.mean(losses)) + +# %% diff --git a/llm-lora_tokenclass.py b/llm-lora_tokenclass.py new file mode 100644 index 0000000000..97d894ff94 --- /dev/null +++ b/llm-lora_tokenclass.py @@ -0,0 +1,174 @@ +#%% +#Trains a model to classify tokens in sentences refering to biology, if a token referes to a biological term +from datasets import load_dataset +from transformers import ( + AutoModelForTokenClassification, + AutoTokenizer, + DataCollatorForTokenClassification, + TrainingArguments, + Trainer, +) +from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType +import evaluate +import torch +import numpy as np + +model_checkpoint = "roberta-large" +lr = 1e-3 +batch_size = 16 +num_epochs = 10 +# %% +bionlp = load_dataset("tner/bionlp2004") +bionlp["train"][0] +# %% +seqeval = evaluate.load("seqeval") +# %% +label_list = [ + "O", + "B-DNA", + "I-DNA", + "B-protein", + "I-protein", + "B-cell_type", + "I-cell_type", + "B-cell_line", + "I-cell_line", + "B-RNA", + "I-RNA", +] + + +def compute_metrics(p): + predictions, labels = p + predictions = np.argmax(predictions, axis=2) + + true_predictions = [ + [label_list[p] for (p, l) in zip(prediction, label) if l != -100] + for prediction, label in zip(predictions, labels) + ] + true_labels = [ + [label_list[l] for (p, l) in zip(prediction, label) if l != -100] + for prediction, label in zip(predictions, labels) + ] + + results = seqeval.compute(predictions=true_predictions, references=true_labels) + return { + "precision": results["overall_precision"], + "recall": results["overall_recall"], + "f1": results["overall_f1"], + "accuracy": results["overall_accuracy"], + } +# %% +path = '/home/oamontoy/workspace/weights-llama-2-7B-chat' +tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True) +# %% +def tokenize_and_align_labels(examples): + tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True) + + labels = [] + for i, label in enumerate(examples[f"tags"]): + word_ids = tokenized_inputs.word_ids(batch_index=i) + previous_word_idx = None + label_ids = [] + for word_idx in word_ids: + if word_idx is None: + label_ids.append(-100) + elif word_idx != previous_word_idx: + label_ids.append(label[word_idx]) + else: + label_ids.append(-100) + previous_word_idx = word_idx + labels.append(label_ids) + + tokenized_inputs["labels"] = labels + return tokenized_inputs +# %% +tokenized_bionlp = bionlp.map(tokenize_and_align_labels, batched=True) +data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer) +# %% +id2label = { + 0: "O", + 1: "B-DNA", + 2: "I-DNA", + 3: "B-protein", + 4: "I-protein", + 5: "B-cell_type", + 6: "I-cell_type", + 7: "B-cell_line", + 8: "I-cell_line", + 9: "B-RNA", + 10: "I-RNA", +} +label2id = { + "O": 0, + "B-DNA": 1, + "I-DNA": 2, + "B-protein": 3, + "I-protein": 4, + "B-cell_type": 5, + "I-cell_type": 6, + "B-cell_line": 7, + "I-cell_line": 8, + "B-RNA": 9, + "I-RNA": 10, +} + +model = AutoModelForTokenClassification.from_pretrained( + model_checkpoint, num_labels=11, id2label=id2label, label2id=label2id +) +# %% +peft_config = LoraConfig( + task_type=TaskType.TOKEN_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all" +) +# %% +model = get_peft_model(model, peft_config) +model.print_trainable_parameters() +# %% +training_args = TrainingArguments( + output_dir="roberta-large-lora-token-classification", + learning_rate=lr, + per_device_train_batch_size=batch_size, + per_device_eval_batch_size=batch_size, + num_train_epochs=num_epochs, + weight_decay=0.01, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, +) +#%% +text = "The activation of IL-2 gene expression and NF-kappa B through CD28 requires reactive oxygen production by 5-lipoxygenase." +inputs = tokenizer(text, return_tensors="pt") +# %% +with torch.no_grad(): + logits = model(**inputs).logits + +tokens = inputs.tokens() +predictions = torch.argmax(logits, dim=2) + +for token, prediction in zip(tokens, predictions[0].numpy()): + print((token, model.config.id2label[prediction])) +# %% +trainer = Trainer( + model=model, + args=training_args, + train_dataset=tokenized_bionlp["train"].select(range(0,1000)), + eval_dataset=tokenized_bionlp["validation"].select(range(0,200)), + tokenizer=tokenizer, + data_collator=data_collator, + compute_metrics=compute_metrics, +) + +trainer.train() +# %% +text = "The activation of IL-2 gene expression and NF-kappa B through CD28 requires reactive oxygen production by 5-lipoxygenase." +inputs = tokenizer(text, return_tensors="pt") +# %% +with torch.no_grad(): + logits = model(**inputs).logits + +tokens = inputs.tokens() +predictions = torch.argmax(logits, dim=2) + +for token, prediction in zip(tokens, predictions[0].numpy()): + print((token, model.config.id2label[prediction])) +# %% diff --git a/openfl-tutorials/Federated_PyTorch_LLM.ipynb b/openfl-tutorials/Federated_PyTorch_LLM.ipynb new file mode 100644 index 0000000000..4ace9f3d30 --- /dev/null +++ b/openfl-tutorials/Federated_PyTorch_LLM.ipynb @@ -0,0 +1,474 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Federated PyTorch TinyImageNet Tutorial" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This notebook is an example of Transfer Learning \n", + "\n", + "Custom DataLoader is used with OpenFL Python API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Install dependencies if not already installed\n", + "!pip install torch torchvision peft transformers sentencepiece huggingface_hub accelerate datasets evaluate seqeval\n", + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Any, Mapping\n", + "import numpy as np\n", + "import openfl.native as fx\n", + "import torch\n", + "import torch as pt\n", + "from accelerate import Accelerator\n", + "from datasets import Dataset, load_dataset, load_metric\n", + "from openfl.federated import PyTorchTaskRunner, TaskRunner\n", + "from openfl.federated.task.runner_pt import change_tags\n", + "from openfl.utilities import Metric, TensorKey\n", + "from openfl.utilities.data_splitters import EqualNumPyDataSplitter\n", + "from peft import LoraConfig, TaskType, get_peft_model\n", + "from peft.utils import get_peft_model_state_dict, set_peft_model_state_dict\n", + "from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss\n", + "from torch.optim import AdamW\n", + "from torch.utils.data import DataLoader\n", + "from tqdm import tqdm\n", + "import torch.nn as nn\n", + "\n", + "from transformers import (AutoConfig, AutoModelForSequenceClassification,\n", + " AutoTokenizer, DataCollatorWithPadding)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "After importing the required packages, the next step is setting up our openfl workspace. To do this, simply run the `fx.init()` command as follows:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Setup default workspace, logging, etc.\n", + "fx.init('torch_cnn_mnist')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we are ready to define our dataset and model to perform federated learning on. The dataset should be composed of a numpy arrayWe start with a simple fully connected model that is trained on the MNIST dataset. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Download the data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def get_glue_mrpc_dataset(tokenizer):\n", + " dataset = load_dataset(\"glue\", \"mrpc\")\n", + "\n", + " def tokenize_function(examples):\n", + " # max_length=None => use the model max length (it's actually the default)\n", + " outputs = tokenizer(\n", + " examples[\"sentence1\"],\n", + " examples[\"sentence2\"],\n", + " truncation=True,\n", + " max_length=None,\n", + " )\n", + " return outputs\n", + "\n", + " tokenized_datasets = dataset.map(\n", + " tokenize_function,\n", + " batched=True,\n", + " remove_columns=[\"idx\", \"sentence1\", \"sentence2\"],\n", + " )\n", + " tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n", + " tokenized_datasets.set_format(\"torch\")\n", + " data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=\"longest\")\n", + " return data_collator, tokenized_datasets\n", + "\n", + "base_model_name = \"roberta-large\"\n", + "padding_side = \"right\"\n", + "tokenizer = AutoTokenizer.from_pretrained(base_model_name, padding_side=padding_side)\n", + "if getattr(tokenizer, \"pad_token_id\") is None:\n", + " tokenizer.pad_token_id = tokenizer.eos_token_id\n", + "data_collator, tokenized_datasets = get_glue_mrpc_dataset(tokenizer)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Describe the dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class GlueMrpc(Dataset):\n", + " \"\"\"\n", + " Has 5.8k pairs of sentences with annotations if the two sentences are equivalent\n", + " \"\"\" \n", + " def get_shape(self):\n", + " \n", + " if not hasattr(self, 'saved_shape'):\n", + " self.saved_shape = max([len(i) for i in self.data['input_ids']])\n", + " return self.saved_shape\n", + "\n", + "train_set = GlueMrpc.from_dict(tokenized_datasets['train'].to_dict())\n", + "valid_set = GlueMrpc.from_dict(tokenized_datasets['test'].to_dict())\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Implement Federated dataset\n", + "We have to implement `split` method" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class GlueMrpcFederatedDataset(DataLoader):\n", + " def __init__(self, train_set, valid_set, batch_size, data_collator=None):\n", + " self.data_splitter = EqualNumPyDataSplitter()\n", + " if isinstance(train_set,Dataset):\n", + " self.train_set = GlueMrpc.from_dict(train_set.to_dict())\n", + " else:\n", + " self.train_set = train_set\n", + " \n", + " if isinstance(valid_set,Dataset):\n", + " self.valid_set = GlueMrpc.from_dict(valid_set.to_dict())\n", + " else:\n", + " self.valid_set = valid_set \n", + " \n", + " self.batch_size = batch_size\n", + " self.data_collator = data_collator\n", + " \n", + " def split(self, num_collaborators):\n", + " train_split = self.data_splitter.split(self.train_set, num_collaborators)\n", + " valid_split = self.data_splitter.split(self.valid_set, num_collaborators)\n", + " return [\n", + " GlueMrpcFederatedDataset(\n", + " self.train_set.select(train_split[i]),\n", + " self.valid_set.select(valid_split[i]),\n", + " self.batch_size\n", + " )\n", + " for i in range(num_collaborators)\n", + " ]\n", + " \n", + " def get_feature_shape(self):\n", + " return self.train_set.get_shape()\n", + " \n", + " def get_train_loader(self, num_batches=None):\n", + " return DataLoader(self.train_set, batch_size=self.batch_size, collate_fn=data_collator)\n", + " \n", + " def get_valid_loader(self):\n", + " return DataLoader(self.valid_set, collate_fn=data_collator)\n", + " \n", + " def get_train_data_size(self):\n", + " return len(self.train_set)\n", + " \n", + " def get_valid_data_size(self):\n", + " return len(self.valid_set)\n", + " \n", + "fl_data = GlueMrpcFederatedDataset(train_set, valid_set, batch_size=32)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Define model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class LLMTaskRunner(PyTorchTaskRunner):\n", + " def __init__(self, base_model_name, data_loader, device=None, metric=None, **kwargs):\n", + " kwargs['data_loader'] = data_loader\n", + " super().__init__(device, **kwargs)\n", + " self.base_model_name = base_model_name\n", + " self.metric = metric\n", + " self._init_model()\n", + " self._init_optimizer()\n", + " \n", + " def _init_model(self):\n", + " model = AutoModelForSequenceClassification.from_pretrained(\n", + " self.base_model_name, return_dict=True)\n", + " peft_config = LoraConfig(task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias=\"all\")\n", + " self.model = get_peft_model(model, peft_config)\n", + " \n", + " def _init_optimizer(self):\n", + " no_decay = [\"bias\", \"LayerNorm.weight\"]\n", + " optimizer_grouped_parameters = [\n", + " {\n", + " \"params\": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)],\n", + " \"weight_decay\": 0.01,\n", + " },\n", + " {\n", + " \"params\": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)],\n", + " \"weight_decay\": 0.0,\n", + " },\n", + " ]\n", + " self.optimizer = AdamW(optimizer_grouped_parameters, lr=0.01)\n", + " \n", + " self.training_round_completed = False\n", + " self.initialize_tensorkeys_for_functions()\n", + " \n", + " def state_dict(self):\n", + " return get_peft_model_state_dict(self.model)\n", + " \n", + " def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):\n", + " return set_peft_model_state_dict(\n", + " self.model, state_dict\n", + " )\n", + " \n", + " def validate(self, col_name, round_num, input_tensor_dict,\n", + " use_tqdm=False, **kwargs):\n", + " \"\"\"Validate.\n", + "\n", + " Run validation of the model on the local data.\n", + "\n", + " Args:\n", + " col_name: Name of the collaborator\n", + " round_num: What round is it\n", + " input_tensor_dict: Required input tensors (for model)\n", + " use_tqdm (bool): Use tqdm to print a progress bar (Default=True)\n", + "\n", + " Returns:\n", + " global_output_dict: Tensors to send back to the aggregator\n", + " local_output_dict: Tensors to maintain in the local TensorDB\n", + "\n", + " \"\"\"\n", + " self.rebuild_model(round_num, input_tensor_dict, validation=True)\n", + " self.model.eval()\n", + " self.model.to(self.device)\n", + " val_score = 0\n", + " total_samples = 0\n", + "\n", + " loader = self.data_loader.get_valid_loader()\n", + " if use_tqdm:\n", + " loader = tqdm(loader, desc='validate')\n", + "\n", + " with pt.no_grad():\n", + " for sample in loader:\n", + " samples = sample['input_ids'].shape[0]\n", + " total_samples += samples\n", + " output = self.model(**sample)\n", + " # get the index of the max log-probability\n", + " logits = output.logits\n", + " predictions = torch.argmax(logits, dim=-1)\n", + " metric.add_batch(predictions=predictions, references=sample['labels'])\n", + " val_score = metric.compute()['accuracy']\n", + "\n", + " origin = col_name\n", + " suffix = 'validate'\n", + " if kwargs['apply'] == 'local':\n", + " suffix += '_local'\n", + " else:\n", + " suffix += '_agg'\n", + " tags = ('metric',)\n", + " tags = change_tags(tags, add_field=suffix)\n", + " # TODO figure out a better way to pass in metric for this pytorch\n", + " # validate function\n", + " output_tensor_dict = {\n", + " TensorKey('acc', origin, round_num, True, tags):\n", + " np.array(val_score)\n", + " }\n", + "\n", + " # Empty list represents metrics that should only be stored locally\n", + " return output_tensor_dict, {}\n", + "\n", + " def train_epoch(self, batch_generator) -> Metric:\n", + " \"\"\"Train single epoch.\n", + "\n", + " Override this function in order to use custom training.\n", + "\n", + " Args:\n", + " batch_generator: Train dataset batch generator. Yields (samples, targets) tuples of\n", + " size = `self.data_loader.batch_size`.\n", + " Returns:\n", + " Metric: An object containing name and np.ndarray value.\n", + " \"\"\"\n", + " losses = []\n", + " for sample in batch_generator:\n", + " self.optimizer.zero_grad()\n", + " output = self.model(**sample)\n", + " loss = output.loss\n", + " loss.backward()\n", + " torch.nn.utils.clip_grad_norm_(self.model.parameters(),1.0)\n", + " self.model.step()\n", + " losses.append(loss.detach().cpu().numpy())\n", + " loss = np.mean(losses)\n", + " if self.model.config.problem_type == \"regression\":\n", + " loss_fct = MSELoss()\n", + " elif self.model.config.problem_type == \"single_label_classification\":\n", + " loss_fct = CrossEntropyLoss()\n", + " elif self.model.config.problem_type == \"multi_label_classification\":\n", + " loss_fct = BCEWithLogitsLoss()\n", + " return Metric(name=loss_fct._get_name(), value=np.array(loss))\n", + " \n", + " \n", + " def save_native(self, filepath, model_state_dict_key='model_state_dict',\n", + " optimizer_state_dict_key='optimizer_state_dict', **kwargs):\n", + " \"\"\"\n", + " Save model and optimizer states in a picked file specified by the \\\n", + " filepath. model_/optimizer_state_dicts are stored in the keys provided. \\\n", + " Uses pt.save().\n", + "\n", + " Args:\n", + " filepath (string) : Path to pickle file to be\n", + " created by pt.save().\n", + " model_state_dict_key (string) : key for model state dict\n", + " in pickled file.\n", + " optimizer_state_dict_key (string) : key for optimizer state\n", + " dict in picked file.\n", + " kwargs : unused\n", + "\n", + " Returns:\n", + " None\n", + " \"\"\"\n", + " pickle_dict = {\n", + " model_state_dict_key: get_peft_model_state_dict(self.model),\n", + " optimizer_state_dict_key: self.optimizer.state_dict()\n", + " }\n", + " pt.save(pickle_dict, filepath)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "num_collaborators = 2\n", + "metric = load_metric('glue', \"mrpc\")\n", + "collaborator_models = [\n", + " LLMTaskRunner(\n", + " base_model_name,\n", + " data_loader=data_slice,\n", + " metric=metric\n", + " )\n", + " for data_slice in fl_data.split(num_collaborators)]\n", + "collaborators = {'one':collaborator_models[0],'two':collaborator_models[1]}#, 'three':collaborator_models[2]}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Original TinyImageNet dataset\n", + "print(f'Original training data size: {len(fl_data.train_set)}')\n", + "print(f'Original validation data size: {len(fl_data.valid_set)}\\n')\n", + "\n", + "#Collaborator one's data\n", + "for i, model in enumerate(collaborator_models):\n", + " print(f'Collaborator {i}\\'s training data size: {len(model.data_loader.train_set)}')\n", + " print(f'Collaborator {i}\\'s validation data size: {len(model.data_loader.valid_set)}\\n')\n", + "\n", + "#Collaborator three's data\n", + "#print(f'Collaborator three\\'s training data size: {len(collaborator_models[2].data_loader.X_train)}')\n", + "#print(f'Collaborator three\\'s validation data size: {len(collaborator_models[2].data_loader.X_valid)}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Run experiment, return trained FederatedModel\n", + "final_fl_model = fx.run_experiment(collaborators,{'aggregator.settings.rounds_to_train':3})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Save final model\n", + "final_fl_model.save_native('final_model.pth')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "llama-env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/openfl-tutorials/llm.py b/openfl-tutorials/llm.py new file mode 100644 index 0000000000..48469e0fb0 --- /dev/null +++ b/openfl-tutorials/llm.py @@ -0,0 +1,306 @@ +# %% + +from typing import Any, Mapping +import numpy as np +import openfl.native as fx +import torch +import torch as pt +from accelerate import Accelerator +from datasets import Dataset, load_dataset, load_metric +from openfl.federated import PyTorchTaskRunner, TaskRunner +from openfl.federated.task.runner_pt import change_tags +from openfl.utilities import Metric, TensorKey +from openfl.utilities.data_splitters import EqualNumPyDataSplitter +from peft import LoraConfig, TaskType, get_peft_model +from peft.utils import get_peft_model_state_dict, set_peft_model_state_dict +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm +import torch.nn as nn + +from transformers import (AutoConfig, AutoModelForSequenceClassification, + AutoTokenizer, DataCollatorWithPadding, Trainer) + + +def get_glue_mrpc_dataset(tokenizer): + dataset = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer( + examples["sentence1"], + examples["sentence2"], + truncation=True, + max_length=None, + ) + return outputs + + tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + tokenized_datasets.set_format("torch") + data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + return data_collator, tokenized_datasets + +class GlueMrpc(Dataset): + """ + Has 5.8k pairs of sentences with annotations if the two sentences are equivalent + """ + def get_shape(self): + + if not hasattr(self, 'saved_shape'): + self.saved_shape = max([len(i) for i in self.data['input_ids']]) + return self.saved_shape + +class GlueMrpcFederatedDataset(DataLoader): + def __init__(self, train_set, valid_set, batch_size, data_collator=None): + self.data_splitter = EqualNumPyDataSplitter() + if isinstance(train_set,Dataset): + self.train_set = GlueMrpc.from_dict(train_set.to_dict()) + else: + self.train_set = train_set + + if isinstance(valid_set,Dataset): + self.valid_set = GlueMrpc.from_dict(valid_set.to_dict()) + else: + self.valid_set = valid_set + + self.batch_size = batch_size + self.data_collator = data_collator + + def split(self, num_collaborators): + train_split = self.data_splitter.split(self.train_set, num_collaborators) + valid_split = self.data_splitter.split(self.valid_set, num_collaborators) + return [ + GlueMrpcFederatedDataset( + self.train_set.select(train_split[i]), + self.valid_set.select(valid_split[i]), + self.batch_size + ) + for i in range(num_collaborators) + ] + + def get_feature_shape(self): + return self.train_set.get_shape() + + def get_train_loader(self, num_batches=None): + return DataLoader(self.train_set, batch_size=self.batch_size, collate_fn=data_collator) + + def get_valid_loader(self): + return DataLoader(self.valid_set, collate_fn=data_collator) + + def get_train_data_size(self): + return len(self.train_set) + + def get_valid_data_size(self): + return len(self.valid_set) + + +# %% +from openfl.federated.task import PyTorchTaskRunner +class LLMTaskRunner(PyTorchTaskRunner): + def __init__(self, base_model_name, data_loader, device=None, metric=None, **kwargs): + kwargs['data_loader'] = data_loader + super().__init__(device, **kwargs) + self.base_model_name = base_model_name + self.metric = metric + self._init_model() + self._init_optimizer() + + def _init_model(self): + model = AutoModelForSequenceClassification.from_pretrained( + self.base_model_name, return_dict=True) + peft_config = LoraConfig(task_type=TaskType.SEQ_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias="all") + self.model = get_peft_model(model, peft_config) + + def _init_optimizer(self): + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [p for n, p in self.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.model.named_parameters() if any(nd in n for nd in no_decay)], + "weight_decay": 0.0, + }, + ] + self.optimizer = AdamW(optimizer_grouped_parameters, lr=0.01) + + self.training_round_completed = False + self.initialize_tensorkeys_for_functions() + + def state_dict(self): + return get_peft_model_state_dict(self.model) + + def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): + return set_peft_model_state_dict( + self.model, state_dict + ) + + def validate(self, col_name, round_num, input_tensor_dict, + use_tqdm=False, **kwargs): + """Validate. + + Run validation of the model on the local data. + + Args: + col_name: Name of the collaborator + round_num: What round is it + input_tensor_dict: Required input tensors (for model) + use_tqdm (bool): Use tqdm to print a progress bar (Default=True) + + Returns: + global_output_dict: Tensors to send back to the aggregator + local_output_dict: Tensors to maintain in the local TensorDB + + """ + self.rebuild_model(round_num, input_tensor_dict, validation=True) + self.model.eval() + self.model.to(self.device) + val_score = 0 + total_samples = 0 + + loader = self.data_loader.get_valid_loader() + if use_tqdm: + loader = tqdm(loader, desc='validate') + + with pt.no_grad(): + for sample in loader: + samples = sample['input_ids'].shape[0] + total_samples += samples + output = self.model(**sample) + # get the index of the max log-probability + logits = output.logits + predictions = torch.argmax(logits, dim=-1) + metric.add_batch(predictions=predictions, references=sample['labels']) + val_score = metric.compute()['accuracy'] + + origin = col_name + suffix = 'validate' + if kwargs['apply'] == 'local': + suffix += '_local' + else: + suffix += '_agg' + tags = ('metric',) + tags = change_tags(tags, add_field=suffix) + # TODO figure out a better way to pass in metric for this pytorch + # validate function + output_tensor_dict = { + TensorKey('acc', origin, round_num, True, tags): + np.array(val_score) + } + + # Empty list represents metrics that should only be stored locally + return output_tensor_dict, {} + + def train_epoch(self, batch_generator) -> Metric: + """Train single epoch. + + Override this function in order to use custom training. + + Args: + batch_generator: Train dataset batch generator. Yields (samples, targets) tuples of + size = `self.data_loader.batch_size`. + Returns: + Metric: An object containing name and np.ndarray value. + """ + losses = [] + for sample in batch_generator: + self.optimizer.zero_grad() + output = self.model(**sample) + loss = output.loss + loss.backward() + self.optimizer.step() + losses.append(loss.detach().cpu().numpy()) + loss = np.mean(losses) + if self.model.config.problem_type == "regression": + loss_fct = MSELoss() + elif self.model.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + elif self.model.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + return Metric(name=loss_fct._get_name(), value=np.array(loss)) + + + def save_native(self, filepath, model_state_dict_key='model_state_dict', + optimizer_state_dict_key='optimizer_state_dict', **kwargs): + """ + Save model and optimizer states in a picked file specified by the \ + filepath. model_/optimizer_state_dicts are stored in the keys provided. \ + Uses pt.save(). + + Args: + filepath (string) : Path to pickle file to be + created by pt.save(). + model_state_dict_key (string) : key for model state dict + in pickled file. + optimizer_state_dict_key (string) : key for optimizer state + dict in picked file. + kwargs : unused + + Returns: + None + """ + pickle_dict = { + model_state_dict_key: get_peft_model_state_dict(self.model), + optimizer_state_dict_key: self.optimizer.state_dict() + } + pt.save(pickle_dict, filepath) + + +# %% +metric = load_metric('glue', "mrpc") +base_model_name = "roberta-large" +padding_side = "right" +tokenizer = AutoTokenizer.from_pretrained(base_model_name, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + +data_collator, tokenized_datasets = get_glue_mrpc_dataset(tokenizer) +#Create a federated model using the pytorch class, lambda optimizer function, and loss function +#tokenized_datasets['train'] = tokenized_datasets['train'].select(range(1000)) +#tokenized_datasets['test'] = tokenized_datasets['test'].select(range(500)) +train_set = GlueMrpc.from_dict(tokenized_datasets['train'].to_dict()) +valid_set = GlueMrpc.from_dict(tokenized_datasets['test'].to_dict()) +fl_data = GlueMrpcFederatedDataset(train_set, valid_set, batch_size=32) + +# %% [markdown] +# The `FederatedModel` object is a wrapper around your Keras, Tensorflow or PyTorch model that makes it compatible with openfl. It provides built in federated training and validation functions that we will see used below. Using it's `setup` function, collaborator models and datasets can be automatically defined for the experiment. + +# %% +# %% +fx.init('torch_cnn_mnist') +num_collaborators = 2 +collaborator_models = [ + LLMTaskRunner( + base_model_name, + data_loader=data_slice, + metric=metric + ) + for data_slice in fl_data.split(num_collaborators)] +collaborators = {'one':collaborator_models[0],'two':collaborator_models[1]}#, 'three':collaborator_models[2]} + +# %% +#Original TinyImageNet dataset +print(f'Original training data size: {len(fl_data.train_set)}') +print(f'Original validation data size: {len(fl_data.valid_set)}\n') + +#Collaborator one's data +for i, model in enumerate(collaborator_models): + print(f'Collaborator {i}\'s training data size: {len(model.data_loader.train_set)}') + print(f'Collaborator {i}\'s validation data size: {len(model.data_loader.valid_set)}\n') + +# %% +#Run experiment, return trained FederatedModel +final_fl_model = fx.run_experiment(collaborators,{'aggregator.settings.rounds_to_train':5}) + +# %% +#Save final model +final_fl_model.save_native('final_model.pth') + +# %% diff --git a/openfl-tutorials/llm_trainer.py b/openfl-tutorials/llm_trainer.py new file mode 100644 index 0000000000..26c0320785 --- /dev/null +++ b/openfl-tutorials/llm_trainer.py @@ -0,0 +1,1055 @@ +# %% + +import warnings +from typing import Any, Dict, List, Mapping, Optional, Union + +import huggingface_hub.utils as hf_hub_utils +import numpy as np +import openfl.native as fx +import torch +import torch as pt +import torch.nn as nn +from accelerate import Accelerator +from datasets import Dataset, load_dataset, load_metric +from openfl.federated import PyTorchTaskRunner, TaskRunner +from openfl.federated.task.runner_pt import change_tags +from openfl.utilities import Metric, TensorKey +from openfl.utilities.data_splitters import EqualNumPyDataSplitter +from peft import LoraConfig, TaskType, get_peft_model +from peft.utils import get_peft_model_state_dict, set_peft_model_state_dict +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.optim import AdamW +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import (AutoModelForSequenceClassification, AutoTokenizer, + DataCollatorWithPadding, Trainer, TrainingArguments, + get_scheduler) + +from transformers import (AutoConfig, AutoModelForSequenceClassification, + AutoTokenizer, DataCollatorWithPadding, Trainer) +from transformers.trainer_utils import (PREFIX_CHECKPOINT_DIR, BestRun, + EvalLoopOutput, EvalPrediction, + FSDPOption, HPSearchBackend, + HubStrategy, IntervalStrategy, + PredictionOutput, + RemoveColumnsCollator, + ShardedDDPOption, TrainerMemoryTracker, + TrainOutput, default_compute_objective, + denumpify_detensorize, + enable_full_determinism, + find_executable_batch_size, + get_last_checkpoint, has_length, + number_of_arguments, seed_worker, + set_seed, speed_metrics) +from transformers.utils import (ADAPTER_CONFIG_NAME, ADAPTER_SAFE_WEIGHTS_NAME, + ADAPTER_WEIGHTS_NAME, CONFIG_NAME, + SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, WEIGHTS_NAME, + PushInProgress, can_return_loss, find_labels, + is_accelerate_available, is_apex_available, + is_bitsandbytes_available, + is_datasets_available, is_in_notebook, + is_ipex_available, is_peft_available, + is_safetensors_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_torch_compile_available, + is_torch_neuroncore_available, + is_torch_tpu_available, logging, strtobool) + +import contextlib +import copy +import functools +import glob +import importlib.metadata +import inspect +import math +import os +import random +import re +import shutil +import sys +import time +import warnings +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union + + +# Integrations must be imported before ML frameworks: +# isort: off +from transformers.integrations import ( + get_reporting_integration_callbacks, + hp_params, + is_fairscale_available, +) + +# isort: on + +import huggingface_hub.utils as hf_hub_utils +import numpy as np +import torch +import torch.distributed as dist +from huggingface_hub import Repository, create_repo, upload_folder +from packaging import version +from torch import nn +from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler + +from transformers import __version__ +from transformers.configuration_utils import PretrainedConfig +from transformers.data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator +from transformers.debug_utils import DebugOption, DebugUnderflowOverflow +from transformers.dependency_versions_check import dep_version_check +from transformers.hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend +from transformers.integrations.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_available +from transformers.modelcard import TrainingSummary +from transformers.modeling_utils import PreTrainedModel, load_sharded_checkpoint, unwrap_model +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, MODEL_MAPPING_NAMES +from transformers.optimization import Adafactor, get_scheduler +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_less_than_1_11 +from transformers.tokenization_utils_base import PreTrainedTokenizerBase +from transformers.trainer_callback import ( + CallbackHandler, + DefaultFlowCallback, + PrinterCallback, + ProgressCallback, + TrainerCallback, + TrainerControl, + TrainerState, +) +from transformers.trainer_pt_utils import ( + DistributedTensorGatherer, + IterableDatasetShard, + LabelSmoother, + LengthGroupedSampler, + SequentialDistributedSampler, + distributed_broadcast_scalars, + distributed_concat, + find_batch_size, + get_dataloader_sampler, + get_model_param_count, + get_module_class_from_name, + get_parameter_names, + nested_concat, + nested_detach, + nested_numpify, + nested_xla_mesh_reduce, + reissue_pt_warnings, + remove_dummy_checkpoint, +) +from transformers.trainer_utils import ( + PREFIX_CHECKPOINT_DIR, + BestRun, + EvalLoopOutput, + EvalPrediction, + FSDPOption, + HPSearchBackend, + HubStrategy, + IntervalStrategy, + PredictionOutput, + RemoveColumnsCollator, + ShardedDDPOption, + TrainerMemoryTracker, + TrainOutput, + default_compute_objective, + denumpify_detensorize, + enable_full_determinism, + find_executable_batch_size, + get_last_checkpoint, + has_length, + number_of_arguments, + seed_worker, + set_seed, + speed_metrics, +) +from transformers.training_args import OptimizerNames, ParallelMode, TrainingArguments +from transformers.utils import ( + ADAPTER_CONFIG_NAME, + ADAPTER_SAFE_WEIGHTS_NAME, + ADAPTER_WEIGHTS_NAME, + CONFIG_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + PushInProgress, + can_return_loss, + find_labels, + is_accelerate_available, + is_apex_available, + is_bitsandbytes_available, + is_datasets_available, + is_in_notebook, + is_ipex_available, + is_peft_available, + is_safetensors_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_torch_compile_available, + is_torch_neuroncore_available, + is_torch_tpu_available, + logging, + strtobool, +) +from transformers.utils.quantization_config import QuantizationMethod +# Name of the files used for checkpointing +TRAINING_ARGS_NAME = "training_args.bin" +TRAINER_STATE_NAME = "trainer_state.json" +OPTIMIZER_NAME = "optimizer.pt" +OPTIMIZER_NAME_BIN = "optimizer.bin" +SCHEDULER_NAME = "scheduler.pt" +SCALER_NAME = "scaler.pt" + +if is_torch_tpu_available(check_device=False): + import torch_xla.core.xla_model as xm + import torch_xla.debug.metrics as met + +if is_apex_available(): + from apex import amp + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") + + from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat +else: + IS_SAGEMAKER_MP_POST_1_10 = False + +if is_accelerate_available(): + from accelerate import Accelerator, skip_first_batches + from accelerate import __version__ as accelerate_version + from accelerate.utils import DistributedDataParallelKwargs, GradientAccumulationPlugin + + if version.parse(accelerate_version) > version.parse("0.20.3"): + from accelerate.utils import ( + load_fsdp_model, + load_fsdp_optimizer, + save_fsdp_model, + save_fsdp_optimizer, + ) + + if is_deepspeed_available(): + from accelerate.utils import DeepSpeedSchedulerWrapper + +logger = logging.get_logger(__name__) +def get_glue_mrpc_dataset(tokenizer): + dataset = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer( + examples["sentence1"], + examples["sentence2"], + truncation=True, + max_length=None, + ) + return outputs + + tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + tokenized_datasets.set_format("torch") + data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + return data_collator, tokenized_datasets + + +class GlueMrpc(Dataset): + """ + Has 5.8k pairs of sentences with annotations if the two sentences are equivalent + """ + + def get_shape(self): + if not hasattr(self, "saved_shape"): + self.saved_shape = max([len(i) for i in self.data["input_ids"]]) + return self.saved_shape + + +class GlueMrpcFederatedDataset(DataLoader): + def __init__(self, train_set, valid_set, batch_size, data_collator=None): + self.data_splitter = EqualNumPyDataSplitter() + if isinstance(train_set, Dataset): + self.train_set = GlueMrpc.from_dict(train_set.to_dict()) + else: + self.train_set = train_set + + if isinstance(valid_set, Dataset): + self.valid_set = GlueMrpc.from_dict(valid_set.to_dict()) + else: + self.valid_set = valid_set + + self.batch_size = batch_size + self.data_collator = data_collator + + def split(self, num_collaborators): + train_split = self.data_splitter.split(self.train_set, num_collaborators) + valid_split = self.data_splitter.split(self.valid_set, num_collaborators) + return [ + GlueMrpcFederatedDataset( + self.train_set.select(train_split[i]), + self.valid_set.select(valid_split[i]), + self.batch_size, + ) + for i in range(num_collaborators) + ] + + def get_feature_shape(self): + return self.train_set.get_shape() + + def get_train_loader(self, num_batches=None): + return DataLoader( + self.train_set, batch_size=self.batch_size, collate_fn=data_collator + ) + + def get_valid_loader(self): + return DataLoader(self.valid_set, collate_fn=data_collator) + + def get_train_data_size(self): + return len(self.train_set) + + def get_valid_data_size(self): + return len(self.valid_set) + + +# %% +from transformers import Trainer + + +class FedTrainer(Trainer): + def setup_training_loop(self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None + ): + + def _inner_training_loop( + self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None + ): + self.accelerator.free_memory() + self._train_batch_size = batch_size + logger.debug(f"Currently training with a batch size of: {self._train_batch_size}") + # Data loader and number of training steps + train_dataloader = self.get_train_dataloader() + + # Setting up training control variables: + # number of training epochs: num_train_epochs + # number of training steps per epoch: num_update_steps_per_epoch + # total number of training steps to execute: max_steps + total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size + + len_dataloader = None + num_train_tokens = None + if has_length(train_dataloader): + len_dataloader = len(train_dataloader) + num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps + num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1) + num_examples = self.num_examples(train_dataloader) + if args.max_steps > 0: + max_steps = args.max_steps + num_train_epochs = args.max_steps // num_update_steps_per_epoch + int( + args.max_steps % num_update_steps_per_epoch > 0 + ) + # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's + # the best we can do. + num_train_samples = args.max_steps * total_train_batch_size + if args.include_tokens_per_second: + num_train_tokens = ( + self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps + ) + else: + max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch) + num_train_epochs = math.ceil(args.num_train_epochs) + num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs + if args.include_tokens_per_second: + num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs + elif args.max_steps > 0: # Rely on max_steps when dataloader does not have a working size + max_steps = args.max_steps + # Setting a very large number of epochs so we go as many times as necessary over the iterator. + num_train_epochs = sys.maxsize + num_update_steps_per_epoch = max_steps + num_examples = total_train_batch_size * args.max_steps + num_train_samples = args.max_steps * total_train_batch_size + if args.include_tokens_per_second: + num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps + else: + raise ValueError( + "args.max_steps must be set to a positive value if dataloader does not have a length, was" + f" {args.max_steps}" + ) + + if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug: + if self.args.n_gpu > 1: + # nn.DataParallel(model) replicates the model, creating new variables and module + # references registered here no longer work on other gpus, breaking the module + raise ValueError( + "Currently --debug underflow_overflow is not supported under DP. Please use DDP" + " (torch.distributed.launch)." + ) + else: + debug_overflow = DebugUnderflowOverflow(self.model) # noqa + + delay_optimizer_creation = ( + self.sharded_ddp is not None + and self.sharded_ddp != ShardedDDPOption.SIMPLE + or is_sagemaker_mp_enabled() + or self.fsdp is not None + or self.is_fsdp_enabled + ) + + # We need to reset the scheduler, as its parameters may be different on subsequent calls + if self._created_lr_scheduler: + self.lr_scheduler = None + self._created_lr_scheduler = False + + if self.is_deepspeed_enabled: + self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps) + + if not delay_optimizer_creation: + self.create_optimizer_and_scheduler(num_training_steps=max_steps) + + self.state = TrainerState() + self.state.is_hyper_param_search = trial is not None + + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + + # Activate gradient checkpointing if needed + if args.gradient_checkpointing: + self.model.gradient_checkpointing_enable() + + model = self._wrap_model(self.model_wrapped) + + # as the model is wrapped, don't use `accelerator.prepare` + # this is for unhandled cases such as + # Fairscale Sharded DDP, FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX + use_accelerator_prepare = True if model is self.model else False + + if delay_optimizer_creation: + if use_accelerator_prepare: + self.model = self.accelerator.prepare(self.model) + self.create_optimizer_and_scheduler(num_training_steps=max_steps) + + # prepare using `accelerator` prepare + if use_accelerator_prepare: + self.model.train() + if hasattr(self.lr_scheduler, "step"): + if self.use_apex: + model = self.accelerator.prepare(self.model) + else: + model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer) + else: + # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config. + model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( + self.model, self.optimizer, self.lr_scheduler + ) + + if self.is_fsdp_enabled: + self.model = self.model_wrapped = model + + # for the rest of this function `model` is the outside model, whether it was wrapped or not + if model is not self.model: + self.model_wrapped = model + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model_wrapped + + # ckpt loading + if resume_from_checkpoint is not None: + if self.is_deepspeed_enabled: + deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint) + elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled: + self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped) + + # Check if saved optimizer or scheduler states exist + self._load_optimizer_and_scheduler(resume_from_checkpoint) + + # important: at this point: + # self.model is the Transformers Model + # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), + # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc. + + # Train! + logger.info("***** Running training *****") + logger.info(f" Num examples = {num_examples:,}") + logger.info(f" Num Epochs = {num_train_epochs:,}") + logger.info(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}") + if self.args.per_device_train_batch_size != self._train_batch_size: + logger.info(f" Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {max_steps:,}") + logger.info(f" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}") + + self.state.epoch = 0 + start_time = time.time() + epochs_trained = 0 + steps_trained_in_current_epoch = 0 + steps_trained_progress_bar = None + + # Check if continuing training from a checkpoint + if resume_from_checkpoint is not None and os.path.isfile( + os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME) + ): + self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) + epochs_trained = self.state.global_step // num_update_steps_per_epoch + if not args.ignore_data_skip: + steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch) + steps_trained_in_current_epoch *= args.gradient_accumulation_steps + else: + steps_trained_in_current_epoch = 0 + + logger.info(" Continuing training from checkpoint, will skip to saved global_step") + logger.info(f" Continuing training from epoch {epochs_trained}") + logger.info(f" Continuing training from global step {self.state.global_step}") + if not args.ignore_data_skip: + logger.info( + f" Will skip the first {epochs_trained} epochs then the first" + f" {steps_trained_in_current_epoch} batches in the first epoch." + ) + + # Update the references + self.callback_handler.model = self.model + self.callback_handler.optimizer = self.optimizer + self.callback_handler.lr_scheduler = self.lr_scheduler + self.callback_handler.train_dataloader = train_dataloader + if self.hp_name is not None and self._trial is not None: + # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial + # parameter to Train when using DDP. + self.state.trial_name = self.hp_name(self._trial) + if trial is not None: + assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial + self.state.trial_params = hp_params(assignments) + else: + self.state.trial_params = None + # This should be the same if the state has been saved but in case the training arguments changed, it's safer + # to set this after the load. + self.state.max_steps = max_steps + self.state.num_train_epochs = num_train_epochs + self.state.is_local_process_zero = self.is_local_process_zero() + self.state.is_world_process_zero = self.is_world_process_zero() + + # tr_loss is a tensor to avoid synchronization of TPUs through .item() + tr_loss = torch.tensor(0.0).to(args.device) + # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses + self._total_loss_scalar = 0.0 + self._globalstep_last_logged = self.state.global_step + model.zero_grad() + + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + # Skip the first epochs_trained epochs to get the random state of the dataloader at the right point. + if not args.ignore_data_skip: + for epoch in range(epochs_trained): + sampler = get_dataloader_sampler(train_dataloader) + is_random_sampler = isinstance(sampler, RandomSampler) + if is_torch_less_than_1_11 or not is_random_sampler: + # We just need to begin an iteration to create the randomization of the sampler. + for _ in train_dataloader: + break + else: + # Otherwise we need to call the whooooole sampler cause there is some random operation added + # AT THE VERY END! + sampler = sampler if sampler is not None else [] + _ = list(sampler) + + total_batched_samples = 0 + for epoch in range(epochs_trained, num_train_epochs): + epoch_iterator = train_dataloader + + # Reset the past mems state at the beginning of each epoch if necessary. + if args.past_index >= 0: + self._past = None + + steps_in_epoch = ( + len(epoch_iterator) + if len_dataloader is not None + else args.max_steps * args.gradient_accumulation_steps + ) + self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) + + if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: + self._load_rng_state(resume_from_checkpoint) + + rng_to_sync = False + steps_skipped = 0 + if steps_trained_in_current_epoch > 0: + epoch_iterator = skip_first_batches(epoch_iterator, steps_trained_in_current_epoch) + steps_skipped = steps_trained_in_current_epoch + steps_trained_in_current_epoch = 0 + rng_to_sync = True + + step = -1 + for step, inputs in enumerate(epoch_iterator): + total_batched_samples += 1 + if rng_to_sync: + self._load_rng_state(resume_from_checkpoint) + rng_to_sync = False + + # Skip past any already trained steps if resuming training + if steps_trained_in_current_epoch > 0: + steps_trained_in_current_epoch -= 1 + if steps_trained_progress_bar is not None: + steps_trained_progress_bar.update(1) + if steps_trained_in_current_epoch == 0: + self._load_rng_state(resume_from_checkpoint) + continue + elif steps_trained_progress_bar is not None: + steps_trained_progress_bar.close() + steps_trained_progress_bar = None + + if step % args.gradient_accumulation_steps == 0: + self.control = self.callback_handler.on_step_begin(args, self.state, self.control) + + with self.accelerator.accumulate(model): + tr_loss_step = self.training_step(model, inputs) + + if ( + args.logging_nan_inf_filter + and not is_torch_tpu_available() + and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)) + ): + # if loss is nan or inf simply add the average of previous logged losses + tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) + else: + tr_loss += tr_loss_step + + self.current_flos += float(self.floating_point_ops(inputs)) + + is_last_step_and_steps_less_than_grad_acc = ( + steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch + ) + + if ( + total_batched_samples % args.gradient_accumulation_steps == 0 + or + # last step in epoch but step is always smaller than gradient_accumulation_steps + is_last_step_and_steps_less_than_grad_acc + ): + # the `or` condition of `is_last_step_and_steps_less_than_grad_acc` is not covered + # in accelerate. So, explicitly enable sync gradients to True in that case. + if is_last_step_and_steps_less_than_grad_acc or ( + version.parse(accelerate_version) <= version.parse("0.20.3") + ): + self.accelerator.gradient_state._set_sync_gradients(True) + + # Gradient clipping + if args.max_grad_norm is not None and args.max_grad_norm > 0: + # deepspeed does its own clipping + + if self.do_grad_scaling: + # Reduce gradients first for XLA + if is_torch_tpu_available(): + gradients = xm._fetch_gradients(self.optimizer) + xm.all_reduce("sum", gradients, scale=1.0 / xm.xrt_world_size()) + # AMP: gradients need unscaling + self.scaler.unscale_(self.optimizer) + + if is_sagemaker_mp_enabled() and args.fp16: + self.optimizer.clip_master_grads(args.max_grad_norm) + elif hasattr(self.optimizer, "clip_grad_norm"): + # Some optimizers (like the sharded optimizer) have a specific way to do gradient clipping + self.optimizer.clip_grad_norm(args.max_grad_norm) + elif hasattr(model, "clip_grad_norm_"): + # Some models (like FullyShardedDDP) have a specific way to do gradient clipping + model.clip_grad_norm_(args.max_grad_norm) + elif self.use_apex: + # Revert to normal clipping otherwise, handling Apex or full precision + nn.utils.clip_grad_norm_( + amp.master_params(self.optimizer), + args.max_grad_norm, + ) + else: + self.accelerator.clip_grad_norm_( + model.parameters(), + args.max_grad_norm, + ) + + # Optimizer step + optimizer_was_run = True + if is_torch_tpu_available(): + if self.do_grad_scaling: + self.scaler.step(self.optimizer) + self.scaler.update() + else: + # tpu-comment: accelerate wrapped optimizers call xm.optimizer_step + self.optimizer.step() + elif self.do_grad_scaling: + scale_before = self.scaler.get_scale() + self.scaler.step(self.optimizer) + self.scaler.update() + scale_after = self.scaler.get_scale() + optimizer_was_run = scale_before <= scale_after + else: + self.optimizer.step() + optimizer_was_run = not self.accelerator.optimizer_step_was_skipped + + if optimizer_was_run: + # Delay optimizer scheduling until metrics are generated + if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + self.lr_scheduler.step() + + model.zero_grad() + self.state.global_step += 1 + self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + + self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval) + else: + self.control = self.callback_handler.on_substep_end(args, self.state, self.control) + + if self.control.should_epoch_stop or self.control.should_training_stop: + break + if step < 0: + logger.warning( + "There seems to be not a single sample in your epoch_iterator, stopping training at step" + f" {self.state.global_step}! This is expected if you're using an IterableDataset and set" + f" num_steps ({max_steps}) higher than the number of available samples." + ) + self.control.should_training_stop = True + + self.control = self.callback_handler.on_epoch_end(args, self.state, self.control) + self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval) + + if DebugOption.TPU_METRICS_DEBUG in self.args.debug: + if is_torch_tpu_available(): + # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.) + xm.master_print(met.metrics_report()) + else: + logger.warning( + "You enabled PyTorch/XLA debug metrics but you don't have a TPU " + "configured. Check your training configuration if this is unexpected." + ) + if self.control.should_training_stop: + break + + if args.past_index and hasattr(self, "_past"): + # Clean the state at the end of training + delattr(self, "_past") + + logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n") + if args.load_best_model_at_end and self.state.best_model_checkpoint is not None: + # Wait for everyone to get here so we are sur the model has been saved by process 0. + if is_torch_tpu_available(): + xm.rendezvous("load_best_model_at_end") + elif args.parallel_mode == ParallelMode.DISTRIBUTED: + dist.barrier() + elif is_sagemaker_mp_enabled(): + smp.barrier() + + self._load_best_model() + + # add remaining tr_loss + self._total_loss_scalar += tr_loss.item() + train_loss = self._total_loss_scalar / self.state.global_step + + metrics = speed_metrics( + "train", + start_time, + num_samples=num_train_samples, + num_steps=self.state.max_steps, + num_tokens=num_train_tokens, + ) + self.store_flos() + metrics["total_flos"] = self.state.total_flos + metrics["train_loss"] = train_loss + + self.is_in_train = False + + self._memory_tracker.stop_and_update_metrics(metrics) + + self.log(metrics) + + run_dir = self._get_output_dir(trial) + checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir) + + # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save. + if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1: + for checkpoint in checkpoints_sorted: + if not os.path.samefile(checkpoint, self.state.best_model_checkpoint): + logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit") + shutil.rmtree(checkpoint) + + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + + # Wait for the checkpoint to be uploaded. + self._finish_current_push() + + return TrainOutput(self.state.global_step, train_loss, metrics) + + + +from openfl.federated.task import PyTorchTaskRunner + + +class LLMTaskRunner(PyTorchTaskRunner): + def __init__( + self, base_model_name, data_loader, device=None, metric=None, collator=None, tokenizer=None, **kwargs + ): + kwargs["data_loader"] = data_loader + super().__init__(device, **kwargs) + self.base_model_name = base_model_name + self.metric = metric + self._init_model() + #self._init_optimizer() + self.LR = 0.01 + self.WEIGHT_DECAY = 0.01 + self.NUM_EPOCHS = 100 + training_args = TrainingArguments( + output_dir="roberta-large-lora-seq1_tests", + learning_rate=self.LR, + per_device_train_batch_size=32, + per_device_eval_batch_size=32, + num_train_epochs=self.NUM_EPOCHS, + weight_decay=self.WEIGHT_DECAY, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, + ) + metric = load_metric('glue', 'mrpc') + def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + self.trainer = Trainer( + model=self.model, + args=training_args, + #train_dataset=tokenized_datasets["train"], + train_dataset=data_loader.train_set, + #eval_dataset=tokenized_datasets["test"], + eval_dataset=data_loader.valid_set, + tokenizer=tokenizer, + compute_metrics=compute_metrics, + data_collator=collator + ) + + def _init_model(self): + model = AutoModelForSequenceClassification.from_pretrained( + self.base_model_name, return_dict=True + ) + peft_config = LoraConfig( + task_type=TaskType.SEQ_CLS, + inference_mode=False, + r=16, + lora_alpha=16, + lora_dropout=0.1, + bias="all", + ) + self.model = get_peft_model(model, peft_config) + + def _init_optimizer(self): + no_decay = ["bias", "LayerNorm.weight"] + optimizer_grouped_parameters = [ + { + "params": [ + p + for n, p in self.model.named_parameters() + if not any(nd in n for nd in no_decay) + ], + "weight_decay": self.WEIGHT_DECAY, + }, + { + "params": [ + p + for n, p in self.model.named_parameters() + if any(nd in n for nd in no_decay) + ], + "weight_decay": 0.0, + }, + ] + self.optimizer = AdamW(optimizer_grouped_parameters, lr=self.LR) + + self.training_round_completed = False + self.initialize_tensorkeys_for_functions() + + def state_dict(self): + return get_peft_model_state_dict(self.model) + + def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): + return set_peft_model_state_dict(self.model, state_dict) + + def validate( + self, col_name, round_num, input_tensor_dict, use_tqdm=False, **kwargs + ): + """Validate. + + Run validation of the model on the local data. + + Args: + col_name: Name of the collaborator + round_num: What round is it + input_tensor_dict: Required input tensors (for model) + use_tqdm (bool): Use tqdm to print a progress bar (Default=True) + + Returns: + global_output_dict: Tensors to send back to the aggregator + local_output_dict: Tensors to maintain in the local TensorDB + + """ + self.rebuild_model(round_num, input_tensor_dict, validation=True) + self.model.eval() + self.model.to(self.device) + val_score = 0 + total_samples = 0 + + loader = self.data_loader.get_valid_loader() + if use_tqdm: + loader = tqdm(loader, desc="validate") + + with pt.no_grad(): + for sample in loader: + samples = sample["input_ids"].shape[0] + total_samples += samples + output = self.model(**sample) + # get the index of the max log-probability + logits = output.logits + predictions = torch.argmax(logits, dim=-1) + metric.add_batch(predictions=predictions, references=sample["labels"]) + val_score = metric.compute()["accuracy"] + + origin = col_name + suffix = "validate" + if kwargs["apply"] == "local": + suffix += "_local" + else: + suffix += "_agg" + tags = ("metric",) + tags = change_tags(tags, add_field=suffix) + # TODO figure out a better way to pass in metric for this pytorch + # validate function + output_tensor_dict = { + TensorKey("acc", origin, round_num, True, tags): np.array(val_score) + } + + # Empty list represents metrics that should only be stored locally + return output_tensor_dict, {} + + def train_epoch(self, batch_generator) -> Metric: + """Train single epoch. + + Override this function in order to use custom training. + + Args: + batch_generator: Train dataset batch generator. Yields (samples, targets) tuples of + size = `self.data_loader.batch_size`. + Returns: + Metric: An object containing name and np.ndarray value. + """ + losses = [] + for sample in batch_generator: + self.optimizer.zero_grad() + output = self.model(**sample) + loss = output.loss + loss.backward() + self.optimizer.step() + losses.append(loss.detach().cpu().numpy()) + loss = np.mean(losses) + if self.model.config.problem_type == "regression": + loss_fct = MSELoss() + elif self.model.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + elif self.model.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + return Metric(name=loss_fct._get_name(), value=np.array(loss)) + + def save_native( + self, + filepath, + model_state_dict_key="model_state_dict", + optimizer_state_dict_key="optimizer_state_dict", + **kwargs, + ): + """ + Save model and optimizer states in a picked file specified by the \ + filepath. model_/optimizer_state_dicts are stored in the keys provided. \ + Uses pt.save(). + + Args: + filepath (string) : Path to pickle file to be + created by pt.save(). + model_state_dict_key (string) : key for model state dict + in pickled file. + optimizer_state_dict_key (string) : key for optimizer state + dict in picked file. + kwargs : unused + + Returns: + None + """ + pickle_dict = { + model_state_dict_key: get_peft_model_state_dict(self.model), + optimizer_state_dict_key: self.optimizer.state_dict(), + } + pt.save(pickle_dict, filepath) + + +# %% +metric = load_metric("glue", "mrpc") +base_model_name = "roberta-large" +padding_side = "right" +tokenizer = AutoTokenizer.from_pretrained(base_model_name, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + +data_collator, tokenized_datasets = get_glue_mrpc_dataset(tokenizer) +# Create a federated model using the pytorch class, lambda optimizer function, and loss function +# tokenized_datasets['train'] = tokenized_datasets['train'].select(range(1000)) +# tokenized_datasets['test'] = tokenized_datasets['test'].select(range(500)) +train_set = GlueMrpc.from_dict(tokenized_datasets["train"].to_dict()) +valid_set = GlueMrpc.from_dict(tokenized_datasets["test"].to_dict()) +fl_data = GlueMrpcFederatedDataset(train_set, valid_set, batch_size=32) + +# %% [markdown] +# The `FederatedModel` object is a wrapper around your Keras, Tensorflow or PyTorch model that makes it compatible with openfl. It provides built in federated training and validation functions that we will see used below. Using it's `setup` function, collaborator models and datasets can be automatically defined for the experiment. + +# %% +# %% +fx.init("torch_cnn_mnist") +num_collaborators = 2 +collaborator_models = [ + LLMTaskRunner(base_model_name, data_loader=data_slice, metric=metric, collator=data_collator, tokenizer=tokenizer) + for data_slice in fl_data.split(num_collaborators) +] +collaborators = { + "one": collaborator_models[0], + "two": collaborator_models[1], +} # , 'three':collaborator_models[2]} + +# %% +# Original TinyImageNet dataset +print(f"Original training data size: {len(fl_data.train_set)}") +print(f"Original validation data size: {len(fl_data.valid_set)}\n") + +# Collaborator one's data +for i, model in enumerate(collaborator_models): + print(f"Collaborator {i}'s training data size: {len(model.data_loader.train_set)}") + print( + f"Collaborator {i}'s validation data size: {len(model.data_loader.valid_set)}\n" + ) + +# %% +# Run experiment, return trained FederatedModel +final_fl_model = fx.run_experiment( + collaborators, {"aggregator.settings.rounds_to_train": 5} +) + +# %% +# Save final model +final_fl_model.save_native("final_model.pth") + +# %% diff --git a/openfl-workspace/torch_llm_sequence_classification/.workspace b/openfl-workspace/torch_llm_sequence_classification/.workspace new file mode 100644 index 0000000000..3c2c5d08b4 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/.workspace @@ -0,0 +1,2 @@ +current_plan_name: default + diff --git a/openfl-workspace/torch_llm_sequence_classification/plan/cols.yaml b/openfl-workspace/torch_llm_sequence_classification/plan/cols.yaml new file mode 100644 index 0000000000..95307de3bc --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/plan/cols.yaml @@ -0,0 +1,5 @@ +# Copyright (C) 2020-2021 Intel Corporation +# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you. + +collaborators: + \ No newline at end of file diff --git a/openfl-workspace/torch_llm_sequence_classification/plan/data.yaml b/openfl-workspace/torch_llm_sequence_classification/plan/data.yaml new file mode 100644 index 0000000000..cc7e6c9595 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/plan/data.yaml @@ -0,0 +1,9 @@ +## Copyright (C) 2020-2021 Intel Corporation +# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you. + +# all keys under 'collaborators' corresponds to a specific colaborator name the corresponding dictionary has data_name, data_path pairs. +# Note that in the mnist case we do not store the data locally, and the data_path is used to pass an integer that helps the data object +# construct the shard of the mnist dataset to be use for this collaborator. + +# collaborator_name ,data_directory_path +one,1 \ No newline at end of file diff --git a/openfl-workspace/torch_llm_sequence_classification/plan/defaults b/openfl-workspace/torch_llm_sequence_classification/plan/defaults new file mode 100644 index 0000000000..fb82f9c5b6 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/plan/defaults @@ -0,0 +1,2 @@ +../../workspace/plan/defaults + diff --git a/openfl-workspace/torch_llm_sequence_classification/plan/plan.yaml b/openfl-workspace/torch_llm_sequence_classification/plan/plan.yaml new file mode 100644 index 0000000000..47ba88b3f2 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/plan/plan.yaml @@ -0,0 +1,45 @@ +# Copyright (C) 2020-2021 Intel Corporation +# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you. + +aggregator : + defaults : plan/defaults/aggregator.yaml + template : openfl.component.Aggregator + settings : + init_state_path : save/torch_cnn_mnist_init.pbuf + best_state_path : save/torch_cnn_mnist_best.pbuf + last_state_path : save/torch_cnn_mnist_last.pbuf + rounds_to_train : 10 + log_metric_callback : + template : src.mnist_utils.write_metric + + +collaborator : + defaults : plan/defaults/collaborator.yaml + template : openfl.component.Collaborator + settings : + delta_updates : false + opt_treatment : RESET + +data_loader : + defaults : plan/defaults/data_loader.yaml + template : src.ptmnist_inmemory.PyTorchMNISTInMemory + settings : + collaborator_count : 2 + data_group_name : mnist + batch_size : 256 + +task_runner : + defaults : plan/defaults/task_runner.yaml + template : src.pt_cnn.PyTorchCNN + +network : + defaults : plan/defaults/network.yaml + +assigner : + defaults : plan/defaults/assigner.yaml + +tasks : + defaults : plan/defaults/tasks_torch.yaml + +compression_pipeline : + defaults : plan/defaults/compression_pipeline.yaml diff --git a/openfl-workspace/torch_llm_sequence_classification/requirements.txt b/openfl-workspace/torch_llm_sequence_classification/requirements.txt new file mode 100644 index 0000000000..2c3839628a --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/requirements.txt @@ -0,0 +1,12 @@ +torch==1.13.1 +torchvision==0.14.1 +tensorboard +wheel>=0.38.0 # not directly required, pinned by Snyk to avoid a vulnerability +peft +transformers +sentencepiece +huggingface_hub +accelerate +datasets +evaluate +seqeval diff --git a/openfl-workspace/torch_llm_sequence_classification/src/__init__.py b/openfl-workspace/torch_llm_sequence_classification/src/__init__.py new file mode 100644 index 0000000000..d5df5b8668 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/src/__init__.py @@ -0,0 +1,3 @@ +# Copyright (C) 2020-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +"""You may copy this file as the starting point of your own model.""" diff --git a/openfl-workspace/torch_llm_sequence_classification/src/llm_utils.py b/openfl-workspace/torch_llm_sequence_classification/src/llm_utils.py new file mode 100644 index 0000000000..032f90849e --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/src/llm_utils.py @@ -0,0 +1,114 @@ +# Copyright (C) 2020-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +import numpy as np +import openfl.native as fx +import torch +from datasets import load_dataset, load_metric +from openfl.federated import PyTorchTaskRunner, TaskRunner +from openfl.utilities import Metric, TensorKey, split_tensor_dict_for_holdouts +from openfl.utilities.data_splitters import EqualNumPyDataSplitter +from peft import LoraConfig, TaskType, get_peft_model +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.utils.data import DataLoader, Dataset, Subset +from tqdm import tqdm + +from transformers import (AutoModelForSequenceClassification, AutoTokenizer, + DataCollatorWithPadding) + +"""You may copy this file as the starting point of your own model.""" + +def get_glue_mrpc_dataset(tokenizer): + dataset = load_dataset("glue", "mrpc") + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer( + examples["sentence1"], + examples["sentence2"], + truncation=True, + max_length=None, + ) + return outputs + + tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + tokenized_datasets.set_format("torch") + data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + return data_collator, tokenized_datasets + + +def _load_raw_datashards(shard_num, collaborator_count, transform=None): + """ + Load the raw data by shard. + + Returns tuples of the dataset shard divided into training and validation. + + Args: + shard_num (int): The shard number to use + collaborator_count (int): The number of collaborators in the federation + transform: torchvision.transforms.Transform to apply to images + + Returns: + 2 tuples: (image, label) of the training, validation dataset + """ + train_data, val_data = ( + datasets.MNIST('data', train=train, download=True, transform=transform) + for train in (True, False) + ) + X_train_tot, y_train_tot = train_data.train_data, train_data.train_labels + X_valid_tot, y_valid_tot = val_data.test_data, val_data.test_labels + + # create the shards + shard_num = int(shard_num) + X_train = X_train_tot[shard_num::collaborator_count].unsqueeze(1).float() + y_train = y_train_tot[shard_num::collaborator_count] + + X_valid = X_valid_tot[shard_num::collaborator_count].unsqueeze(1).float() + y_valid = y_valid_tot[shard_num::collaborator_count] + + return (X_train, y_train), (X_valid, y_valid) + + +def load_mnist_shard(shard_num, collaborator_count, + categorical=False, channels_last=True, **kwargs): + """ + Load the MNIST dataset. + + Args: + shard_num (int): The shard to use from the dataset + collaborator_count (int): The number of collaborators in the + federation + categorical (bool): True = convert the labels to one-hot encoded + vectors (Default = True) + channels_last (bool): True = The input images have the channels + last (Default = True) + **kwargs: Additional parameters to pass to the function + + Returns: + list: The input shape + int: The number of classes + numpy.ndarray: The training data + numpy.ndarray: The training labels + numpy.ndarray: The validation data + numpy.ndarray: The validation labels + """ + num_classes = 10 + + (X_train, y_train), (X_valid, y_valid) = _load_raw_datashards( + shard_num, collaborator_count, transform=transforms.ToTensor()) + + logger.info(f'MNIST > X_train Shape : {X_train.shape}') + logger.info(f'MNIST > y_train Shape : {y_train.shape}') + logger.info(f'MNIST > Train Samples : {X_train.shape[0]}') + logger.info(f'MNIST > Valid Samples : {X_valid.shape[0]}') + + if categorical: + # convert class vectors to binary class matrices + y_train = one_hot(y_train, num_classes) + y_valid = one_hot(y_valid, num_classes) + + return num_classes, X_train, y_train, X_valid, y_valid diff --git a/openfl-workspace/torch_llm_sequence_classification/src/pt_llm.py b/openfl-workspace/torch_llm_sequence_classification/src/pt_llm.py new file mode 100644 index 0000000000..9dc2d60afb --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/src/pt_llm.py @@ -0,0 +1,208 @@ +# Copyright (C) 2020-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""You may copy this file as the starting point of your own model.""" +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import tqdm + +from openfl.federated import PyTorchTaskRunner +from openfl.utilities import TensorKey + + +def cross_entropy(output, target): + """Binary cross-entropy metric. + + Args: + output: The mode prediction + target: The target (ground truth label) + + Returns: + Binary cross-entropy with logits + + """ + return F.cross_entropy(input=output, target=target) + + +class PyTorchCNN(PyTorchTaskRunner): + """Simple CNN for classification.""" + + def __init__(self, device='cpu', **kwargs): + """Initialize. + + Args: + data: The data loader class + device: The hardware device to use for training (Default = "cpu") + **kwargs: Additional arguments to pass to the function + + """ + super().__init__(device=device, **kwargs) + + self.num_classes = self.data_loader.num_classes + self.init_network(device=self.device, **kwargs) + self._init_optimizer() + self.loss_fn = cross_entropy + self.initialize_tensorkeys_for_functions() + + def _init_optimizer(self): + """Initialize the optimizer.""" + self.optimizer = optim.Adam(self.parameters(), lr=1e-4) + + def init_network(self, + device, + print_model=True, + pool_sqrkernel_size=2, + conv_sqrkernel_size=5, + conv1_channels_out=20, + conv2_channels_out=50, + fc2_insize=500, + **kwargs): + """Create the network (model). + + Args: + device: The hardware device to use for training + print_model (bool): Print the model topology (Default=True) + pool_sqrkernel_size (int): Max pooling kernel size (Default=2), + assumes square 2x2 + conv_sqrkernel_size (int): Convolutional filter size (Default=5), + assumes square 5x5 + conv1_channels_out (int): Number of filters in first + convolutional layer (Default=20) + conv2_channels_out: Number of filters in second convolutional + layer (Default=50) + fc2_insize (int): Number of neurons in the + fully-connected layer (Default = 500) + **kwargs: Additional arguments to pass to the function + + FIXME: We are tracking only side lengths (rather than + length and width) as we are assuming square + shapes for feature and kernels. + In order that all of the input and activation components are + used (not cut off), we rely on a criterion: appropriate integers + are divisible so that all casting to int perfomed below does no + rounding (i.e. all int casting simply converts a float with '0' + in the decimal part to an int.) + + (Note this criterion held for the original input sizes considered + for this model: 28x28 and 32x32 when used with the default values + above) + """ + self.pool_sqrkernel_size = pool_sqrkernel_size + channel = self.data_loader.get_feature_shape()[0] # (channel, dim1, dim2) + self.conv1 = nn.Conv2d(channel, conv1_channels_out, conv_sqrkernel_size, 1) + + # perform some calculations to track the size of the single channel activations + # channels are first for pytorch + conv1_sqrsize_in = self.feature_shape[-1] + conv1_sqrsize_out = conv1_sqrsize_in - (conv_sqrkernel_size - 1) + # a pool operation happens after conv1 out + # (note dependence on 'forward' function below) + conv2_sqrsize_in = int(conv1_sqrsize_out / pool_sqrkernel_size) + + self.conv2 = nn.Conv2d(conv1_channels_out, conv2_channels_out, conv_sqrkernel_size, 1) + + # more tracking of single channel activation size + conv2_sqrsize_out = conv2_sqrsize_in - (conv_sqrkernel_size - 1) + # a pool operation happens after conv2 out + # (note dependence on 'forward' function below) + l0 = int(conv2_sqrsize_out / pool_sqrkernel_size) + self.fc1_insize = l0 * l0 * conv2_channels_out + self.fc1 = nn.Linear(self.fc1_insize, fc2_insize) + self.fc2 = nn.Linear(fc2_insize, self.num_classes) + if print_model: + print(self) + self.to(device) + + def forward(self, x): + """Forward pass of the model. + + Args: + x: Data input to the model for the forward pass + """ + x = F.relu(self.conv1(x)) + pl = self.pool_sqrkernel_size + x = F.max_pool2d(x, pl, pl) + x = F.relu(self.conv2(x)) + x = F.max_pool2d(x, pl, pl) + x = x.view(-1, self.fc1_insize) + x = F.relu(self.fc1(x)) + x = self.fc2(x) + return x + + def validate(self, col_name, round_num, input_tensor_dict, use_tqdm=False, **kwargs): + """Validate. + + Run validation of the model on the local data. + + Args: + col_name: Name of the collaborator + round_num: What round is it + input_tensor_dict: Required input tensors (for model) + use_tqdm (bool): Use tqdm to print a progress bar (Default=True) + + Returns: + global_output_dict: Tensors to send back to the aggregator + local_output_dict: Tensors to maintain in the local TensorDB + + """ + self.rebuild_model(round_num, input_tensor_dict, validation=True) + self.eval() + val_score = 0 + total_samples = 0 + + loader = self.data_loader.get_valid_loader() + if use_tqdm: + loader = tqdm.tqdm(loader, desc='validate') + + with torch.no_grad(): + for data, target in loader: + samples = target.shape[0] + total_samples += samples + data, target = torch.tensor(data).to( + self.device), torch.tensor(target).to( + self.device, dtype=torch.int64) + output = self(data) + # get the index of the max log-probability + pred = output.argmax(dim=1) + val_score += pred.eq(target).sum().cpu().numpy() + + origin = col_name + suffix = 'validate' + if kwargs['apply'] == 'local': + suffix += '_local' + else: + suffix += '_agg' + tags = ('metric', suffix) + # TODO figure out a better way to pass + # in metric for this pytorch validate function + output_tensor_dict = { + TensorKey('acc', origin, round_num, True, tags): + np.array(val_score / total_samples) + } + + # Empty list represents metrics that should only be stored locally + return output_tensor_dict, {} + + def reset_opt_vars(self): + """Reset optimizer variables. + + Resets the optimizer state variables. + + """ + self._init_optimizer() + + def save_native(self, filepath): + """ + Save model in a picked file specified by the filepath. + Uses torch.save(). + + Args: + filepath (string) : Path to pickle file to be + created by pt.save(). + Returns: + None + """ + torch.save(self, filepath) diff --git a/openfl-workspace/torch_llm_sequence_classification/src/ptllm_inmemory.py b/openfl-workspace/torch_llm_sequence_classification/src/ptllm_inmemory.py new file mode 100644 index 0000000000..6ff9d6c5d1 --- /dev/null +++ b/openfl-workspace/torch_llm_sequence_classification/src/ptllm_inmemory.py @@ -0,0 +1,41 @@ +# Copyright (C) 2020-2021 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""You may copy this file as the starting point of your own model.""" + +from openfl.federated import PyTorchDataLoader +from .mnist_utils import load_mnist_shard + + +class PyTorchMNISTInMemory(PyTorchDataLoader): + """PyTorch data loader for MNIST dataset.""" + + def __init__(self, data_path, batch_size, **kwargs): + """Instantiate the data object. + + Args: + data_path: The file path to the data + batch_size: The batch size of the data loader + **kwargs: Additional arguments, passed to super + init and load_mnist_shard + """ + super().__init__(batch_size, **kwargs) + + # TODO: We should be downloading the dataset shard into a directory + # TODO: There needs to be a method to ask how many collaborators and + # what index/rank is this collaborator. + # Then we have a way to automatically shard based on rank and size + # of collaborator list. + + num_classes, X_train, y_train, X_valid, y_valid = load_mnist_shard( + shard_num=int(data_path), **kwargs) + + self.X_train = X_train + self.y_train = y_train + self.train_loader = self.get_train_loader() + + self.X_valid = X_valid + self.y_valid = y_valid + self.val_loader = self.get_valid_loader() + + self.num_classes = num_classes diff --git a/rebuild_model.py b/rebuild_model.py new file mode 100644 index 0000000000..4421d745f4 --- /dev/null +++ b/rebuild_model.py @@ -0,0 +1,176 @@ +# %% +import numpy as np +from datasets import load_dataset, load_metric +from peft import LoraConfig, TaskType, get_peft_model, LoraModel + + +from transformers import ( + AutoModelForSequenceClassification, + AutoTokenizer, + DataCollatorWithPadding, + Trainer, + TrainingArguments, +) +from transformers.utils import find_adapter_config_file +from peft import PeftConfig, PeftType +from peft.utils.save_and_load import get_peft_model_state_dict +import transformers.modeling_utils +from peft import AutoPeftModelForTokenClassification, PeftModelForSequenceClassification +from peft.utils import set_peft_model_state_dict, get_peft_model_state_dict + + +# %% +# trained model and tokenizer +model_checkpoint = "/home/oamontoy/workspace/saved_model" +tokenizer_name = "roberta-large" +padding_side = "right" +tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + +task = "mrpc" +metric = load_metric("glue", task) + + +def compute_metrics(eval_pred): + predictions, labels = eval_pred + predictions = np.argmax(predictions, axis=1) + return metric.compute(predictions=predictions, references=labels) + + +BATCH_SIZE = 32 +NUM_EPOCHS = 5 +LR = 1e-3 +WEIGHT_DECAY = 0.01 + + +def dataset_function(tokenizer, task="mrpc"): + dataset = load_dataset("glue", task) + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer( + examples["sentence1"], + examples["sentence2"], + truncation=True, + max_length=None, + ) + return outputs + + tokenized_datasets = dataset.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + tokenized_datasets.set_format("torch") + data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="longest") + return data_collator, tokenized_datasets + + +data_collator, tokenized_datasets = dataset_function(tokenizer, task) + +# %% +# get initial model + +model = AutoModelForSequenceClassification.from_pretrained( + tokenizer_name, return_dict=True +) +model = PeftModelForSequenceClassification.from_pretrained(model, + model_checkpoint +) +# %% +# test performance +training_args = TrainingArguments( + output_dir="roberta-large-lora-seq", + learning_rate=LR, + per_device_train_batch_size=BATCH_SIZE, + per_device_eval_batch_size=BATCH_SIZE, + num_train_epochs=NUM_EPOCHS, + weight_decay=WEIGHT_DECAY, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, +) +trainer = Trainer( + model=model, + args=training_args, + train_dataset=tokenized_datasets["train"].select(range(500)), + eval_dataset=tokenized_datasets["test"].select(range(200)), + tokenizer=tokenizer, + compute_metrics=compute_metrics, + data_collator=data_collator, +) +# %% +print(trainer.evaluate()) +# %% +lora_params = get_peft_model_state_dict(model) +tokenizer_name = "roberta-large" +tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, padding_side=padding_side) +if getattr(tokenizer, "pad_token_id") is None: + tokenizer.pad_token_id = tokenizer.eos_token_id +model_re = AutoModelForSequenceClassification.from_pretrained( + tokenizer_name, return_dict=True +) +loaded_peft_config = PeftConfig.from_pretrained(model_checkpoint) +model_re = get_peft_model(model_re, loaded_peft_config) +# %% + + +def set_params( + model, + processed_adapter_state_dict, + adapter_name="default", + device_map="auto", + max_memory=None, + offload_folder=None, + offload_index=None, +): + model._hf_peft_config_loaded = True + incompatible_keys = set_peft_model_state_dict( + model, processed_adapter_state_dict, adapter_name + ) + model._dispatch_accelerate_model( + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_index=offload_index, + ) +set_params(model_re, lora_params, 'default') +print(model_re.peft_config["default"]) +# %% + + +# %% +data_collator, tokenized_datasets = dataset_function(tokenizer, task) +training_args = TrainingArguments( + output_dir="roberta-large-lora-seq", + learning_rate=LR, + per_device_train_batch_size=BATCH_SIZE, + per_device_eval_batch_size=BATCH_SIZE, + num_train_epochs=NUM_EPOCHS, + weight_decay=WEIGHT_DECAY, + evaluation_strategy="epoch", + save_strategy="epoch", + load_best_model_at_end=True, +) +trainer = Trainer( + model=model_re, + args=training_args, + train_dataset=tokenized_datasets["train"].select(range(500)), + eval_dataset=tokenized_datasets["test"].select(range(200)), + tokenizer=tokenizer, + compute_metrics=compute_metrics, + data_collator=data_collator, +) + +# %% +print(trainer.evaluate()) +#%% + +p = tokenized_datasets["train"][0]['input_ids'].unsqueeze(0) +model_re(p) +# %% +model(tokenized_datasets["train"][0]['input_ids'].unsqueeze(0)) +# %%