Build the model.
(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True)
| 182 | verbose=True, write_to_tensorboard=not args.skip_train) |
| 183 | |
| 184 | def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True): |
| 185 | """Build the model.""" |
| 186 | args = get_args() |
| 187 | args.model_type = model_type |
| 188 | |
| 189 | # Build model. |
| 190 | if mpu.get_pipeline_model_parallel_world_size() > 1 and \ |
| 191 | args.virtual_pipeline_model_parallel_size is not None: |
| 192 | assert model_type != ModelType.encoder_and_decoder, \ |
| 193 | "Interleaved schedule not supported for model with both encoder and decoder" |
| 194 | model = [] |
| 195 | for i in range(args.virtual_pipeline_model_parallel_size): |
| 196 | mpu.set_virtual_pipeline_model_parallel_rank(i) |
| 197 | # Set pre_process and post_process only after virtual rank is set. |
| 198 | pre_process = mpu.is_pipeline_first_stage() |
| 199 | post_process = mpu.is_pipeline_last_stage() |
| 200 | this_model = model_provider_func( |
| 201 | pre_process=pre_process, |
| 202 | post_process=post_process |
| 203 | ) |
| 204 | this_model.model_type = model_type |
| 205 | model.append(this_model) |
| 206 | else: |
| 207 | pre_process = mpu.is_pipeline_first_stage() |
| 208 | post_process = mpu.is_pipeline_last_stage() |
| 209 | add_encoder = True |
| 210 | add_decoder = True |
| 211 | if model_type == ModelType.encoder_and_decoder: |
| 212 | if mpu.get_pipeline_model_parallel_world_size() > 1: |
| 213 | assert args.pipeline_model_parallel_split_rank is not None, \ |
| 214 | "Split rank needs to be specified for model with both encoder and decoder" |
| 215 | rank = mpu.get_pipeline_model_parallel_rank() |
| 216 | split_rank = args.pipeline_model_parallel_split_rank |
| 217 | world_size = mpu.get_pipeline_model_parallel_world_size() |
| 218 | pre_process = rank == 0 or rank == split_rank |
| 219 | post_process = (rank == (split_rank - 1)) or ( |
| 220 | rank == (world_size - 1)) |
| 221 | add_encoder = mpu.is_pipeline_stage_before_split() |
| 222 | add_decoder = mpu.is_pipeline_stage_after_split() |
| 223 | model = model_provider_func( |
| 224 | pre_process=pre_process, |
| 225 | post_process=post_process, |
| 226 | add_encoder=add_encoder, |
| 227 | add_decoder=add_decoder) |
| 228 | else: |
| 229 | model = model_provider_func( |
| 230 | pre_process=pre_process, |
| 231 | post_process=post_process |
| 232 | ) |
| 233 | model.model_type = model_type |
| 234 | |
| 235 | if not isinstance(model, list): |
| 236 | model = [model] |
| 237 | |
| 238 | # Disallow training and inference with Transformer Engine |
| 239 | # for non-GPT models |
| 240 | #args.allow_transformer_engine = all([type(m) == GPTModel for m in model]) |
| 241 | args.allow_transformer_engine = True |