improvements

Lightning-AI · May 8, 2024 · 88634cc · 88634cc
1 parent 4bc9c51
commit 88634cc
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 49 deletions.
diff --git a/litgpt/external/galore.py b/litgpt/external/galore.py
@@ -232,7 +232,6 @@ def get_galore_params(model):
         if not any(target_key in module_name for target_key in target_modules_list):
             continue
 
-        print('enable GaLore for weights in module: ', module_name)
         galore_params.append(module.weight)
     id_galore_params = [id(p) for p in galore_params]
     # make parameters without "rank" to another group
@@ -377,7 +376,7 @@ def __init__(
         eps: float = 1e-6,
         weight_decay: float = 0.0,
         correct_bias: bool = True,
-        no_deprecation_warning: bool = False,
+        no_deprecation_warning: bool = True,
     ):
         if not no_deprecation_warning:
             warnings.warn(

diff --git a/litgpt/finetune/adapter.py b/litgpt/finetune/adapter.py
@@ -29,7 +29,6 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
-    get_linear_nonlinear_params,
     init_out_dir,
     load_checkpoint,
     num_parameters,
@@ -164,16 +163,13 @@ def main(
             optimizer_cls = torch.optim.AdamW
 
     elif optim.optimizer in ("galore_adamw", "galore_adamw_8bit"):
-        if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
-            raise ValueError("The combinatiomn of QLoRA and GaLore is currently not supported.")
+        from litgpt.external.galore import get_galore_params
 
-        linear_params, nonlinear_params = get_linear_nonlinear_params(model)
-        # Currently apply galore to all parameters;
-        # we could add options to target specific layers for AdamW and GaLore later
+        regular_params, galore_params = get_galore_params(model)
         trainable_params = [
-            {'params': nonlinear_params},
+            {'params': regular_params},
             {
-                'params': linear_params,
+                'params': galore_params,
                 'rank': optim.galore_r,
                 'update_proj_gap': optim.galore_update_proj_gap,
                 'scale': optim.galore_scale,

diff --git a/litgpt/finetune/adapter_v2.py b/litgpt/finetune/adapter_v2.py
@@ -29,7 +29,6 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
-    get_linear_nonlinear_params,
     init_out_dir,
     load_checkpoint,
     num_parameters,
@@ -165,16 +164,13 @@ def main(
             optimizer_cls = torch.optim.AdamW
 
     elif optim.optimizer in ("galore_adamw", "galore_adamw_8bit"):
-        if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
-            raise ValueError("The combinatiomn of QLoRA and GaLore is currently not supported.")
+        from litgpt.external.galore import get_galore_params
 
-        linear_params, nonlinear_params = get_linear_nonlinear_params(model)
-        # Currently apply galore to all parameters;
-        # we could add options to target specific layers for AdamW and GaLore later
+        regular_params, galore_params = get_galore_params(model)
         trainable_params = [
-            {'params': nonlinear_params},
+            {'params': regular_params},
             {
-                'params': linear_params,
+                'params': galore_params,
                 'rank': optim.galore_r,
                 'update_proj_gap': optim.galore_update_proj_gap,
                 'scale': optim.galore_scale,

diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py
@@ -30,7 +30,6 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
-    get_linear_nonlinear_params,
     load_checkpoint,
     init_out_dir,
     num_parameters,
@@ -200,17 +199,17 @@ def main(
         if isinstance(fabric.strategy.precision, BitsandbytesPrecision):
             raise ValueError("The combination of QLoRA and GaLore is currently not supported.")
 
-        linear_params, nonlinear_params = get_linear_nonlinear_params(model)
-        # Currently apply galore to all parameters;
-        # we could add options to target specific layers for AdamW and GaLore later
+        from litgpt.external.galore import get_galore_params
+
+        regular_params, galore_params = get_galore_params(model)
         trainable_params = [
-            {'params': nonlinear_params},
+            {'params': regular_params},
             {
-                'params': linear_params,
+                'params': galore_params,
                 'rank': optim.galore_r,
-                'update_proj_gap':  optim.galore_update_proj_gap,
-                'scale':  optim.galore_scale,
-                'proj_type':  optim.galore_proj_type
+                'update_proj_gap': optim.galore_update_proj_gap,
+                'scale': optim.galore_scale,
+                'proj_type': optim.galore_proj_type
             }
         ]
         if optim.optimizer == "galore_adamw_8bit":

diff --git a/litgpt/pretrain.py b/litgpt/pretrain.py
@@ -30,7 +30,6 @@
     chunked_cross_entropy,
     copy_config_files,
     get_default_supported_precision,
-    get_linear_nonlinear_params,
     init_out_dir,
     num_parameters,
     parse_devices,
@@ -190,13 +189,13 @@ def main(
         optimizer_cls = torch.optim.AdamW
 
     elif optim.optimizer in ("galore_adamw", "galore_adamw_8bit"):
-        linear_params, nonlinear_params = get_linear_nonlinear_params(model)
-        # Currently apply galore to all parameters;
-        # we could add options to target specific layers for AdamW and GaLore later
+        from litgpt.external.galore import get_galore_params
+
+        regular_params, galore_params = get_galore_params(model)
         trainable_params = [
-            {'params': nonlinear_params},
+            {'params': regular_params},
             {
-                'params': linear_params,
+                'params': galore_params,
                 'rank': optim.galore_r,
                 'update_proj_gap': optim.galore_update_proj_gap,
                 'scale': optim.galore_scale,

diff --git a/litgpt/utils.py b/litgpt/utils.py
@@ -485,19 +485,3 @@ def choose_logger(
     if logger_name == "wandb":
         return WandbLogger(project=name, resume=resume, **kwargs)
     raise ValueError(f"`--logger_name={logger_name}` is not a valid option. Choose from 'csv', 'tensorboard', 'wandb'.")
-
-
-## REMOVE
-def get_linear_nonlinear_params(model):
-    linear_params = []
-    nonlinear_params = []
-    for module in model.modules():
-        if isinstance(module, torch.nn.Linear):
-            linear_params.extend(list(module.parameters()))
-        else:
-
-            nonlinear_params.extend(list(module.parameters()))
-    linear_params = list(set(linear_params))
-    nonlinear_params = list(set(nonlinear_params) - set(linear_params))
-    return linear_params, nonlinear_params
-