configuration_bitnet.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. # Copyright 2025 The BitNet Team and The HuggingFace Inc. team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. """BitNet model configuration"""
  14. from huggingface_hub.dataclasses import strict
  15. from ...configuration_utils import PreTrainedConfig
  16. from ...modeling_rope_utils import RopeParameters
  17. from ...utils import auto_docstring
  18. @auto_docstring(checkpoint="microsoft/bitnet-b1.58-2B-4T")
  19. @strict
  20. class BitNetConfig(PreTrainedConfig):
  21. r"""
  22. ```python
  23. >>> from transformers import BitNetModel, BitNetConfig
  24. >>> # Initializing a BitNet style configuration
  25. >>> configuration = BitNetConfig()
  26. >>> # Initializing a model from the BitNet style configuration
  27. >>> model = BitNetModel(configuration)
  28. >>> # Accessing the model configuration
  29. >>> configuration = model.config
  30. ```
  31. """
  32. model_type = "bitnet"
  33. keys_to_ignore_at_inference = ["past_key_values"]
  34. default_theta = 500000.0
  35. vocab_size: int = 128256
  36. hidden_size: int = 2560
  37. intermediate_size: int = 6912
  38. num_hidden_layers: int = 30
  39. num_attention_heads: int = 20
  40. num_key_value_heads: int | None = 5
  41. hidden_act: str = "relu2"
  42. max_position_embeddings: int = 2048
  43. initializer_range: float = 0.02
  44. rms_norm_eps: float = 1e-5
  45. use_cache: bool = True
  46. pad_token_id: int | None = None
  47. bos_token_id: int | None = 128000
  48. eos_token_id: int | list[int] | None = 128001
  49. tie_word_embeddings: bool = False
  50. attention_bias: bool = False
  51. attention_dropout: float | int | None = 0.0
  52. rope_parameters: RopeParameters | dict | None = None
  53. def __post_init__(self, **kwargs):
  54. # for backward compatibility
  55. if self.num_key_value_heads is None:
  56. self.num_key_value_heads = self.num_attention_heads
  57. super().__post_init__(**kwargs)
  58. __all__ = ["BitNetConfig"]