[Colossal-Llama-2] Add finetuning Colossal-Llama-2 example (#4878)

* Add finetuning Colossal-Llama-2 example

* Add finetuning Colossal-Llama-2 example 2

* Add finetuning Colossal-Llama-2 example and support NEFTuning

* Add inference example and refine neftune

* Modify readme file

* update the imports

---------
Co-authored-by: Xu Yuanchen <yuanchen.xu00@gmail.com>
Co-authored-by: Camille Zhong <44392324+Camille7777@users.noreply.github.com>

applications/Colossal-LLaMA-2/README.md

@@ -11,7 +11,10 @@
   - [Performance Evaluation](#performance-evaluation)
   - [Examples](#examples)
   - [Training Logs](#training-logs)
-  - [Import from Transformers (Inference)](#import-from-transformers-inference)
+  - [Inference](#inference)
+    - [Import from HuggingFace](#import-from-huggingface)
+    - [Import from Modelscope](#import-from-modelscope)
+    - [Quick Start](#quick-start)
 - [Usage](#usage)
   - [Install](#install)
     - [0. Pre-requisite](#0-pre-requisite)
@@ -21,8 +24,14 @@
     - [1. Init Tokenizer Preparation](#1-init-tokenizer-preparation)
     - [2. Init Model Preparation](#2-init-model-preparation)
     - [3. Data Preparation](#3-data-preparation)
+        - [3.1 Data for Pretraining](#31-data-for-pretraining)
+        - [3.2 Data for Supervised Fine-tuning](#32-data-for-supervised-fine-tuning)
     - [4. Command Line Arguments for Training](#4-command-line-arguments-for-training)
+        - [4.1 Arguments for Pretraining](#41-arguments-for-pretraining)
+        - [4.2 Arguments for Supervised Fine-tuning](#42-arguments-for-supervised-fine-tuning)
     - [5. Running Command](#5-running-command)
+        - [5.1 Command for Pretraining](#51-command-for-pretraining)
+        - [5.2 Command for Supervised Fine-tuning](#52-command-for-supervised-fine-tuning)
 - [Technical Insights](#technical-insights)
   - [Data](#data)
   - [Tokenizer](#tokenizer)
@@ -117,7 +126,8 @@ We also recorded the training logs for the experiment
 <img src="https://github.com/hpcaitech/public_assets/blob/main/applications/colossal-llama-2/trainingLossByTokens.jpeg?raw=true" width=600/>
 </p>
 
-### Import from Transformers (Inference)
+### Inference
+#### Import from HuggingFace
 To load Colossal-LLaMA-2-7B-base model using Transformers, use the following code:
 ```Python
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -135,6 +145,7 @@ pred = model.generate(**inputs,
 print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)[len(input):])
 ```
 
+#### Import from Modelscope
 You can also load our model using modelscope, use the following code:
 ```Python
 from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download
@@ -153,6 +164,30 @@ print(tokenizer.decode(output.cpu()[0], skip_special_tokens=True)[len(input):])
 ```
 You can download model weights from [🤗HuggingFace](https://huggingface.co/hpcai-tech/Colossal-LLaMA-2-7b-base) or [👾Modelscope](https://modelscope.cn/models/colossalai/Colossal-LLaMA-2-7b-base/summary).
 
+#### Quick Start
+You can run [`inference_example.py`](inference_example.py) to quickly start the inference of our base model by loading model weights from HF.
+
+Command to run the script:
+```bash
+python inference_example.py \
+    --model_path "<HF_REPO_NAME_OR_LOCAL_PATH_TO_MODEL>" \
+    --device "cuda:0" \
+    --max_new_tokens 512 \
+    --do_sample True \
+    --temperature 0.3 \
+    --top_k 50 \
+    --top_p 0.95 \
+    --input_txt "YOUR_PROMPT_OR_QUESTION" 
+```
+Here is details about CLI arguments:
+* Model path: `--model_path`. HF repo name or local path of the model.
+* Device: `--device`. Set the device.
+* Max new tokens: `--max_new_tokens`. Set maximum numbers of tokens to generate, ignoring the number of tokens in the prompt.
+* Do sample: `--do_sample`. Set whether or not to use sampling.
+* Temperature: `--temperature`. Set temperature value.
+* Top_k: `--top_k`. Set top_k value for top-k-filtering.
+* Top_p: `--top_p`. Set top_p value for generation.
+* Input_txt: `--input_txt`. The prompt string input to the model.
 ## Usage
 ### Install
 
@@ -218,6 +253,8 @@ Here is details about CLI arguments:
 ❗️**Important**: Once you initialize the new model checkpoint, copy your new tokenizer files (`special_tokens_map.json`, `tokenizer.model` and `tokenizer_config.json`) to your new model folder.
 
 #### 3. Data Preparation
+
+##### 3.1 Data for Pretraining
 Raw data should be formatted as `jsonl` format. Each data point should have the following fields:
 * `source` (str, compulsory): This part is ignored when calculating loss. Default can be empty.
 * `target` (str, compulsory): Loss will be calculated.
@@ -250,7 +287,31 @@ Here is details about CLI arguments:
 * Max length: `max_length`. Max length of spliced samples. Default value is 4096.
 * Number of bins for each category: `num_spliced_dataset_bins`. Number of bins for each category, used for bucket-based training.
 
+##### 3.2 Data for Supervised Fine-tuning
+We prepare data for supervised fine-tuning in a similar way. The main difference lies in the data format. Each data point should have the following field:
+* `messages` (list, compulsory): This part consists of a conversation between a human and assistant. The length of `messages` can vary and only content from `assistant` is used for calculating loss.
+
+Examples:
+```JSON
+{"messages": [{"from": "human", "content": "What are the three primary colors?"}, {"from": "assistant", "content": "The three primary colors are red, blue, and yellow."}]}
+{"messages": [{"from": "human", "content": "解释个人电脑和服务器之间的区别。"}, {"from": "assistant", "content": "个人电脑和服务器是两种不同类型的计算机系统，它们的主要区别在于用途、硬件配置和性能。 个人电脑，顾名思义，是为个人使用而设计的计算机。它们通常用于日常的工作、娱乐和学习，可以运行各种各样的应用程序和游戏。个人电脑的硬件配置一般是按照标准配置来设计的，不过也可以根据个人需求进行定制。 而服务器是为了满足大量用户的需求而设计的计算机系统，它们通常用于为用户提供各种网络服务，如网站、电子邮件和文件传输等。服务器通常需要高性能的硬件配置，并且可以承受高负载和长时间的运行。由于服务器需要支持大量用户的访问，它们通常配备多核处理器、大容量内存和大容量硬盘驱动器，以提高系统的运行速度和稳定性。 总之，个人电脑和服务器之间的主要区别在于它们的用途、硬件配置和性能。个人电脑用于个人使用，而服务器用于支持大量用户的访问。服务器的硬件配置通常比个人电脑更高，以保证系统的性能和稳定性。"}]}
+```
+
+Command to convert jsonl dataset to arrow format is similar to the command in [3.1 Data for Pretraining](#31-data-for-pretraining). In `prepare_sft_dataset.py`, we don't concatenate different data samples.
+```
+python prepare_sft_dataset.py.py \
+    --data_input_dirs "<JOSNL_DIR_1>,<JOSNL_DIR_2>,<JOSNL_DIR_3>" \
+    --tokenizer_dir "<TOKENIZER_DIR>" \
+    --data_cache_dir "jsonl_to_arrow_cache" \
+    --data_jsonl_output_dir "spliced_tokenized_output_jsonl" \
+    --data_arrow_output_dir "spliced_tokenized_output_arrow" \
+    --max_length 4096 \
+    --num_spliced_dataset_bins 10
+```
+
 #### 4. Command Line Arguments for Training
+
+##### 4.1 Arguments for Pretraining
 You can use `colossalai run` to launch multi-nodes training:
 ```bash
 colossalai run --nproc_per_node YOUR_GPU_PER_NODE --hostfile YOUR_HOST_FILE \
@@ -288,7 +349,16 @@ Here is details about CLI arguments:
 * Tensor parallelism size: `--tp`. TP size for 3d Parallelism. The default value is 1.
 * Zero stage: `--zero`. Zero stage for 3d Parallelism. The default value is 1.
 
+##### 4.2 Arguments for Supervised Fine-tuning
+We add support for gradient accumulation and NEFTuning for supervised fine-tuning and thus there are two more arguments apart from the arguments listed in [4.1 Arguments for Pretraining](#41-arguments-for-pretraining).
+
+Here is details about CLI arguments:
+* Accumulation steps: `--accumulation_steps`. The default value is `8`.
+* NEFTuning: `--use_neft`. The default value is `False`. It can help improve the performance of chat models.
+
 #### 5. Running Command
+
+##### 5.1 Command for Pretraining
 An [example bash](train.example.sh) is also provided for the experiment. Here is the steps to run the experiment:
 * Create your own hostfile: `cp hostfile.example hostfile`.
 * Create your own bash: `cp train.example.sh train.sh`.
@@ -310,6 +380,10 @@ declare -a dataset=(
     "<DIR_2>/part-00000"
 )
 ```
+
+##### 5.2 Command for Supervised Fine-tuning
+An [example bash](train_sft.example.sh) is provided. The only difference with the command for pretraining is the two arguments (`--accumulation_steps` and `--use_neft`) in the script. You can refer to [4.2 Arguments for Supervised Fine-tuning](#42-arguments-for-supervised-fine-tuning) for more details.
+
 ## Technical Insights
 In order to enhance LLaMA-2's capabilities for understanding and generating Chinese content, The [Colossal-AI](https://github.com/hpcaitech/ColossalAI) team proposes the continuation of pre-training the LLaMA-2 model using both Chinese and English corpora. The overall pipeline can be described as follows:
 
@@ -416,3 +490,11 @@ Applying the above process to perform knowledge transfer in any field allows for
     year={2023}
 }
 ```
+```bibtex
+@article{jain2023neftune,
+    title={NEFTune: Noisy Embeddings Improve Instruction Finetuning},
+    author={Jain, Neel and Chiang, Ping-yeh and Wen, Yuxin and Kirchenbauer, John and Chu, Hong-Min and Somepalli, Gowthami and Bartoldson, Brian R and Kailkhura, Bhavya and Schwarzschild, Avi and Saha, Aniruddha and others},
+    journal={arXiv preprint arXiv:2310.05914},
+    year={2023}
+}
+```

applications/Colossal-LLaMA-2/colossal_llama2/dataset/conversation.py

+#    Copyright 2023 lm-sys@FastChat
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import dataclasses
+from enum import Enum, auto
+from typing import List
+
+
+class SeparatorStyle(Enum):
+    ADD_BOS_EOS_TOKEN = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle
+    seps: List[str]
+
+    def clear(self):
+        self.messages = []
+
+    def get_prompt(self, length: int = None):
+        if length is None:
+            length = len(self.messages)
+
+        if self.sep_style == SeparatorStyle.ADD_BOS_EOS_TOKEN:
+            ret = self.system
+            for role, message in self.messages[0:length]:
+                if message:
+                    ret += role + ": " + self.seps[0] + message + self.seps[1]
+                else:
+                    ret += role + ": " + self.seps[0]
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def save_prompt(self):
+        if self.sep_style == SeparatorStyle.ADD_BOS_EOS_TOKEN:
+            ret = self.system
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + self.seps[0] + message + self.seps[1] + "\n"
+                else:
+                    ret += role + ": " + self.seps[0]
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            seps=self.seps,
+        )
+
+    def dict(self):
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "seps": self.seps,
+        }
+
+
+conv = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+    roles=("Human", "Assistant"),
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.ADD_BOS_EOS_TOKEN,
+    seps=["<s>", "</s>"],
+)
+
+default_conversation = conv

applications/Colossal-LLaMA-2/colossal_llama2/dataset/spliced_and_tokenized_dataset.py

@@ -4,22 +4,29 @@
 Splicing multiple pre-tokenized sequence data points
 """
 
+import bisect
 import random
 import warnings
 from copy import deepcopy
-from datasets import dataset_dict
-from typing import Any, Callable, Dict, Iterable, List, Union, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Tuple, Union
 
+from datasets import dataset_dict
 from torch.utils.data import ConcatDataset, Dataset, IterableDataset
 from transformers.models.llama.tokenization_llama import LlamaTokenizer
 from transformers.tokenization_utils import PreTrainedTokenizer
 
+from colossalai.logging import get_dist_logger
+
+from .conversation import Conversation, default_conversation
+
+logger = get_dist_logger()
+
 IGNORE_INDEX = -100
 
 DSType = Union[Dataset, ConcatDataset, dataset_dict.Dataset]
 
 
-def supervised_tokenize(
+def supervised_tokenize_pretrain(
     data_point: Dict[str, str], tokenizer: LlamaTokenizer, ignore_index: int = None, max_length: int = 4096
 ) -> Dict[str, Union[int, str, List[int]]]:
     """
@@ -62,6 +69,121 @@ def supervised_tokenize(
     )
 
 
+def supervised_tokenize_sft(
+    data_point: Dict[str, str],
+    tokenizer: LlamaTokenizer,
+    conversation_template: Conversation = default_conversation,
+    ignore_index: int = None,
+    max_length: int = 4096,
+) -> Dict[str, Union[int, str, List[int]]]:
+    """
+    A tokenization function to tokenize an original supervised data point as following:
+        {"messages": [{"from": "human", "content": "xxx"}, {"from": "assistant", "content": "xxx"}]}
+    """
+    assert tokenizer.add_bos_token is False and tokenizer.add_eos_token is False, (
+        "Initially set `tokenizer.add_bos_token` and `tokenizer.add_eos_token` to False, "
+        "add <bos> and <eos> manually later"
+    )
+
+    assert (
+        tokenizer.bos_token == conversation_template.seps[0] and tokenizer.eos_token == conversation_template.seps[1]
+    ), "`bos_token` and `eos_token` should be the same with `conversation_template.seps`."
+
+    if ignore_index is None:
+        ignore_index = IGNORE_INDEX
+
+    messages = data_point["messages"]
+    template = deepcopy(conversation_template)
+    template.messages = []
+
+    for mess in messages:
+        from_str = mess["from"]
+        if from_str.lower() == "human":
+            from_str = template.roles[0]
+        elif from_str.lower() == "assistant":
+            from_str = template.roles[1]
+        else:
+            raise ValueError(f"Unsupported role {from_str.lower()}")
+
+        template.append_message(from_str, mess["content"])
+
+    if len(template.messages) % 2 != 0:
+        template.messages = template.messages[0:-1]
+
+    # `target_turn_index` is the number of turns which exceeds `max_length - 1` for the first time.
+    turns = [i for i in range(1, len(messages) // 2 + 1)]
+    target_turn_index = bisect.bisect_right(
+        turns,
+        max_length - 1,
+        key=lambda x: len(tokenizer([template.get_prompt(2 * x)], add_special_tokens=False)["input_ids"][0]),
+    )
+
+    # The tokenized length for first turn already exceeds `max_length - 1`.
+    if target_turn_index - 1 < 0:
+        return dict(
+            input_ids=None,
+            labels=None,
+            inputs_decode=None,
+            labels_decode=None,
+            seq_length=None,
+            seq_category=None,
+        )
+
+    target_turn = turns[target_turn_index - 1]
+    prompt = template.get_prompt(2 * target_turn)
+    tokenized = tokenizer([prompt], add_special_tokens=False)["input_ids"][0]
+
+    template.messages = template.messages[0 : 2 * target_turn]
+
+    starts = []
+    ends = []
+    gpt_bos = False if template.messages[0][0] == template.roles[0] else True
+    gpt_eos = False if template.messages[0][0] == template.roles[0] else True
+
+    for i, token_id in enumerate(tokenized):
+        if token_id == tokenizer.bos_token_id:
+            if gpt_bos:
+                starts.append(i)
+            gpt_bos = not gpt_bos
+        elif token_id == tokenizer.eos_token_id:
+            if gpt_eos:
+                ends.append(i)
+            gpt_eos = not gpt_eos
+
+    if len(starts) != target_turn or len(ends) != target_turn:
+        logger.info(
+            "Please check whether the tokenizer add additional `bos_token` and `eos_token`.\n\nOr the original message contains `bos_token` or `eos_token`."
+        )
+        return dict(
+            input_ids=None,
+            labels=None,
+            inputs_decode=None,
+            labels_decode=None,
+            seq_length=None,
+            seq_category=None,
+        )
+
+    tokenized = [tokenizer.bos_token_id] + tokenized
+    labels = [ignore_index] * len(tokenized)
+    for start, end in zip(starts, ends):
+        labels[start + 1 : end + 2] = tokenized[start + 1 : end + 2]
+
+    labels_decode = deepcopy(labels)
+    for i, z in enumerate(labels_decode):
+        if z == ignore_index:
+            labels_decode[i] = tokenizer.unk_token_id
+
+    # `inputs_decode` and `labels_decode` can be used to check whether the tokenization method is true.
+    return dict(
+        input_ids=tokenized,
+        labels=labels,
+        inputs_decode=tokenizer.decode(tokenized),
+        labels_decode=tokenizer.decode(labels_decode),
+        seq_length=len(tokenized),
+        seq_category=data_point["category"] if "category" in data_point else "None",
+    )
+
+
 class ClosedToConstantLengthSplicedDataset(IterableDataset):
     """
     Define an iterable dataset that returns a (close to) constant length data point spliced from multiple
@@ -169,12 +291,7 @@ class ClosedToConstantLengthSplicedDataset(IterableDataset):
                     spliced_labels.extend(seq_labels)
             # For residual spliced data point at the end of the data set
             if self.infinite is False and more_data_points is False and len(spliced_input_ids) > 0:
-                examples.append(
-                    {
-                        self.input_ids_field: spliced_input_ids,
-                        self.labels_field: spliced_labels
-                    }
-                )
+                examples.append({self.input_ids_field: spliced_input_ids, self.labels_field: spliced_labels})
             if self.shuffle:
                 random.shuffle(examples)
             for spliced_data_point in examples:

applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py

+#    Copyright 2023 The Hugging Face team
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import torch
+
+
+def unwrap(model):
+    return model.unwrap().module
+
+
+def neftune_post_forward_hook(module, input, output):
+    """
+    Implements the NEFTune forward pass for the model using forward hooks. Note this works only for torch.nn.Embedding
+    layers. This method is slightly adapted from the original source code that can be found here:
+    https://github.com/neelsjain/NEFTune Simply add it to your model as follows:
+    ```python
+    model = ...
+    model.embed_tokens.neftune_noise_alpha = 0.1
+    model.embed_tokens.register_forward_hook(neftune_post_forward_hook)
+    ```
+    Args:
+        module (`torch.nn.Module`):
+            The embedding module where the hook is attached. Note that you need to set `module.neftune_noise_alpha` to
+            the desired noise alpha value.
+        input (`torch.Tensor`):
+            The input tensor to the model.
+        output (`torch.Tensor`):
+            The output tensor of the model (i.e. the embeddings).
+    """
+    if module.training:
+        dims = torch.tensor(output.size(1) * output.size(2))
+        mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)
+        output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
+    return output
+
+
+def activate_neftune(model, neftune_noise_alpha=0.1):
+    r"""
+    Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper:
+    https://arxiv.org/abs/2310.05914
+    """
+    embeddings = unwrap(model).get_input_embeddings()
+
+    embeddings.neftune_noise_alpha = neftune_noise_alpha
+    hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook)
+    neftune_hook_handle = hook_handle
+
+    return model, neftune_hook_handle
+
+
+def deactivate_neftune(model, neftune_hook_handle):
+    """
+    Deactivates the neftune method. Make sure to call `_activate_neftune` first.
+    """
+    embeddings = unwrap(model).get_input_embeddings()
+
+    neftune_hook_handle.remove()
+    del embeddings.neftune_noise_alpha

applications/Colossal-LLaMA-2/inference_example.py

+import argparse
+import os
+
+import torch
+from colossalai.logging import get_dist_logger
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+logger = get_dist_logger()
+
+
+def load_model(model_path, device="cuda", **kwargs):
+    logger.info(
+        "Please check whether the tokenizer and model weights are properly stored in the same folder."
+    )
+    model = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
+    model.to(device)
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+    except OSError:
+        raise ImportError("Tokenizer not found. Please check if the tokenizer exists or the model path is correct.")
+
+    return model, tokenizer
+
+
+@torch.inference_mode()
+def generate(args):
+    model, tokenizer = load_model(model_path=args.model_path, device=args.device)
+
+    BASE_INFERENCE_SUFFIX = "\n\n->\n\n"
+    input_txt = f"{args.input_txt}{BASE_INFERENCE_SUFFIX}"
+
+    inputs = tokenizer(args.input_txt, return_tensors='pt').to(args.device)
+    output = model.generate(**inputs,
+                            max_new_tokens=args.max_new_tokens,
+                            do_sample=args.do_sample,
+                            temperature=args.temperature,
+                            top_k=args.top_k,
+                            top_p=args.top_p,
+                            num_return_sequences=1)
+    response = tokenizer.decode(output.cpu()[0], skip_special_tokens=True)[len(input_txt):]
+    logger.info(f"Question: {input_txt} \n\n Answer: \n{response}")
+    return response
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Colossal-LLaMA-2 inference Process.")
+    parser.add_argument('--model_path', type=str, default="hpcai-tech/Colossal-LLaMA-2-7b-base", help="HF repo name or local path of the model")
+    parser.add_argument('--device', type=str, default="cuda:0", help="Set the device")
+    parser.add_argument('--max_new_tokens', type=int, default=512, help=" Set maximum numbers of tokens to generate, ignoring the number of tokens in the prompt")
+    parser.add_argument('--do_sample', type=bool, default=True, help="Set whether or not to use sampling")
+    parser.add_argument('--temperature', type=float, default=0.3, help="Set temperature value")
+    parser.add_argument('--top_k', type=int, default=50, help="Set top_k value for top-k-filtering")
+    parser.add_argument('--top_p', type=int, default=0.95, help="Set top_p value for generation")
+    parser.add_argument('--input_txt', type=str, default="明月松间照，", help="The prompt input to the model")
+    args = parser.parse_args()
+    generate(args)
\ No newline at end of file

applications/Colossal-LLaMA-2/prepare_pretrain_dataset.py

@@ -11,14 +11,14 @@ import os
 import time
 from multiprocessing import cpu_count
 
+from colossal_llama2.dataset.spliced_and_tokenized_dataset import (
+    ClosedToConstantLengthSplicedDataset,
+    supervised_tokenize_pretrain,
+)
 from datasets import dataset_dict, load_dataset
 from transformers.models.llama.tokenization_llama import LlamaTokenizer
 
 from colossalai.logging import get_dist_logger
-from colossal_llama2.dataset.spliced_and_tokenized_dataset import (
-    supervised_tokenize,
-    ClosedToConstantLengthSplicedDataset,
-)
 
 logger = get_dist_logger()
 
@@ -104,7 +104,7 @@ def main():
         assert isinstance(dataset, dataset_dict.Dataset)
         logger.info(f"Start to process part-{index}/{len(list_dataset)} of all original datasets.")
         dataset = dataset.map(
-            function=supervised_tokenize,
+            function=supervised_tokenize_pretrain,
             fn_kwargs={"tokenizer": tokenizer, "max_length": args.max_length},
             keep_in_memory=False,
             num_proc=min(len(dataset), cpu_count()),
@@ -149,5 +149,5 @@ def main():
         spliced_dataset.save_to_disk(dataset_path=output_arrow_path, num_proc=min(len(spliced_dataset), cpu_count()))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

applications/Colossal-LLaMA-2/prepare_sft_dataset.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Prepare sft dataset for fine-tuning
+"""
+
+import argparse
+import json
+import math
+import os
+from multiprocessing import cpu_count
+
+from colossal_llama2.dataset.conversation import default_conversation
+from colossal_llama2.dataset.spliced_and_tokenized_dataset import supervised_tokenize_sft
+from datasets import dataset_dict, load_dataset
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+
+from colossalai.logging import get_dist_logger
+
+logger = get_dist_logger()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_input_dirs",
+        type=str,
+        required=True,
+        default=None,
+        help="Comma(i.e., ',') separated list of all data directories containing `.jsonl` data files.",
+    )
+    parser.add_argument(
+        "--tokenizer_dir", type=str, required=True, default=None, help="A directory containing the tokenizer"
+    )
+    parser.add_argument("--data_cache_dir", type=str, default="cache", help="Data cache directory")
+    parser.add_argument(
+        "--data_jsonl_output_dir",
+        type=str,
+        default="jsonl_output",
+        help="Output directory of spliced dataset with jsonl format",
+    )
+    parser.add_argument(
+        "--data_arrow_output_dir",
+        type=str,
+        default="arrow_output",
+        help="Output directory of spliced dataset with arrow format",
+    )
+    parser.add_argument("--max_length", type=int, default=4096, help="Max length of each spliced tokenized sequence")
+    parser.add_argument("--num_spliced_dataset_bins", type=int, default=10, help="Number of spliced dataset bins")
+    args = parser.parse_args()
+
+    if args.num_spliced_dataset_bins >= 100000:
+        raise ValueError("Too many spliced divisions, must be smaller than 100000")
+
+    assert not os.path.exists(args.data_cache_dir), f"Find existed data cache dir {args.data_cache_dir}"
+    assert not os.path.exists(
+        args.data_jsonl_output_dir
+    ), f"Find existed jsonl data output dir {args.data_jsonl_output_dir}"
+    assert not os.path.exists(
+        args.data_arrow_output_dir
+    ), f"Find existed arrow data output dir {args.data_arrow_output_dir}"
+    os.makedirs(args.data_jsonl_output_dir)
+    os.makedirs(args.data_arrow_output_dir)
+
+    # Prepare to all input datasets
+    input_data_paths = []
+    input_data_dirs = args.data_input_dirs.split(",")
+    for ds_dir in input_data_dirs:
+        ds_dir = os.path.abspath(ds_dir)
+        assert os.path.exists(ds_dir), f"Not find data dir {ds_dir}"
+        ds_files = [name for name in os.listdir(ds_dir) if name.endswith(".jsonl")]
+        ds_paths = [os.path.join(ds_dir, name) for name in ds_files]
+        input_data_paths.extend(ds_paths)
+
+    # Prepare to data splitting.
+    train_splits = []
+    split_interval = math.ceil(100 / args.num_spliced_dataset_bins)
+    for i in range(0, 100, split_interval):
+        start = i
+        end = i + split_interval
+        if end > 100:
+            end = 100
+        train_splits.append(f"train[{start}%:{end}%]")
+
+    # Prepare to the tokenizer.
+    tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_dir)
+    tokenizer.add_bos_token = False
+    tokenizer.add_eos_token = False
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.unk_token
+
+    list_dataset = load_dataset(
+        path="json",
+        data_files=input_data_paths,
+        cache_dir=os.path.join(args.data_cache_dir, "raw"),
+        keep_in_memory=False,
+        split=train_splits,
+        num_proc=cpu_count(),
+    )
+    for index, dataset in enumerate(list_dataset):
+        assert isinstance(dataset, dataset_dict.Dataset)
+        logger.info(f"Start to process part-{index}/{len(list_dataset)} of all original datasets.")
+        dataset = dataset.map(
+            function=supervised_tokenize_sft,
+            fn_kwargs={
+                "tokenizer": tokenizer,
+                "conversation_template": default_conversation,
+                "max_length": args.max_length,
+            },
+            keep_in_memory=False,
+            num_proc=min(len(dataset), cpu_count()),
+        )
+
+        dataset = dataset.filter(lambda data: data["labels"] is not None)
+        dataset = dataset.sort(column_names=("seq_category", "seq_length"), reverse=False, keep_in_memory=False)
+
+        # We don't concatenate data samples here.
+        spliced_dataset = dataset
+        # Save each jsonl spliced dataset.
+        output_index = "0" * (5 - len(str(index))) + str(index)
+        output_name = f"part-{output_index}"
+        output_jsonl_path = os.path.join(args.data_jsonl_output_dir, output_name + ".jsonl")
+        # st = time.time()
+        with open(file=output_jsonl_path, mode="w", encoding="utf-8") as fp_writer:
+            spliced_count = 0
+            for spliced_data_point in spliced_dataset:
+                if spliced_count % 500 == 0:
+                    logger.info(f"processing {spliced_count} spliced data points for {fp_writer.name}")
+                spliced_count += 1
+                fp_writer.write(json.dumps(spliced_data_point, ensure_ascii=False) + "\n")
+
+        # Save each arrow spliced dataset
+        output_arrow_path = os.path.join(args.data_arrow_output_dir, output_name)
+        logger.info(f"Start to save {output_arrow_path}")
+        spliced_dataset = load_dataset(
+            path="json",
+            data_files=[output_jsonl_path],
+            cache_dir=os.path.join(args.data_cache_dir, "spliced_and_tokenized"),
+            keep_in_memory=False,
+            num_proc=cpu_count(),
+            split="train",
+        )
+        spliced_dataset.save_to_disk(dataset_path=output_arrow_path, num_proc=min(len(spliced_dataset), cpu_count()))
+
+
+if __name__ == "__main__":
+    main()

applications/Colossal-LLaMA-2/train_sft.example.sh

+#!/bin/bash
+
+# NCCL IB environment variables
+export NCCL_IB_HCA=mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1
+export NCCL_IB_DISABLE=0
+export NCCL_SOCKET_IFNAME=eth0
+export NCCL_IB_GID_INDEX=3
+export NCCL_IB_TIMEOUT=23
+export NCCL_IB_RETRY_CNT=7
+export OMP_NUM_THREADS=8
+
+PROJECT_NAME=""
+PARENT_SAVE_DIR=""
+PARENT_TENSORBOARD_DIR=""
+PARENT_CONFIG_FILE=""
+PRETRAINED_MODEL_PATH=""
+
+declare -a dataset=(
+    "PATH TO THE DATASET"
+)
+
+TIMESTAMP=$(date +%Y-%m-%d-%H-%M-%S)
+FULL_PROJECT_NAME="${PROJECT_NAME}-${TIMESTAMP}"
+SAVE_DIR="${PARENT_SAVE_DIR}${FULL_PROJECT_NAME}"
+TENSORBOARD_DIR="${PARENT_TENSORBOARD_DIR}${FULL_PROJECT_NAME}"
+CONFIG_FILE="${PARENT_CONFIG_FILE}${FULL_PROJECT_NAME}.json"
+
+colossalai run --nproc_per_node 8 --hostfile hostfile --master_port 30013 train_sft.py \
+    --pretrained $PRETRAINED_MODEL_PATH \
+    --dataset ${dataset[@]} \
+    --plugin "zero2" \
+    --save_interval 400 \
+    --save_dir $SAVE_DIR \
+    --tensorboard_dir $TENSORBOARD_DIR \
+    --config_file $CONFIG_FILE \
+    --num_epochs 1 \
+    --accumulation_steps 8 \
+    --micro_batch_size 8 \
+    --lr 5e-5 \
+    --mixed_precision "bf16" \
+    --grad_clip 1.0 \
+    --weight_decay 0.01 \
+    --warmup_steps 100 \
+    --use_grad_checkpoint \
+    --use_flash_attn \
+    --use_neft \

applications/Colossal-LLaMA-2/train_sft.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Supervised fine-tuning of Colossal-LLaMA-2-base developed by Colossal-AI Team
+"""
+
+import argparse
+import json
+import os
+import resource
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+from colossal_llama2.dataset.loader import (
+    DataCollatorForSupervisedDataset,
+    StatefulDistributedSampler,
+    load_tokenized_dataset,
+    setup_distributed_dataloader,
+)
+from colossal_llama2.utils.ckpt_io import load_checkpoint, save_checkpoint
+from colossal_llama2.utils.flash_attention_patch import replace_with_flash_attention
+from colossal_llama2.utils.froze import freeze_non_embeds_parameters
+from colossal_llama2.utils.neftune_patch import activate_neftune, deactivate_neftune
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer
+
+import colossalai
+from colossalai.booster import Booster
+from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin
+from colossalai.cluster import DistCoordinator
+from colossalai.lazy import LazyInitContext
+from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
+from colossalai.nn.optimizer import HybridAdam
+from colossalai.utils import get_current_device
+
+
+def get_model_numel(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+def format_numel_str(numel: int) -> str:
+    B = 1024**3
+    M = 1024**2
+    K = 1024
+    if numel >= B:
+        return f"{numel / B:.2f} B"
+    elif numel >= M:
+        return f"{numel / M:.2f} M"
+    elif numel >= K:
+        return f"{numel / K:.2f} K"
+    else:
+        return f"{numel}"
+
+
+def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor:
+    dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
+    tensor.div_(dist.get_world_size())
+    return tensor
+
+
+def main() -> None:
+    # ==============================
+    # Parse Arguments
+    # ==============================
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained",
+        type=str,
+        default=None,
+        help="Address of the pre-trained modeling",
+    )
+    parser.add_argument("--dataset", nargs="+", default=[])
+    parser.add_argument(
+        "--plugin",
+        type=str,
+        default="gemini",
+        choices=["gemini", "gemini_auto", "zero2", "zero2_cpu", "3d"],
+        help="Choose which plugin to use",
+    )
+    parser.add_argument("--load_checkpoint", type=str, default=None, help="Load checkpoint")
+    parser.add_argument("--save_interval", type=int, default=1000, help="Save interval")
+    parser.add_argument("--save_dir", type=str, default="checkpoint_dir", help="Checkpoint directory")
+    parser.add_argument("--tensorboard_dir", type=str, default="logs_dir", help="Tensorboard directory")
+    parser.add_argument("--config_file", type=str, default="config_file", help="Config file")
+    parser.add_argument("--num_epochs", type=int, default=1, help="Number of training epochs")
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="Number of accumulation steps")
+    parser.add_argument("--micro_batch_size", type=int, default=2, help="Batch size of each process")
+    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
+    parser.add_argument("--max_length", type=int, default=4096, help="Model max length")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="fp16",
+        choices=["fp16", "bf16"],
+        help="Mixed precision",
+    )
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping value")
+    parser.add_argument("--weight_decay", type=float, default=0.1, help="Weight decay")
+    parser.add_argument("--warmup_steps", type=int, default=None, help="Warmup steps")
+    parser.add_argument(
+        "--use_grad_checkpoint",
+        action="store_true",
+        default=False,
+        help="Use gradient checkpointing",
+    )
+    parser.add_argument(
+        "--use_flash_attn",
+        action="store_true",
+        default=False,
+        help="Use flash-attention",
+    )
+    parser.add_argument(
+        "--use_neft",
+        action="store_true",
+        default=False,
+        help="Use NEFTune",
+    )
+    parser.add_argument(
+        "--freeze_non_embeds_params",
+        action="store_true",
+        default=False,
+        help="Freeze non embeddings parameters",
+    )
+    parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--zero", type=int, default=1)
+    args = parser.parse_args()
+
+    with open(args.config_file, "w") as f:
+        json.dump(args.__dict__, f, indent=4)
+
+    # ==============================
+    # Initialize Distributed Training
+    # ==============================
+    colossalai.launch_from_torch({})
+    coordinator = DistCoordinator()
+
+    # ==============================
+    # Initialize Tensorboard
+    # ==============================
+    if coordinator.is_master():
+        os.makedirs(args.tensorboard_dir, exist_ok=True)
+        writer = SummaryWriter(args.tensorboard_dir)
+
+    # ==============================
+    # Initialize Booster
+    # ==============================
+    if args.plugin == "gemini":
+        plugin = GeminiPlugin(
+            precision=args.mixed_precision,
+            initial_scale=2**16,
+            max_norm=args.grad_clip,
+        )
+    elif args.plugin == "gemini_auto":
+        plugin = GeminiPlugin(
+            precision=args.mixed_precision,
+            placement_policy="auto",
+            initial_scale=2**16,
+            max_norm=args.grad_clip,
+        )
+    elif args.plugin == "zero2":
+        plugin = LowLevelZeroPlugin(
+            stage=2,
+            precision=args.mixed_precision,
+            initial_scale=2**16,
+            max_norm=args.grad_clip,
+        )
+    elif args.plugin == "zero2_cpu":
+        plugin = LowLevelZeroPlugin(
+            stage=2,
+            precision=args.mixed_precision,
+            initial_scale=2**16,
+            cpu_offload=True,
+            max_norm=args.grad_clip,
+        )
+    elif args.plugin == "3d":
+        plugin = HybridParallelPlugin(
+            tp_size=args.tp,
+            pp_size=1,
+            zero_stage=args.zero,
+            max_norm=args.grad_clip,
+            precision=args.mixed_precision,
+        )
+    else:
+        raise ValueError(f"Unknown plugin {args.plugin}")
+
+    booster = Booster(plugin=plugin)
+
+    # ======================================================
+    # Initialize Tokenizer, Dataset, Collator and Dataloader
+    # ======================================================
+    tokenizer = LlamaTokenizer.from_pretrained(args.pretrained)
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.add_bos_token = False
+    tokenizer.add_eos_token = False
+
+    coordinator.print_on_master(f"Configuration file will be saved at: {args.config_file}")
+    coordinator.print_on_master(f"Tensorboard logs will be saved at: {args.tensorboard_dir}")
+    coordinator.print_on_master(f"Model checkpoint will be saved at: {args.save_dir}")
+
+    coordinator.print_on_master(f"Load dataset: {args.dataset}")
+
+    dataset = load_tokenized_dataset(dataset_paths=args.dataset, mode="train")
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer, max_length=args.max_length)
+    dataloader = setup_distributed_dataloader(
+        dataset=dataset,
+        batch_size=args.micro_batch_size,
+        shuffle=True,
+        drop_last=True,
+        collate_fn=data_collator,
+    )
+    coordinator.print_on_master(
+        f"Max CUDA memory after data loader: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB"
+    )
+
+    # ======================================================
+    # Initialize Model, Objective, Optimizer and LR Scheduler
+    # ======================================================
+    init_ctx = (
+        LazyInitContext(default_device=get_current_device()) if isinstance(plugin, (GeminiPlugin,)) else nullcontext()
+    )
+    with init_ctx:
+        model = LlamaForCausalLM(LlamaConfig.from_pretrained(args.pretrained))
+        # Freeze part of parameters.
+        if args.freeze_non_embeds_params:
+            freeze_non_embeds_parameters(model=model)
+
+    if args.use_grad_checkpoint:
+        model.gradient_checkpointing_enable()
+        coordinator.print_on_master(msg="Gradient checkpointing enabled successfully")
+    if args.use_flash_attn:
+        replace_with_flash_attention(model=model)
+        coordinator.print_on_master(msg="Flash-attention enabled successfully")
+
+    model_numel = get_model_numel(model)
+    coordinator.print_on_master(f"Model params: {format_numel_str(model_numel)}")
+
+    optimizer = HybridAdam(
+        model_params=filter(lambda p: p.requires_grad, model.parameters())
+        if args.freeze_non_embeds_params
+        else model.parameters(),
+        lr=args.lr,
+        betas=(0.9, 0.95),
+        weight_decay=args.weight_decay,
+        adamw_mode=True,
+    )
+
+    if args.warmup_steps is None:
+        args.warmup_steps = int(args.num_epochs * 0.025 * (len(dataloader) // args.accumulation_steps))
+        coordinator.print_on_master(f"Warmup steps is set to {args.warmup_steps}")
+
+    lr_scheduler = CosineAnnealingWarmupLR(
+        optimizer=optimizer,
+        total_steps=args.num_epochs * (len(dataloader) // args.accumulation_steps),
+        warmup_steps=args.warmup_steps,
+        eta_min=0.1 * args.lr,
+    )
+
+    # Flash attention will be disabled because it does NOT support fp32.
+    default_dtype = torch.float16 if args.mixed_precision == "fp16" else torch.bfloat16
+    torch.set_default_dtype(default_dtype)
+    model, optimizer, _, dataloader, lr_scheduler = booster.boost(
+        model=model,
+        optimizer=optimizer,
+        lr_scheduler=lr_scheduler,
+        dataloader=dataloader,
+    )
+
+    torch.set_default_dtype(torch.float)
+
+    if args.load_checkpoint is None:
+        coordinator.print_on_master(f"Load pretrained model checkpoint from {args.pretrained}")
+        booster.load_model(model, args.pretrained, strict=False)
+
+    coordinator.print_on_master(f"Booster init max CUDA memory: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB")
+    coordinator.print_on_master(
+        f"Booster init max CPU memory: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024:.2f} MB"
+    )
+
+    start_epoch = 0
+    start_step = 0
+    sampler_start_idx = 0
+    if args.load_checkpoint is not None:
+        if "modeling" in args.load_checkpoint:
+            coordinator.print_on_master(f"Continued pretrain from checkpoint {args.load_checkpoint}")
+            booster.load_model(model, args.load_checkpoint)
+        else:
+            coordinator.print_on_master(f"Load model checkpoint from {args.load_checkpoint}")
+            start_epoch, start_step, sampler_start_idx = load_checkpoint(
+                load_dir=args.load_checkpoint,
+                booster=booster,
+                model=model,
+                optimizer=optimizer,
+                lr_scheduler=lr_scheduler,
+            )
+            coordinator.print_on_master(
+                f"Loaded checkpoint {args.load_checkpoint} at epoch {start_epoch} step {start_step}"
+            )
+            coordinator.print_on_master(f"Loaded sample at index {sampler_start_idx}")
+
+        coordinator.print_on_master(
+            f"Checkpoint loaded max CUDA memory: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB"
+        )
+        coordinator.print_on_master(
+            f"Checkpoint loaded CUDA memory: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB"
+        )
+        coordinator.print_on_master(
+            f"Checkpoint loaded max CPU memory: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024:.2f} MB"
+        )
+
+    if args.use_neft:
+        coordinator.print_on_master("Activate NEFTune.")
+        model, handle = activate_neftune(model)
+
+    num_steps_per_epoch = len(dataloader) // args.accumulation_steps
+    # If resume training, set the sampler start index to the correct value
+    assert isinstance(dataloader.sampler, StatefulDistributedSampler)
+    dataloader.sampler.set_start_index(start_index=sampler_start_idx)
+
+    for epoch in range(start_epoch, args.num_epochs):
+        dataloader.sampler.set_epoch(epoch=epoch)
+        pbar = tqdm(desc=f"Epoch {epoch}", disable=not coordinator.is_master(), total=num_steps_per_epoch)
+        total_loss = torch.tensor(0.0).to(torch.cuda.current_device())
+        for step, batch in enumerate(dataloader):
+            batch = {k: v.to(get_current_device()) for k, v in batch.items() if isinstance(v, torch.Tensor)}
+
+            batch_output = model(**batch)
+
+            loss = batch_output.loss / args.accumulation_steps
+            total_loss += loss.item()
+
+            booster.backward(loss=loss, optimizer=optimizer)
+
+            if (step + 1) % args.accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+                all_reduce_mean(tensor=total_loss)
+                pbar.set_postfix({"Loss": f"{total_loss.item():.4f}"})
+                if coordinator.is_master():
+                    global_step = (epoch * num_steps_per_epoch) + (step + 1) // args.accumulation_steps
+                    writer.add_scalar(tag="Loss", scalar_value=total_loss.item(), global_step=global_step)
+                    writer.add_scalar(
+                        tag="Learning Rate",
+                        scalar_value=lr_scheduler.get_last_lr()[0],
+                        global_step=global_step,
+                    )
+                total_loss.fill_(0.0)
+                pbar.update()
+            # Save modeling.
+
+            if (args.save_interval > 0 and (step + 1) % (args.save_interval * args.accumulation_steps) == 0) or (
+                step + 1
+            ) == len(dataloader):
+                coordinator.print_on_master("\nStart saving model checkpoint with running states")
+
+                if args.use_neft:
+                    coordinator.print_on_master("Deactivate NEFTune before saving model.")
+                    deactivate_neftune(model, handle)
+
+                save_checkpoint(
+                    save_dir=args.save_dir,
+                    booster=booster,
+                    model=model,
+                    optimizer=optimizer,
+                    lr_scheduler=lr_scheduler,
+                    epoch=epoch,
+                    step=step + 1,
+                    batch_size=args.micro_batch_size,
+                    coordinator=coordinator,
+                )
+                coordinator.print_on_master(
+                    f"Saved checkpoint at epoch {epoch} step {step + 1} at folder {args.save_dir}"
+                )
+
+                if args.use_neft:
+                    coordinator.print_on_master("Activate NEFTune.")
+                    model, handle = activate_neftune(model)
+
+            # Delete CUDA cache.
+            # del batch, batch_labels, batch_output, loss
+            torch.cuda.empty_cache()
+
+        # the continue epochs are not resumed, so we need to reset the sampler start index and start step
+        dataloader.sampler.set_start_index(start_index=0)
+        start_step = 0
+
+    if args.use_neft:
+        coordinator.print_on_master("Deactivate NEFTune.")
+        deactivate_neftune(model, handle)
+
+    # Final save.
+    coordinator.print_on_master("Start saving final model checkpoint")
+    booster.save_model(model, os.path.join(args.save_dir, "modeling"), shard=True)
+    coordinator.print_on_master(f"Saved final model checkpoint at epoch {epoch} at folder {args.save_dir}")
+
+    coordinator.print_on_master(f"Max CUDA memory usage: {torch.cuda.max_memory_allocated()/1024**2:.2f} MB")
+
+
+if __name__ == "__main__":
+    main()