| 971 | |
| 972 | |
| 973 | def _add_mixed_precision_args(parser): |
| 974 | group = parser.add_argument_group(title="mixed precision") |
| 975 | |
| 976 | group.add_argument("--fp16", action="store_true", help="Run model in fp16 mode.") |
| 977 | group.add_argument("--ln-fp16", action="store_true", help="Run layernorm in fp16 mode.") |
| 978 | group.add_argument( |
| 979 | "--bf16", action="store_true", help="Run model in bfloat16 mode." |
| 980 | ) |
| 981 | group.add_argument( |
| 982 | "--loss-scale", |
| 983 | type=float, |
| 984 | default=None, |
| 985 | help="Static loss scaling, positive power of 2 " |
| 986 | "values can improve fp16 convergence. If None, dynamic" |
| 987 | "loss scaling is used.", |
| 988 | ) |
| 989 | group.add_argument( |
| 990 | "--initial-loss-scale", |
| 991 | type=float, |
| 992 | default=2 ** 32, |
| 993 | help="Initial loss-scale for dynamic loss scaling.", |
| 994 | ) |
| 995 | group.add_argument( |
| 996 | "--min-loss-scale", |
| 997 | type=float, |
| 998 | default=1.0, |
| 999 | help="Minimum loss scale for dynamic loss scale.", |
| 1000 | ) |
| 1001 | group.add_argument( |
| 1002 | "--loss-scale-window", |
| 1003 | type=float, |
| 1004 | default=1000, |
| 1005 | help="Window over which to raise/lower dynamic scale.", |
| 1006 | ) |
| 1007 | group.add_argument( |
| 1008 | "--hysteresis", type=int, default=2, help="hysteresis for dynamic loss scaling" |
| 1009 | ) |
| 1010 | group.add_argument( |
| 1011 | "--fp32-residual-connection", |
| 1012 | action="store_true", |
| 1013 | help="Move residual connections to fp32.", |
| 1014 | ) |
| 1015 | group.add_argument('--apply-query-key-layer-scaling', action='store_true', |
| 1016 | help='Scale Q * K^T by 1 / layer-number. If this flag ' |
| 1017 | 'is set, then it will automatically set ' |
| 1018 | 'attention-softmax-in-fp32 to true') |
| 1019 | group.add_argument( |
| 1020 | "--attention-softmax-in-fp32", |
| 1021 | action="store_true", |
| 1022 | help="Run attention masking and softmax in fp32. " |
| 1023 | "This flag is ignored unless " |
| 1024 | "--no-query-key-layer-scaling is specified.", |
| 1025 | ) |
| 1026 | group.add_argument( |
| 1027 | "--accumulate-allreduce-grads-in-fp32", |
| 1028 | action="store_true", |
| 1029 | help="Gradient accumulation and all-reduce in fp32.", |
| 1030 | ) |