| 1038 | |
| 1039 | |
| 1040 | def _add_distributed_args(parser): |
| 1041 | group = parser.add_argument_group(title="distributed") |
| 1042 | |
| 1043 | group.add_argument( |
| 1044 | "--tensor-model-parallel-size", |
| 1045 | type=int, |
| 1046 | default=1, |
| 1047 | help="Degree of tensor model parallelism.", |
| 1048 | ) |
| 1049 | group.add_argument( |
| 1050 | "--pipeline-model-parallel-size", |
| 1051 | type=int, |
| 1052 | default=1, |
| 1053 | help="Degree of pipeline model parallelism.", |
| 1054 | ) |
| 1055 | group.add_argument( |
| 1056 | "--model-parallel-size", |
| 1057 | type=int, |
| 1058 | default=None, |
| 1059 | help="Old model parallel argument, do not use. Use " |
| 1060 | "--tensor-model-parallel-size instead.", |
| 1061 | ) |
| 1062 | group.add_argument( |
| 1063 | "--num-layers-per-virtual-pipeline-stage", |
| 1064 | type=int, |
| 1065 | default=None, |
| 1066 | help="Number of layers per virtual pipeline stage", |
| 1067 | ) |
| 1068 | group.add_argument( |
| 1069 | "--distributed-backend", |
| 1070 | default="nccl", |
| 1071 | choices=["nccl", "gloo"], |
| 1072 | help="Which backend to use for distributed training.", |
| 1073 | ) |
| 1074 | group.add_argument( |
| 1075 | "--DDP-impl", |
| 1076 | default="local", |
| 1077 | choices=["local", "torch"], |
| 1078 | help="which DistributedDataParallel implementation " "to use.", |
| 1079 | ) |
| 1080 | group.add_argument( |
| 1081 | "--use-contiguous-buffers-in-ddp", |
| 1082 | action="store_true", |
| 1083 | help="If set, use contiguous buffer in DDP. Note that " |
| 1084 | "this option only works woth local DDP.", |
| 1085 | ) |
| 1086 | group.add_argument( |
| 1087 | "--no-scatter-gather-tensors-in-pipeline", |
| 1088 | action="store_false", |
| 1089 | help="Use scatter/gather to optimize communication of tensors in pipeline", |
| 1090 | dest="scatter_gather_tensors_in_pipeline", |
| 1091 | ) |
| 1092 | group.add_argument( |
| 1093 | "--local_rank", |
| 1094 | type=int, |
| 1095 | default=None, |
| 1096 | help="local rank passed from distributed launcher.", |
| 1097 | ) |