Skip to main content

Ctrl+K

Megatron-LLM 0.1.0 documentation

User guide
API

User guide
API

Section Navigation

megatron.arguments
megatron.checkpointing
megatron.dist_signal_handler
megatron.global_vars
megatron.indexer
- megatron.indexer.IndexBuilder
megatron.initialize
megatron.memory
megatron.microbatches
megatron.optimizer_param_scheduler
- megatron.optimizer_param_scheduler.OptimizerParamScheduler
megatron.p2p_communication
megatron.schedules
megatron.text_generation_server
- megatron.text_generation_server.MegatronGenerate
- megatron.text_generation_server.MegatronServer
megatron.timers
megatron.training
megatron.utils
megatron.wandb_logger

megatron.core.parallel_state
megatron.core.utils

megatron.core.tensor_parallel.cross_entropy
- megatron.core.tensor_parallel.cross_entropy.vocab_parallel_cross_entropy
- megatron.core.tensor_parallel.cross_entropy.vocab_parallel_max_indices
megatron.core.tensor_parallel.data
- megatron.core.tensor_parallel.data.broadcast_data
megatron.core.tensor_parallel.layers
megatron.core.tensor_parallel.mappings
megatron.core.tensor_parallel.random
megatron.core.tensor_parallel.utils

megatron.data.autoaugment
- megatron.data.autoaugment.ImageNetPolicy
- megatron.data.autoaugment.SubPolicy
megatron.data.blendable_dataset
- megatron.data.blendable_dataset.BlendableDataset
megatron.data.gpt_dataset
megatron.data.image_folder
megatron.data.realm_dataset_utils
megatron.data.bert_dataset
megatron.data.data_samplers
megatron.data.indexed_dataset
megatron.data.orqa_wiki_dataset
megatron.data.realm_index
megatron.data.biencoder_dataset_utils
megatron.data.dataset_utils
megatron.data.ict_dataset
megatron.data.t5_dataset

megatron.model.bert_model
megatron.model.biencoder_model
megatron.model.classification
- megatron.model.classification.Classification
megatron.model.distributed
megatron.model.enums
megatron.model.falcon_model
- megatron.model.falcon_model.FalconModel
megatron.model.fused_bias_gelu
- megatron.model.fused_bias_gelu.GeLUFunction
megatron.model.fused_layer_norm
megatron.model.fused_softmax
megatron.model.glu_activations
megatron.model.gpt_model
- megatron.model.gpt_model.GPTModel
- megatron.model.gpt_model.post_language_model_processing
megatron.model.language_model
megatron.model.llama_model
- megatron.model.llama_model.LlamaModel
megatron.model.module
megatron.model.multiple_choice
- megatron.model.multiple_choice.MultipleChoice
megatron.model.positional_embeddings
megatron.model.t5_model
megatron.model.transformer
megatron.model.utils

megatron.optimizer.clip_grads
- megatron.optimizer.clip_grads.clip_grad_norm_fp32
- megatron.optimizer.clip_grads.count_zeros_fp32
megatron.optimizer.distrib_optimizer
- megatron.optimizer.distrib_optimizer.DistributedOptimizer
- megatron.optimizer.distrib_optimizer.Range
megatron.optimizer.grad_scaler
megatron.optimizer.optimizer

megatron.text_generation.api
megatron.text_generation.beam_utils
- megatron.text_generation.beam_utils.BeamHypotheses
megatron.text_generation.communication
megatron.text_generation.forward_step
- megatron.text_generation.forward_step.ForwardStep
- megatron.text_generation.forward_step.InferenceParams
megatron.text_generation.generation
megatron.text_generation.sampling
megatron.text_generation.tokenization
- megatron.text_generation.tokenization.detokenize_generations
- megatron.text_generation.tokenization.tokenize_prompts

megatron.tokenizer.bert_tokenization
megatron.tokenizer.gpt2_tokenization
megatron.tokenizer.tokenizer
- megatron.tokenizer.tokenizer.AbstractTokenizer
- megatron.tokenizer.tokenizer.build_tokenizer

API
megatron.schedules
megatron.sch...

megatron.schedules.get_tensor_shapes#

megatron.schedules.get_tensor_shapes(rank, model_type: ModelType)#

previous

megatron.schedules.get_forward_backward_func

next

megatron.schedules.recv_backward

On this page

get_tensor_shapes()

© Copyright 2023, Alejandro Hernández Cano, Matteo Pagliardini, Kyle Matoba, Amirkeivan Mohtashami, Olivia Simin Fan, Axel Marmet, Deniz Bayazit, Igor Krawczuk, Zeming Chen, Francesco Salvi, Antoine Bosselut, Martin Jaggi.

Created using Sphinx 7.1.0.

Built with the PyData Sphinx Theme 0.15.2.