(
self,
config: TransformerConfig,
transformer_layer_spec: ModuleSpec,
vocab_size: int,
max_sequence_length: int,
pre_process: bool = True,
post_process: bool = True,
fp16_lm_cross_entropy: bool = False,
parallel_output: bool = True,
share_embeddings_and_output_weights: bool = False,
position_embedding_type: Literal['learned_absolute', 'rope'] = 'learned_absolute',
rotary_percent: float = 1.0,
rotary_base: int = 10000,
seq_len_interpolation_factor: Optional[float] = None,
)
| 47 | """ |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | config: TransformerConfig, |
| 52 | transformer_layer_spec: ModuleSpec, |
| 53 | vocab_size: int, |
| 54 | max_sequence_length: int, |
| 55 | pre_process: bool = True, |
| 56 | post_process: bool = True, |
| 57 | fp16_lm_cross_entropy: bool = False, |
| 58 | parallel_output: bool = True, |
| 59 | share_embeddings_and_output_weights: bool = False, |
| 60 | position_embedding_type: Literal['learned_absolute', 'rope'] = 'learned_absolute', |
| 61 | rotary_percent: float = 1.0, |
| 62 | rotary_base: int = 10000, |
| 63 | seq_len_interpolation_factor: Optional[float] = None, |
| 64 | ) -> None: |
| 65 | super().__init__(config=config) |
| 66 | |
| 67 | self.transformer_layer_spec: ModuleSpec = transformer_layer_spec |
| 68 | self.vocab_size = vocab_size |
| 69 | self.max_sequence_length = max_sequence_length |
| 70 | self.pre_process = pre_process |
| 71 | self.post_process = post_process |
| 72 | self.fp16_lm_cross_entropy = fp16_lm_cross_entropy |
| 73 | self.parallel_output = parallel_output |
| 74 | self.share_embeddings_and_output_weights = share_embeddings_and_output_weights |
| 75 | self.position_embedding_type = position_embedding_type |
| 76 | |
| 77 | # megatron core pipelining currently depends on model type |
| 78 | # TODO: remove this dependency ? |
| 79 | self.model_type = ModelType.encoder_or_decoder |
| 80 | |
| 81 | # These 2 attributes are needed for TensorRT-LLM export. |
| 82 | self.max_position_embeddings = max_sequence_length |
| 83 | self.rotary_percent = rotary_percent |
| 84 | |
| 85 | if self.pre_process: |
| 86 | self.embedding = LanguageModelEmbedding( |
| 87 | config=self.config, |
| 88 | vocab_size=self.vocab_size, |
| 89 | max_sequence_length=self.max_sequence_length, |
| 90 | position_embedding_type=position_embedding_type, |
| 91 | ) |
| 92 | |
| 93 | if self.position_embedding_type == 'rope': |
| 94 | self.rotary_pos_emb = RotaryEmbedding( |
| 95 | kv_channels=self.config.kv_channels, |
| 96 | rotary_percent=rotary_percent, |
| 97 | rotary_interleaved=self.config.rotary_interleaved, |
| 98 | seq_len_interpolation_factor=seq_len_interpolation_factor, |
| 99 | rotary_base=rotary_base, |
| 100 | ) |
| 101 | |
| 102 | # Transformer. |
| 103 | self.decoder = TransformerBlock( |
| 104 | config=self.config, |
| 105 | spec=transformer_layer_spec, |
| 106 | pre_process=self.pre_process, |
nothing calls this directly
no test coverage detected