Build the model.
(model_provider_func)
| 248 | |
| 249 | |
| 250 | def get_model(model_provider_func): |
| 251 | """Build the model.""" |
| 252 | args = get_args() |
| 253 | |
| 254 | # Build model. |
| 255 | if ( |
| 256 | mpu.get_pipeline_model_parallel_world_size() > 1 |
| 257 | and args.virtual_pipeline_model_parallel_size is not None |
| 258 | ): |
| 259 | model = [] |
| 260 | for i in range(args.virtual_pipeline_model_parallel_size): |
| 261 | mpu.set_virtual_pipeline_model_parallel_rank(i) |
| 262 | # Set pre_process and post_process only after virtual rank is set. |
| 263 | pre_process = mpu.is_pipeline_first_stage() |
| 264 | post_process = mpu.is_pipeline_last_stage() |
| 265 | this_model = model_provider_func( |
| 266 | pre_process=pre_process, post_process=post_process |
| 267 | ) |
| 268 | model.append(this_model) |
| 269 | else: |
| 270 | pre_process = mpu.is_pipeline_first_stage() |
| 271 | post_process = mpu.is_pipeline_last_stage() |
| 272 | model = model_provider_func(pre_process=pre_process, post_process=post_process) |
| 273 | |
| 274 | if not isinstance(model, list): |
| 275 | model = [model] |
| 276 | |
| 277 | # Set tensor model parallel attributes if not set. |
| 278 | # Only parameters that are already tensor model parallel have these |
| 279 | # attributes set for them. We should make sure the default attributes |
| 280 | # are set for all params so the optimizer can use them. |
| 281 | for model_module in model: |
| 282 | for param in model_module.parameters(): |
| 283 | mpu.set_defaults_if_not_set_tensor_model_parallel_attributes(param) |
| 284 | |
| 285 | # Print number of parameters. |
| 286 | if mpu.get_data_parallel_rank() == 0: |
| 287 | print( |
| 288 | " > number of parameters on (tensor, pipeline) " |
| 289 | "model parallel rank ({}, {}): {}".format( |
| 290 | mpu.get_tensor_model_parallel_rank(), |
| 291 | mpu.get_pipeline_model_parallel_rank(), |
| 292 | sum( |
| 293 | [ |
| 294 | sum( |
| 295 | [ |
| 296 | p.ds_numel if hasattr(p, "ds_id") else p.nelement() |
| 297 | for p in model_module.parameters() |
| 298 | ] |
| 299 | ) |
| 300 | for model_module in model |
| 301 | ] |
| 302 | ), |
| 303 | ), |
| 304 | flush=True, |
| 305 | ) |
| 306 | |
| 307 | if args.deepspeed: |