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

Reset scaler state #1999

Merged
merged 2 commits into from
Feb 24, 2023
Merged

Reset scaler state #1999

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aaf5b0e
reset scaler
mvpatel2000 Feb 24, 2023
8eeab0b
Merge branch 'dev' into mvpatel2000/reset-optimizer
mvpatel2000 Feb 24, 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
7 changes: 6 additions & 1 deletion composer/trainer/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import re
import time
import warnings
from collections import defaultdict
from copy import deepcopy
from pathlib import Path
from typing import Any, Callable, ContextManager, Dict, Iterable, List, Optional, Sequence, TextIO, Tuple, Union, cast
Expand All @@ -24,7 +25,7 @@
import torch.distributed
import torch.nn as nn
import torch.utils.data
from torch.cuda.amp.grad_scaler import GradScaler
from torch.cuda.amp.grad_scaler import GradScaler, _refresh_per_optimizer_state
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
from torchmetrics import Metric
Expand Down Expand Up @@ -257,6 +258,8 @@ def _adjust_grad_accum(state: State, device_batch_size: int):
del state.loss
for optimizer in state.optimizers:
optimizer.zero_grad(set_to_none=True)
if state.scaler is not None:
state.scaler._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
torch.cuda.empty_cache()


Expand Down Expand Up @@ -285,6 +288,8 @@ def _adjust_device_train_microbatch_size(state: State):
del state.loss
for optimizer in state.optimizers:
optimizer.zero_grad(set_to_none=True)
if state.scaler is not None:
state.scaler._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
torch.cuda.empty_cache()


Expand Down