| 377 | |
| 378 | |
| 379 | def get_optimizer(config, params_to_optimize, use_8BitAdam, use_deepspeed: bool = False): |
| 380 | |
| 381 | optimizer = config["optimizer"] |
| 382 | use_8bit_adam = use_8BitAdam |
| 383 | |
| 384 | |
| 385 | # Use DeepSpeed optimzer |
| 386 | if use_deepspeed: |
| 387 | from accelerate.utils import DummyOptim |
| 388 | |
| 389 | return DummyOptim( |
| 390 | params_to_optimize, |
| 391 | lr = config["learning_rate"], |
| 392 | betas = (config["adam_beta1"], config["adam_beta2"]), |
| 393 | eps = config["adam_epsilon"], |
| 394 | weight_decay = config["adam_weight_decay"], |
| 395 | ) |
| 396 | |
| 397 | # Optimizer creation |
| 398 | supported_optimizers = ["adam", "adamw", "prodigy"] |
| 399 | if optimizer not in supported_optimizers: |
| 400 | logger.warning( |
| 401 | f"Unsupported choice of optimizer: {optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW" |
| 402 | ) |
| 403 | optimizer = "adamw" |
| 404 | |
| 405 | if use_8bit_adam and optimizer.lower() not in ["adam", "adamw"]: |
| 406 | logger.warning( |
| 407 | f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was " |
| 408 | f"set to {optimizer.lower()}" |
| 409 | ) |
| 410 | |
| 411 | if use_8bit_adam: |
| 412 | try: |
| 413 | from bitsandbytes.optim import AdamW8bit |
| 414 | # from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit |
| 415 | |
| 416 | except ImportError: |
| 417 | raise ImportError( |
| 418 | "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`." |
| 419 | ) |
| 420 | |
| 421 | if optimizer.lower() == "adamw": # Most cases we should use AdamW with 8Bit adam to accelerate |
| 422 | optimizer_class = AdamW8bit if use_8bit_adam else torch.optim.AdamW |
| 423 | |
| 424 | optimizer = optimizer_class( |
| 425 | params_to_optimize, |
| 426 | betas = (config["adam_beta1"], config["adam_beta2"]), |
| 427 | eps = config["adam_epsilon"], |
| 428 | weight_decay = config["adam_weight_decay"], |
| 429 | ) |
| 430 | |
| 431 | elif optimizer.lower() == "adam": |
| 432 | optimizer_class = bnb.optim.Adam8bit if use_8bit_adam else torch.optim.Adam |
| 433 | |
| 434 | optimizer = optimizer_class( |
| 435 | params_to_optimize, |
| 436 | betas = (config["adam_beta1"], config["adam_beta2"]), |