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.

• rope_scaling (bool, optional) – Toggle RoPE scaling.

• rope_scaling_factor (float) – RoPE scaling factor. Default 8.

• scatter_embedding_sequence_parallel (bool, optional) – Whether embeddings should be scattered across sequence parallel region or not. Defaults to True.

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

• pg_collection (ProcessGroupCollection) – Model communication process groups

build_schedule_plan(input_ids: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, decoder_input: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, extra_block_kwargs: Optional[dict] = None, runtime_gather_output: Optional[bool] = None, inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, loss_mask: Optional[torch.Tensor] = None)

Builds a computation schedule plan for the model.

This function creates a schedule plan for a model chunk, including preprocessing, transformer layers, and postprocessing. The schedule plan is used to optimize computation and memory usage in distributed environments.

Parameters

• input_ids (Tensor) – Input token IDs.

• position_ids (Tensor) – Position IDs.

• attention_mask (Tensor) – Attention mask.

• decoder_input (Tensor, optional) – Decoder input tensor. Defaults to None.

• labels (Tensor, optional) – Labels for loss computation. Defaults to None.

• inference_context (BaseInferenceContext, optional) – Inference context. Defaults to None.

• packed_seq_params (PackedSeqParams, optional) – Parameters for packed sequences. Defaults to None.

• extra_block_kwargs (dict, optional) – Additional keyword arguments for blocks. Defaults to None.

• runtime_gather_output (Optional[bool], optional) – Whether to gather output at runtime. Defaults to None.

• inference_params (InferenceParams, optional) – Parameters for inference. Defaults to None.

• loss_mask (Optional[Tensor], optional) – Loss mask. Defaults to None.

Returns

The model chunk schedule plan.

Return type

TransformerModelChunkSchedulePlan

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_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, extra_block_kwargs: Optional[dict] = None, runtime_gather_output: Optional[bool] = None, *, inference_params: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, loss_mask: Optional[torch.Tensor] = 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

Parameters

runtime_gather_output (bool) – Gather output at runtime. Default None means parallel_output arg in the constructor will be used.

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 = (), metadata: Optional[Dict] = None) → megatron.core.dist_checkpointing.mapping.ShardedStateDict

Sharded state dict implementation for GPTModel backward-compatibility.

Removing extra state. Tie word embeddings and output layer in mtp process stage.

Parameters

• prefix (str) – Module name prefix.

• sharded_offsets (tuple) – PP related offsets, expected to be empty at this module level.

• metadata (Optional[Dict]) – metadata controlling sharded state dict creation.

Returns

sharded state dict for the GPTModel

Return type

ShardedStateDict

shared_embedding_or_output_weight() → torch.Tensor

Gets the embedding weight or output logit weights when share input embedding and output weights set to True or when use Multi-Token Prediction (MTP) feature.

Returns

During pre processing or MTP process it returns the input embeddings weight. Otherwise, during post processing it returns the final output layers weight.

Return type

Tensor