Megatron-LM

Megatron-LM permite o treinamento de grandes modelos de linguagem transformer em escala. Ele fornece paralelismo de modelo baseado em tensor, pipeline e sequência para o pré-treinamento de Modelos de Linguagem baseados em transformer, como GPT (Apenas Decodificador), BERT (Apenas Codificador) e T5 (Codificador-Decodificador). Para informações detalhadas e como as coisas funcionam nos bastidores, consulte o repositório no github.

O que está integrado?

O Accelerate integra o seguinte recurso do Megatron-LM para habilitar o pré-treinamento/ajuste fino em larga escala de modelos BERT (Encoder), GPT (Decoder) ou T5 (Encoder e Decoder):

a. Paralelismo de Tensors (TP): reduz a pegada de memória sem muita comunicação adicional nos ranks intra-nó. Cada tensor é dividido em vários blocos, com cada shard residindo em uma GPU separada. A cada etapa, o mesmo mini-lote de dados é processado independentemente e em paralelo por cada shard, seguido de sincronização entre todas as GPUs (operação all-reduce). Em uma camada de transformer simples, isso resulta em 2 all-reduces no caminho de forward e 2 no caminho de backward. Para mais detalhes, consulte o artigo de pesquisa Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism e esta seção do blogpost The Technology Behind BLOOM Training.

b. Paralelismo de Pipeline (PP): reduz a pegada de memória e permite treinamento em larga escala por meio da paralelização entre nós. Reduz o "bubble" do PP ingênuo por meio do agendamento PipeDream-Flush/1F1B e do agendamento Interleaved 1F1B. As camadas são distribuídas uniformemente entre os estágios de PP. Por exemplo, se um modelo tem 24 camadas e temos 4 GPUs para paralelismo de pipeline, cada GPU terá 6 camadas (24/4). Para mais detalhes sobre agendamentos para reduzir o tempo ocioso do PP, consulte o artigo de pesquisa Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM e esta seção do post de blog The Technology Behind BLOOM Training.

c. Paralelismo de Sequência (SP): reduz a pegada de memória sem nenhuma comunicação adicional. Aplicável apenas ao usar TP. Ele reduz a memória de ativação necessária, pois impede que as mesmas cópias fiquem nos ranks de paralelismo de tensor após o all-reduce, substituindo-as por reduce-scatter, e a operação no-op seria substituída por all-gather. Como all-reduce = reduce-scatter + all-gather, isso economiza uma grande quantidade de memória de ativação sem custo adicional de comunicação. Em termos simples, ele divide as saídas de cada camada do transformer ao longo da dimensão da sequência, por exemplo, se o comprimento da sequência for 1024 e o tamanho do TP for 4, cada GPU terá 256 tokens (1024/4) para cada amostra. Isso aumenta o tamanho do lote que pode ser suportado para o treinamento. Para mais detalhes, consulte o artigo de pesquisa Reducing Activation Recomputation in Large Transformer Models.

d. Paralelismo de Dados (DP) via Otimizador Distribuído: reduz a pegada de memória ao dividir os estados do otimizador e os gradientes entre os ranks de DP (em oposição ao método tradicional de replicar o estado do otimizador entre os ranks de paralelismo de dados). Por exemplo, ao usar o otimizador Adam com treinamento em precisão mista, cada parâmetro consome 12 bytes de memória. Isso é distribuído igualmente entre as GPUs, ou seja, cada parâmetro consome 3 bytes (12/4) se tivermos 4 GPUs. Para mais detalhes, consulte o artigo de pesquisa ZeRO: Memory Optimizations Toward Training Trillion Parameter Models e a seção seguinte do blog The Technology Behind BLOOM Training.

e. Expert Parallelism (EP) O paralelismo de especialistas no Megatron-LM é usado para camadas Mixture-of-Experts (MoE), onde existem muitos “especialistas” (pequenas redes feed-forward), mas apenas alguns são ativados para cada token. Em vez de colocar todos os especialistas em cada GPU, o Megatron distribui diferentes especialistas em diferentes GPUs — esse é o paralelismo de especialistas. Durante o treinamento, os tokens são roteados para as GPUs que hospedam seus especialistas selecionados, são computados lá e, em seguida, enviados de volta, reduzindo o custo de memória. Frequentemente, ele é combinado com paralelismo de tensor/pipeline para modelos em larga escala. f. Full Activation Recomputation: Reduz significativamente a pegada de memória das ativações por meio de checkpointing inteligente de ativações. Ele não armazena ativações que ocupam grande quantidade de memória, mas que são rápidas de recomputar, alcançando assim um ótimo compromisso entre memória e recomputação. Por exemplo, para o GPT-3, isso resulta em uma redução de 70% na memória necessária para ativações, ao custo de apenas 2,7% de sobrecarga de FLOPs para a recomputação das ativações. Para mais detalhes, consulte o artigo de pesquisa Reducing Activation Recomputation in Large Transformer Models.

g. Kernels Fundidos: Softmax fundido, Layer Norm fundido em precisão mista e acúmulo de gradientes fundido ao cálculo do gradiente de pesos da camada linear. Adição fundida de GeLU compilada pelo JIT do PyTorch e de Bias+Dropout+Residual.

h. Suporte a conjuntos de dados indexados: formato binário eficiente de conjuntos de dados para treinamento em larga escala. Suporte ao arquivo de índice mmap, cached e ao formato de carregador lazy.

i. Redimensionamento de checkpoints e interoperabilidade: Utilitário para redimensionar checkpoints do Megatron-LM com tamanhos variáveis de paralelismo de tensor e pipeline para os checkpoints fragmentados do Transformers, amplamente suportados por uma variedade de ferramentas, como Accelerate Big Model Inference, Megatron-DeepSpeed Inference, entre outras. Também há suporte para converter checkpoints fragmentados do Transformers em checkpoints do Megatron-LM com tamanhos variáveis de paralelismo de tensor e pipeline para treinamento em larga escala.

Pré-requisitos

Você precisará instalar as versões mais recentes de pytorch, cuda, nccl e NVIDIA APEX, além da biblioteca nltk. Consulte a documentação para mais detalhes. Outra forma de configurar o ambiente é obter um Container NVIDIA PyTorch que já vem com todas as instalações necessárias do NGC.

Abaixo está um método passo a passo para configurar o ambiente conda:

1. Crie um ambiente virtual

conda create --name ml

2. Pressupondo que a máquina tenha o CUDA 11.3 instalado, a instalação da versão GPU correspondente do PyTorch

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

3. Instalar Nvidia APEX

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
cd ..

4. Instalando o Megatron-LM

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout 9a1c0d05c992c8a241da384ab27dce2021bb56dd
you need to manually move gpt_builders.py to megatron/training and update
include = [
    "megatron.core", 
    "megatron.core.*",
    "megatron.training",
    "megatron.training.*",
    "megatron.legacy",
    "megatron.legacy.*",
]
in pyproject.toml file to unblock yourself from using Megatron
pip install --no-use-pep517 -e .

Preparar checkpoint do Megaton-LM

Se você deseja fazer o ajuste fino de um modelo, certifique-se de ter um checkpoint no formato torch dist pronto. Se você tiver acesso apenas ao modelo no huggingface, considere convertê-lo para um checkpoint no formato torch dist aceitável pelo Megatron. Um exemplo é usar o script do slime, tomando os modelos GLM como exemplo:

source /your/path/to/slime/scripts/models/glm4.5-355B-A32B.sh
srun torchrun --nproc-per-node 8 \
   /your/path/to/slime/tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /your/path/to/huggingface/models/GLM4.5-355B-A32B \
    --save /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist

Após a conversão, verifique: 1. em /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist: altere o conteúdo de latest_checkpointed_iteration.txt de release para 0 e renomeie o diretório release para iter_0000000; 2: na configuração, certifique-se de que megatron_lm_no_load_optim esteja definido como true, para que nenhum estado de otimizador seja necessário.

Plugin do Accelerate para Megatron-LM

Recursos importantes são suportados diretamente por meio do comando accelerate config. Um exemplo das perguntas correspondentes para o uso de recursos do Megatron-LM é mostrado abaixo:

:~$ accelerate config --config_file "megatron_gpt_config.yaml"
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU): 2
How many different machines will you use (use more than 1 for multi-node training)? [1]: 
Do you want to use DeepSpeed? [yes/NO]: 
Do you want to use FullyShardedDataParallel? [yes/NO]: 
Do you want to use Megatron-LM ? [yes/NO]: yes
What is the Tensor Parallelism degree/size? [1]:2
Do you want to enable Sequence Parallelism? [YES/no]: 
What is the Pipeline Parallelism degree/size? [1]:2
What is the number of micro-batches? [1]:2
Do you want to enable selective activation recomputation? [YES/no]: 
Do you want to use distributed optimizer which shards optimizer state and gradients across data parallel ranks? [YES/no]: 
What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
How many GPU(s) should be used for distributed training? [1]:4
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16

A configuração resultante é mostrada abaixo:

~$ cat megatron_gpt_config.yaml 
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: MEGATRON_LM
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config:
  megatron_lm_gradient_clipping: 1.0
  megatron_lm_num_micro_batches: 2
  megatron_lm_pp_degree: 2
  megatron_lm_recompute_activations: true
  megatron_lm_sequence_parallelism: true
  megatron_lm_tp_degree: 2
  megatron_lm_use_distributed_optimizer: true
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false

Vamos usar o exemplo de pré-treinamento do GPT. As alterações mínimas necessárias no run_clm_no_trainer.py oficial para usar o Megatron-LM são as seguintes:

1. Como o Megatron-LM usa sua própria implementação de Optimizer, é necessário usar o scheduler correspondente compatível com ele. Dessa forma, o suporte é apenas para o scheduler do Megatron-LM. O usuário precisará criar accelerate.utils.MegatronLMDummyScheduler. Um exemplo é dado abaixo:

from accelerate.utils import MegatronLMDummyScheduler

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    lr_scheduler = MegatronLMDummyScheduler(
        optimizer=optimizer,
        total_num_steps=args.max_train_steps,
        warmup_num_steps=args.num_warmup_steps,
    )
else:
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
    )

2. Obter os detalhes do tamanho total do lote agora exige levar em consideração os tamanhos de paralelismo de tensor e de pipeline. Um exemplo de como obter o tamanho total efetivo do lote é mostrado abaixo:

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
else:
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

3. Ao usar o Megatron-LM, as perdas já são médias calculadas em todo o grupo de paralelismo de dados

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses.append(loss)
else:
    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    losses = torch.tensor(losses)
else:
    losses = torch.cat(losses)

4. Para o Megatron-LM, precisamos salvar o modelo usando accelerator.save_state

if accelerator.distributed_type == DistributedType.MEGATRON_LM:
    accelerator.save_state(args.output_dir)
else:
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(
        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
    )

Pronto! Estamos prontos para começar 🚀. Você pode encontrar o script de exemplo na pasta de exemplos, no caminho accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py. Vamos executá-lo para a arquitetura de modelo gpt-large usando 4 GPUs A100-80GB.

accelerate launch --config_file megatron_gpt_config.yaml \
examples/by_feature/megatron_lm_gpt_pretraining.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--block_size 1024 \
--learning_rate 5e-5 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 5 \
--with_tracking \
--report_to "wandb" \
--output_dir "awesome_model"

Abaixo estão alguns trechos importantes dos logs de saída:

Loading extension module fused_dense_cuda...
>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
 > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
Building gpt model in the pre-training mode.
The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
Preparing dataloader
Preparing dataloader
Preparing model
 > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
 > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
 > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
Preparing optimizer
Preparing scheduler
> learning rate decay style: linear
10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]

Há um grande número de outras opções/recursos que podem ser configurados usando accelerate.utils.MegatronLMPlugin.

Recursos avançados para aproveitar a escrita de etapas de treinamento personalizadas e Conjuntos de Dados Indexados do Megatron-LM

Para aproveitar mais recursos, consulte os detalhes abaixo.

1. Abaixo está um exemplo das alterações necessárias para personalizar o Train Step ao usar Megatron-LM. Você implementará o accelerate.utils.AbstractTrainStep ou herdará de seus filhos correspondentes accelerate.utils.GPTTrainStep, accelerate.utils.BertTrainStep ou accelerate.utils.T5TrainStep.

from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group


# Custom loss function for the Megatron model
class GPTTrainStepWithCustomLoss(GPTTrainStep):
    def __init__(self, megatron_args, **kwargs):
        super().__init__(megatron_args)
        self.kwargs = kwargs

    def get_loss_func(self):
        def loss_func(inputs, loss_mask, output_tensor):
            batch_size, seq_length = output_tensor.shape
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = losses.view(-1) * loss_mask

            # Resize and average loss per sample
            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
            loss_per_sample = loss_per_sample / loss_mask_per_sample

            # Calculate and scale weighting
            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
            weights = 1.0 + self.kwargs["alpha"] * weights
            # Calculate weighted average
            weighted_loss = (loss_per_sample * weights).mean()

            # Reduce loss across data parallel groups
            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])

            return weighted_loss, {"lm loss": averaged_loss[0]}

        return loss_func

    def get_forward_step_func(self):
        def forward_step(data_iterator, model):
            """Forward step."""
            # Get the batch.
            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)

            return output_tensor, partial(self.loss_func, tokens, loss_mask)

        return forward_step


def main():
    # Custom loss function for the Megatron model
    keytoken_ids = []
    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
    for keyword in keywords:
        ids = tokenizer([keyword]).input_ids[0]
        if len(ids) == 1:
            keytoken_ids.append(ids[0])
    accelerator.print(f"Keytoken ids: {keytoken_ids}")
    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
        "keytoken_ids": keytoken_ids,
        "alpha": 0.25,
    }

2. Para usar os conjuntos de dados do Megatron-LM, são necessárias algumas alterações adicionais. Os dataloaders para esses conjuntos de dados estão disponíveis apenas no rank 0 de cada grupo de paralelismo de tensor. Dessa forma, existem ranks em que o dataloader não estará disponível, o que exige ajustes no loop de treinamento. O fato de ser possível fazer tudo isso demonstra o quão flexível e extensível é o Accelerate. As alterações necessárias são as seguintes.

a. Para conjuntos de dados indexados do Megatron-LM, precisamos usar MegatronLMDummyDataLoader e passar os argumentos de conjunto de dados necessários para ele, como data_path, seq_length, etc. Veja aqui a lista de argumentos disponíveis.

from accelerate.utils import MegatronLMDummyDataLoader

megatron_dataloader_config = {
    "data_path": args.data_path,
    "splits_string": args.splits_string,
    "seq_length": args.block_size,
    "micro_batch_size": args.per_device_train_batch_size,
}
megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True

b. megatron_dataloader é repetido 3 vezes para obter os dataloaders de treinamento, validação e teste, conforme as proporções de args.splits_string

model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
)

c. Alterações nos loops de treinamento e avaliação, pois o dataloader está disponível apenas nos ranks 0 do tensor parallel Portanto, precisamos iterar apenas se o dataloader não for None, caso contrário, fornecer um dicionário vazio Dessa forma, fazemos o loop usando um loop while e interrompemos quando completed_steps é igual a args.max_train_steps. Isso é semelhante à configuração do Megatron-LM, na qual o usuário precisa fornecer max_train_steps ao usar conjuntos de dados indexados do Megaton-LM. Isso mostra como o Accelerate é flexível e extensível.

while completed_steps < args.max_train_steps:
    model.train()
    batch = next(train_dataloader) if train_dataloader is not None else {}
    outputs = model(**batch)
    loss = outputs.loss
    ...

    if completed_steps % eval_interval == 0:
        eval_completed_steps = 0
        losses = []
        while eval_completed_steps < eval_iters:
            model.eval()
            with torch.no_grad():
                batch = next(eval_dataloader) if eval_dataloader is not None else {}
                outputs = model(**batch)

Utilitário para remodelação de checkpoints e interoperabilidade

1. Os scripts para isso estão presentes na biblioteca Transformers, sob os respectivos modelos. Atualmente, está disponível para o modelo GPT checkpoint_reshaping_and_interoperability.py

2. Abaixo está um exemplo de conversão de checkpoint do Megatron-LM para checkpoint fragmentado universal do Transformers.

python checkpoint_reshaping_and_interoperability.py \
--convert_checkpoint_from_megatron_to_transformers \
--load_path "gpt/iter_0005000" \
--save_path "gpt/trfs_checkpoint" \
--max_shard_size "200MB" \
--tokenizer_name "gpt2" \
--print-checkpoint-structure

3. Conversão de checkpoint de transformers para megatron com tp_size=2, pp_size=2 e dp_size=2.

python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
--load_path "gpt/trfs_checkpoint" \
--save_path "gpt/megatron_lm_checkpoint" \
--target_tensor_model_parallel_size 2 \
--target_pipeline_model_parallel_size 2 \
--target_data_parallel_size 2 \
--target_params_dtype "bf16" \
--make_vocab_size_divisible_by 128 \
--use_distributed_optimizer \
--print-checkpoint-structure

Os modelos GPT do Megatron-LM suportam o retorno de logits e a função megatron_generate para geração de texto

1. Retornar logits exige definir require_logits=True em MegatronLMPlugin, conforme mostrado abaixo. Eles estarão disponíveis no último estágio do pipeline.

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

2. Método megatron_generate para o modelo GPT do Megatron-LM: este método usa Tensor e Pipeline Parallelism para concluir gerações para um lote de entradas ao usar greedy com ou sem amostragem top_k/top_p e para entradas de prompt individuais ao usar decodificação com beam search. Apenas um subconjunto dos recursos de generate do transformers é suportado. Isso ajudará a usar modelos grandes via tensor e pipeline parallelism para geração (já faz key-value caching e usa kernels fundidos por padrão). Isso requer que o tamanho de data parallel seja 1, que sequence parallelism e activation checkpointing estejam desativados. Também requer especificar o caminho para o arquivo de vocabulário e o arquivo de merges do tokenizer. O exemplo abaixo mostra como configurar e usar o método megatron_generate para o modelo GPT do Megatron-LM.

# specifying tokenizer's vocab and merges file
vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

# inference using `megatron_generate` functionality
tokenizer.pad_token = tokenizer.eos_token
max_new_tokens = 64
batch_texts = [
    "Are you human?",
    "The purpose of life is",
    "The arsenal was constructed at the request of",
    "How are you doing these days?",
]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)

# top-p sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_p=0.8,
    top_p_decay=0.5,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# top-k sampling
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    top_k=50,
    temperature=0.9,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# adding `bos` token at the start
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

# beam search => only takes single prompt
batch_texts = ["The purpose of life is"]
batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
generated_tokens = model.megatron_generate(
    batch_encodings["input_ids"],
    batch_encodings["attention_mask"],
    max_new_tokens=max_new_tokens,
    num_beams=20,
    length_penalty=1.5,
)
decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
accelerator.print(decoded_preds)

3. Um exemplo de ponta a ponta do uso do método megatron_generate para o modelo GPT do Megatron-LM está disponível em megatron_gpt2_generation.py com o arquivo de configuração megatron_lm_gpt_generate_config.yaml. O script bash com o comando accelerate launch está disponível em megatron_lm_gpt_generate.sh. Os logs de saída do script estão disponíveis em megatron_lm_gpt_generate.log.

Suporte a embeddings posicionais ROPE e ALiBi e a Multi-Query Attention

1. Para a atenção ROPE/ALiBi, passe position_embedding_type com ("absolute" | "rotary" | "alibi") para MegatronLMPlugin conforme mostrado abaixo.

other_megatron_args = {"position_embedding_type": "alibi"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

2. Para Multi-Query Attention, passe attention_head_type com ("multihead" | "multiquery") para MegatronLMPlugin conforme mostrado abaixo.

other_megatron_args = {"attention_head_type": "multiquery"}
megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)

Observações

1. Suporta modelos GPT2, Megatron-BERT e T5 do Transformers. Isso abrange as classes de modelos apenas Decoder, apenas Encoder e Encoder-Decoder.

2. Apenas a perda é retornada da passagem direta do modelo, pois há uma interação bastante complexa entre paralelismo de pipeline, de tensor e de dados nos bastidores. A chamada model(**batch_data) retorna a(s) perda(s) média(s) entre os ranks de paralelismo de dados. Isso é adequado para a maioria dos casos em que os trabalhos de pré-treino são executados usando os recursos do Megatron-LM e você pode calcular facilmente a perplexity usando a perda. Para o modelo GPT, é suportado o retorno de logits em adição à(s) perda(s). Esses logits não são agregados entre os ranks de paralelismo de dados. Use accelerator.utils.gather_across_data_parallel_groups para agregar os logits entre os ranks de paralelismo de dados. Esses logits, juntamente com os rótulos, podem ser usados para calcular diversas métricas de desempenho.

3. O processo principal é o último rank, pois as perdas/logits estão disponíveis no último estágio do pipeline. accelerator.is_main_process e accelerator.is_local_main_process retornam True para o último rank ao usar a integração com Megatron-LM.

4. Na chamada de accelerator.prepare, um modelo Megatron-LM correspondente a um modelo Transformers específico é criado com pesos aleatórios. Use accelerator.load_state para carregar o checkpoint do Megatron-LM com as partições TP, PP e DP correspondentes.

5. Atualmente, o suporte para remodelação de checkpoints e interoperabilidade está disponível apenas para GPT. Em breve, será estendido para BERT e T5.

6. gradient_accumulation_steps precisa ser 1. Ao usar Megatron-LM, micro lotes no paralelismo de pipeline são sinônimos de acumulação de gradiente.

7. Ao usar o Megatron-LM, utilize accelerator.save_state e accelerator.load_state para salvar e carregar checkpoints.

8. Abaixo estão os mapeamentos das arquiteturas de modelos do Megatron-LM para as arquiteturas de modelos equivalentes do transformers. Somente essas arquiteturas de modelos do transformers são suportadas.

a. Megatron-LM BertModel : modelos do transformers com megatron-bert no tipo de modelo do config, por exemplo, MegatronBERT

b. Megatron-LM GPTModel : modelos do transformers com gpt2 no tipo de modelo da config, por exemplo, OpenAI GPT2

c. Megatron-LM T5Model : modelos do transformers com t5 no tipo de modelo da config, por exemplo, T5 e MT5