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