tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

fix LISA by ensuring params are not frozen during __init__

Browse Source
This commit is contained in:
Aman Karmani
2024-04-01 06:57:28 +00:00
parent b357c93f23
commit 6185cd5227
1 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/axolotl/utils/callbacks/lisa.py
View File

@@ -16,7 +16,7 @@ from transformers import TrainerCallback
if TYPE_CHECKING: if TYPE_CHECKING:
from axolotl.core.trainer_builder import AxolotlTrainer from axolotl.core.trainer_builder import AxolotlTrainer
LOG = logging.getLogger("axolotl.callbacks") LOG = logging.getLogger("axolotl.callbacks.lisa")
def lisa_callback_factory(trainer: "AxolotlTrainer"): def lisa_callback_factory(trainer: "AxolotlTrainer"):
@@ -37,20 +37,22 @@ def lisa_callback_factory(trainer: "AxolotlTrainer"):
self.total_layers = len( self.total_layers = len(
reduce(getattr, self.layers_attribute.split("."), self.trainer.model) reduce(getattr, self.layers_attribute.split("."), self.trainer.model)
) )
self.freeze_all_layers(True)
self.active_layers_indices = [] self.active_layers_indices = []
def freeze_all_layers(self, summarize=False): layers = reduce(
getattr, self.layers_attribute.split("."), self.trainer.model
)
LOG.info(
f"LISA will activate {self.n_layers}/{len(layers)} layers ({self.n_layers*100/len(layers)}%) every {self.step_interval} steps"
)
def freeze_all_layers(self):
layers = reduce( layers = reduce(
getattr, self.layers_attribute.split("."), self.trainer.model getattr, self.layers_attribute.split("."), self.trainer.model
) )
for layer in layers: for layer in layers:
for param in layer.parameters(): for param in layer.parameters():
param.requires_grad = False param.requires_grad = False
if summarize:
LOG.info(
f"Freezing {len(layers)} layers; will activate {self.n_layers} layers ({self.n_layers*100/len(layers)}%) every {self.step_interval} steps"
)
def on_step_begin( def on_step_begin(
self, args, state, control, **kwargs self, args, state, control, **kwargs