| 77 | |
| 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 |
| 110 | for k, m in self.mixins.items(): |
| 111 | if mixin_names is None or k in mixin_names: |
| 112 | m.reinit(self) |
| 113 | |
| 114 | def add_mixin(self, name, new_mixin, reinit=False): |
| 115 | assert name not in self.mixins |
| 116 | assert isinstance(new_mixin, BaseMixin) |
| 117 | |
| 118 | self.mixins[name] = new_mixin # will auto-register parameters |
| 119 | object.__setattr__(new_mixin, 'transformer', self.transformer) # cannot use pytorch set_attr |
| 120 | |
| 121 | self.collect_hooks_() |
| 122 | if reinit: |
| 123 | new_mixin.reinit(self) # also pass current mixins |
| 124 | |
| 125 | def del_mixin(self, name): |
| 126 | assert name in self.mixins |
| 127 | del self.mixins[name] |
| 128 | self.collect_hooks_() |
| 129 | |
| 130 | def get_mixin(self, name): |
| 131 | return self.mixins[name] |
| 132 | |
| 133 | def forward(self, *args, **kwargs): |
| 134 | # update hooks as the current model (overrided forwards) |
| 135 | # Attention! the transformer might be shared by multiple models |
| 136 | self.transformer.hooks.clear() |
no outgoing calls