doc/tutorial/fedmkt/fedmkt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Federated Tuning With FedMKT methods in FATE-LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will demonstrate how to efficiently train federated large language models using the FATE-LLM framework. In FATE-LLM, we introduce the \"FedMKT\" module, specifically designed for federated learning with large language models. FedMKT introduces a novel\n",
    "federated mutual knowledge transfer framework that enables effective knowledge transfer between an LLM deployed on the server and SLMs residing on clients.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Algorithm is based on paper [\"FedMKT: Federated Mutual Knowledge Transfer for Large and SmallLanguage Models\"](https://arxiv.org/pdf/2406.02224), We integrate its code into the FATE-LLM framework.  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiments\n",
    "\n",
    "Chapter List: \n",
    "* settings\n",
    "  1. DataSet: ARC-Challenge\n",
    "  2. Models Use in \"FEDMKT\" Paper\n",
    "  3. Prepare Optimal Vocabulary Mapping Tables\n",
    "  4. Training LLMs with Lora\n",
    "* experiment examples:\n",
    "  1. Running FEDMKT With Launcher (Experimential Using): 4-SLMs\n",
    "  2. Running FEDMKT With Launcher (Experimential Using): 1-SLM (One To One)\n",
    "  3. Running FEDMKT With Launcher (Experimential Using): 1-SLM And SLM Trains Only (LLM2SLM)\n",
    "  4. Running FEDMKT With Launcher (Experimential Using): 4-SLMs Homogeneous SFT\n",
    "  5. Running FEDMKT with Pipeline (Industrial Using)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset: ARC-Challenge\n",
    "\n",
    "ARC-Challenge is a dataset of 7,787 genuine grade-school level, multiple-choice science questions, assembled to encourage research in advanced question-answering. \n",
    "\n",
    "You can refer to following link for more details about [ARC-Challange](https://huggingface.co/datasets/allenai/ai2_arc)\n",
    "\n",
    "In this section, we will download ARC-Challenge dataset from huggingface and splits it into five parts, part \"common\" for public dataset and other parts for slms(opt2, gpt2, llama, opt)'s  training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "\n",
    "\n",
    "data = datasets.load_dataset(\"ai2_arc\", \"ARC-Challenge\", download_mode=\"force_redownload\", ignore_verifications=True)\n",
    "train_data = data.pop(\"train\")\n",
    "\n",
    "seed=123\n",
    "n = train_data.shape[0]\n",
    "client_num = 4\n",
    "process_data_output_dir = \"\" # processed data saved directory should be specified, it will be used in later.\n",
    "\n",
    "client_data_num = n // (client_num + 1)\n",
    "\n",
    "for i in range(client_num):\n",
    "    splits = train_data.train_test_split(train_size=client_data_num, shuffle=True, seed=seed)\n",
    "    client_name = f\"client_{i}\"\n",
    "    data[client_name] = splits[\"train\"]\n",
    "    train_data = splits[\"test\"]\n",
    "\n",
    "if train_data.shape[0] == client_data_num:\n",
    "    data[\"common\"] = train_data\n",
    "else:\n",
    "    data[\"common\"] = train_data.train_test_split(\n",
    "        train_size=client_data_num, shuffle=True, seed=args.seed\n",
    "    )[\"train\"]\n",
    "\n",
    "data.save_to_disk(process_data_output_dir)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Models Use In \"FEDMKT\" Paper\n",
    "\n",
    "LLM: [Llama-2-7B](https://huggingface.co/meta-llama/Llama-2-7b-hf)  \n",
    "SLM-0: [opt-1.3b](https://huggingface.co/facebook/opt-1.3b)  \n",
    "SLM-1: [gpt2-xlarge](https://huggingface.co/openai-community/gpt2-xl)  \n",
    "SLM-2: [Llama-1.3b](https://huggingface.co/princeton-nlp/Sheared-LLaMA-1.3B)  \n",
    "SLM-3: [bloom-1.1B](https://huggingface.co/bigscience/bloom-1b1)\n",
    "\n",
    "Users should download the models from huggingface before the following steps and saved them in local directories, as models are too big, redownload them cost too much times.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# replaoce the names of models to local save directories\n",
    "llm_pretrained_path = \"llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMA-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare Optimal Vocabulary Mapping Tables\n",
    "\n",
    "To use \"FEDMKT\" for federated knowledge transfer, we need to build pptimal vocabulary mapping tables first.\n",
    "In paper of \"FEDMKT\", it has One LLM and four SLMs, so we need to build eight pptimal vocabulary mapping tables. For each paired of (LLM, SLM), two tables should be built as co-training are needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_llm.algo.fedmkt.token_alignment.vocab_mapping import get_vocab_mappings\n",
    "\n",
    "\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\" # replace this to actually paths\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "for idx, (llm_pretrained, slm_pretrained) in enumerate(llm_slm_pairs):\n",
    "    slm_to_llm_vocab_mapping_path = slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_path = llm_to_slm_vocab_mapping_paths[idx]\n",
    "    _ = get_vocab_mappings(slm_pretrained, llm_pretrained, slm_to_llm_vocab_mapping_paths[idx], num_processors=16)\n",
    "    _ = get_vocab_mappings(llm_pretrained, slm_pretrained, llm_to_slm_vocab_mapping_paths[idx], num_processors=16)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training LLMs with Lora\n",
    "\n",
    "In this section, We will introduce the lora configs use in five models listed in paper: one LLM (Llama-2-7B), four SLMs(opt-1.3B, gpt2-xlarge, Llama-1.3B, bloom-1.1B)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LLM models with peft is located on fate_llm/model_zoo, we will give a guide to use them. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init LLm Llama-2-7B's Lora Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init SLMs Lora Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "\n",
    "def get_slm_conf(slm_idx):\n",
    "    slm_pretrained_path = slm_pretrained_paths[slm_idx]\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running FEDMKT With Launcher (Experimential Using): 4-SLMs\n",
    "\n",
    "Using launcher to startup is mainly for experimential. Before running this section, make sure that [FATE-LLM Standalone](https://github.com/FederatedAI/FATE-LLM?tab=readme-ov-file#standalone-deployment) has been deployed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Global Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "process_data_output_dir = \"\"\n",
    "llm_pretrained_path = \"Llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMa-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\"\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "#### all variables has been defined above\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "\n",
    "global_epochs = 1\n",
    "batch_size=4\n",
    "llm_lr = 3e-5\n",
    "slm_lrs = [3e-5, 3e-4, 3e-5, 3e-5, 3e-5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init FEDMKTLLM Runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "In this Section, we will introduce how to initialize \"FEDMKTLLM\" object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step1: Initialize LLM With LoraConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, TaskType\n",
    "from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTLLM\n",
    "from fate.ml.nn.homo.fedavg import FedAVGArguments\n",
    "from fate_llm.dataset.qa_dataset import QaDataset\n",
    "from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "from transformers import AutoConfig\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    ")\n",
    "\n",
    "model = LLaMa(\n",
    "    pretrained_path=llm_pretrained_path,\n",
    "    peft_type=\"LoraConfig\",\n",
    "    peft_config=lora_config.to_dict(),\n",
    "    torch_dtype=\"bfloat16\"    \n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step2: Specify Public Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pub_data = QaDataset(tokenizer_name_or_path=llm_pretrained_path,\n",
    "                     dataset_name=\"arc_challenge\",\n",
    "                     data_part=\"common\",\n",
    "                     seq_max_len=512,\n",
    "                     need_preprocess=True)\n",
    "pub_data.load(process_data_output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step3: Initialize FEDMKT Training Args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = FedMKTTrainingArguments(\n",
    "    global_epochs=global_epochs,\n",
    "    per_device_train_batch_size=1,\n",
    "    gradient_accumulation_steps=batch_size,\n",
    "    learning_rate=llm_lr,\n",
    "    output_dir=\"./\",\n",
    "    dataloader_num_workers=4,\n",
    "    remove_unused_columns=False,\n",
    "    warmup_ratio=0.008,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    optim=\"adamw_torch\",\n",
    "    adam_beta1=0.9,\n",
    "    adam_beta2=0.95,\n",
    "    weight_decay=0.1,\n",
    "    max_grad_norm=1.0,\n",
    "    use_cpu=False,\n",
    "    vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size, # pay attention to this, \n",
    "                                                                           # vocab_size must be specified to avoid dimension mismatch \n",
    "                                                                           # of tokenizer's vocab_size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step4: Initialize Other Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fed_args = FedAVGArguments(\n",
    "    aggregate_strategy='epoch',\n",
    "    aggregate_freq=1\n",
    ")\n",
    "\n",
    "slm_to_llm_vocab_mapping = []\n",
    "for path in slm_to_llm_vocab_mapping_paths:\n",
    "    with open(path, \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "        slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "slm_tokenizers = [get_tokenizer(slm_pretrained_path) for slm_pretrained_path in slm_pretrained_paths]\n",
    "tokenizer = get_tokenizer(llm_pretrained_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step5: New FEDMKTLLM Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = FedMKTLLM(\n",
    "    ctx=ctx,\n",
    "    model=model,\n",
    "    training_args=training_args,\n",
    "    fed_args=fed_args,\n",
    "    train_set=pub_data,\n",
    "    tokenizer=tokenizer,\n",
    "    slm_tokenizers=slm_tokenizers,\n",
    "    slm_to_llm_vocab_mappings=slm_to_llm_vocab_mapping,\n",
    "    save_trainable_weights_only=True, # save lora weights only\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Step6: Training And Save Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()\n",
    "trainer.save_model(output_dir=\"fill the path to save llm finetuning result\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Init FEDMKTSLM Runner\n",
    "\n",
    "FEDMKTSLM Runner is a slightly different of FEDMKTLLM Runner, we only introduce different variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Import SLMs you need to run, here we choose four Slms Using In Original Paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "from peft import LoraConfig, TaskType    \n",
    "from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "from fate_llm.model_zoo.pellm.gpt2 import GPT2CLM\n",
    "from fate_llm.model_zoo.pellm.opt import OPT\n",
    "from fate_llm.model_zoo.pellm.bloom import Bloom\n",
    "from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTSLM\n",
    "from fate_llm.dataset.qa_dataset import QaDataset\n",
    "from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "from transformers import AutoConfig\n",
    "\n",
    "slm_idx = 0\n",
    "\n",
    "slm_model_class = [\n",
    "    OPT,\n",
    "    GPT2CLM,\n",
    "    LLaMa,\n",
    "    Bloom\n",
    "]\n",
    "    \n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "    target_modules=slm_lora_target_modules[slm_idx]\n",
    ")\n",
    "\n",
    "model = slm_model_class[slm_idx](\n",
    "    pretrained_path=slm_pretrained_paths[slm_idx],\n",
    "    peft_type=\"LoraConfig\",\n",
    "    peft_config=lora_config.to_dict(),\n",
    "    torch_dtype=\"bfloat16\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Specify Private Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "priv_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                      dataset_name=\"arc_challenge\",\n",
    "                      data_part=f\"client_{slm_idx}\",\n",
    "                      seq_max_len=512,\n",
    "                      need_preprocess=True)\n",
    "priv_data.load(process_data_output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Other Variables "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = get_tokenizer(slm_pretrained_paths[slm_idx])\n",
    "\n",
    "import json\n",
    "with open(llm_to_slm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "    vocab_mapping = json.loads(fin.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### New FEDMKTSLM Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = FedMKTSLM(\n",
    "    ctx=ctx,\n",
    "    model=model,\n",
    "    training_args=training_args,\n",
    "    fed_args=fed_args,\n",
    "    pub_train_set=pub_data,\n",
    "    priv_train_set=priv_data,\n",
    "    tokenizer=tokenizer,\n",
    "    save_trainable_weights_only=True, # save lora weights only\n",
    "    llm_tokenizer=get_tokenizer(llm_pretrained_path), # different with LLM setting\n",
    "    llm_to_slm_vocab_mapping=vocab_mapping, # different with LLM setting\n",
    "    data_collator=transformers.DataCollatorForSeq2Seq(tokenizer) # use to train private dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete Code To DO SFT With 4 SLMs\n",
    "\n",
    "Please paste the code in \"fedmkt_4_slms.py\" and execute it with the following command"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "python fedmkt_4_slms.py --parties guest:9999 host:9999 host:10000 host:10001 arbiter:9999 --log_level INFO\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fedmkt_4_slms.py\n",
    "\n",
    "import os\n",
    "\n",
    "from fate.arch import Context\n",
    "from fate.arch.launchers.multiprocess_launcher import launch\n",
    "import json\n",
    "\n",
    "process_data_output_dir = \"\"\n",
    "llm_pretrained_path = \"Llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMa-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\"\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "\n",
    "global_epochs = 5\n",
    "batch_size=4\n",
    "llm_lr = 3e-5\n",
    "slm_lrs = [3e-5, 3e-4, 3e-5, 3e-5, 3e-5]\n",
    "\n",
    "llm_model_saved_directory = \"./models/fedmkt_4_slms_llm_model\"\n",
    "slm_models_saved_directory = [\n",
    "    \"./models/fedmkt_4_slms_slm_0\", \n",
    "    \"./models/fedmkt_4_slms_slm_1\", \n",
    "    \"./models/fedmkt_4_slms_slm_2\", \n",
    "    \"./models/fedmkt_4_slms_slm_3\"\n",
    "]\n",
    "\n",
    "\n",
    "def train_llm(ctx):\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTLLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    "    )\n",
    "\n",
    "    model = LLaMa(\n",
    "        pretrained_path=llm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=llm_pretrained_path,\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=llm_lr,\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size,\n",
    "    )\n",
    "\n",
    "    slm_to_llm_vocab_mapping = []\n",
    "    for path in slm_to_llm_vocab_mapping_paths:\n",
    "        with open(path, \"r\") as fin:\n",
    "            vocab_mapping = json.loads(fin.read())\n",
    "            slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "    slm_tokenizers = [get_tokenizer(slm_pretrained_path) for slm_pretrained_path in slm_pretrained_paths]\n",
    "\n",
    "    tokenizer = get_tokenizer(llm_pretrained_path)\n",
    "    trainer = FedMKTLLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        train_set=pub_data,\n",
    "        tokenizer=tokenizer,\n",
    "        slm_tokenizers=slm_tokenizers,\n",
    "        slm_to_llm_vocab_mappings=slm_to_llm_vocab_mapping,\n",
    "        save_trainable_weights_only=True,\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(llm_model_saved_directory)\n",
    "\n",
    "\n",
    "def train_slm(ctx, slm_idx):\n",
    "    import transformers\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.model_zoo.pellm.gpt2 import GPT2CLM\n",
    "    from fate_llm.model_zoo.pellm.opt import OPT\n",
    "    from fate_llm.model_zoo.pellm.bloom import Bloom\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTSLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    slm_model_class = [\n",
    "        OPT,\n",
    "        GPT2CLM,\n",
    "        LLaMa,\n",
    "        Bloom\n",
    "    ]\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )\n",
    "\n",
    "    model = slm_model_class[slm_idx](\n",
    "        pretrained_path=slm_pretrained_paths[slm_idx],\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    priv_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                          dataset_name=\"arc_challenge\",\n",
    "                          data_part=f\"client_{slm_idx}\",\n",
    "                          seq_max_len=512,\n",
    "                          need_preprocess=True)\n",
    "    priv_data.load(process_data_output_dir)\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=slm_lrs[slm_idx],\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(slm_pretrained_paths[slm_idx]).vocab_size,\n",
    "    )\n",
    "\n",
    "    tokenizer = get_tokenizer(slm_pretrained_paths[slm_idx])\n",
    "\n",
    "    import json\n",
    "    with open(llm_to_slm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "\n",
    "    trainer = FedMKTSLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        pub_train_set=pub_data,\n",
    "        priv_train_set=priv_data,\n",
    "        tokenizer=tokenizer,\n",
    "        save_trainable_weights_only=True,\n",
    "        llm_tokenizer=get_tokenizer(llm_pretrained_path),\n",
    "        llm_to_slm_vocab_mapping=vocab_mapping,\n",
    "        data_collator=transformers.DataCollatorForSeq2Seq(tokenizer)\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(slm_models_saved_directory[slm_idx])\n",
    "\n",
    "\n",
    "def run(ctx: Context):\n",
    "    if ctx.is_on_arbiter:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "        train_llm(ctx)\n",
    "    elif ctx.is_on_guest:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "        train_slm(ctx, slm_idx=0)\n",
    "    else:\n",
    "        if ctx.local.party[1] == \"9999\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "            slm_idx = 1\n",
    "        elif ctx.local.party[1] == \"10000\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "            slm_idx = 2\n",
    "        elif ctx.local.party[1] == \"10001\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "            slm_idx = 3\n",
    "        else:\n",
    "            raise ValueError(f\"party_id={ctx.local.party[1]} is illegal\")\n",
    "\n",
    "        train_slm(ctx, slm_idx=slm_idx)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    launch(run)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running FEDMKT With Launcher (Experimential Using): 1-SLM (One To One)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actually, a slightly modifications from 4-SLMs running code are enough to do sft with single clients, it will be listed in below sections, we take SLM-0(OPT) as an example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Only Use Single Optimal Vocabulary Mapping Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slm_idx = 0\n",
    "slm_to_llm_vocab_mapping = []\n",
    "with open(slm_to_llm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "    vocab_mapping = json.loads(fin.read())\n",
    "    slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "slm_tokenizers = [get_tokenizer(slm_pretrained_paths[slm_idx])]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete Code To DO SFT With 1 SLM\n",
    "\n",
    "Please paste the code in \"fedmkt_1_slm.py\" and execute it with the following command"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "python fedmkt_1_slm.py --parties guest:9999 arbiter:9999 --log_level INFO\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fedmkt_1_slm.py\n",
    "\n",
    "import os\n",
    "\n",
    "from fate.arch import Context\n",
    "from fate.arch.launchers.multiprocess_launcher import launch\n",
    "import json\n",
    "\n",
    "process_data_output_dir = \"\"\n",
    "llm_pretrained_path = \"Llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMa-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\"\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "\n",
    "global_epochs = 5\n",
    "batch_size = 4\n",
    "llm_lr = 3e-5\n",
    "slm_lrs = [3e-5]\n",
    "\n",
    "llm_model_saved_directory = \"./models/fedmkt_single_slm_llm\"\n",
    "slm_models_saved_directory = [\n",
    "    \"./models/fedmkt_single_slm_opt\",\n",
    "]\n",
    "\n",
    "\n",
    "def train_llm(ctx, slm_idx):\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTLLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    "    )\n",
    "\n",
    "    model = LLaMa(\n",
    "        pretrained_path=llm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=llm_pretrained_path,\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=llm_lr,\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size,\n",
    "    )\n",
    "\n",
    "    slm_to_llm_vocab_mapping = []\n",
    "    with open(slm_to_llm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "        slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "    slm_tokenizers = [get_tokenizer(slm_pretrained_paths[slm_idx])]\n",
    "\n",
    "    tokenizer = get_tokenizer(llm_pretrained_path)\n",
    "    trainer = FedMKTLLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        train_set=pub_data,\n",
    "        tokenizer=tokenizer,\n",
    "        slm_tokenizers=slm_tokenizers,\n",
    "        slm_to_llm_vocab_mappings=slm_to_llm_vocab_mapping,\n",
    "        save_trainable_weights_only=True,\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(llm_model_saved_directory)\n",
    "\n",
    "\n",
    "def train_slm(ctx, slm_idx):\n",
    "    import transformers\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.model_zoo.pellm.gpt2 import GPT2CLM\n",
    "    from fate_llm.model_zoo.pellm.opt import OPT\n",
    "    from fate_llm.model_zoo.pellm.bloom import Bloom\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTSLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    slm_model_class = [\n",
    "        OPT,\n",
    "        GPT2CLM,\n",
    "        LLaMa,\n",
    "        Bloom\n",
    "    ]\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )\n",
    "\n",
    "    model = slm_model_class[slm_idx](\n",
    "        pretrained_path=slm_pretrained_paths[slm_idx],\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    priv_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                          dataset_name=\"arc_challenge\",\n",
    "                          data_part=f\"client_{slm_idx}\",\n",
    "                          seq_max_len=512,\n",
    "                          need_preprocess=True)\n",
    "    priv_data.load(process_data_output_dir)\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=slm_lrs[slm_idx],\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(slm_pretrained_paths[slm_idx]).vocab_size,\n",
    "    )\n",
    "\n",
    "    tokenizer = get_tokenizer(slm_pretrained_paths[slm_idx])\n",
    "\n",
    "    import json\n",
    "    with open(llm_to_slm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "\n",
    "    trainer = FedMKTSLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        pub_train_set=pub_data,\n",
    "        priv_train_set=priv_data,\n",
    "        tokenizer=tokenizer,\n",
    "        save_trainable_weights_only=True,\n",
    "        llm_tokenizer=get_tokenizer(llm_pretrained_path),\n",
    "        llm_to_slm_vocab_mapping=vocab_mapping,\n",
    "        data_collator=transformers.DataCollatorForSeq2Seq(tokenizer)\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(slm_models_saved_directory[slm_idx])\n",
    "\n",
    "\n",
    "def run(ctx: Context):\n",
    "    if ctx.is_on_arbiter:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "        train_llm(ctx, slm_idx=0)\n",
    "    else:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "        train_slm(ctx, slm_idx=0)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    launch(run)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running FEDMKT With Launcher (Experimential Using): 1-SLM And SLM Trains Only (LLM2SLM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we introduce how to do SFT using FEDMKT algorithm, with only single SLM are trained, but without LLM training, means that SLM distill knowlege from LLM only, not co-training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Difference With Section \"Running FEDMKT With Launcher (Experimential Using): 1-SLMs\"\n",
    "\n",
    "Add llm_training=False to fedmkt_training_args to both LLM and LLM is enough!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete Code To DO SFT With 1 SLM And SLM Trains Only\n",
    "\n",
    "Please paste the code in \"fedmkt_llm_to_slm.py\" and execute it with the following command"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "python fedmkt_llm_to_slm.py --parties guest:9999 arbiter:9999 --log_level INFO\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fedmkt_llm_to_slm.py\n",
    "\n",
    "import os\n",
    "\n",
    "from fate.arch import Context\n",
    "from fate.arch.launchers.multiprocess_launcher import launch\n",
    "import json\n",
    "\n",
    "process_data_output_dir = \"\"\n",
    "llm_pretrained_path = \"Llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMa-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\"\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "\n",
    "global_epochs = 5\n",
    "batch_size = 4\n",
    "llm_lr = 3e-5\n",
    "slm_lrs = [3e-5]\n",
    "\n",
    "slm_models_saved_directory = [\n",
    "    \"./models/fedmkt_llm_to_slm_opt\",\n",
    "]\n",
    "\n",
    "\n",
    "def train_llm(ctx, slm_idx):\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTLLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    "    )\n",
    "\n",
    "    model = LLaMa(\n",
    "        pretrained_path=llm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=llm_pretrained_path,\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=llm_lr,\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size,\n",
    "        llm_training=False\n",
    "    )\n",
    "\n",
    "    slm_to_llm_vocab_mapping = []\n",
    "    with open(slm_to_llm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "        slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "    slm_tokenizers = [get_tokenizer(slm_pretrained_paths[slm_idx])]\n",
    "\n",
    "    tokenizer = get_tokenizer(llm_pretrained_path)\n",
    "    trainer = FedMKTLLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        train_set=pub_data,\n",
    "        tokenizer=tokenizer,\n",
    "        slm_tokenizers=slm_tokenizers,\n",
    "        slm_to_llm_vocab_mappings=slm_to_llm_vocab_mapping,\n",
    "        save_trainable_weights_only=True,\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "\n",
    "\n",
    "def train_slm(ctx, slm_idx):\n",
    "    import transformers\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.model_zoo.pellm.gpt2 import GPT2CLM\n",
    "    from fate_llm.model_zoo.pellm.opt import OPT\n",
    "    from fate_llm.model_zoo.pellm.bloom import Bloom\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTSLM\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    slm_model_class = [\n",
    "        OPT,\n",
    "        GPT2CLM,\n",
    "        LLaMa,\n",
    "        Bloom\n",
    "    ]\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.1,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )\n",
    "\n",
    "    model = slm_model_class[slm_idx](\n",
    "        pretrained_path=slm_pretrained_paths[slm_idx],\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    priv_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                          dataset_name=\"arc_challenge\",\n",
    "                          data_part=f\"client_{slm_idx}\",\n",
    "                          seq_max_len=512,\n",
    "                          need_preprocess=True)\n",
    "    priv_data.load(process_data_output_dir)\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=slm_lrs[slm_idx],\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(slm_pretrained_paths[slm_idx]).vocab_size,\n",
    "        llm_training=False\n",
    "    )\n",
    "\n",
    "    tokenizer = get_tokenizer(slm_pretrained_paths[slm_idx])\n",
    "\n",
    "    import json\n",
    "    with open(llm_to_slm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "\n",
    "    trainer = FedMKTSLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        pub_train_set=pub_data,\n",
    "        priv_train_set=priv_data,\n",
    "        tokenizer=tokenizer,\n",
    "        save_trainable_weights_only=True,\n",
    "        llm_tokenizer=get_tokenizer(llm_pretrained_path),\n",
    "        llm_to_slm_vocab_mapping=vocab_mapping,\n",
    "        data_collator=transformers.DataCollatorForSeq2Seq(tokenizer)\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(slm_models_saved_directory[slm_idx])\n",
    "\n",
    "\n",
    "def run(ctx: Context):\n",
    "    if ctx.is_on_arbiter:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "        train_llm(ctx, slm_idx=0)\n",
    "    else:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "        train_slm(ctx, slm_idx=0)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    launch(run)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running FEDMKT With Launcher (Experimential Using): 4-SLMs Homogeneous SFT\n",
    "\n",
    "To run homogeneous experiments, two steps are needed.\n",
    "1. add post_fedavg=True to fedmkt_training_args to both LLM and LLM is enough!\n",
    "2. add fed_args to FEDMKTLLM/FEDMKTSLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialze fed args\n",
    "from fate.ml.nn.homo.fedavg import FedAVGArguments\n",
    "\n",
    "fed_args = FedAVGArguments(\n",
    "    aggregate_strategy='epoch',\n",
    "    aggregate_freq=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Complete Code To DO SFT With 4-SLMs Homogeneous SFT\n",
    "\n",
    "Please paste the code in \"fedmkt_4_slms_homo.py\" and execute it with the following command"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "python fedmkt_4_slms_homo.py --parties guest:9999 host:9999 host:10000 host:10001 arbiter:9999 --log_level INFO\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fedmkt_4_slms_homo.py\n",
    "\n",
    "import os\n",
    "\n",
    "from fate.arch import Context\n",
    "from fate.arch.launchers.multiprocess_launcher import launch\n",
    "import json\n",
    "\n",
    "process_data_output_dir = \"\"\n",
    "llm_pretrained_path = \"Llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"opt-1.3b\"\n",
    "slm_2_pretrained_path = \"opt-1.3b\"\n",
    "slm_3_pretrained_path = \"opt-1.3b\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\"\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\"] * 4\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\"] * 4\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path] * 4\n",
    "slm_lora_target_modules = [[\"q_proj\", \"v_proj\"]] * 4\n",
    "\n",
    "global_epochs = 5\n",
    "batch_size = 4\n",
    "llm_lr = 3e-5\n",
    "slm_lrs = [3e-5, 3e-5, 3e-5, 3e-5, 3e-5]\n",
    "\n",
    "llm_model_saved_directory = \"./models/fedmkt_homo_4_slms_llm_model\"\n",
    "slm_models_saved_directory = [\n",
    "    \"./models/fedmkt_homo_4_slms_slm_0\",\n",
    "]\n",
    "\n",
    "\n",
    "def train_llm(ctx):\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.llama import LLaMa\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTLLM\n",
    "    from fate.ml.nn.homo.fedavg import FedAVGArguments\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1,\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    "    )\n",
    "\n",
    "    model = LLaMa(\n",
    "        pretrained_path=llm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=llm_pretrained_path,\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=llm_lr,\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size,\n",
    "        post_fedavg=True, # difference\n",
    "    )\n",
    "\n",
    "    # difference\n",
    "    fed_args = FedAVGArguments(\n",
    "        aggregate_strategy='epoch',\n",
    "        aggregate_freq=1\n",
    "    )\n",
    "\n",
    "    slm_to_llm_vocab_mapping = []\n",
    "    for path in slm_to_llm_vocab_mapping_paths:\n",
    "        with open(path, \"r\") as fin:\n",
    "            vocab_mapping = json.loads(fin.read())\n",
    "            slm_to_llm_vocab_mapping.append(vocab_mapping)\n",
    "\n",
    "    slm_tokenizers = [get_tokenizer(slm_pretrained_path) for slm_pretrained_path in slm_pretrained_paths]\n",
    "\n",
    "    tokenizer = get_tokenizer(llm_pretrained_path)\n",
    "    trainer = FedMKTLLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args,\n",
    "        fed_args=fed_args, # difference\n",
    "        train_set=pub_data,\n",
    "        tokenizer=tokenizer,\n",
    "        slm_tokenizers=slm_tokenizers,\n",
    "        slm_to_llm_vocab_mappings=slm_to_llm_vocab_mapping,\n",
    "        save_trainable_weights_only=True,\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    trainer.save_model(llm_model_saved_directory)\n",
    "\n",
    "\n",
    "def train_slm(ctx, slm_idx):\n",
    "    import transformers\n",
    "    from peft import LoraConfig, TaskType\n",
    "    from fate_llm.model_zoo.pellm.opt import OPT\n",
    "    from fate_llm.algo.fedmkt import FedMKTTrainingArguments, FedMKTSLM\n",
    "    from fate.ml.nn.homo.fedavg import FedAVGArguments\n",
    "    from fate_llm.dataset.qa_dataset import QaDataset\n",
    "    from fate_llm.data.tokenizers.cust_tokenizer import get_tokenizer\n",
    "    from transformers import AutoConfig\n",
    "\n",
    "    slm_model_class = [OPT] * 4\n",
    "\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )\n",
    "\n",
    "    model = slm_model_class[slm_idx](\n",
    "        pretrained_path=slm_pretrained_paths[slm_idx],\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    priv_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                          dataset_name=\"arc_challenge\",\n",
    "                          data_part=f\"client_{slm_idx}\",\n",
    "                          seq_max_len=512,\n",
    "                          need_preprocess=True)\n",
    "    priv_data.load(process_data_output_dir)\n",
    "\n",
    "    pub_data = QaDataset(tokenizer_name_or_path=slm_pretrained_paths[slm_idx],\n",
    "                         dataset_name=\"arc_challenge\",\n",
    "                         data_part=\"common\",\n",
    "                         seq_max_len=512,\n",
    "                         need_preprocess=True)\n",
    "    pub_data.load(process_data_output_dir)\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=global_epochs,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=batch_size,\n",
    "        learning_rate=slm_lrs[slm_idx],\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(slm_pretrained_paths[slm_idx]).vocab_size,\n",
    "        post_fedavg=True, # difference\n",
    "    )\n",
    "\n",
    "    # difference\n",
    "    fed_args = FedAVGArguments(\n",
    "        aggregate_strategy='epoch',\n",
    "        aggregate_freq=1\n",
    "    )\n",
    "\n",
    "    tokenizer = get_tokenizer(slm_pretrained_paths[slm_idx])\n",
    "\n",
    "    import json\n",
    "    with open(llm_to_slm_vocab_mapping_paths[slm_idx], \"r\") as fin:\n",
    "        vocab_mapping = json.loads(fin.read())\n",
    "\n",
    "    trainer = FedMKTSLM(\n",
    "        ctx=ctx,\n",
    "        model=model,\n",
    "        training_args=training_args, \n",
    "        fed_args=fed_args, # difference\n",
    "        pub_train_set=pub_data,\n",
    "        priv_train_set=priv_data,\n",
    "        tokenizer=tokenizer,\n",
    "        save_trainable_weights_only=True,\n",
    "        llm_tokenizer=get_tokenizer(llm_pretrained_path),\n",
    "        llm_to_slm_vocab_mapping=vocab_mapping,\n",
    "        data_collator=transformers.DataCollatorForSeq2Seq(tokenizer)\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "    if slm_idx == 0:\n",
    "        trainer.save_model(slm_models_saved_directory[slm_idx])\n",
    "\n",
    "\n",
    "def run(ctx: Context):\n",
    "    if ctx.is_on_arbiter:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "        train_llm(ctx)\n",
    "    elif ctx.is_on_guest:\n",
    "        os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "        train_slm(ctx, slm_idx=0)\n",
    "    else:\n",
    "        if ctx.local.party[1] == \"9999\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "            slm_idx = 1\n",
    "        elif ctx.local.party[1] == \"10000\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "            slm_idx = 2\n",
    "        elif ctx.local.party[1] == \"10001\":\n",
    "            os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "            slm_idx = 3\n",
    "        else:\n",
    "            raise ValueError(f\"party_id={ctx.local.party[1]} is illegal\")\n",
    "\n",
    "        train_slm(ctx, slm_idx=slm_idx)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    launch(run)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running FEDMKT with Pipeline (Industrial Using)\n",
    "\n",
    "Please make sure that [FATE-LLM Cluster](https://github.com/FederatedAI/FATE/wiki/Download#llm%E9%83%A8%E7%BD%B2%E5%8C%85) has been deployed, ensure that multiple machines has been deployed in FATE-LLM Cluster mode, past the following code to test_fedmkt_4_slms.py, the execute \"python test_fedmkt_4_slms.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fate_client.pipeline.components.fate.homo_nn import HomoNN, get_config_of_fedmkt_runner\n",
    "from fate_client.pipeline.components.fate.nn.algo_params import FedMKTTrainingArguments, FedAVGArguments\n",
    "from fate_client.pipeline.components.fate.nn.loader import LLMModelLoader, LLMDatasetLoader, LLMDataFuncLoader\n",
    "from peft import LoraConfig, TaskType\n",
    "from fate_client.pipeline import FateFlowPipeline\n",
    "from fate_client.pipeline.components.fate.reader import Reader\n",
    "from transformers import AutoConfig\n",
    "\n",
    "guest = '9999' # replace this party id to actual guest party id in your enviroment\n",
    "host = ['9999', '10000', '10001'] # replace host party ids in your enviroment\n",
    "arbiter = '9999' # replace this party id to actual arbiter party id in your enviroment\n",
    "\n",
    "\n",
    "process_data_output_dir = \"\" # replace this to actual process_data_output_dir\n",
    "# replaoce the names of models to local save directories\n",
    "llm_pretrained_path = \"llama-2-7b-hf\"\n",
    "slm_0_pretrained_path = \"opt-1.3b\"\n",
    "slm_1_pretrained_path = \"gpt2-xl\"\n",
    "slm_2_pretrained_path = \"Sheared-LLaMA-1.3B\"\n",
    "slm_3_pretrained_path = \"bloom-1b1\"\n",
    "llm_slm_pairs = [\n",
    "    (llm_pretrained_path, slm_0_pretrained_path),\n",
    "    (llm_pretrained_path, slm_1_pretrained_path),\n",
    "    (llm_pretrained_path, slm_2_pretrained_path),\n",
    "    (llm_pretrained_path, slm_3_pretrained_path)\n",
    "]\n",
    "\n",
    "vocab_mapping_directory = \"\" # reploace this to actual voacb_mapping_directory\n",
    "\n",
    "slm_to_llm_vocab_mapping_paths = [\"opt_to_llama.json\", \"gpt2_to_llama.json\", \"llama_small_to_llama.json\", \"bloom_to_llama.json\"]\n",
    "llm_to_slm_vocab_mapping_paths = [\"llama_to_opt.json\", \"llama_to_gpt2.json\", \"llama_to_llama_small\", \"llama_to_bloom.json\"]\n",
    "\n",
    "for idx in range(4):\n",
    "    slm_to_llm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + slm_to_llm_vocab_mapping_paths[idx]\n",
    "    llm_to_slm_vocab_mapping_paths[idx] = vocab_mapping_directory + \"/\" + llm_to_slm_vocab_mapping_paths[idx]\n",
    "\n",
    "slm_pretrained_paths = [slm_0_pretrained_path, slm_1_pretrained_path, slm_2_pretrained_path, slm_3_pretrained_path]\n",
    "slm_lora_target_modules = [\n",
    "    [\"q_proj\", \"v_proj\"],\n",
    "    [\"c_attn\"],\n",
    "    ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    [\"query_key_value\"]\n",
    "]\n",
    "slm_models = [\n",
    "    (\"pellm.opt\", \"OPT\"),\n",
    "    (\"pellm.gpt2\", \"GPT2CLM\"),\n",
    "    (\"pellm.llama\", \"LLaMa\"),\n",
    "    (\"pellm.bloom\", \"Bloom\")\n",
    "]\n",
    "\n",
    "\n",
    "def get_llm_conf():\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=16, lora_dropout=0.05,\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj']\n",
    "    )\n",
    "    lora_config.target_modules = list(lora_config.target_modules)\n",
    "\n",
    "    llm_model = LLMModelLoader(\n",
    "        \"pellm.llama\",\n",
    "        \"LLaMa\",\n",
    "        pretrained_path=llm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "        torch_dtype=\"bfloat16\"\n",
    "    )\n",
    "\n",
    "    pub_dataset = LLMDatasetLoader(\n",
    "        \"qa_dataset\",\n",
    "        \"QaDataset\",\n",
    "        tokenizer_name_or_path=llm_pretrained_path,\n",
    "        need_preprocess=True,\n",
    "        dataset_name=\"arc_challenge\",\n",
    "        data_part=\"common\",\n",
    "        seq_max_len=512\n",
    "    )\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=5,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=4,\n",
    "        learning_rate=3e-5,\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=AutoConfig.from_pretrained(llm_pretrained_path).vocab_size,\n",
    "    )\n",
    "\n",
    "    fed_args = FedAVGArguments(\n",
    "        aggregate_strategy='epoch',\n",
    "        aggregate_freq=1\n",
    "    )\n",
    "\n",
    "    tokenizer = LLMDataFuncLoader(\n",
    "        \"tokenizers.cust_tokenizer\",\n",
    "        \"get_tokenizer\",\n",
    "        tokenizer_name_or_path=llm_pretrained_path\n",
    "    )\n",
    "\n",
    "    slm_tokenizers = list()\n",
    "    for slm_pretrained_path in slm_pretrained_paths:\n",
    "        slm_tokenizers.append(\n",
    "            LLMDataFuncLoader(\"tokenizers.cust_tokenizer\", \"get_tokenizer\", tokenizer_name_or_path=slm_pretrained_path)\n",
    "        )\n",
    "\n",
    "    return get_config_of_fedmkt_runner(\n",
    "        model=llm_model,\n",
    "        training_args=training_args,\n",
    "        fed_args=fed_args,\n",
    "        pub_dataset=pub_dataset,\n",
    "        tokenizer=tokenizer,\n",
    "        slm_tokenizers=slm_tokenizers,\n",
    "        slm_to_llm_vocab_mapping_paths=slm_to_llm_vocab_mapping_paths,\n",
    "        pub_dataset_path=process_data_output_dir,\n",
    "        save_trainable_weights_only=True,\n",
    "    )\n",
    "\n",
    "\n",
    "def get_slm_conf(slm_idx):\n",
    "    slm_pretrained_path = slm_pretrained_paths[slm_idx]\n",
    "    lora_config = LoraConfig(\n",
    "        task_type=TaskType.CAUSAL_LM,\n",
    "        inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1,\n",
    "        target_modules=slm_lora_target_modules[slm_idx]\n",
    "    )\n",
    "    lora_config.target_modules = list(lora_config.target_modules)\n",
    "    llm_to_slm_vocab_mapping = llm_to_slm_vocab_mapping_paths[slm_idx]\n",
    "\n",
    "    slm_model = LLMModelLoader(\n",
    "        slm_models[slm_idx][0],\n",
    "        slm_models[slm_idx][1],\n",
    "        pretrained_path=slm_pretrained_path,\n",
    "        peft_type=\"LoraConfig\",\n",
    "        peft_config=lora_config.to_dict(),\n",
    "    )\n",
    "    vocab_size = AutoConfig.from_pretrained(slm_pretrained_path).vocab_size\n",
    "\n",
    "    pub_dataset = LLMDatasetLoader(\n",
    "        \"qa_dataset\",\n",
    "        \"QaDataset\",\n",
    "        tokenizer_name_or_path=slm_pretrained_path,\n",
    "        need_preprocess=True,\n",
    "        dataset_name=\"arc_challenge\",\n",
    "        data_part=\"common\",\n",
    "        seq_max_len=512\n",
    "    )\n",
    "\n",
    "    priv_dataset = LLMDatasetLoader(\n",
    "        \"qa_dataset\",\n",
    "        \"QaDataset\",\n",
    "        tokenizer_name_or_path=slm_pretrained_path,\n",
    "        need_preprocess=True,\n",
    "        dataset_name=\"arc_challenge\",\n",
    "        data_part=\"client_0\",\n",
    "        seq_max_len=512\n",
    "    )\n",
    "\n",
    "    training_args = FedMKTTrainingArguments(\n",
    "        global_epochs=5,\n",
    "        per_device_train_batch_size=1,\n",
    "        gradient_accumulation_steps=4,\n",
    "        learning_rate=3e-5 if slm_idx != 1 else 3e-4\n",
    "        output_dir=\"./\",\n",
    "        dataloader_num_workers=4,\n",
    "        remove_unused_columns=False,\n",
    "        warmup_ratio=0.008,\n",
    "        lr_scheduler_type=\"cosine\",\n",
    "        optim=\"adamw_torch\",\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.95,\n",
    "        weight_decay=0.1,\n",
    "        max_grad_norm=1.0,\n",
    "        use_cpu=False,\n",
    "        vocab_size=vocab_size,\n",
    "        # post_fedavg=True,\n",
    "        # llm_training=False,\n",
    "    )\n",
    "\n",
    "    fed_args = FedAVGArguments(\n",
    "        aggregate_strategy='epoch',\n",
    "        aggregate_freq=1\n",
    "    )\n",
    "\n",
    "    tokenizer = LLMDataFuncLoader(\n",
    "        \"tokenizers.cust_tokenizer\",\n",
    "        \"get_tokenizer\",\n",
    "        tokenizer_name_or_path=slm_pretrained_path\n",
    "    )\n",
    "\n",
    "    llm_tokenizer = LLMDataFuncLoader(\n",
    "        \"tokenizers.cust_tokenizer\", \"get_tokenizer\", tokenizer_name_or_path=llm_pretrained_path\n",
    "    )\n",
    "\n",
    "    data_collator = LLMDataFuncLoader(module_name='data_collator.cust_data_collator',\n",
    "                                      item_name='get_seq2seq_data_collator', tokenizer_name_or_path=slm_pretrained_path)\n",
    "\n",
    "    return get_config_of_fedmkt_runner(\n",
    "        model=slm_model,\n",
    "        training_args=training_args,\n",
    "        fed_args=fed_args,\n",
    "        pub_dataset=pub_dataset,\n",
    "        priv_dataset=priv_dataset,\n",
    "        tokenizer=tokenizer,\n",
    "        llm_tokenizer=llm_tokenizer,\n",
    "        llm_to_slm_vocab_mapping_path=llm_to_slm_vocab_mapping,\n",
    "        pub_dataset_path=process_data_output_dir,\n",
    "        save_trainable_weights_only=True,\n",
    "        data_collator=data_collator\n",
    "    )\n",
    "\n",
    "\n",
    "pipeline = FateFlowPipeline().set_parties(guest=guest, arbiter=arbiter, host=host)\n",
    "pipeline.bind_local_path(path=process_data_output_dir, namespace=\"experiment\", name=\"arc_challenge\")\n",
    "\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=\"arc_challenge\"\n",
    ")\n",
    "reader_0.hosts[[0, 1, 2]].task_parameters(\n",
    "    namespace=\"experiment\",\n",
    "    name=\"arc_challenge\"\n",
    ")\n",
    "\n",
    "\n",
    "homo_nn_0 = HomoNN(\n",
    "    'nn_0',\n",
    "    train_data=reader_0.outputs[\"output_data\"],\n",
    "    runner_module=\"fedmkt_runner\",\n",
    "    runner_class=\"FedMKTRunner\",\n",
    ")\n",
    "\n",
    "homo_nn_0.arbiter.task_parameters(\n",
    "    runner_conf=get_llm_conf()\n",
    ")\n",
    "\n",
    "homo_nn_0.guest.task_parameters(\n",
    "    runner_conf=get_slm_conf(slm_idx=0)\n",
    ")\n",
    "\n",
    "for idx in range(3):\n",
    "    homo_nn_0.hosts[idx].task_parameters(\n",
    "        runner_conf=get_slm_conf(slm_idx=idx + 1)\n",
    "    )\n",
    "\n",
    "homo_nn_0.guest.conf.set(\"launcher_name\", \"deepspeed\") # tell schedule engine to run task with deepspeed\n",
    "homo_nn_0.hosts[[0, 1, 2]].conf.set(\"launcher_name\", \"deepspeed\") # tell schedule engine to run task with deepspeed\n",
    "homo_nn_0.arbiter.conf.set(\"launcher_name\", \"deepspeed\") # tell schedule engine to run task with deepspeed\n",
    "\n",
    "pipeline.add_tasks([reader_0, homo_nn_0])\n",
    "pipeline.conf.set(\"task\", dict(engine_run={\"cores\": 1})) # the number of gpus of each party\n",
    "\n",
    "pipeline.compile()\n",
    "pipeline.fit()\n",
    "\n"
   ]
  }
 ],
 "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}