Cleanup CosineLRScheduler and change defaults (#1487)

Summary:
Pull Request resolved: fairinternal/fairseq-py#1487

Here's the code for CosineLRScheduler that I used as a reference: https://github.com/pytorch/fairseq/blob/577e4fa78a295fd7cd3ee7e9fd4b936ca800ebea/fairseq/optim/lr_scheduler/cosine_lr_schedul

In the reference:
- `warmup_init_lr` defaults to `args.lr[0]`
- `warmup_end_lr` defaults to `args.max_lr`
- `min_lr` defaults to `args.lr[0]`  (note that there's also a `args.min_lr` option defined in the global fairseq config, but this is unused by the cosine scheduler)
- `max_lr` is a required option

This diff removes `max_lr` and replaces it with `lr[0]` to be more
consistent with other LR schedulers. We then add an explicit `min_lr`
option to the Config.

Test Plan: Imported from OSS

Reviewed By: alexeib

Differential Revision: D25342180

Pulled By: myleott

fbshipit-source-id: 61281666e68839da8efc4714c2ce8c49dc4c8e6e

Author: Myle Ott

examples/language_model/README.adaptive_inputs.md

@@ -19,8 +19,8 @@ fairseq-train --task language_modeling \
     data-bin/wikitext-103 \
     --save-dir checkpoints/transformer_wikitext-103 \
     --arch transformer_lm_wiki103 \
-    --max-update 286000 --max-lr 1.0 --t-mult 2 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 \
-    --warmup-updates 16000 --warmup-init-lr 1e-07 --stop-min-lr 1e-09 --optimizer nag --lr 0.0001 --clip-norm 0.1 \
+    --max-update 286000 --lr 1.0 --t-mult 2 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 \
+    --warmup-updates 16000 --warmup-init-lr 1e-07 --stop-min-lr 1e-09 --optimizer nag --min-lr 0.0001 --clip-norm 0.1 \
     --criterion adaptive_loss --max-tokens 3072 --update-freq 3 --tokens-per-sample 3072 --seed 1 \
     --sample-break-mode none --skip-invalid-size-inputs-valid-test --ddp-backend=no_c10d
 ```

examples/pay_less_attention_paper/README.md

@@ -140,7 +140,7 @@ python -m torch.distributed.launch --nproc_per_node 8 $(which fairseq-train) \
     --stop-min-lr 1e-09 --update-freq 16 --attention-dropout 0.1 --keep-last-epochs 10 \
     --ddp-backend=no_c10d --max-tokens 3584 \
     --lr-scheduler cosine --warmup-init-lr 1e-7 --warmup-updates 10000 \
-    --lr-shrink 1 --max-lr 0.001 --lr 1e-7 --warmup-init-lr 1e-07 \
+    --lr-shrink 1 --lr 0.001 --min-lr 1e-7 --warmup-init-lr 1e-07 \
     --t-mult 1 --lr-period-updates 20000 \
     --arch lightconv_wmt_en_de_big --save-dir $SAVE \
     --dropout 0.3 --attention-dropout 0.1 --weight-dropout 0.1 \
@@ -165,7 +165,7 @@ python -m torch.distributed.launch --nproc_per_node 8 $(which fairseq-train) \
     --stop-min-lr 1e-09 --update-freq 16 --attention-dropout 0.1 --keep-last-epochs 10 \
     --ddp-backend=no_c10d --max-tokens 3584 \
     --lr-scheduler cosine --warmup-init-lr 1e-7 --warmup-updates 10000 \
-    --lr-shrink 1 --max-lr 0.001 --lr 1e-7 --warmup-init-lr 1e-07 \
+    --lr-shrink 1 --lr 0.001 --min-lr 1e-7 --warmup-init-lr 1e-07 \
     --t-mult 1 --lr-period-updates 70000 \
     --arch lightconv_wmt_en_fr_big --save-dir $SAVE \
     --dropout 0.1 --attention-dropout 0.1 --weight-dropout 0.1 \

examples/quant_noise/README.md

@@ -208,7 +208,7 @@ fairseq-train --task language_modeling /path/to/wikitext-103/data \
     --ddp-backend no_c10d \
     --decoder-attention-heads 8 --decoder-embed-dim 1024 --decoder-ffn-embed-dim 4096 --decoder-input-dim 1024 \
     --decoder-layers 16 --decoder-normalize-before --decoder-output-dim 1024 \
-    --lr 0.0001 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 --max-lr 1.0 --t-mult 2.0 \
+    --min-lr 0.0001 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 --lr 1.0 --t-mult 2.0 \
     --max-tokens 3072 --tokens-per-sample 3072 --momentum 0.99 --optimizer nag \
     --sample-break-mode none --update-freq 3 \
     --warmup-init-lr 1e-07 --warmup-updates 16000 \
@@ -269,7 +269,7 @@ fairseq-train --task language_modeling /path/to/wikitext-103/data \
     --ddp-backend no_c10d \
     --decoder-attention-heads 8 --decoder-embed-dim 1024 --decoder-ffn-embed-dim 4096 --decoder-input-dim 1024 --decoder-layers 16 --decoder-normalize-before --decoder-output-dim 1024 \
     --fp16 --keep-last-epochs -1 \
-    --lr 0.0001 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 --max-lr 0.05 --stop-min-lr 1e-09 \
+    --min-lr 0.0001 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75 --lr 0.05 --stop-min-lr 1e-09 \
     --max-tokens 2944  --tokens-per-sample 2944\
     --momentum 0.99 --no-epoch-checkpoints --no-progress-bar --optimizer nag --required-batch-size-multiple 8 \
     --sample-break-mode none --t-mult 2.0 --skip-invalid-size-inputs-valid-test \

examples/truncated_bptt/README.md

@@ -37,7 +37,7 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 fairseq-train \
     --arch transformer_xl --n-layer 16 --d-model 410 --n-head 10 \
     --d-head 41 --d-inner 2100 --dropout 0.1 --dropatt 0.0 --mem-len 150 \
     --optimizer adam --clip-norm 0.25 \
-    --lr-scheduler cosine --warmup-updates 0 --lr 0.0 --max-lr 0.00025  \
+    --lr-scheduler cosine --warmup-updates 0 --min-lr 0.0 --lr 0.00025  \
     --log-format json --log-interval 25 \
     --fp16
 ```

examples/wav2vec/README.md

@@ -186,7 +186,7 @@ $ python examples/wav2vec/wav2vec_manifest.py /path/to/waves --dest /manifest/pa
 
 ```
 $ python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 --save-interval 1 --no-epoch-checkpoints \
---arch wav2vec --task audio_pretraining --lr 1e-06 --stop-min-lr 1e-09 --optimizer adam --max-lr 0.005 --lr-scheduler cosine \
+--arch wav2vec --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 --optimizer adam --lr 0.005 --lr-scheduler cosine \
 --conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1)] \
 --conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
 --skip-connections-agg --residual-scale 0.5 --log-compression --warmup-updates 500 --warmup-init-lr 1e-07 --criterion wav2vec --num-negatives 10 \
@@ -244,8 +244,8 @@ $ python examples/wav2vec/wav2vec_manifest.py /path/to/waves --dest /manifest/pa
 
 ```
 $ python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 \
---save-interval 1 --no-epoch-checkpoints --arch wav2vec --task audio_pretraining --lr 1e-06 --stop-min-lr 1e-09 \
---optimizer adam --max-lr 1e-05 --lr-scheduler cosine \
+--save-interval 1 --no-epoch-checkpoints --arch wav2vec --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 \
+--optimizer adam --lr 1e-05 --lr-scheduler cosine \
 --conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1), (512, 1, 1)] \
 --conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
 --activation gelu --offset auto --skip-connections-agg --residual-scale 0.5 \

fairseq/optim/lr_scheduler/cosine_lr_scheduler.py

@@ -26,9 +26,11 @@ class CosineLRScheduleConfig(FairseqDataclass):
             "help": "initial learning rate during warmup phase; default is cfg.lr"
         },
     )
-    max_lr: float = field(
-        default=1.0, metadata={"help": "max learning rate, must be more than cfg.lr"}
+    lr: List[float] = field(
+        default=II("optimization.lr"),
+        metadata={"help": "max learning rate, must be more than cfg.min_lr"},
     )
+    min_lr: float = field(default=0.0, metadata={"help": "min learning rate"})
     t_mult: float = field(
         default=1.0, metadata={"help": "factor to grow the length of each period"}
     )
@@ -38,7 +40,7 @@ class CosineLRScheduleConfig(FairseqDataclass):
     lr_shrink: float = field(
         default=0.1, metadata={"help": "shrink factor for annealing"}
     )
-    lr: List[float] = II("optimization.lr")
+    # This is not required, but is for convenience in inferring lr_period_updates
     max_update: int = II("optimization.max_update")
 
 
@@ -50,7 +52,7 @@ class CosineLRSchedule(FairseqLRScheduler):
 
     We also support a warmup phase where we linearly increase the learning rate
     from some initial learning rate (``--warmup-init-lr``) until the configured
-    max learning rate (``--max-lr``).
+    max learning rate (``--lr``).
 
     During warmup::
 
@@ -59,7 +61,7 @@ class CosineLRSchedule(FairseqLRScheduler):
 
     After warmup::
 
-      lr = lr_min + 0.5*(lr_max - lr_min)*(1 + cos(t_curr / t_i))
+      lr = cfg.min_lr + 0.5*(cfg.lr - cfg.min_lr)*(1 + cos(t_curr / t_i))
 
     where ``t_curr`` is current percentage of updates within the current period
     range and ``t_i`` is the current period range, which is scaled by ``t_mul``
@@ -74,23 +76,21 @@ def __init__(self, cfg: CosineLRScheduleConfig, fairseq_optimizer):
                 f" Consider --lr-scheduler=fixed instead. ({cfg.lr})"
             )
 
-        warmup_end_lr = cfg.max_lr
-        lr = cfg.lr[0] if isinstance(cfg.lr, Collection) else cfg.lr
-        if cfg.warmup_init_lr < 0:
-            cfg.warmup_init_lr = lr
+        self.max_lr = cfg.lr[0] if isinstance(cfg.lr, Collection) else cfg.lr
+        assert (
+            self.max_lr > cfg.min_lr
+        ), f"max_lr (={cfg.lr}) must be more than min_lr (={cfg.min_lr})"
 
-        # default min_lr=-1 -> cosine anneale to lr=0.0
-        # otherwise pick min_lr from config
-        self.min_lr = cfg.min_lr if cfg.min_lr > 0.0 else 0.0
-        self.max_lr = lr
-        assert self.max_lr > self.min_lr, "max_lr must be more than lr"
+        warmup_end_lr = self.max_lr
+        if cfg.warmup_init_lr < 0:
+            cfg.warmup_init_lr = cfg.min_lr
 
         self.t_mult = cfg.t_mult
         self.period = cfg.lr_period_updates
 
         if self.period <= 0:
             assert (
-                cfg.max_update >= 0
+                cfg.max_update > 0
             ), "Either --max_update or --lr-period-updates must be set"
             self.period = cfg.max_update - cfg.warmup_updates
 
@@ -136,7 +136,7 @@ def step_update(self, num_updates):
                 t_curr = curr_updates - (self.period * i)
 
             lr_shrink = self.lr_shrink ** i
-            min_lr = self.min_lr * lr_shrink
+            min_lr = self.cfg.min_lr * lr_shrink
             max_lr = self.max_lr * lr_shrink
 
             self.lr = min_lr + 0.5 * (max_lr - min_lr) * (