| 587 | _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.causal_mask"] |
| 588 | |
| 589 | def __init__(self, config): |
| 590 | super().__init__(config) |
| 591 | if not hasattr(config, 'wbits'): |
| 592 | config.wbits = 32 |
| 593 | config.groupsize = 128 |
| 594 | |
| 595 | if config.wbits not in [4, 8, 32]: |
| 596 | logger.warning(f'Specify `wbits` with 4, 8 or 32 to load the model. ') |
| 597 | if config.wbits in [4, 8]: |
| 598 | def noop(*args, **kwargs): |
| 599 | pass |
| 600 | torch.nn.init.kaiming_uniform_ = noop |
| 601 | torch.nn.init.uniform_ = noop |
| 602 | torch.nn.init.normal_ = noop |
| 603 | |
| 604 | torch.set_default_dtype(torch.half) |
| 605 | transformers.modeling_utils._init_weights = False |
| 606 | torch.set_default_dtype(torch.half) |
| 607 | self.transformer = MossModel(config) |
| 608 | self.lm_head = nn.Linear(config.n_embd, config.vocab_size) |
| 609 | if config.wbits in [4, 8]: |
| 610 | torch.set_default_dtype(torch.float) |
| 611 | transformers.modeling_utils._init_weights = True |
| 612 | self.quantize(config.wbits, config.groupsize) |
| 613 | # Initialize weights and apply final processing |
| 614 | self.post_init() |
| 615 | |
| 616 | def get_output_embeddings(self): |
| 617 | return self.lm_head |