DOCSEARCHbyAlgolia
Submit Docs

Adan: Adaptive Nesterov Momentum Algorithm

repository·main·Indexed 21 days ago

https://github.com/sail-sg/adan

An Adaptive Nesterov Momentum Algorithm designed for faster optimization of deep learning models across LLMs, Vision, NLP, and Diffusion models. The library provides integration guides for MAE, timm, DreamFusion, and Megatron-LM, as well as a fused kernel implementation for improved GPU performance via multi-tensor and single-tensor access patterns.

Tokens
11K
Snippets
28
Records
37
Agent score
74%

What's inside Adan

  1. How fused Adan access patterns affect performance and memory

    main

    The fused and foreach parameters in the Adan optimizer change how the GPU kernel interacts with model parameters:

    Multi-tensor access (foreach=True, fused=True)

    • Mechanism: All layer parameters are passed to the kernel simultaneously. The kernel internally iterates through the layers.
    • Pros: Significantly reduces kernel launch overhead by requiring only one kernel start.
    • Cons: Theoretically increases peak memory usage, though benchmarking shows the increase is typically insignificant.

    Single-tensor access (foreach=False, fused=True)

    • Mechanism: A Python for loop traverses each layer, launching a separate kernel for each layer's gradient calculation.
    • Pros: Theoretically reduces peak memory usage by only accessing one layer at a time.
    • Cons: Introduces significant kernel launch overhead due to multiple starts.

  2. Integrate Adan into timm for vision tasks

    main

    To use the Adan optimizer within the timm (PyTorch Image Models) framework, you must perform two main steps: adding the required hyper-parameters to your training script and modifying the optimizer factory to recognize and instantiate Adan.

    Step 1: Add Adan Hyper-parameters

    Add the following arguments to your train.py argument parser to support Adan-specific configurations:

    • --max-grad-norm: L2 norm threshold for gradient clipping (default: 0.0, no clipping).
    • --weight-decay: Weight decay value, similar to AdamW (default: 0.02).
    • --opt-eps: Optimizer epsilon to prevent division by zero (default: None, uses 1e-8 in Adan).
    • --opt-betas: Optimizer betas (default: None, uses [0.98, 0.92, 0.99] in Adan).
    • --no-prox: If set, performs weight decay like AdamW. If unset (default), uses the proximal update rule described in the Adan paper. Note: The paper uses no-prox=False.
    • --bias-decay: If set, performs weight decay on bias terms, batch normalization (BN), and other 1D parameters. By default, these are filtered out in timm.
    parser.add_argument('--max-grad-norm', type=float, default=0.0, help='if the l2 norm is large than this hyper-parameter, then we clip the gradient  (default: 0.0, no gradient clip)')
parser.add_argument('--weight-decay', type=float, default=0.02,  help='weight decay, similar one used in AdamW (default: 0.02)')
parser.add_argument('--opt-eps', default=None, type=float, metavar='EPSILON', help='optimizer epsilon to avoid the bad case where second-order moment is zero (default: None, use opt default 1e-8 in adan)')
parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA', help='optimizer betas in Adan (default: None, use opt default [0.98, 0.92, 0.99] in Adan)')
parser.add_argument('--no-prox', action='store_true', default=False, help='whether perform weight decay like AdamW (default=False)')
parser.add_argument('--bias-decay', action='store_true', default=False, help='Perform the weight decay on bias term (default=False)')

  3. Train ViT with Setting I

    main

    Use Setting I for training Vision Transformers (ViT) following prevalent ResNet training settings. This configuration is suitable for models like deit_small_patch16_224.

    Key hyperparameters for this setting:

    • --opt Adan
    • --lr 1.5e-2
    • --opt-betas 0.98 0.92 0.99
    • --warmup-epochs 60
    • --mixup 0.2 and --cutmix 1.0
    python -m torch.distributed.launch --nproc_per_node=8 ./train.py \"    --data-dir ${IMAGENET_DIR}   \"    --model deit_small_patch16_224 \"    --sched cosine -j 10 \"    --epochs ${EPOCH} --weight-decay 0.02 \"    --opt Adan \" --lr 1.5e-2  --opt-betas 0.98 0.92 0.99 \"    --opt-eps 1e-8 --max-grad-norm 0.0 \"    --warmup-lr 1e-8 --min-lr 1.0e-08 \"    -b 256 --amp \"    --aug-repeats 0 \"    --warmup-epochs 60 \"    --aa rand-m7-mstd0.5-inc1 \"    --smoothing 0.1 \"    --remode pixel \"    --reprob 0.0 \"    --bce \"    --drop 0.0 --drop-path 0.05 \"    --mixup 0.2 --cutmix 1.0 \"    --output ${OUT_DIR} \"    --experiment ${EXP_DIR}

  4. Prepare and binarize datasets for pre-training

    main

    Before pre-training, you must download the dataset and convert it into a binarized format compatible with the training pipeline.

    1. Download: Run the provided download script.
    2. Binarize: Use tools/preprocess_data.py to process the JSON dataset. You will need to provide the input file, output prefix, vocabulary, and tokenizer settings (e.g., GPT2BPETokenizer).
    # Step 1: Download
python ./download_dataset.py

# Step 2: Binarize
python tools/preprocess_data.py \
      --input stack_python.json \
      --output-prefix codegpt \
      --vocab checkpoints/gpt2-adan/tokenizer/vocab.json \
      --json-key content \
      --dataset-impl mmap \
      --workers 16 \
      --chunk-size 25 \
      --tokenizer-type GPT2BPETokenizer \
      --merge-file checkpoints/gpt2-adan/tokenizer/merges.txt \
      --append-eod

  5. Fine-tune BERT-base on GLUE tasks

    main

    Fine-tuning involves three main steps: downloading GLUE data, preprocessing it, and running the fine-tuning script.

    1. Download: Use the provided download_glue_data.py script.
    2. Preprocess: Run ./examples/roberta/preprocess_GLUE_tasks.sh with the target task name (e.g., RTE, SST-2, or ALL).
    3. Fine-tune: Execute the acc_test.py script with the appropriate task configuration.

    Required arguments for acc_test.py:

    • --avg_num: Number of repetitions.
    • --data_path: Path to raw GLUE task data.
    • --bin_path: Path to binarized GLUE task data.
    • --pre_path: Path to the pre-trained BERT model checkpoint.
    • --finetune_path: Directory to save/load fine-tuned models.
    • --task: The specific task configuration name (e.g., rte-adan).
    TASK=RTE;

python  path/to/fairseq/examples/roberta/config/finetuning/acc_test.py --avg_num 1 \
--data_path /path/to/fairseq/GLUE/glue_data/$TASK \
--bin_path /path/to/fairseq/GLUE/$TASK-bin \
--pre_path /path/to/fairseq/bert-adan/checkpoint_best.pt \
--finetune_path /path/to/fairseq/bert-fintune/adan/$TASK/ \
--task rte-adan

  6. Configure the MAE environment for Adan

    main

    To ensure compatibility with the Adan experiments for MAE, use the following package versions or the provided Docker image.

    Required Package Versions:

    • torch: 1.7.1+cu110
    • torchvision: 0.8.2+cu110
    • timm: 0.4.5
    • torchaudio: 0.7.2

    Docker Image: xyxie/adan-image:mae (available on Docker Hub).

  7. Train ResNet-50

    main

    Use the following command to train ResNet-50 using the default settings for ResNets.

    Note on Learning Rate (--lr):

    • Use 3e-2 for 100 epochs.
    • Use 1.5e-2 for 200 or 300 epochs.

    Key hyperparameters:

    • --opt Adan
    • --opt-betas 0.98 0.92 0.99
    • --bias-decay
    • --bce
    python -m torch.distributed.launch --nproc_per_node=8 ./train.py \"    --data-dir ${IMAGENET_DIR} \"    --model resnet50 \"    --sched cosine -j 8 \"    --epochs ${EPOCH} --weight-decay .02 \"    --opt Adan \" --lr ${LR}  --opt-betas 0.98 0.92 0.99 \"    --opt-eps 1e-8 --max-grad-norm 5.0 \"    --warmup-lr 1e-9 --min-lr 1.0e-05 --bias-decay \"    -b 256 --amp \"    --aug-repeats 0 \"    --warmup-epochs 60 \"    --aa rand-m7-mstd0.5-inc1 \"    --smoothing 0.0 \"    --remode pixel \"    --crop-pct 0.95 \"    --reprob 0.0 \"    --bce \"    --drop 0.0 --drop-path 0.05 \"    --mixup 0.1 --cutmix 1.0 \"    --output ${OUT_DIR} \"    --experiment ${EXP_DIR}

  8. Integrate Adan into MAE training and fine-tuning scripts

    main

    To use the Adan optimizer with the official MAE repository, you must perform two integration steps: adding the necessary command-line arguments and updating the optimizer instantiation logic.

    Step 1: Add CLI Arguments

    Add the following arguments to main_pretrain.py and main_finetune.py using argparse:

    • --use-adan: Boolean flag to enable Adan. Defaults to False (uses AdamW).
    • --max-grad-norm: Float. If the L2 norm exceeds this value, the gradient is clipped. Default 0.0 (no clipping).
    • --opt-eps: Float. Optimizer epsilon to prevent division by zero. If None, Adan defaults to 1e-8.
    • --opt-betas: List of floats. Optimizer betas for Adan. If None, Adan defaults to [0.98, 0.92, 0.99].

    Step 2: Update Optimizer Creation

    Replace the vanilla optimizer creation logic with the following pattern. Note that following timm conventions, weight decay should be set to 0.0 for bias and norm layers when using Adan.

    # following timm: set wd as 0 for bias and norm layers
param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
if args.use_adan:
  if args.bias_decay:
    param = model_without_ddp.parameters() 
  else: 
    param = param_groups
    args.weight_decay = 0.0
    optimizer = Adan(param, weight_decay=args.weight_decay, 
                     lr=args.lr, betas=args.opt_betas, 
                     eps = args.opt_eps, max_grad_norm=args.max_grad_norm)
  else:
    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))

  9. Install dependencies for Adan Optimizer fused kernel

    main

    To use the fused Adan kernel, ensure the following dependencies are installed and compatible:

    1. Libtorch/PyTorch: ATen is required. (Tested/passed on PyTorch 1.13.1)
    2. CUDA Toolkit: (Tested/passed on CUDA 11.6+)
    3. ninja: Required for the build process.

  10. Configure Adan optimizer arguments in Megatron-LM

    main

    To enable Adan via command line arguments, add the following parameters to Megatron-LM/megatron/arguments.py. Note that beta3 is specific to Adan and is not used by the standard Adam optimizer.

    # Add these to Megatron-LM/megatron/arguments.py
group.add_argument('--adan-beta1', type=float, default=0.98,
                   help='First coefficient for computing running averages ' 
                   'of gradient and its square')
group.add_argument('--adan-beta2', type=float, default=0.92,
                   help='Second coefficient for computing running averages ' 
                   'of gradient and its square')
group.add_argument('--adan-beta3', type=float, default=0.99,
                   help='Second coefficient for computing running averages ' 
                   'of gradient and its square')
group.add_argument('--adan-eps', type=float, default=1e-08,
                   help='Term added to the denominator to improve' 
                   'numerical stability')
group.add_argument('--optimizer', type=str, default='adam',
                   choices=['adam', 'sgd', 'adan'])

  11. Integrate Adan optimizer into Megatron-LM

    main

    To use the Adan optimizer within a Megatron-LM environment, you must manually integrate the Adan class and register its specific hyperparameters.

    1. Add the optimizer implementation: Place adan.py into Megatron-LM/megatron/optimizer/adan.py.
    2. Register the optimizer: In Megatron-LM/megatron/optimizer/__init__.py, import Adan and add logic to initialize it when args.optimizer == 'adan' is selected. The Adan constructor requires betas as a tuple of three values: (adan_beta1, adan_beta2, adan_beta3).
    from .adan import Adan

# Inside the optimizer selection logic
elif args.optimizer == 'adan':
  optimizer = Adan(param_groups, lr=args.lr, weight_decay=args.weight_decay,
                   betas=(args.adan_beta1, args.adan_beta2, args.adan_beta3),
                   eps=args.adan_eps)

  12. Train ResNet-101

    main

    Use the following command to train ResNet-101.

    Note on Learning Rate (--lr):

    • Use 1e-2 for 100 epochs.
    • Use 1.5e-2 for 200 or 300 epochs.

    Key hyperparameters:

    • --opt Adan
    • --opt-betas 0.98 0.92 0.99
    • --bias-decay
    • --bce-loss
    python -m torch.distributed.launch --nproc_per_node=8 train.py \" \"    --data-dir ${IMAGENET_DIR} \"    --model resnet101 \"    --sched cosine -j 8 \"    --epochs 300 --weight-decay .02 \"    --lr 1.5e-2  --warmup-lr 1e-9 --min-lr 1.0e-05 \"    -b 256 --amp --opt adan --opt-betas 0.98 0.92 0.99 --opt-eps 1e-8 \"    --max-grad-norm 5 \"    --bias-decay \"    --aug-repeats 0 \"    --warmup-epochs 90 \"    --aa rand-m7-mstd0.5-inc1 \"    --smoothing 0.0 \"    --remode pixel \" \"    --bce-loss \"    --crop-pct 0.95 \"    --reprob 0.0 \"    --drop 0.0 --drop-path 0.2 \"    --mixup 0.1 --cutmix 1.0 \"    --output ${OUT_DIR} \"    --experiment ${EXP_DIR}