MTF dataset and packing

[thomasw21]

No description provided.

[Muennighoff]

Why does the last segment id have to be padding?

[thomasw21]

You're right, I'll fix them.

[thomasw21]

Done in c74dbb7

[Muennighoff]

Still need to cast them to torch tensors / directly use torch tensors instead of numpy

[thomasw21]

Doing it in the next one.

[Muennighoff]

Note: We will check this when preprocessing

[thomasw21]

I think a good opportunity to have the EOS token that we didn't have from pretrianing is adding one here.

{INPUT} {TARGET} <EOS>

Not sure we should not have a between input and target
WDYT @TevenLeScao @Muennighoff @lintangsutawika

[Muennighoff]

Good idea!
I think we should add an EOS token only after each target

[thomasw21]

Yeah agree on only after target.

[TevenLeScao]

Agreed

[thomasw21]

I haven't thought about it very clearly, but essentially this mean that the entire dataset can be random accessed which isn't possible if you pack things greedily (the problem is given an index, can you know which original sample you need to add in the batches ....) So I'm refactoring this to put this in another dataset ... Sorry about this @Muennighoff you were right.

[stas00]

FYI, tests/test_dataloaders.py::TestDataLoading::test_mlm_dataset fails.

I've added a skip to not run it by CI. #302
Please fix the test at your convenience.

I partially fixed it while at it - the test couldn't have possibly ever worked as the args were wrong

After the partial fix so that it actually runs it fails on CI and on my desktop as well:

>               train_ds, valid_ds, test_ds = mlm_dataset.build_train_valid_test_datasets(
                    data_prefix=args.data_path,
                    data_impl=args.data_impl,
                    splits_string=args.split,
                    # TODO @thomasw21 figure how that value works
                    train_valid_test_num_samples=train_val_test_num_samples,
                    sequence_length=args.seq_length,
                    noise_density=args.noise_density,
                    mean_noise_span_length=args.mean_noise_span_length,
                    seed=args.seed,
                    skip_warmup=(not args.mmap_warmup)
                )

tests/test_dataloaders.py:90: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
megatron/data/mlm_dataset.py:25: in build_train_valid_test_datasets
    return _build_train_valid_test_datasets(
megatron/data/mlm_dataset.py:263: in _build_train_valid_test_datasets
    train_dataset = build_dataset(0, 'train')
megatron/data/mlm_dataset.py:250: in build_dataset
    dataset = MLMDataset(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <megatron.data.mlm_dataset.MLMDataset object at 0x7fc94710c070>, name = 'train'
indexed_dataset = <megatron.data.indexed_dataset.MMapIndexedDataset object at 0x7fc94710c880>
documents = array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
        13,  14,  15,  16,  17,  18,  19,  20,...9, 930, 931, 932, 933, 934, 935,
       936, 937, 938, 939, 940, 941, 942, 943, 944, 945, 946, 947],
      dtype=int32)
data_prefix = '/mnt/nvme0/code/huggingface/Megatron-DeepSpeed-master/tests/data/gpt2/meg-gpt2-openwebtext_text_document'
sequence_length = 512, num_samples = 40000, seed = 1234, noise_density = 0.15, mean_noise_span_length = 3

    def __init__(
        self,
        name,
        indexed_dataset,
        documents,
        data_prefix,
        sequence_length,
        num_samples,
        seed,
        noise_density=0.15,
        mean_noise_span_length=3
    ):
    
        # Params to store.
        self.name = name
        self.seed = seed
        self.sequence_length = sequence_length
    
        # Dataset.
        self.indexed_dataset = indexed_dataset
    
        self.noise_density = noise_density
        self.mean_noise_span_length = mean_noise_span_length
        # T5-like span masked language modeling will fuse consecutively masked tokens to a single sentinel token.
        # To ensure that the input length is `sequence_length`, we need to increase the maximum length
        # according to `noise_density` and `mean_noise_span_length`. We can also define the label length accordingly.
        number_of_raw_tokens, inputs_length, targets_length, num_noise_spans = compute_input_and_target_lengths(
            sequence_length=self.sequence_length,
            noise_density=self.noise_density,
            mean_noise_span_length=self.mean_noise_span_length
        )
        self.inputs_length = inputs_length
        # In order to compute loss, we need an extra token at the end.
        self.number_of_raw_tokens = number_of_raw_tokens + 1
        self.targets_length = targets_length + 1
        self.num_noise_spans = num_noise_spans
    
        # Build the samples mapping.
        self._gpt_dataset = GPTDataset(
            name=self.name,
            data_prefix=data_prefix,
            documents=documents,
            indexed_dataset=self.indexed_dataset,
            num_samples=num_samples,
            # -1 because GPTDataset will return `seq_length + 1` sequences.
            seq_length=number_of_raw_tokens - 1,
            seed=seed
        )
    
        # Vocab stuff.
        tokenizer = get_tokenizer()
        self.sep_id = tokenizer.sep
        self.sentinel_token_ids = tokenizer.additional_special_tokens_ids
        assert self.sep_id is not None, "MLM dataset requires tokenizer to have a <sep> token"
        assert len(self.sentinel_token_ids) > 0, "Provide the argument --vocab-extra-ids 100 to the script"
        assert len(self.sentinel_token_ids) >= self.num_noise_spans, "Not enough sentinel tokens, please add more"
    
        args = get_args()
        if hasattr(args, "encoder_seq_length") and args.encoder_seq_length is not None:
            # T5 style
>           assert self.inputs_length == args.encoder_seq_length
E           AssertionError

megatron/data/mlm_dataset.py:332: AssertionError