storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-18 00:33:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Low-bit memory optimizers, CpuOffloadOptimizer, Memory Reports (#3)

Browse Source 
* update

* update

* update

* update

* update

* update

* update

* update

* update
This commit is contained in:
Aryan
2024-10-04 18:18:54 +05:30
committed by GitHub
parent 512a8a6854
commit d81209bf49
11 changed files with 434 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/args.py
+20 -2
View File
@@ -301,13 +301,21 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
"--optimizer",
type=lambda s: s.lower(),
default="adam",
choices=["adam", "adamw", "prodigy"],
choices=["adam", "adamw", "prodigy", "came"],
help=("The optimizer type to use."),
)
parser.add_argument(
"--use_8bit",
action="store_true",
help="Whether or not to use 8-bit optimizers from `bitsandbytes`. Ignored if incompatible optimzer selected.",
help="Whether or not to use 8-bit optimizers from `bitsandbytes` or `bitsandbytes`.",
)
parser.add_argument(
"--use_4bit",
action="store_true",
help="Whether or not to use 4-bit optimizers from `torchao`.",
)
parser.add_argument(
"--use_torchao", action="store_true", help="Whether or not to use the `torchao` backend for optimizers."
)
parser.add_argument(
"--beta1",
@@ -355,6 +363,16 @@ def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
action="store_true",
help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage.",
)
parser.add_argument(
"--use_cpu_offload_optimizer",
action="store_true",
help="Whether or not to use the CPUOffloadOptimizer from TorchAO to perform optimization step and maintain parameters on the CPU.",
)
parser.add_argument(
"--offload_gradients",
action="store_true",
help="Whether or not to offload the gradients to CPU when using the CPUOffloadOptimizer from TorchAO.",
)
def _get_configuration_args(parser: argparse.ArgumentParser) -> None: