Skip to content
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

Training stability: Continue training even if a data batch was hit which causes OOM #113

Closed
wants to merge 4 commits into from
+21 −2
Closed

Training stability: Continue training even if a data batch was hit which causes OOM #113

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fe5e1ef
continue training even if a broken data batch was hit
RuntimeRacer May 2, 2023
234c723
also perform grad scaling in pre-run OOM check.
RuntimeRacer May 2, 2023
2541aed
try explicit cleanup of everything related to the batch to prevent tr…
RuntimeRacer May 3, 2023
f9292b2
avoid error in case we hit OOM already on forward pass
RuntimeRacer May 4, 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
23 changes: 21 additions & 2 deletions valle/bin/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -692,8 +692,23 @@ def train_one_epoch(

set_batch_count(model, params.batch_idx_train)
except: # noqa
# Save the broken batch
logging.warning(f"Hit a broken batch of training data. Cut ID: {batch['utt_id']} Text: {batch['text']} - Skipping...")
display_and_save_batch(batch, params=params)
raise
# Clean up batch data from Memory and GPU
del batch["text_tokens"]
del batch["text_tokens_lens"]
del batch["audio_features"]
del batch["audio_features_lens"]
del batch
try:
del loss
del loss_info
except UnboundLocalError:
pass
torch.cuda.empty_cache()
# Continue training
continue

if params.average_period > 0:
if (
Expand Down Expand Up @@ -1101,6 +1116,10 @@ def scan_pessimistic_batches_for_oom(
elif params.dtype in ["float16", "fp16"]:
dtype = torch.float16

scaler = GradScaler(
enabled=(params.dtype in ["fp16", "float16"]), init_scale=1.0
)

for criterion, cuts in batches.items():
batch = train_dl.dataset[cuts]
try:
Expand All @@ -1111,7 +1130,7 @@ def scan_pessimistic_batches_for_oom(
batch=batch,
is_training=True,
)
loss.backward()
scaler.scale(loss).backward()
optimizer.zero_grad()
except Exception as e:
if "CUDA out of memory" in str(e):
Expand Down