| 804 | |
| 805 | |
| 806 | def _add_learning_rate_args(parser): |
| 807 | group = parser.add_argument_group(title="learning rate") |
| 808 | |
| 809 | group.add_argument( |
| 810 | "--lr", |
| 811 | type=float, |
| 812 | default=None, |
| 813 | help="Initial learning rate. Depending on decay style " |
| 814 | "and initial warmup, the learing rate at each " |
| 815 | "iteration would be different.", |
| 816 | ) |
| 817 | group.add_argument( |
| 818 | "--lr-decay-style", |
| 819 | type=str, |
| 820 | default="linear", |
| 821 | choices=["constant", "linear", "cosine"], |
| 822 | help="Learning rate decay function.", |
| 823 | ) |
| 824 | group.add_argument( |
| 825 | "--lr-decay-iters", |
| 826 | type=int, |
| 827 | default=None, |
| 828 | help="number of iterations to decay learning rate over," |
| 829 | " If None defaults to `--train-iters`", |
| 830 | ) |
| 831 | group.add_argument( |
| 832 | "--lr-decay-samples", |
| 833 | type=int, |
| 834 | default=None, |
| 835 | help="number of samples to decay learning rate over," |
| 836 | " If None defaults to `--train-samples`", |
| 837 | ) |
| 838 | group.add_argument( |
| 839 | "--lr-decay-tokens", |
| 840 | type=int, |
| 841 | default=None, |
| 842 | help="number of tokens to decay learning rate over," |
| 843 | " If not None will override iter/sample-based decay", |
| 844 | ) |
| 845 | group.add_argument( |
| 846 | "--lr-warmup-fraction", |
| 847 | type=float, |
| 848 | default=None, |
| 849 | help="fraction of lr-warmup-(iters/samples) to use " "for warmup (as a float)", |
| 850 | ) |
| 851 | group.add_argument( |
| 852 | "--lr-warmup-iters", |
| 853 | type=int, |
| 854 | default=0, |
| 855 | help="number of iterations to linearly warmup " "learning rate over.", |
| 856 | ) |
| 857 | group.add_argument( |
| 858 | "--lr-warmup-samples", |
| 859 | type=int, |
| 860 | default=0, |
| 861 | help="number of samples to linearly warmup " "learning rate over.", |
| 862 | ) |
| 863 | group.add_argument( |