Skip to content

[Fix] Hanging for Fully Randomized Bucketing #4348

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
titu1994 merged 4 commits into from
Jun 9, 2022
Merged

[Fix] Hanging for Fully Randomized Bucketing #4348

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2da5a9e
Update container to 22.05 (#4329)
ericharper Jun 7, 2022
b390f1e
Merge r1.9.0 main (#4331)
ericharper Jun 7, 2022
c1ced75
fix full_randn bucket hang
stevehuang52 Jun 7, 2022
7503d47
remove unused variables
stevehuang52 Jun 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 6 additions & 26 deletions nemo/collections/asr/modules/rnnt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -845,8 +845,6 @@ def forward(
)

losses = []
wer_numer_list = []
wer_denom_list = []
batch_size = int(encoder_outputs.size(0)) # actual batch size

# Iterate over batch using fused_batch_size steps
Expand Down Expand Up @@ -914,31 +912,14 @@ def forward(
else:
losses = None

# Compute WER for sub batch
# Update WER for sub batch
if compute_wer:
sub_enc = sub_enc.transpose(1, 2) # [B, T, D] -> [B, D, T]
sub_enc = sub_enc.detach()
sub_transcripts = sub_transcripts.detach()

original_log_prediction = self.wer.log_prediction
if original_log_prediction and batch_idx == 0:
self.wer.log_prediction = True
else:
self.wer.log_prediction = False

# Compute the wer (with logging for just 1st sub-batch)
# Update WER on each process without syncing
self.wer.update(sub_enc, sub_enc_lens, sub_transcripts, sub_transcript_lens)
wer, wer_num, wer_denom = self.wer.compute()
self.wer.reset()

wer_numer_list.append(wer_num)
wer_denom_list.append(wer_denom)

# Reset logging default
self.wer.log_prediction = original_log_prediction

else:
wer = None

del sub_enc, sub_transcripts, sub_enc_lens, sub_transcript_lens

Expand All @@ -951,12 +932,11 @@ def forward(

# Collect sub batch wer results
if compute_wer:
wer_num = torch.tensor(wer_numer_list, dtype=torch.long)
wer_denom = torch.tensor(wer_denom_list, dtype=torch.long)

wer_num = wer_num.sum() # global sum of correct words/chars
wer_denom = wer_denom.sum() # global sum of all words/chars
# Sync and all_reduce on all processes, compute global WER
wer, wer_num, wer_denom = self.wer.compute()
self.wer.reset()
else:
wer = None
wer_num = None
wer_denom = None

Expand Down