(self, args, transformer=None, params_dtype=torch.float, **kwargs)
| 78 | |
| 79 | class BaseModel(torch.nn.Module, metaclass=MetaModel): |
| 80 | def __init__(self, args, transformer=None, params_dtype=torch.float, **kwargs): |
| 81 | super(BaseModel, self).__init__() |
| 82 | self.mixins = torch.nn.ModuleDict() |
| 83 | self.collect_hooks_() |
| 84 | if transformer is not None: |
| 85 | self.transformer = transformer |
| 86 | else: |
| 87 | # check if model-only mode |
| 88 | from sat.arguments import _simple_init |
| 89 | success = _simple_init(model_parallel_size=args.model_parallel_size, seed=args.seed if hasattr(args, 'seed') else 1234) |
| 90 | |
| 91 | args_dict = {k: (getattr(args, v[0]) if hasattr(args, v[0]) else v[1]) for k, v in ARGS_DEFAULT.items()} |
| 92 | |
| 93 | self.transformer = BaseTransformer( |
| 94 | num_layers=args.num_layers, |
| 95 | vocab_size=args.vocab_size, |
| 96 | hidden_size=args.hidden_size, |
| 97 | num_attention_heads=args.num_attention_heads, |
| 98 | max_sequence_length=args.max_sequence_length, |
| 99 | layernorm_order=args.layernorm_order, |
| 100 | **args_dict, |
| 101 | hooks=self.hooks, |
| 102 | params_dtype=params_dtype, |
| 103 | skip_init=args.skip_init, |
| 104 | device=torch.cuda.current_device() if hasattr(args, 'use_gpu_initialization') and args.use_gpu_initialization else torch.device('cpu'), |
| 105 | **kwargs |
| 106 | ) |
| 107 | |
| 108 | def reinit(self, mixin_names=None): # will be called when loading model, None means all |
| 109 | # if some mixins are loaded, overrides this function |
nothing calls this directly
no test coverage detected