doc/tutorial/offsite_tuning/Offsite_tuning_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2345e19-83eb-4196-9606-74658c8fbdc5",
   "metadata": {},
   "source": [
    "# Offsite-tuning Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f1d728c-09e1-418e-8d80-53dd0ec467b1",
   "metadata": {},
   "source": [
    "In this tutorial, we'll focus on how to leverage Offsite-Tuning framework in FATE-LLM-2.0 to fine-tune your LLM. You'll learn how to:\n",
    "\n",
    "1. Define models, including main models(which are at server side and will offer adapters and emulators) and submodel(which are at client side and will load adapters and emulators for local fine-tuning) compatible with Offsite-Tuning framework.\n",
    "2. Get hands-on experience with the Offsite-Tuning trainer.\n",
    "3. Define configurations for advanced setup(Using Deepspeed, offsite-tuning + federation) through FATE-pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31432345-5cce-4efa-9a9b-844f997f14ad",
   "metadata": {},
   "source": [
    "## Introduction of Offsite-tuning\n",
    "\n",
    "Offsite-Tuning is a novel approach designed for the efficient and privacy-preserving adaptation of large foundational models for specific downstream tasks. The framework allows data owners to fine-tune models locally without uploading sensitive data to the LLM owner's servers. Specifically, the LLM owner sends a lightweight \"Adapter\" and a lossy compressed \"Emulator\" to the data owner. Using these smaller components, the data owner can then fine-tune the model solely on their private data. The Adapter, once fine-tuned, is returned to the model owner and integrated back into the large model to enhance its performance on the specific dataset.\n",
    "\n",
    "Offsite-Tuning addresses the challenge of unequal distribution of computational power and data. It allows thLLMel owner to enhance the model's capabilities without direct access to private data, while also enabling data owners who may not have the resources to train a full-scale model to fine-tune a portion of it using less computational power. This mutually beneficial arrangement accommodates both parties involve.\n",
    "\n",
    "Beyond the standard two-party setup involving the model owner and the data ownin FATE-LLM, er, Offsite-Tunframework ing is also extendable to scenarios with multiple data owners. FATE supports multi-party Offsite-Tuning, allowing multiple data owners to fine-tune and aggregate their Adapters locally, further enhancing the flexibility and applicability of this framewrFor more details of Offsite-tuning, please refer to the [original paper](https://arxiv.org/pdf/2302.04870.pdf).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e7ac467-e5df-4bf3-8571-0a477ab4612d",
   "metadata": {},
   "source": [
    "## Preliminary\n",
    "\n",
    "We strongly recommend you finish reading our NN tutorial to get familiar with Model and Dataset customizations: [NN Tutorials](https://github.com/FederatedAI/FATE/blob/master/doc/2.0/fate/components/pipeline_nn_cutomization_tutorial.md)\n",
    "\n",
    "In this tutorial, we assume that you have deploy the codes of FATE(including fateflow & fate-client) & FATE-LLM-2.0. You can add python path so that you can run codes in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f33516e8-0d28-4c97-bc38-ba28d60acf37",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "your_path_to_fate_python = 'xxx/fate/fate/python'\n",
    "sys.path.append(your_path_to_fate_python)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f2fc794",
   "metadata": {},
   "source": [
    "If you install FATE & FATE-LLM-2.0 via pip, you can directly use the following codes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7309281b-5956-4158-9256-d6db230e086d",
   "metadata": {},
   "source": [
    "## Define Main Model and Sub Model\n",
    "\n",
    "Main models are at server side and will provides weights of adapters and emulators to client sides, while Sub Models are at client side and will load adapters and emulators for local fine-tuning. In this chapter we will take a standard GPT2 as the example and show you how to quickly develop main model class and sub model class for offsite-tuning.\n",
    "\n",
    "### Base Classes and Interfaces\n",
    "\n",
    "The base classes for the Main and Sub Models are OffsiteTuningMainModel and OffsiteTuningSubModel, respectively. To build your own models upon these base classes, you need to:\n",
    "\n",
    "1. Implement three key interfaces: get_base_model, get_model_transformer_blocks, and forward. The get_base_model interface should return the full Main or Sub Model. Meanwhile, the get_model_transformer_blocks function should return a ModuleList of all transformer blocks present in your language model, enabling the extraction of emulators and adapters from these blocks. Finally, you're required to implement the forward process for model inference.\n",
    "\n",
    "2. Supply the parameters emulator_layer_num, adapter_top_layer_num, and adapter_bottom_layer_num to the parent class. This allows the framework to automatically generate the top and bottom adapters as well as the dropout emulator for you. Specifically, the top adapters are taken from the top of the transformer blocks, while the bottom adapters are taken from the bottom. The emulator uses a dropout emulator consistent with the paper's specifications. Once the adapter layers are removed, the emulator is formed by selecting transformer blocks at fixed intervals and finally stack them to make a dropout emulator.\n",
    "\n",
    "Our framework will automatically detect the emulator and adapters of a main model, and send them to clients. Clients' models them load the weights of emulators and adapters to get trainable models.\n",
    "\n",
    "### Example\n",
    "\n",
    "Let us take a look of our built-in GPT-2 model. It will be easy for you to build main models and sub models based on the framework. Please notice that the GPT2LMHeadSubModel's base model is intialized from a GPTConfig, that is to say, it's weights are random and need to load pretrained weights from server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8611c115-0321-458f-b190-49dcb127a653",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.model_zoo.offsite_tuning.offsite_tuning_model import OffsiteTuningSubModel, OffsiteTuningMainModel\n",
    "from transformers import GPT2LMHeadModel, GPT2Config\n",
    "from torch import nn\n",
    "import torch as t\n",
    "\n",
    "\n",
    "class GPT2LMHeadMainModel(OffsiteTuningMainModel):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            model_name_or_path,\n",
    "            emulator_layer_num: int,\n",
    "            adapter_top_layer_num: int = 2,\n",
    "            adapter_bottom_layer_num: int = 2):\n",
    "\n",
    "        self.model_name_or_path = model_name_or_path\n",
    "        super().__init__(\n",
    "            emulator_layer_num,\n",
    "            adapter_top_layer_num,\n",
    "            adapter_bottom_layer_num)\n",
    "\n",
    "    def get_base_model(self):\n",
    "        return GPT2LMHeadModel.from_pretrained(self.model_name_or_path)\n",
    "\n",
    "    def get_model_transformer_blocks(self, model: GPT2LMHeadModel):\n",
    "        return model.transformer.h\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(**x)\n",
    "\n",
    "class GPT2LMHeadSubModel(OffsiteTuningSubModel):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            model_name_or_path,\n",
    "            emulator_layer_num: int,\n",
    "            adapter_top_layer_num: int = 2,\n",
    "            adapter_bottom_layer_num: int = 2,\n",
    "            fp16_mix_precision=False,\n",
    "            partial_weight_decay=None):\n",
    "\n",
    "        self.model_name_or_path = model_name_or_path\n",
    "        self.emulator_layer_num = emulator_layer_num\n",
    "        self.adapter_top_layer_num = adapter_top_layer_num\n",
    "        self.adapter_bottom_layer_num = adapter_bottom_layer_num\n",
    "        super().__init__(\n",
    "            emulator_layer_num,\n",
    "            adapter_top_layer_num,\n",
    "            adapter_bottom_layer_num,\n",
    "            fp16_mix_precision)\n",
    "        self.partial_weight_decay = partial_weight_decay\n",
    "\n",
    "    def get_base_model(self):\n",
    "        total_layer_num = self.emulator_layer_num + \\\n",
    "            self.adapter_top_layer_num + self.adapter_bottom_layer_num\n",
    "        config = GPT2Config.from_pretrained(self.model_name_or_path)\n",
    "        config.num_hidden_layers = total_layer_num\n",
    "        # initialize a model without pretrained weights\n",
    "        return GPT2LMHeadModel(config)\n",
    "\n",
    "    def get_model_transformer_blocks(self, model: GPT2LMHeadModel):\n",
    "        return model.transformer.h\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.model(**x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd1f63f-afa7-4f09-a67e-63812ddcd801",
   "metadata": {},
   "source": [
    "We can define a server side model and a client side model that can work together in the offsite-tuning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04870e76-11cc-4d79-a09e-b6fd16ed2f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_main = GPT2LMHeadMainModel('gpt2', 4, 2, 2)\n",
    "model_sub = GPT2LMHeadSubModel('gpt2', 4, 2, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19d34937-b4ae-436e-b4ea-1620fb80bed4",
   "metadata": {},
   "source": [
    "### Share additional parameters with clients\n",
    "\n",
    "Additionally, beyond the weights of emulators and adapters, you may also want to share other model parameters, such as embedding weights, with your client partners. To achieve this, you'll need to implement two more interfaces: get_additional_param_state_dict and load_additional_param_state_dict for both the Main and Sub Models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "189fce0e-8e4d-4368-8e14-907b30ce0a49",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_additional_param_state_dict(self):\n",
    "    # get parameter of additional parameter\n",
    "    model = self.model\n",
    "    param_dict = {\n",
    "        'wte': model.transformer.wte,\n",
    "        'wpe': model.transformer.wpe,\n",
    "        'last_ln_f': model.transformer.ln_f\n",
    "    }\n",
    "\n",
    "    addition_weights = self.get_numpy_state_dict(param_dict)\n",
    "\n",
    "    wte = addition_weights.pop('wte')\n",
    "    wte_dict = split_numpy_array(wte, 10, 'wte')\n",
    "    wpe = addition_weights.pop('wpe')\n",
    "    wpe_dict = split_numpy_array(wpe, 10, 'wpe')\n",
    "    addition_weights.update(wte_dict)\n",
    "    addition_weights.update(wpe_dict)\n",
    "    return addition_weights\n",
    "\n",
    "def load_additional_param_state_dict(self, submodel_weights: dict):\n",
    "    # load additional weights:\n",
    "    model = self.model\n",
    "    param_dict = {\n",
    "        'wte': model.transformer.wte,\n",
    "        'wpe': model.transformer.wpe,\n",
    "        'last_ln_f': model.transformer.ln_f\n",
    "    }\n",
    "\n",
    "    new_submodel_weight = {}\n",
    "    new_submodel_weight['last_ln_f'] = submodel_weights['last_ln_f']\n",
    "    wte_dict, wpe_dict = {}, {}\n",
    "    for k, v in submodel_weights.items():\n",
    "        if 'wte' in k:\n",
    "            wte_dict[k] = v\n",
    "        if 'wpe' in k:\n",
    "            wpe_dict[k] = v\n",
    "    wte = recover_numpy_array(wte_dict, 'wte')\n",
    "    wpe = recover_numpy_array(wpe_dict, 'wpe')\n",
    "    new_submodel_weight['wte'] = wte\n",
    "    new_submodel_weight['wpe'] = wpe\n",
    "\n",
    "    self.load_numpy_state_dict(param_dict, new_submodel_weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59d9aa6a-80e9-4130-8af1-c7d2bd0fbba3",
   "metadata": {},
   "source": [
    "From these codes we can see that we use 'split_numpy_array, recover_numpy_array' to cut embedding weights into pieces and recover them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda6f5e3-d05a-4cdf-afd4-affbc162fce4",
   "metadata": {},
   "source": [
    "## Submit a Offsite-tuning Task - A QA Task Sample with GPT2\n",
    "\n",
    "Now we are going to show you how to run a 2 party(server & client) offsite-tuning task using the GPT-2 model defined above. Before we submit the task we need to prepare the QA dataset.\n",
    "\n",
    "### Prepare QA Dataset - Sciq\n",
    "\n",
    "In this example, we use sciq dataset. You can use tools provided in our qa_dataset.py to tokenize the sciq dataset and save the tokenized result. **Remember to modify the save_path to your own path.** For the sake of simplicity, in this tutorial, for every party we only use this dataset to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "84f6947e-f0a3-4a42-9549-a9776a15b66d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.dataset.qa_dataset import tokenize_qa_dataset\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer_name_or_path = 'gpt2'\n",
    "tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)\n",
    "\n",
    "if 'llama' in tokenizer_name_or_path:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, unk_token=\"<unk>\",  bos_token=\"<s>\", eos_token=\"</s>\", add_eos_token=True)   \n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "else:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)\n",
    "if 'gpt2' in tokenizer_name_or_path:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "import os\n",
    "# bind data path to name & namespace\n",
    "save_path = 'xxxx/sciq'\n",
    "rs = tokenize_qa_dataset('sciq', tokenizer, save_path, seq_max_len=600)  # we save the cache dataset to the fate root folder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adabe89a-37be-4c64-bd83-4f8c8b80096f",
   "metadata": {},
   "source": [
    "We can use our built-in QA dataset to load tokenized dataset, to see if everything is working correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6500c2ba-bc39-4db4-b2ea-947fb09c334e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.dataset.qa_dataset import QaDataset\n",
    "\n",
    "ds = QaDataset(tokenizer_name_or_path=tokenizer_name_or_path)\n",
    "ds.load(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d6f62b60-eed0-4bd0-874e-ae3feeebb120",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11679\n",
      "600\n"
     ]
    }
   ],
   "source": [
    "print(len(ds))  # train set length\n",
    "print(ds[0]['input_ids'].__len__()) # first sample length"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0609c63d-35a4-43bc-bd4b-f1c61adea587",
   "metadata": {},
   "source": [
    "## Submit a Task\n",
    "\n",
    "Now the model and the dataset is prepared! We can submit a training task. In the FATE-2.0, you can define your pipeline in a much easier manner.\n",
    "\n",
    "After we submit the task below, the following process will occur: The server and client each initialize their respective models. The server extracts shared parameters and sends them to the client. The client then loads these parameters and conducts training on a miniaturized GPT-2 model composed of an emulator and adapter on SciqP \n",
    "\n",
    "If you are not familiar with trainer configuration, please refer to [NN Tutorials](https://github.com/FederatedAI/FATE/blob/master/doc/2.0/fate/components/pipeline_nn_cutomization_tutorial.md).\n",
    "\n",
    " Upon completion of the training, the client sends the adapter parameters back to the server. Since we are directly using Hugging Face's LMHeadGPT2, there's no need to supply a loss function. Simply inputting the preprocessed data and labels into the model will calculate the correct loss and proceed with gradient descent\n",
    "\n",
    "One thing to pay special attention to is that Offsite-Tuning differs from FedAvg within FATE. In Offsite-Tuning, the server (the arbiter role) needs to initialize the model. Therefore, please refer to the example below and set the runner conf separately for the client and the server.\n",
    "\n",
    "To make this a quick demo, we only select 100 samples from the origin qa datset, see 'select_num=100' in the LLMDatasetLoader."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261dfb43",
   "metadata": {},
   "source": [
    "### Bind Dataset Path with Name & Namespace\n",
    "\n",
    "Plase execute the following code to bind the dataset path with name & namespace. Remember to modify the path to your own dataset save path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc1e82b",
   "metadata": {},
   "outputs": [],
   "source": [
    "! flow table bind --namespace experiment --name sciq --path YOUR_SAVE_PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e8c5ff4",
   "metadata": {},
   "source": [
    "### Pipeline codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c9113d10-c3e7-4875-9502-ce46aa0b86b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<fate_client.pipeline.pipeline.FateFlowPipeline at 0x7fc69aa33a00>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_conf_of_ot_runner\n",
    "from fate_client.pipeline.components.fate.nn.algo_params import Seq2SeqTrainingArguments, FedAVGArguments\n",
    "from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMDataFuncLoader\n",
    "from fate_client.pipeline.components.fate.nn.torch.base import Sequential\n",
    "from fate_client.pipeline.components.fate.nn.torch import nn\n",
    "\n",
    "\n",
    "guest = '9999'\n",
    "host = '9999'\n",
    "arbiter = '9999'\n",
    "\n",
    "pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter)\n",
    "pipeline.set_site_party_id('9999')\n",
    "reader_0 = Reader(\"reader_0\", runtime_parties=dict(guest=guest))\n",
    "reader_0.guest.task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"sciq\"\n",
    ")\n",
    "\n",
    "client_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadSubModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=4,\n",
    "    adapter_top_layer_num=1,\n",
    "    adapter_bottom_layer_num=1\n",
    ")\n",
    "\n",
    "server_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadMainModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=4,\n",
    "    adapter_top_layer_num=1,\n",
    "    adapter_bottom_layer_num=1  \n",
    ")\n",
    "\n",
    "train_args = Seq2SeqTrainingArguments(\n",
    "    per_device_train_batch_size=1,\n",
    "    learning_rate=5e-5,\n",
    "    disable_tqdm=False,\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=10,\n",
    "    logging_strategy='steps',\n",
    "    use_cpu=False\n",
    ")\n",
    "\n",
    "dataset = LLMDatasetLoader(\n",
    "    module_name='qa_dataset', item_name='QaDataset',\n",
    "    tokenizer_name_or_path='gpt2',\n",
    "    select_num=100\n",
    ")\n",
    "\n",
    "data_collator = LLMDataFuncLoader(module_name='data_collator.cust_data_collator', item_name='get_seq2seq_data_collator', tokenizer_name_or_path='gpt2')\n",
    "\n",
    "client_conf = get_conf_of_ot_runner(\n",
    "    model=client_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=False\n",
    ")\n",
    "\n",
    "server_conf = get_conf_of_ot_runner(\n",
    "    model=server_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=False\n",
    ")\n",
    "\n",
    "homo_nn_0 = HomoNN(\n",
    "    'nn_0',\n",
    "    train_data=reader_0.outputs[\"output_data\"],\n",
    "    runner_module=\"offsite_tuning_runner\",\n",
    "    runner_class=\"OTRunner\"\n",
    ")\n",
    "\n",
    "homo_nn_0.guest.task_parameters(runner_conf=client_conf)\n",
    "homo_nn_0.arbiter.task_parameters(runner_conf=server_conf)\n",
    "pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "pipeline.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e97c2823",
   "metadata": {},
   "source": [
    "You can try to initialize your models, datasets to check if they can be loaded correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "872817e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPT2LMHeadSubModel(\n",
      "  (model): GPT2LMHeadModel(\n",
      "    (transformer): GPT2Model(\n",
      "      (wte): Embedding(50257, 768)\n",
      "      (wpe): Embedding(1024, 768)\n",
      "      (drop): Dropout(p=0.1, inplace=False)\n",
      "      (h): ModuleList(\n",
      "        (0): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (1): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (2): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (3): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (4): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (5): GPT2Block(\n",
      "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (attn): GPT2Attention(\n",
      "            (c_attn): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "          (mlp): GPT2MLP(\n",
      "            (c_fc): Conv1D()\n",
      "            (c_proj): Conv1D()\n",
      "            (act): NewGELUActivation()\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "      (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
      "  )\n",
      "  (emulator): ModuleList(\n",
      "    (0): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "    (1): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "    (2): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "    (3): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (adapter_bottom): ModuleList(\n",
      "    (0): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (adapter_top): ModuleList(\n",
      "    (0): GPT2Block(\n",
      "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (attn): GPT2Attention(\n",
      "        (c_attn): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): GPT2MLP(\n",
      "        (c_fc): Conv1D()\n",
      "        (c_proj): Conv1D()\n",
      "        (act): NewGELUActivation()\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n",
      "**********\n",
      "<fate_llm.dataset.qa_dataset.QaDataset object at 0x7fc724fdfd00>\n",
      "**********\n",
      "DataCollatorForSeq2Seq(tokenizer=GPT2TokenizerFast(name_or_path='gpt2', vocab_size=50257, model_max_length=1024, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<|endoftext|>', 'eos_token': '<|endoftext|>', 'unk_token': '<|endoftext|>', 'pad_token': '<|endoftext|>'}, clean_up_tokenization_spaces=True),  added_tokens_decoder={\n",
      "\t50256: AddedToken(\"<|endoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
      "}, model=None, padding=True, max_length=None, pad_to_multiple_of=None, label_pad_token_id=-100, return_tensors='pt')\n"
     ]
    }
   ],
   "source": [
    "print(client_model())\n",
    "print('*' * 10)\n",
    "print(dataset())\n",
    "print('*' * 10)\n",
    "print(data_collator())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "898c3491",
   "metadata": {},
   "source": [
    "Seems that everything is ready! Now we can submit the task. Submit the code below to submit your task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "74497742-4030-4a7a-a13e-2c020da47cd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b33b2e2b-3b53-4881-8db6-a67e1293e88b",
   "metadata": {},
   "source": [
    "## Add Deepspeed Setting\n",
    "\n",
    "By simply adding a ds_config, we can run our task with a deepspeed backend. If you have deployed eggroll envoironment, you can submmit the task with deepspeed to eggroll accelerate your training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6e8f063b-263c-4ba5-b2ba-98a86ce38b94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pipeline.backend.pipeline.PipeLine at 0x7f8002385e50>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_conf_of_ot_runner\n",
    "from fate_client.pipeline.components.fate.nn.algo_params import Seq2SeqTrainingArguments, FedAVGArguments\n",
    "from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMDataFuncLoader\n",
    "from peft import LoraConfig, TaskType\n",
    "from transformers.modeling_utils import unwrap_model\n",
    "\n",
    "\n",
    "guest = '10000'\n",
    "host = '10000'\n",
    "arbiter = '10000'\n",
    "\n",
    "# pipeline = FateFlowPipeline().set_parties(guest=guest, host=host, arbiter=arbiter)\n",
    "pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter)\n",
    "\n",
    "reader_0 = Reader(\"reader_0\", runtime_parties=dict(guest=guest))\n",
    "reader_0.guest.task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"sciq\"\n",
    ")\n",
    "\n",
    "client_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadSubModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=18,\n",
    "    adapter_top_layer_num=2,\n",
    "    adapter_bottom_layer_num=2\n",
    ")\n",
    "\n",
    "server_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadMainModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=18,\n",
    "    adapter_top_layer_num=2,\n",
    "    adapter_bottom_layer_num=2  \n",
    ")\n",
    "\n",
    "dataset = LLMDatasetLoader(\n",
    "    module_name='qa_dataset', item_name='QaDataset',\n",
    "    tokenizer_name_or_path='gpt2',\n",
    "    select_num=100\n",
    ")\n",
    "\n",
    "data_collator = LLMDataFuncLoader(module_name='data_collator.cust_data_collator', item_name='get_seq2seq_data_collator', tokenizer_name_or_path='gpt2')\n",
    "\n",
    "batch_size = 1\n",
    "lr = 5e-5\n",
    "ds_config = {\n",
    "    \"train_micro_batch_size_per_gpu\": batch_size,\n",
    "    \"optimizer\": {\n",
    "        \"type\": \"Adam\",\n",
    "        \"params\": {\n",
    "            \"lr\": lr,\n",
    "            \"torch_adam\": True,\n",
    "            \"adam_w_mode\": False\n",
    "        }\n",
    "    },\n",
    "    \"fp16\": {\n",
    "        \"enabled\": True\n",
    "    },\n",
    "    \"gradient_accumulation_steps\": 1,\n",
    "    \"zero_optimization\": {\n",
    "        \"stage\": 2,\n",
    "        \"allgather_partitions\": True,\n",
    "        \"allgather_bucket_size\": 1e8,\n",
    "        \"overlap_comm\": True,\n",
    "        \"reduce_scatter\": True,\n",
    "        \"reduce_bucket_size\": 1e8,\n",
    "        \"contiguous_gradients\": True,\n",
    "        \"offload_optimizer\": {\n",
    "            \"device\": \"cpu\"\n",
    "        },\n",
    "        \"offload_param\": {\n",
    "            \"device\": \"cpu\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "train_args = Seq2SeqTrainingArguments(\n",
    "    per_device_train_batch_size=1,\n",
    "    learning_rate=5e-5,\n",
    "    disable_tqdm=False,\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=10,\n",
    "    logging_strategy='steps',\n",
    "    dataloader_num_workers=4,\n",
    "    use_cpu=False,\n",
    "    deepspeed=ds_config,  # Add deepspeed config here\n",
    "    remove_unused_columns=False,\n",
    "    fp16=True\n",
    ")\n",
    "\n",
    "client_conf = get_conf_of_ot_runner(\n",
    "    model=client_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=False,\n",
    ")\n",
    "\n",
    "server_conf = get_conf_of_ot_runner(\n",
    "    model=server_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=False\n",
    ")\n",
    "\n",
    "\n",
    "homo_nn_0 = HomoNN(\n",
    "    'nn_0',\n",
    "    train_data=reader_0.outputs[\"output_data\"],\n",
    "    runner_module=\"offsite_tuning_runner\",\n",
    "    runner_class=\"OTRunner\"\n",
    ")\n",
    "\n",
    "homo_nn_0.guest.task_parameters(runner_conf=client_conf)\n",
    "homo_nn_0.arbiter.task_parameters(runner_conf=server_conf)\n",
    "\n",
    "# if you have deployed eggroll, you can add this line to submit your job to eggroll\n",
    "homo_nn_0.guest.conf.set(\"launcher_name\", \"deepspeed\")\n",
    "\n",
    "pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "pipeline.conf.set(\"task\", dict(engine_run={\"cores\": 4}))\n",
    "pipeline.compile()\n",
    "pipeline.fit()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97249681-c3a3-43bd-8167-7ae3f4e1616b",
   "metadata": {},
   "source": [
    "## Offsite-tuning + Multi Client Federation\n",
    "\n",
    "\n",
    "The Offsite-Tuning + FedAVG federation is configured based on the standard Offsite-Tuning. In this situation, you need to add data input & configurations for all clients. And do remember to add 'aggregate_model=True' for client & server conf so that model federation will be conducted during the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdbdc60c-a948-4be3-bba6-519d8640b0a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_conf_of_ot_runner\n",
    "from fate_client.pipeline.components.fate.nn.algo_params import Seq2SeqTrainingArguments, FedAVGArguments\n",
    "from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMCustFuncLoader\n",
    "from peft import LoraConfig, TaskType\n",
    "\n",
    "\n",
    "guest = '10000'\n",
    "host = '10000'\n",
    "arbiter = '10000'\n",
    "\n",
    "pipeline = FateFlowPipeline().set_parties(guest=guest, host=host, arbiter=arbiter)\n",
    "\n",
    "reader_0 = Reader(\"reader_0\", runtime_parties=dict(guest=guest, host=host))\n",
    "reader_0.guest.task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"sciq\"\n",
    ")\n",
    "reader_0.hosts[0].task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"sciq\"\n",
    ")\n",
    "\n",
    "client_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadSubModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=4,\n",
    "    adapter_top_layer_num=1,\n",
    "    adapter_bottom_layer_num=1\n",
    ")\n",
    "\n",
    "server_model = LLMModelLoader(\n",
    "    module_name='offsite_tuning.gpt2', item_name='GPT2LMHeadMainModel',\n",
    "    model_name_or_path='gpt2',\n",
    "    emulator_layer_num=4,\n",
    "    adapter_top_layer_num=1,\n",
    "    adapter_bottom_layer_num=1  \n",
    ")\n",
    "\n",
    "dataset = LLMDatasetLoader(\n",
    "    module_name='qa_dataset', item_name='QaDataset',\n",
    "    tokenizer_name_or_path='gpt2',\n",
    "    select_num=100\n",
    ")\n",
    "\n",
    "data_collator = LLMCustFuncLoader(module_name='cust_data_collator', item_name='get_seq2seq_tokenizer', model_path='gpt2')\n",
    "\n",
    "train_args = Seq2SeqTrainingArguments(\n",
    "    per_device_train_batch_size=1,\n",
    "    learning_rate=5e-5,\n",
    "    disable_tqdm=False,\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=10,\n",
    "    logging_strategy='steps',\n",
    "    dataloader_num_workers=4\n",
    ")\n",
    "\n",
    "client_conf = get_conf_of_ot_runner(\n",
    "    model=client_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=True\n",
    ")\n",
    "\n",
    "server_conf = get_conf_of_ot_runner(\n",
    "    model=server_model,\n",
    "    dataset=dataset,\n",
    "    data_collator=data_collator,\n",
    "    training_args=train_args,\n",
    "    fed_args=FedAVGArguments(),\n",
    "    aggregate_model=True\n",
    ")\n",
    "\n",
    "homo_nn_0 = HomoNN(\n",
    "    'nn_0',\n",
    "    train_data=reader_0.outputs[\"output_data\"],\n",
    "    runner_module=\"offsite_tuning_runner\",\n",
    "    runner_class=\"OTRunner\"\n",
    ")\n",
    "\n",
    "homo_nn_0.guest.task_parameters(runner_conf=client_conf)\n",
    "homo_nn_0.hosts[0].task_parameters(runner_conf=client_conf)\n",
    "homo_nn_0.arbiter.task_parameters(runner_conf=server_conf)\n",
    "\n",
    "pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "\n",
    "pipeline.compile()\n",
    "pipeline.fit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}