"Assertion `srcIndex < srcSelectDimSize` failed" showed when I tried to train using my own script

I'm trying to train vicuna-7B model with this script

class TrainDataset(torch.utils.data.Dataset):

    def __init__(self, data_path: str, tokenizer: transformers.PreTrainedTokenizer, max_length: int = 0):
        abs_data_path = os.path.abspath(data_path)
        self.data = self.get_data(abs_data_path)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index: int) -> torch.Tensor:

        items = {'input_ids': self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(self.data[index]))}
        items['attention_mask'] = [1]*len(items['input_ids'])
        if self.max_length:
            items['input_ids'] = items['input_ids'][:self.max_length]
            items = self.tokenizer.pad(items, padding='max_length', max_length=self.max_length)

        return {k: torch.tensor(v, dtype=(torch.long if k != 'attention_mask' else torch.bool)) for k,v in items.items()}

    def get_data(self, data_path: str) -> List[str]:
        
        with open(data_path) as f:
            str_data = f.readlines()
        
        return str_data

def load_pair(model_type, model_path) -> Tuple[transformers.PreTrainedModel, transformers.PreTrainedTokenizer]:

    if model_type not in MODEL_TOKEN_PAIR:
        raise KeyError(f'{model_type} model type is not supported')

    model_cls, token_cls = MODEL_TOKEN_PAIR[model_type]

    tokenizer: transformers.PreTrainedTokenizer = token_cls.from_pretrained(model_path, use_fast=False)
    if tokenizer.pad_token is None:
        try:
            tokenizer.convert_tokens_to_ids('<pad>')
            pad_token = '<pad>'
        except NotImplementedError:
            pad_token = tokenizer.convert_ids_to_tokens(0)
        tokenizer.add_special_tokens({'pad_token': pad_token})
    
    model: transformers.PreTrainedModel = model_cls.from_pretrained(model_path)

    return model, tokenizer

def main():

    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    model_args, data_args, training_args = parse_args(parser)

    mlflow.set_tracking_uri(training_args.mlflow_url)
    os.environ['MLFLOW_EXPERIMENT_NAME'] = training_args.experiment_name

    mlflowCallback = transformers.integrations.MLflowCallback()

    model, tokenizer = load_pair(model_args.model_type, model_args.model_name_or_path)

    train_dataset = TrainDataset(data_args.data_path, tokenizer, training_args.model_max_length)

    data_collator = transformers.DataCollatorForLanguageModeling(
        tokenizer=tokenizer, mlm=False
    )

    trainer = transformers.Trainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        callbacks=[EarlyStopping(training_args.max_epoch_without_progress), mlflowCallback]
    )

    trainer.train()

After I check, the difference between my TrainDataset with SupervisedDataset is my dataset did not contain labels key. But, it should be fixed by data_collator and the return value is same. Then I ran my script, error like in #199 here showed up.

The weird thing is, if I run train.py from this repo, the training is running smoothly. Is there any problem with my script or I should do some special preprocessing in my dataset?

Start by comparing TrainDataset.getitem, data_collator, and the repository's train.py and SupervisedDataset, then reproduce the failure at trainer.train(). Check how the returned input_ids, attention_mask, and missing labels differ before collation; done means identifying the preprocessing mismatch and confirming training runs without the assertion.