Pretrain data processing

1. Raw data source

First, prepare raw data into data_1.txt, data_2.txt, …

In the raw text files, use line breaks to separate different datas, i.e., each line has one data (for example, a long document). Notice that there may be line breaks inside data, they should be replaced with unique identifier, such as <n>, and replaced back in step 2.

Examples:

1The Lord of the Rings<n>The Lord of the Rings is an epic high-fantasy novel by English author and scholar J. R. R. Tolkien. Set in Middle-earth, intended to be Earth at some distant time in the past, the story began as a sequel to Tolkien's 1937 children's book The Hobbit, but eventually developed into a much larger work. Written in stages between 1937 and 1949, The Lord of the Rings is one of the best-selling books ever written, with over 150 million copies sold.<n>The title refers to the story's main antagonist, the Dark Lord Sauron, who in an earlier age created the One Ring to rule the other Rings of Power given to Men, Dwarves, and Elves, in his campaign to conquer all of Middle-earth. From homely beginnings in the Shire, a hobbit land reminiscent of the English countryside, the story ranges across Middle-earth, following the quest to destroy the One Ring mainly through the eyes of the hobbits Frodo, Sam, Merry and Pippin.
2The Little Prince<n>I thought that I was rich, with a flower that was unique in all the world; and all I had was a common rose. A common rose...<n>To me, you are still nothing more than a little boy who is just like a hundred thousand other little boys. And I have no need of you. And you, on your part, have no need of me. To you, I am nothing more than a fox like a hundred thousand other foxes. But if you tame me, then we shall need each other. To me, you will be unique in all the world. To you, I shall be unique in all the world.<n>The wheat fields have nothing to say to me. And that is sad. But you have hair that is the color of gold. Think how wonderful that will be when you have tamed me! The grain, which is also golden, will bring me back the thought of you. And I shall love to listen to the wind in the wheat.
3...

Important: Since Streaming data processing and training are used in subsequent steps, data should be pre-shuffled before putting into data.txt sequentially.

2. Tokenize

  • Implement an Encoder, which gives it a line of input data and it returns you the tokenized result.

  • Mapping all datas with Encoder, with the help of multiprocessing

  • We also provide a tools called indexed_dataset, which compress the tokenized data into binary format.

There is an example code in model_center/tools/preprocess_cpm1_lm.py, we simplify it and made it into a example like the following:

import multiprocessing
from model_center.tools import indexed_dataset

# 1. Implement an Encoder, which gives it a line of input data and it returns you the tokenized result.
class Encoder(object): 
    def initializer(self):
        Encoder.tokenizer = YourTokenizer()

    def encode(self, line):
        data = line.strip().replace("<n>", "\n") # replace back line break symbol

        doc_ids = Encoder.tokenizer.encode(data)

        max_length = 512 # model's maximum input length

        pieces = []
        while i < len(doc_ids): # split document into chunks with model's maximum length
            piece = doc_ids[i:i+max_length]
            if len(piece) < 32: # drop too short chunks
                break
            i += max_length

            pieces.append(piece)

        return pieces

if __name__ == '__main__':
    # assumes that there are 100 raw data files, named `data_1.txt` to `data_100.txt`
    for ith in range(1, 101):
        fin = open(f"data_{ith}.txt", "r", encoding="utf-8")

        # encoder use the tokenizer to encode data
        encoder = Encoder()

        # 2. Mapping all datas with Encoder, with the help of multiprocessing
        pool = multiprocessing.Pool(processes=64, initializer=encoder.initializer)
        encoded_docs = pool.imap_unordered(encoder.encode, fin, chunksize=10)

        # 3. tool `indexed_dataset` compress the tokenized data into binary format `bin_file`
        # it will also generate another small `idx_file` for saving meta information in order to decode `bin_file`.
        bin_file = os.path.join("path/to/your/tokenized_output_folder/", f"tokenized_{ith}.bin")
        idx_file = os.path.join("path/to/your/tokenized_output_folder/", f"tokenized_{ith}.idx")
    
        binary_builder = indexed_dataset.make_builder(bin_file, impl="mmap", dtype=np.uint16)

        # put tokenized data into binary_builder
        for pieces in encoded_docs:
            for doc_ids in pieces:
                binary_builder.add_item(torch.IntTensor(doc_ids))

        # finish compressing tokenized data into `bin_file`, and generate meta information into `idx_file`
        binary_builder.finalize(idx_file)

        # close multiproceessing mapping
        pool.close()

3. Self-Supervised Dataset

We provide a tool model_center.dataset.DistributedMMapIndexedDataset for loading those compressed binary files and do extra processing, for example, a commonly used way is to randomly mask a span and ask the model to regenerate it.

It is important to note that, tokenization is an rather slow operation, it may become a bottleneck if tokenization is perform while the model is being trained. Thus, we put tokenization into data pre-processing stage (step 2), which is run beforehand.

However, We do not calculate things like attention_mask and save them into binary files because it will lead to multiplication of space occupation, which is not affordable for the large amount of data needed for training large language models. Putting those extra information that are less time consuming into __getitem__ function of pytorch Dataset, space will be saved since the space will be recycled when the next batch arrives.

from model_center.dataset import DistributedMMapIndexedDataset

class YourDataset(torch.utils.data.Dataset):
    def __init__(self, ctx : MMapIndexedDataset):
        self.ctx = ctx

    def __len__(self):
        return len(self.ctx)
    
    def __get_item_data(self, input_ids): # do extra processing
        length = input_ids.shape[0]

        # randomly mask a span in range [lef, rig)
        lef = random.randint(input_length)
        rig = random.randint(lef, input_length)

        # pretrain objective: regenerate the masked span and ignore other positions
        targets = np.full((input_length), -100) # -100 as ignore_index
        targets[lef:rig] = input_ids[lef:rig]

        # calculate attention_mask, to tell model which positions are visible
        attention_mask = (np.arange((input_length)) < lef) | (np.arange((input_length)) >= rig)
        return input_ids, targets, input_length, attention_mask

    def __getitem__(self, ith):
        input_ids = self.ctx[ith] # get the i-th data from DistributedMMapIndexedDataset
        return self.__get_item_data(ctx) # do extra processing and return

if __name__ == '__main__':
    dataset = YourDataset(
        DistributedMMapIndexedDataset("path/to/your/tokenized_output_folder", "tokenized", bmt.rank(), bmt.world_size()), 
        # the second argument "tokenized" is the common prefix of your tokenized file name,
        # here we assumes that they are called "tokenized_1.bin", "tokenized_2.idx", etc.
    )