Prodigy Optimizer Documentation

repository·main·Indexed 19 days ago

https://github.com/konstmish/prodigy

An adaptive, parameter-free optimizer implemented in PyTorch designed to automatically estimate the optimal learning rate. Includes guidance on installation via the prodigyopt package, configuration of parameters like d_coef and slice_p, scheduler recommendations using CosineAnnealingLR, and specific stability settings for training diffusion models.

Tokens
553
Snippets
3
Records
4
Agent score
17%

What's inside Prodigy

  1. Configure schedulers with Prodigy

    main

    When using a scheduler with Prodigy, it is recommended to use no scheduler or CosineAnnealingLR.

    If using CosineAnnealingLR, set T_max to the total number of steps to avoid unwanted restarts. If you do use restarts, set safeguard_warmup=True to mitigate issues with the optimizer overestimating the learning rate during the initial phase.

    # Recommended: Cosine Annealing without restarts
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps)
  2. Training Diffusion Models with Prodigy

    main

    For training diffusion models, the following configuration is recommended for better stability:

    • Set safeguard_warmup=True.
    • Set use_bias_correction=True.
    • Set weight_decay=0.01.
    • Optionally set betas=(0.9, 0.99).

    If the model fails to train and the estimated learning rate d remains too small, consider increasing d0 to 1e-5 or 1e-4.

  3. Initialize the Prodigy optimizer

    main

    Prodigy is a PyTorch optimizer designed to be parameter-free. You should typically use lr=1. and adjust the learning rate scale using d_coef.

    Key arguments:

    • params: The model parameters (e.g., net.parameters()).
    • lr: Recommended to keep at 1. (default).
    • weight_decay: Standard values to try are 0 (default), 0.001, 0.01, or 0.1.
    • slice_p: Controls memory consumption. Use 1 (default) for maximum accuracy or 11 for a better trade-off between accuracy and memory efficiency.
    • decouple: If True (default), uses weight decay like AdamW. If False, uses standard $\ell_2$ regularization like Adam.
    • d_coef: Adjusts the estimated learning rate. Values $>1$ (e.g., 2 or 10) force a larger estimate; values $<1$ (e.g., 0.5 or 0.1) force a smaller estimate.
    from prodigyopt import Prodigy
    
    # Example initialization
    opt = Prodigy(net.parameters(), lr=1., weight_decay=0.01, slice_p=1)