Skip to content

Gradscaler flags #2281

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
mravanelli merged 9 commits into from
Dec 13, 2023
Merged

Gradscaler flags #2281

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
f720eab
add flags for gradscaler
Adel-Moumen Nov 30, 2023
de7e0b8
add check_loss_isfinite
Adel-Moumen Nov 30, 2023
18a7021
update dict
Adel-Moumen Nov 30, 2023
1faa548
typo
Adel-Moumen Nov 30, 2023
65612e2
remove default
Adel-Moumen Nov 30, 2023
1d2d16d
better message
Adel-Moumen Nov 30, 2023
7ffb517
fix pre-commit
Adel-Moumen Nov 30, 2023
744ae6d
remove checks
Adel-Moumen Dec 1, 2023
9c32cab
remove new arguments
mravanelli Dec 13, 2023
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: 32 additions & 0 deletions speechbrain/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1133,6 +1133,7 @@ def fit_batch(self, batch):
scaled_loss = self.scaler.scale(
loss / self.grad_accumulation_factor
)
self.check_loss_isfinite(scaled_loss)

@TParcollet TParcollet Nov 30, 2023

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checking for all parameters to be finite is redundant with the GradScaler. It already does this. Also it can be crazy expensive for very large large models. Checking if the loss is not finite makes sense, not the parameters.

@asumagic asumagic Dec 1, 2023 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The grad scaler does this but it does not care for how long the parameters have gone non-finite. An idea I suggested yesterday was to occasionally check the gradscaler scale for insane values with a patience mechanism, as I've sometimes seen the scale vanish or explode when issues occurred.
That does induce a CPU-GPU sync though, but I'm not sure how often we have one in the first place.

@Adel-Moumen Adel-Moumen Dec 1, 2023

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with @TParcollet. I will remove the part where I'm checking the NaNs/inf in the weights and will only cover the loss part. BTW, the function was also intended for other non-GradScaler use cases like fp32.

scaled_loss.backward()

if should_step:
Expand All @@ -1141,6 +1142,37 @@ def fit_batch(self, batch):
self.on_fit_batch_end(batch, outputs, loss, should_step)
return loss.detach().cpu()

def check_loss_isfinite(self, loss):
"""Check if the loss is finite.

If the loss is not finite, log a helpful message and increment the `nonfinite_count`.
If the `nonfinite_count` exceeds the `--nonfinite_patience` threshold, stop the training
and raise an error.

This check is particularly useful when the loss becomes NaN or inf, while the
parameters and gradients remain finite. It helps prevent getting stuck in an
infinite loop during training.

Arguments
---------
loss : tensor
The loss tensor after ``backward()`` has been called but
before the optimizers ``step()``.
"""
if not torch.isfinite(loss):
self.nonfinite_count += 1

# Check if patience is exhausted
if self.nonfinite_count > self.nonfinite_patience:
raise ValueError(
"Loss is not finite and patience is exhausted. "
"To debug, wrap `fit()` with "
"autograd's `detect_anomaly()`, e.g.\n\nwith "
"torch.autograd.detect_anomaly():\n\tbrain.fit(...)"
)
else:
logger.warning("Patience not yet exhausted.")

def check_gradients(self):
""" Checks if the gradients are finite. If not, it will emit a warning and set them to zero."""
for param in self.modules.parameters():
Expand Down