Fixing bug in Megatron BERT when loss mask is all zeros

[shanmugamr1992]

Signed-off-by: Shanmugam Ramasamy [email protected]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Collection: [Note which collection this PR will affect]

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[yidong72]

LGTM

[MaximumEntropy]

Hey! Sorry to block merging, but I think this might be a symptom of a bug at the data level and we should try to get to the bottom of that and not handle it this way.

[yidong72]

Hey! Sorry to block merging, but I think this might be a symptom of a bug at the data level and we should try to get to the bottom of that and not handle it this way.

It related to the dataset @shanmugamr1992 is using. I didn't see this behavior in my biomegatron dataset before.

[shanmugamr1992]

@yidong72 @MaximumEntropy So this is not about my dataset. The number of tokens to be padded is calculated as (total_tokens. * masking_probability) . So if total tokens is 200 and masking probability is 0.15, then 30 random tokens are padded. However if the number of tokens is small (Say 4 or 5) meaning remaining of the 512 are just zero padded, we have this case that there is nothing to be predicted. (Loss mask is all zeros).
So this can be handled in two ways :

1. Change the code that says minimum 1 token needs to be masked (However I think it might throw some weird behaviour when we only have 1 token in a batch)
2. We can assign the loss for that batch to be zero , which seems quite right, since the actual loss for that batch, when there is nothing to predict is zero.
   I tried both approaches. Approach 1 threw Nan's after 20k or something. I can run it locally and debug, or we can go ahead with approach 2 , which is not that hacky.
   NOTE : This is a very weird corner case that will happen only when you set the microbatchsize to be 1. So wont really affect users as much.