Build the model.
(model_provider_func)
| 133 | |
| 134 | |
| 135 | def get_model(model_provider_func): |
| 136 | """Build the model.""" |
| 137 | args = get_args() |
| 138 | |
| 139 | # Build model on cpu. |
| 140 | model = model_provider_func() |
| 141 | |
| 142 | if args.deepspeed: |
| 143 | # DeepSpeed handles CUDA, FP16, and DDP components. |
| 144 | return model |
| 145 | |
| 146 | # GPU allocation. |
| 147 | model.cuda(torch.cuda.current_device()) |
| 148 | |
| 149 | # Fp16 conversion. |
| 150 | if args.fp16: |
| 151 | model = FP16_Module(model) |
| 152 | |
| 153 | # Wrap model for distributed training.""" |
| 154 | if args.DDP_impl == 'torch': |
| 155 | i = torch.cuda.current_device() |
| 156 | model = torchDDP(model, device_ids=[i], output_device=i, |
| 157 | process_group=mpu.get_data_parallel_group()) |
| 158 | return model |
| 159 | if args.DDP_impl == 'local': |
| 160 | model = LocalDDP(model) |
| 161 | return model |
| 162 | |
| 163 | raise NotImplementedError('Unknown DDP implementation specified: {}. ' |
| 164 | 'Exiting.'.format(args.DDP_impl)) |
| 165 | |
| 166 | |
| 167 | def get_optimizer(model): |
no test coverage detected