| 213 | |
| 214 | @classmethod |
| 215 | def from_pretrained(cls, name, args=None, *, home_path=None, url=None, prefix='', build_only=False, use_node_group=True, overwrite_args={}, **kwargs): |
| 216 | if build_only or 'model_parallel_size' not in overwrite_args: |
| 217 | return cls.from_pretrained_base(name, args=args, home_path=home_path, url=url, prefix=prefix, build_only=build_only, overwrite_args=overwrite_args, **kwargs) |
| 218 | else: |
| 219 | new_model_parallel_size = overwrite_args['model_parallel_size'] |
| 220 | if new_model_parallel_size != 1 or new_model_parallel_size == 1 and args.model_parallel_size == 1: |
| 221 | model, model_args = cls.from_pretrained_base(name, args=args, home_path=home_path, url=url, prefix=prefix, build_only=True, overwrite_args=overwrite_args, **kwargs) |
| 222 | local_rank = get_node_rank() if use_node_group else get_model_parallel_rank() |
| 223 | world_size = torch.distributed.get_world_size() |
| 224 | assert world_size % new_model_parallel_size == 0, "world size should be a multiplier of new model_parallel_size." |
| 225 | destroy_model_parallel() |
| 226 | initialize_model_parallel(1) |
| 227 | if local_rank == 0: |
| 228 | args.skip_init = True |
| 229 | args.use_gpu_initialization = False |
| 230 | args.device = 'cpu' |
| 231 | overwrite_args.pop('model_parallel_size') |
| 232 | model_full, args_ = cls.from_pretrained_base(name, args=args, home_path=home_path, url=url, prefix=prefix, build_only=False, overwrite_args=overwrite_args, **kwargs) |
| 233 | if args_.model_parallel_size != 1: |
| 234 | raise Exception("We do not support overwriting model_parallel_size when original model_parallel_size != 1. Try merging the model using `from_pretrained(xxx,overwrite_args={'model_parallel_size':1})` first if you still want to change model_parallel_size!") |
| 235 | if hasattr(args, 'mode') and args.mode == 'inference': # For multi-node inference, we should prevent rank 0 eagerly printing some info. |
| 236 | torch.distributed.barrier() |
| 237 | destroy_model_parallel() |
| 238 | initialize_model_parallel(new_model_parallel_size) |
| 239 | if local_rank == 0: |
| 240 | mp_split_model_rank0(model, model_full, use_node_group=use_node_group) |
| 241 | del model_full |
| 242 | else: |
| 243 | mp_split_model_receive(model, use_node_group=use_node_group) |
| 244 | reset_random_seed(6) |
| 245 | else: |
| 246 | overwrite_args.pop('model_parallel_size') |
| 247 | model, model_args = cls.from_pretrained_base(name, args=args, home_path=home_path, url=url, prefix=prefix, build_only=False, overwrite_args=overwrite_args, **kwargs) |
| 248 | rank = torch.distributed.get_rank() |
| 249 | world_size = torch.distributed.get_world_size() |
| 250 | assert world_size == model_args.model_parallel_size, "world size should be equal to model_parallel_size." |
| 251 | destroy_model_parallel() |
| 252 | initialize_model_parallel(1) |
| 253 | if rank == 0: |
| 254 | args.use_gpu_initialization = False |
| 255 | args.device = 'cpu' |
| 256 | overwrite_args['model_parallel_size'] = 1 |
| 257 | model_full, args_ = cls.from_pretrained_base(name, args=args, home_path=home_path, url=url, prefix=prefix, build_only=True, overwrite_args=overwrite_args, **kwargs) |
| 258 | torch.distributed.barrier() |
| 259 | destroy_model_parallel() |
| 260 | initialize_model_parallel(model_args.model_parallel_size) |
| 261 | if rank == 0: |
| 262 | mp_merge_model_rank0(model, model_full) |
| 263 | model, model_args = model_full, args_ |
| 264 | else: |
| 265 | mp_merge_model_send(model) |
| 266 | model_args.model_parallel_size = 1 |
| 267 | destroy_model_parallel() |
| 268 | initialize_model_parallel(1) |
| 269 | return model, model_args |
| 270 | |
| 271 | @classmethod |
| 272 | def list_avail_args(cls, print=True): |