| 811 | LLAMA_START_DOCSTRING, |
| 812 | ) |
| 813 | class LlamaPreTrainedModel(PreTrainedModel): |
| 814 | config_class = LlamaConfig |
| 815 | base_model_prefix = "model" |
| 816 | supports_gradient_checkpointing = True |
| 817 | _no_split_modules = ["LlamaDecoderLayer"] |
| 818 | _skip_keys_device_placement = ["past_key_values"] |
| 819 | _supports_flash_attn_2 = True |
| 820 | _supports_sdpa = True |
| 821 | _supports_cache_class = True |
| 822 | _supports_quantized_cache = True |
| 823 | _supports_static_cache = True |
| 824 | |
| 825 | def _init_weights(self, module): |
| 826 | std = self.config.initializer_range |
| 827 | if isinstance(module, nn.Linear): |
| 828 | module.weight.data.normal_(mean=0.0, std=std) |
| 829 | if module.bias is not None: |
| 830 | module.bias.data.zero_() |
| 831 | elif isinstance(module, nn.Embedding): |
| 832 | module.weight.data.normal_(mean=0.0, std=std) |
| 833 | if module.padding_idx is not None: |
| 834 | module.weight.data[module.padding_idx].zero_() |
| 835 | |
| 836 | |
| 837 | LLAMA_INPUTS_DOCSTRING = r""" |
nothing calls this directly
no outgoing calls
no test coverage detected