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

Jump to bottom

Adam optim support #9

Merged
merged 2 commits into from
May 16, 2024
Merged

Adam optim support #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f31c78c
adam optim support added, explictly mask exp_avg_sq buffer
mklasby May 16, 2024
fe89efe
adding defaults as param to dst algos so lop can hook grad accum with…
mklasby May 16, 2024
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
15 changes: 10 additions & 5 deletions sparsimony/dst/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,11 @@


class DSTMixin(ABC):
_OPTIM_REG = {optim.SGD: "momentum_buffer", optim.AdamW: "exp_avg"}
_OPTIM_REG = {
optim.SGD: ["momentum_buffer"],
optim.AdamW: ["exp_avg", "exp_avg_sq"],
optim.Adam: ["exp_avg", "exp_avg_sq"],
}

def __init__(self, optimizer: torch.optim.Optimizer, *args, **kwargs):
if type(optimizer) not in self._OPTIM_REG:
Expand Down Expand Up @@ -174,10 +178,11 @@ def _momentum_zero_wrapper():
if config["sparsity"] == 0:
continue
original_param = get_original_tensor(**config)
state_kw = self._OPTIM_REG[type(self.optimizer)]
if state_kw in self.optimizer.state[original_param]:
mask = get_mask(**config)
self.optimizer.state[original_param][state_kw] *= mask
state_kw_list = self._OPTIM_REG[type(self.optimizer)]
mask = get_mask(**config)
for state_kw in state_kw_list:
if state_kw in self.optimizer.state[original_param]:
self.optimizer.state[original_param][state_kw] *= mask

self.optimizer.step = _momentum_zero_wrapper

Expand Down
6 changes: 4 additions & 2 deletions sparsimony/dst/rigl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, Dict, Any
import torch
import torch.nn as nn
from torch.ao.pruning.sparsifier.base_sparsifier import BaseSparsifier
Expand All @@ -22,6 +22,7 @@ def __init__(
scheduler: BaseScheduler,
distribution: BaseDistribution,
optimizer: torch.optim.Optimizer,
defaults: Optional[Dict[str, Any]] = None,
sparsity: float = 0.5,
grown_weights_init: float = 0.0,
init_method: Optional[str] = "grad_flow",
Expand All @@ -31,7 +32,8 @@ def __init__(
self.sparsity = sparsity
self.grown_weights_init = grown_weights_init
self.init_method = init_method
defaults = dict(parametrization=FakeSparsityDenseGradBuffer)
if defaults is None:
defaults = dict(parametrization=FakeSparsityDenseGradBuffer)
super().__init__(optimizer=optimizer, defaults=defaults)

def prune_mask(
Expand Down
6 changes: 4 additions & 2 deletions sparsimony/dst/set.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, Dict, Any
import torch
import torch.nn as nn
from torch.ao.pruning.sparsifier.base_sparsifier import BaseSparsifier
Expand Down Expand Up @@ -27,6 +27,7 @@ def __init__(
scheduler: BaseScheduler,
distribution: BaseDistribution,
optimizer: torch.optim.Optimizer,
defaults: Optional[Dict[str, Any]] = None,
sparsity: float = 0.5,
grown_weights_init: float = 0.0,
init_method: Optional[str] = "grad_flow",
Expand All @@ -36,7 +37,8 @@ def __init__(
self.sparsity = sparsity
self.grown_weights_init = grown_weights_init
self.init_method = init_method
defaults = dict()
if defaults is None:
defaults = dict()
super().__init__(optimizer=optimizer, defaults=defaults)

def prune_mask(
Expand Down
Loading