models.gpt package

This is the implementation of the popular GPT model. It supports several features like model parallelization (Tensor Parallel, Pipeline Parallel, Data Parallel) , mixture of experts, FP8 , Distributed optimizer etc. We are constantly adding new features. So be on the lookout or raise an issue if you want to have something added.

class core.models.gpt.gpt_model.GPTModel(*args: Any, **kwargs: Any)

Bases: megatron.core.models.common.language_module.language_module.LanguageModule

GPT Transformer language model.

Parameters

• config (TransformerConfig) – Transformer config

• transformer_layer_spec (ModuleSpec) – Specifies module to use for transformer layers

• vocab_size (int) – Vocabulary size

• max_sequence_length (int) – maximum size of sequence. This is used for positional embedding

• pre_process (bool, optional) – Include embedding layer (used with pipeline parallelism). Defaults to True.

• post_process (bool, optional) – Include an output layer (used with pipeline parallelism). Defaults to True.

• fp16_lm_cross_entropy (bool, optional) – Defaults to False.

• parallel_output (bool, optional) – Do not gather the outputs, keep them split across tensor parallel ranks. Defaults to True.

• share_embeddings_and_output_weights (bool, optional) – When True, input embeddings and output logit weights are shared. Defaults to False.

• position_embedding_type (Literal[learned_absolute,rope], optional) – Position embedding type.. Defaults to ‘learned_absolute’.

• 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.

• rotary_base (int, optional) – Base period for rotary position embeddings. Ignored unless position_embedding_type is ‘rope’. Defaults to 10000.

• 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.

forward(input_ids: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, decoder_input: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, inference_params: Optional[megatron.core.InferenceParams] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, extra_block_kwargs: Optional[dict] = None) → torch.Tensor

Forward function of the GPT Model This function passes the input tensors through the embedding layer, and then the decoeder and finally into the post processing layer (optional).

It either returns the Loss values if labels are given or the final hidden units

set_input_tensor(input_tensor: torch.Tensor) → None

Sets input tensor to the model.

See megatron.model.transformer.set_input_tensor()

Parameters

input_tensor (Tensor) – Sets the input tensor for the model.

sharded_state_dict(prefix: str = '', sharded_offsets: tuple = ()) → megatron.core.dist_checkpointing.mapping.ShardedStateDict