Mamba language model. Args: config (TransformerConfig): Transformer config mamba_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types vocab_size (int): Vocabulary size max_sequence_length (int): maximum size of sequence. This is used
| 14 | |
| 15 | |
| 16 | class MambaModel(LanguageModule): |
| 17 | """Mamba language model. |
| 18 | |
| 19 | Args: |
| 20 | config (TransformerConfig): Transformer config |
| 21 | mamba_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types |
| 22 | vocab_size (int): Vocabulary size |
| 23 | max_sequence_length (int): maximum size of sequence. This is used for positional embedding |
| 24 | pre_process (bool, optional): Include embedding layer (used with pipeline parallelism). Defaults to True. |
| 25 | mamba_ssm_ngroups (int, optional): Specifies the number of groups to use. The default value is 8, as in the NVIDIA Mamba2 (pure and hybrid) 8b. However, in the original Mamba2 paper, the checkpoints use a setting of 1. Defaults to 8. |
| 26 | hybrid_attention_ratio (float, optional): The target ratio of attention layers to total layers |
| 27 | hybrid_mlp_ratio (float, optional): The target ratio of mlp layers to total layers |
| 28 | hybrid_override_pattern (str, optional): The hybrid layer pattern to override with |
| 29 | post_process (bool, optional): Include an output layer (used with pipeline parallelism). Defaults to True. |
| 30 | fp16_lm_cross_entropy (bool, optional): Defaults to False. |
| 31 | parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True. |
| 32 | share_embeddings_and_output_weights (bool, optional): When True, input embeddings and output logit weights are shared. Defaults to False. |
| 33 | position_embedding_type (Literal[learned_absolute,rope,none], optional): Position embedding type. Defaults to 'none'. |
| 34 | rotary_percent (float, optional): Percent of rotary dimension to use for rotary position embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0. |
| 35 | rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 10000. |
| 36 | seq_len_interpolation_factor (Optional[float], optional): scale of linearly interpolating RoPE for longer sequences. The value must be a float larger than 1.0. Defaults to None. |
| 37 | """ |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | config: TransformerConfig, |
| 42 | mamba_stack_spec: ModuleSpec, |
| 43 | vocab_size: int, |
| 44 | max_sequence_length: int, |
| 45 | mamba_ssm_ngroups: int = 8, |
| 46 | pre_process: bool = True, |
| 47 | hybrid_attention_ratio: float = 0.0, |
| 48 | hybrid_mlp_ratio: float = 0.0, |
| 49 | hybrid_override_pattern: str = None, |
| 50 | post_process: bool = True, |
| 51 | fp16_lm_cross_entropy: bool = False, |
| 52 | parallel_output: bool = True, |
| 53 | share_embeddings_and_output_weights: bool = False, |
| 54 | # Mamba with no attention has no need for position embeddings, so none is default |
| 55 | position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'none', |
| 56 | rotary_percent: float = 1.0, |
| 57 | rotary_base: int = 10000, |
| 58 | seq_len_interpolation_factor: Optional[float] = None, |
| 59 | ) -> None: |
| 60 | super().__init__(config=config) |
| 61 | |
| 62 | self.mamba_stack_spec: ModuleSpec = mamba_stack_spec |
| 63 | self.vocab_size = vocab_size |
| 64 | self.max_sequence_length = max_sequence_length |
| 65 | self.mamba_ssm_ngroups = mamba_ssm_ngroups |
| 66 | self.pre_process = pre_process |
| 67 | self.hybrid_attention_ratio = hybrid_attention_ratio |
| 68 | self.hybrid_mlp_ratio = hybrid_mlp_ratio |
| 69 | self.hybrid_override_pattern = hybrid_override_pattern |
| 70 | self.post_process = post_process |
| 71 | self.fp16_lm_cross_entropy = fp16_lm_cross_entropy |
| 72 | self.parallel_output = parallel_output |
| 73 | self.share_embeddings_and_output_weights = share_embeddings_and_output_weights |
no outgoing calls
no test coverage detected