| 348 | |
| 349 | |
| 350 | def get_optimizer(config, params_to_optimize, use_8BitAdam, use_deepspeed: bool = False): |
| 351 | |
| 352 | optimizer = config["optimizer"] |
| 353 | use_8bit_adam = use_8BitAdam |
| 354 | |
| 355 | |
| 356 | # Use DeepSpeed optimzer |
| 357 | if use_deepspeed: |
| 358 | from accelerate.utils import DummyOptim |
| 359 | |
| 360 | return DummyOptim( |
| 361 | params_to_optimize, |
| 362 | lr = config["learning_rate"], |
| 363 | betas = (config["adam_beta1"], config["adam_beta2"], config["adam_beta3"]), |
| 364 | eps = config["adam_epsilon"], |
| 365 | weight_decay = config["adam_weight_decay"], |
| 366 | ) |
| 367 | |
| 368 | # Optimizer creation |
| 369 | supported_optimizers = ["adam", "adamw", "prodigy"] |
| 370 | if optimizer not in supported_optimizers: |
| 371 | logger.warning( |
| 372 | f"Unsupported choice of optimizer: {optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW" |
| 373 | ) |
| 374 | optimizer = "adamw" |
| 375 | |
| 376 | if use_8bit_adam and optimizer.lower() not in ["adam", "adamw"]: |
| 377 | logger.warning( |
| 378 | f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was " |
| 379 | f"set to {optimizer.lower()}" |
| 380 | ) |
| 381 | |
| 382 | if use_8bit_adam: |
| 383 | try: |
| 384 | from bitsandbytes.optim import AdamW8bit |
| 385 | # from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit |
| 386 | |
| 387 | except ImportError: |
| 388 | raise ImportError( |
| 389 | "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." |
| 390 | ) |
| 391 | |
| 392 | if optimizer.lower() == "adamw": # Most cases we should use AdamW with 8Bit adam to accelerate |
| 393 | optimizer_class = AdamW8bit if use_8bit_adam else torch.optim.AdamW |
| 394 | |
| 395 | optimizer = optimizer_class( |
| 396 | params_to_optimize, |
| 397 | betas = (config["adam_beta1"], config["adam_beta2"]), |
| 398 | eps = config["adam_epsilon"], |
| 399 | weight_decay = config["adam_weight_decay"], |
| 400 | ) |
| 401 | |
| 402 | elif optimizer.lower() == "adam": |
| 403 | optimizer_class = bnb.optim.Adam8bit if use_8bit_adam else torch.optim.Adam |
| 404 | |
| 405 | optimizer = optimizer_class( |
| 406 | params_to_optimize, |
| 407 | betas = (config["adam_beta1"], config["adam_beta2"]), |