GPTSFTChatDataset loss_mask becomes all False when prompt length > max_seq_length

[shengyangs]

Describe the bug

In the GPTSFTChatDataset, if the first prompt length exceeds max_seq_length, all following turns are truncated out. Then the loss_mask becomes all False for the example.

https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py#L359

This is problematic because, when the loss_mask is all False, the loss of the MegatronGPTModel will be nan.

https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py#L1015

Steps/Code to reproduce bug

This can be reproduced by passing into GPTSFTChatDataset an example whose first turn prompt has > 2048 tokens when the max_seq_length=2048. Then use the GPTSFTChatDataset in a supervised fine-tuning job (e.g., train_gpt_sft.py in NeMo-Aligner)

Expected behavior

In the collate_fn function, check if the loss_masks of all examples are not empty. If not, raise an error.

[odelalleau]

I ran into the same issue => hacked the code to zero out the loss on problematic micro-batches: odelalleau@b383e6a

(obviously not a proper fix but can be useful to get unblocked)

[shengyangs]

Yeah, this will also fix the problem.

[github-actions]

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[shengyangs]

This issue should keep active since the same problem has been encountered by others as well.

[github-actions]

This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[github-actions]

This issue was closed because it has been inactive for 7 days since being marked as stale.