| 896 | |
| 897 | |
| 898 | def _add_checkpointing_args(parser): |
| 899 | group = parser.add_argument_group(title="checkpointing") |
| 900 | |
| 901 | group.add_argument( |
| 902 | "--save", |
| 903 | type=str, |
| 904 | default=None, |
| 905 | help="Output directory to save checkpoints to.", |
| 906 | ) |
| 907 | group.add_argument( |
| 908 | "--save-interval", |
| 909 | type=int, |
| 910 | default=None, |
| 911 | help="Number of iterations between checkpoint saves.", |
| 912 | ) |
| 913 | group.add_argument( |
| 914 | "--no-save-optim", |
| 915 | action="store_true", |
| 916 | default=None, |
| 917 | help="Do not save current optimizer.", |
| 918 | ) |
| 919 | group.add_argument( |
| 920 | "--no-save-rng", |
| 921 | action="store_true", |
| 922 | default=None, |
| 923 | help="Do not save current rng state.", |
| 924 | ) |
| 925 | group.add_argument( |
| 926 | "--load", |
| 927 | type=str, |
| 928 | default=None, |
| 929 | help="Directory containing a model checkpoint.", |
| 930 | ) |
| 931 | group.add_argument( |
| 932 | "--low-memory-load", |
| 933 | action="store_true", |
| 934 | default=None, |
| 935 | help="Load model checkpoint in low memory mode." |
| 936 | "On each machine, workers load the checkpoint one at a time." |
| 937 | ) |
| 938 | group.add_argument( |
| 939 | "--dist-timeout", |
| 940 | type=int, |
| 941 | default=30, |
| 942 | help="Timeout for Pytorch Distributed backend (in minutes).", |
| 943 | ) |
| 944 | group.add_argument( |
| 945 | "--load-state", |
| 946 | type=str, |
| 947 | default=None, |
| 948 | help="Start training from a existing model state.", |
| 949 | ) |
| 950 | group.add_argument( |
| 951 | "--no-load-optim", |
| 952 | action="store_true", |
| 953 | default=None, |
| 954 | help="Do not load optimizer when loading checkpoint.", |
| 955 | ) |