Build the model.
(args, model_cls, **kwargs)
| 394 | return model, model_args |
| 395 | |
| 396 | def get_model(args, model_cls, **kwargs): |
| 397 | """Build the model.""" |
| 398 | import torch |
| 399 | from sat.helpers import print_rank0,print_all |
| 400 | from sat import mpu |
| 401 | |
| 402 | print_rank0(f'building {model_cls.__name__} model ...') |
| 403 | if 'params_dtype' not in kwargs: |
| 404 | if hasattr(args, 'fp16') and args.fp16: |
| 405 | params_dtype = torch.half |
| 406 | elif hasattr(args, 'bf16') and args.bf16: |
| 407 | params_dtype = torch.bfloat16 |
| 408 | else: |
| 409 | params_dtype = torch.float32 |
| 410 | else: |
| 411 | # pop params_dtype from kwargs |
| 412 | params_dtype = kwargs.pop('params_dtype') |
| 413 | |
| 414 | from sat.helpers import check_if_zero3 |
| 415 | if check_if_zero3(args): |
| 416 | import deepspeed |
| 417 | with deepspeed.zero.Init(): |
| 418 | model = model_cls(args, params_dtype=params_dtype, **kwargs) |
| 419 | else: |
| 420 | model = model_cls(args, params_dtype=params_dtype, **kwargs) |
| 421 | |
| 422 | if mpu.get_data_parallel_rank() == 0: |
| 423 | print_all(' > number of parameters on model parallel rank {}: {}'.format( |
| 424 | mpu.get_model_parallel_rank(), |
| 425 | sum([p.nelement() for p in model.parameters()])), flush=True) |
| 426 | |
| 427 | if hasattr(args, 'fp16') and args.fp16: |
| 428 | model.half() |
| 429 | elif hasattr(args, 'bf16') and args.bf16: |
| 430 | model.bfloat16() |
| 431 | |
| 432 | try: # TODO: is this useful? |
| 433 | if not hasattr(args, 'device'): |
| 434 | args.device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu' |
| 435 | model = model.to(args.device) |
| 436 | except Exception as e: |
| 437 | print_all(e) |
| 438 | |
| 439 | return model |
no test coverage detected