Como usar o Rapidformer para otimizar o treinamento de modelos Transformer no PyTorch - Platform For AI: Central de Documentação do Alibaba Cloud

Pré-requisitos

Antes de começar, verifique se você já:

Instale a imagem de runtime do Rapidformer. Consulte Install a PAI-Megatron-Patch runtime image
Revisou a referência de parâmetros do Rapidformer. Consulte Parameter settings guide
Revisou a referência da API do Rapidformer. Consulte Rapidformer API

Escolha um método de aceleração

Use esta tabela para selecionar o método adequado antes de prosseguir.

	Caixa preta	Caixa branca
Funcionamento	Passe os nomes do modelo e do conjunto de dados para a CLI do Rapidformer. Nenhuma alteração no código é necessária.	Crie uma subclasse de um template do Rapidformer (`Finetuner` ou `PreTrainer`), implemente os métodos obrigatórios e invoque via CLI com `--user-script`.
Quando usar	O modelo e o conjunto de dados estão registrados no Hugging Face ou Megatron e não há necessidade de personalizar o loop de treinamento.	É necessário controlar o carregamento de dados, a arquitetura do modelo ou a lógica do forward pass — ou você usa um modelo Megatron totalmente personalizado.
Personalização	Nenhuma	Controle total sobre dados, modelo, otimizador e forward pass
Recursos de aceleração	Treinamento de precisão mista, otimização de grafo ONNX Runtime, otimização de memória ZeRO, FSDP	Mesmos recursos da caixa preta, além de paralelismo de tensor e paralelismo de pipeline (apenas Megatron)

Para programas com Trainers personalizados existentes, o Rapidformer oferece aceleração limitada e pouco invasiva — apenas otimizador Apex, particionamento de estado do modelo e otimização de grafo computacional. O treinamento de precisão mista exige modificações significativas. Use o método de template de caixa branca para obter cobertura completa de aceleração.

Referência de flags de aceleração

Todos os exemplos de CLI neste documento usam combinações destas flags. Consulte esta tabela para entender a função de cada flag e quando ativá-la.

Flag	Técnica de aceleração	Descrição	Combinar com
`--mixed-precision`	Treinamento de precisão mista	Treina em FP16 para reduzir o uso de memória e acelerar a computação	Qualquer outra flag
`--onnx-runtime-training`	Otimização de grafo computacional	Usa o ONNX Runtime para otimizar o grafo de computação	Qualquer outra flag
`--zero-1-memory-optimization`	ZeRO Estágio 1 — particionamento de estado do otimizador	Particiona os estados do otimizador entre GPUs	`--mixed-precision`
`--zero-2-memory-optimization`	ZeRO Estágio 2 — particionamento de otimizador + gradiente	Particiona estados do otimizador e gradientes entre GPUs	`--mixed-precision`
`--zero-3-memory-optimization`	ZeRO Estágio 3 — particionamento completo do estado do modelo	Particiona estados do otimizador, gradientes e parâmetros do modelo	`--mixed-precision`
`--fsdp-memory-optimization`	FSDP — Fully Sharded Data Parallel	Fragmenta os estados do modelo usando PyTorch FSDP	`--mixed-precision`
`--tensor-model-parallel-size <N>`	Paralelismo de tensor	Divide operadores entre N GPUs	Apenas pré-treinamento Megatron
`--pipeline-model-parallel-size <N>`	Paralelismo de pipeline	Divide as camadas do modelo entre N GPUs	Apenas pré-treinamento Megatron
`--checkpoint-activations`	Checkpointing de gradiente	Recalcula ativações durante o backward pass para economizar VRAM	Qualquer outra flag
`--data-impl mmap`	Dados mapeados em memória	Lê o conjunto de dados do disco usando arquivos mapeados em memória, reduzindo a sobrecarga de carregamento	Apenas pré-treinamento

Para a referência completa de parâmetros, consulte Parameter settings guide.

Caixa preta: Ajuste fino com Hugging Face

Não requer código. Registre seu conjunto de dados e modelo no Hugging Face e invoque a CLI. Este exemplo ativa treinamento de precisão mista, otimização de grafo ONNX Runtime e particionamento de estado do otimizador ZeRO Estágio 1.

Registre seu conjunto de dados no Hugging Face ou use um existente. Passe o conjunto de dados para o Rapidformer usando --dataset-name.
- Register a Hugging Face dataset
- Browse existing Hugging Face datasets
Registre seu modelo no Hugging Face ou use um existente. Passe o modelo para o Rapidformer usando --pretrained-model-name-or-path.
- Register a Hugging Face model
- Browse existing Hugging Face models

Execute a CLI do Rapidformer para iniciar o treinamento.

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MASTER_ADDR=localhost
export MASTER_PORT=6010
export NNODES=1
export NODE_RANK=0

rapidformer --task sequence_classification \
            --pretrained-model-name-or-path 'bert-base-cased' \
            --data-path glue \
            --data-name mrpc \
            --epochs 3 \
            --micro-batch-size 16 \
            --global-batch-size 64 \
            --lr 2e-5 \
            --lr-decay-style linear \
            --lr-warmup-iters 100 \
            --weight-decay 1e-2 \
            --clip-grad 1.0 \
            --seed 42 \
            --mixed-precision \
            --onnx-runtime-training \
            --zero-1-memory-optimization

Caixa preta: Pré-treinamento com Hugging Face

Não requer código. Crie um conjunto de dados mmap, registre um modelo e invoque a CLI. Este exemplo ativa carregamento de dados mapeados em memória, treinamento de precisão mista, otimização de grafo ONNX Runtime e particionamento de estado do modelo FSDP.

Crie um conjunto de dados no formato mmap para pré-treinamento usando o Megatron data processing script.

python preprocess_data.py \
  --input book_wiki_owtv2_small.json \
  --output-prefix gpt_small \
  --vocab gpt2-vocab.json \
  --dataset-impl mmap \
  --tokenizer-type GPT2BPETokenizer \
  --merge-file gpt2-merges.txt \
  --append-eod

Registre seu modelo no Hugging Face ou use um existente. Passe o modelo para o Rapidformer usando --pretrained-model-name-or-path.
- Register a Hugging Face model
- Browse existing Hugging Face models

Execute a CLI do Rapidformer para iniciar o treinamento.

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MASTER_ADDR=localhost
export MASTER_PORT=6010
export NNODES=1
export NODE_RANK=0

rapidformer --task pretraining \
            --pretrained-model-name-or-path 'bert-base-uncased' \
            --num-layers 12 \
            --hidden-size 768 \
            --num-attention-heads 12 \
            --micro-batch-size 16 \
            --global-batch-size 128 \
            --seq-length 512 \
            --tokenizer-type BertWordPieceLowerCase \
            --max-position-embeddings 512 \
            --train-iters 100 \
            --data-path book_wiki_owtv2_small_text_sentence \
            --vocab-file bert-en-uncased-vocab.txt \
            --data-impl mmap \
            --split 980,20 \
            --lr 1e-3 \
            --lr-decay-style linear \
            --min-lr 0.0 \
            --lr-decay-iters 2000 \
            --weight-decay 1e-2 \
            --clip-grad 1.0 \
            --lr-warmup-fraction .01 \
            --mixed-precision \
            --onnx-runtime-training \
            --fsdp-memory-optimization

Caixa branca: Ajuste fino com Hugging Face usando o template Finetuner

Crie uma subclasse de Finetuner e implemente quatro métodos. O Rapidformer gerencia automaticamente o treinamento distribuído, precisão mista, ZeRO e FSDP — seu código define apenas o carregamento de dados, construção do modelo, lógica do forward pass e métricas de avaliação.

Estrutura do template

A classe base Finetuner exige quatro métodos:

Método	Entrada	Saída	Finalidade
`train_valid_test_datasets_provider`	Nenhum	`train_dataset, valid_dataset, test_dataset, collate_fn`	Criar conjuntos de dados de treino, validação e teste
`model_optimizer_lr_scheduler_provider`	Nenhum	`model, optimizer, lr_scheduler`	Construir o modelo, otimizador e agendador de taxa de aprendizado
`run_forward_step`	`batch_or_iterator, model`	`loss`	Definir a lógica do forward pass
`run_compute_metrics`	`model, eval_dataloader`	`metric`	Calcular métricas de avaliação (apenas ajuste fino)

Consulte a Rapidformer API para visualizar as assinaturas completas dos métodos e detalhes dos tipos de retorno.

Etapas de implementação

Prepare o conjunto de dados e o modelo seguindo a mesma abordagem descrita em Black box: Hugging Face fine-tuning. Em seguida:

Importe as interfaces do Rapidformer e do Hugging Face.

from transformers.easytexmier import AutoConfig, BertForSequenceClassification
from datasets import load_dataset, load_metric
from rapidformer import RapidformerEngine
from rapidformer import get_args
from rapidformer import get_logger
from rapidformer import get_timers
from rapidformer import Finetuner
from rapidformer import Pretrainer
from rapidformer import build_train_valid_test_datasets_for_huggingface

Implemente os quatro métodos na sua subclasse de Finetuner.

class MyFintuner(Finetuner):
    def __init__(self, engine):
        super().__init__(engine=engine)

    def train_valid_test_datasets_provider(self):
        tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

        def tokenize_function(examples):
            # max_length=None uses the model's maximum length
            outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
            return outputs

        datasets = load_dataset(args.dataset_path, args.dataset_name)
        tokenized_datasets = datasets.map(
            tokenize_function,
            batched=True,
            remove_columns=["idx", "sentence1", "sentence2"],
        )
        tokenized_datasets.rename_column_("label", "labels")

        train_dataset = tokenized_datasets["train"]
        valid_dataset = tokenized_datasets['validation']
        test_dataset = tokenized_datasets['test']

        def collate_fn(examples):
            return tokenizer.pad(examples, padding="longest", return_tensors="pt")

        return train_dataset, valid_dataset, test_dataset, collate_fn

    def model_optimizer_lr_scheduler_provider(self):
        args = get_args()
        model = BertForSequenceClassification.from_pretrained(args.load)
        return model, None, None

    def run_forward_step(self, batch, model):
        output_tensor = model(**batch)
        return output_tensor.loss

    def run_compute_metrics(self, model, eval_dataloader):
        model = model[0]
        metric = load_metric(args.dataset_path, args.dataset_name)
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                outputs = model(**batch)
            predictions = outputs.logits.argmax(dim=-1)
            metric.add_batch(
                predictions=self.gather(predictions),
                references=self.gather(batch["labels"]),
            )
        eval_metric = metric.compute()
        return eval_metric

Inicialize o engine, crie um trainer e salve como rapidformer_finetune_huggingface_bert_trainer.py.
```
engine = RapidformerEngine()
trainer = MyFintuner(engine=engine)
trainer.train()
```

Execute via CLI, apontando --user-script para o arquivo criado na etapa 3.

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MASTER_ADDR=localhost
export MASTER_PORT=6010
export NNODES=1
export NODE_RANK=0

rapidformer --user-script rapidformer_finetune_huggingface_bert_trainer.py \
            --task sequence_classification \
            --pretrained-model-name-or-path 'bert-base-cased' \
            --data-path glue \
            --data-name mrpc \
            --epochs 3 \
            --micro-batch-size 16 \
            --global-batch-size 16 \
            --lr 2e-5 \
            --lr-decay-style linear \
            --lr-warmup-iters 100 \
            --weight-decay 1e-2 \
            --clip-grad 1.0 \
            --mixed-precision \
            --zero-3-memory-optimization \
            --onnx-runtime-training

Caixa branca: Pré-treinamento com Hugging Face usando o template PreTrainer

Crie uma subclasse de PreTrainer e implemente três métodos. Diferentemente do Finetuner, o pré-treinamento lê dados via iterador; portanto, importe mpu para broadcast de paralelismo de dados.

Estrutura do template

A classe base PreTrainer exige três métodos:

Método	Entrada	Saída	Finalidade
`train_valid_test_datasets_provider`	`train_val_test_num_samples`	`train_ds, valid_ds, test_ds`	Criar conjuntos de dados de treino, validação e teste
`model_optimizer_lr_scheduler_provider`	Nenhum	`model, optimizer, lr_scheduler`	Construir o modelo, otimizador e agendador de taxa de aprendizado
`run_forward_step`	`data_iterator, model`	`loss`	Definir a lógica do forward pass (baseado em iterador)

Consulte White box: Hugging Face fine-tuning with the Finetuner template para detalhes sobre os tipos de retorno dos métodos.

Etapas de implementação

Prepare o conjunto de dados e o modelo conforme descrito em Black box: Hugging Face fine-tuning. Em seguida:

Importe as interfaces do Rapidformer e do Hugging Face.

O pré-treinamento lê dados via iterador. Importe mpu para fazer broadcast dos dados entre workers paralelos.

from megatron import mpu
from transformers import BertConfig, BertForPreTraining
from rapidformer import RapidformerEngine, get_args, PreTrainer
from rapidformer import build_train_valid_test_datasets_for_huggingface

Implemente os três métodos na sua subclasse de PreTrainer.

class MyBertPreTrainer(PreTrainer):

    def __init__(self, engine):
        super().__init__(engine=engine)

    def train_valid_test_datasets_provider(self, train_val_test_num_samples):
        args = get_args()
        train_ds, valid_ds, test_ds = build_train_valid_test_datasets_for_huggingface(
            data_prefix=args.data_path,
            data_impl=args.data_impl,
            splits_string=args.split,
            train_valid_test_num_samples=train_val_test_num_samples,
            max_seq_length=args.seq_length,
            masked_lm_prob=args.mask_prob,
            short_seq_prob=args.short_seq_prob,
            seed=args.seed,
            skip_warmup=(not args.mmap_warmup),
            binary_head=True)
        return train_ds, valid_ds, test_ds

    def model_optimizer_lr_scheduler_provider(self):
        args = get_args()
        model = AutoModelForPreTraining.from_pretrained(args.pretrained_model_name_or_path)
        return model, None, None

    def run_forward_step(self, data_iterator, model):
        keys = ['input_ids', 'attention_mask', 'token_type_ids', 'labels', 'next_sentence_label']
        datatype = torch.int64

        data = next(data_iterator) if data_iterator is not None else None
        data_b = mpu.broadcast_data(keys, data, datatype)

        input_ids = data_b['input_ids'].long()
        attention_mask = data_b['attention_mask'].long()
        token_type_ids = data_b['token_type_ids'].long()
        labels = data_b['labels'].long()
        next_sentence_label = data_b['next_sentence_label'].long()

        output_tensor = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            labels=labels,
            next_sentence_label=next_sentence_label)

        return output_tensor['loss']

Inicialize o engine, crie um trainer e salve como rapidformer_pretrain_huggingface_bert_trainer.py.
```
engine = RapidformerEngine()
trainer = MyBertPreTrainer(engine=engine)
trainer.train()
```

Execute via CLI, apontando --user-script para o arquivo criado na etapa 3.

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MASTER_ADDR=localhost
export MASTER_PORT=6010
export NNODES=1
export NODE_RANK=0

DATA_PATH=book_wiki_owtv2_small_text_sentence

rapidformer --user-script rapidformer_pretrain_huggingface_bert_trainer.py \
            --pretrained-model-name-or-path 'bert-base-uncased' \
            --num-layers 12 \
            --hidden-size 768 \
            --num-attention-heads 12 \
            --micro-batch-size 16 \
            --global-batch-size 64 \
            --seq-length 512 \
            --tokenizer-type BertWordPieceLowerCase \
            --max-position-embeddings 512 \
            --train-iters 100 \
            --data-path $DATA_PATH \
            --vocab-file bert-en-uncased-vocab.txt \
            --data-impl mmap \
            --split 980,20 \
            --lr 1e-3 \
            --lr-decay-style linear \
            --weight-decay 1e-2 \
            --clip-grad 1.0 \
            --lr-warmup-fraction .01 \
            --zero-3-memory-optimization \
            --onnx-runtime-training \
            --mixed-precision

Caixa branca: Ajuste fino com Hugging Face usando um Trainer personalizado

Para códigos de treinamento existentes do Hugging Face que usam um Trainer personalizado, adicione as APIs do Rapidformer incrementalmente. Esta seção mostra as duas alterações mais impactantes: paralelismo de dados e um otimizador mais rápido.

Este método fornece aceleração limitada — apenas otimizador Apex, particionamento de estado do modelo e otimização de grafo computacional. O treinamento de precisão mista exige modificações significativas no seu loop de treinamento. Para aceleração completa, use o Finetuner template method .

Adicionar paralelismo de dados

Substitua o DataLoader padrão por finetuner.build_data_loader. Este loader suporta paralelismo de dados e move automaticamente os lotes para a GPU, então remova batch.to(device) do seu loop de treinamento.

+ from rapidformer import RapidformerEngine, Finetuner
+ engine = RapidformerEngine()
+ finetuner = Finetuner(engine=engine)

- train_dataloader = DataLoader(tokenized_datasets["train"])
- eval_dataloader = DataLoader(tokenized_datasets["train"])

+ train_dataloader = finetuner.build_data_loader(tokenized_datasets["train"])
+ eval_dataloader = finetuner.build_data_loader(tokenized_datasets["validation"])

Usar o otimizador Apex Fused Adam

Substitua AdamW pelo otimizador Apex Fused Adam do Rapidformer. Chame engine.compose para encapsular o modelo, o otimizador e o agendador de taxa de aprendizado juntos.

- optimizer = AdamW(params=model.parameters(), lr=args.lr, correct_bias=True)
- lr_scheduler = get_linear_schedule_with_warmup(
-     optimizer=optimizer,
-     num_warmup_steps=args.lr_warmup_iters,
-     num_training_steps=args.train_iters
- )

+ lr_scheduler = partial(
+     get_linear_schedule_with_warmup,
+     num_warmup_steps=args.lr_warmup_iters,
+     num_training_steps=args.train_iters
+ )

+ model, optimizer, lr_scheduler = engine.compose(
+     model_obj=model,
+     lr_scheduler_fn=lr_scheduler)

Combinar o otimizador Apex com treinamento de precisão mista em um loop de Trainer personalizado exige alternar o modelo para FP16 e adicionar escalonamento de perda — uma modificação significativa. O template Finetuner integra automaticamente paralelismo de dados, Apex, treinamento de precisão mista do PyTorch, precisão mista do otimizador Megatron e otimização de VRAM do FairScale e DeepSpeed.

Caixa branca: Pré-treinamento Megatron com o template PreTrainer

Para modelos Megatron totalmente personalizados — sem Hugging Face ou Data/Model Hub — implemente todos os três métodos do PreTrainer com lógica personalizada. Conjuntos de dados devem herdar de torch.utils.data.Dataset; modelos devem herdar de torch.nn.Module.

Este método também ativa paralelismo de tensor e paralelismo de pipeline, disponíveis apenas para modelos Megatron.

Crie um conjunto de dados no formato mmap usando o Megatron data processing script.

python preprocess_data.py \
  --input /apsarapangu/disk2/jerry.lp/pretrain_datasets/en/book_wiki_owtv2_small.json \
  --output-prefix /apsarapangu/disk2/jerry.lp/pretrain_datasets/en/gpt_small \
  --vocab gpt2-vocab.json \
  --dataset-impl mmap \
  --tokenizer-type GPT2BPETokenizer \
  --merge-file gpt2-merges.txt \
  --append-eod

Implemente train_valid_test_datasets_provider com lógica personalizada de criação de conjunto de dados.

from rapidformer import RapidformerEngine, get_args, PreTrainer

class MegatronGPTPreTrainer(PreTrainer):
    def __init__(self, engine):
        super().__init__(engine=engine)

    def train_valid_test_datasets_provider(self, train_val_test_num_samples):
        args = get_args()
        train_ds, valid_ds, test_ds = build_train_valid_test_datasets(
            data_prefix=args.data_path,
            data_impl=args.data_impl,
            splits_string=args.split,
            train_valid_test_num_samples=train_val_test_num_samples,
            seq_length=args.seq_length,
            seed=args.seed,
            skip_warmup=(not args.mmap_warmup))
        return train_ds, valid_ds, test_ds

Implemente model_optimizer_lr_scheduler_provider com seu modelo personalizado.

from rapidformer import RapidformerEngine, get_args, PreTrainer
from yourmodel import GPTModel

class MegatronGPTPreTrainer(PreTrainer):
    def __init__(self, engine):
        super().__init__(engine=engine)

    def model_optimizer_lr_scheduler_provider(self):
        model = GPTModel()
        return model, None, None

Implemente run_forward_step com sua lógica personalizada de forward pass.

from rapidformer import RapidformerEngine, get_args, PreTrainer

class MyGPTPreTrainer(PreTrainer):
    def __init__(self, engine):
        super().__init__(engine=engine)

    def run_forward_step(self, data_iterator, model):
        """Forward step."""
        args = get_args()
        tokenizer = get_tokenizer()

        keys = ['text']
        datatype = torch.int64

        data = next(data_iterator) if data_iterator is not None else None
        data_b = mpu.broadcast_data(keys, data, datatype)

        tokens_ = data_b['text'].long()
        labels = tokens_[:, 1:].contiguous()
        tokens = tokens_[:, :-1].contiguous()

        attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
            tokens,
            tokenizer.eod,
            args.reset_position_ids,
            args.reset_attention_mask,
            args.eod_mask_loss)

        output_tensor = model(tokens, position_ids, attention_mask, labels=labels)

        losses = output_tensor.float()
        loss_mask = loss_mask.view(-1).float()
        loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()

        return loss

Inicialize o engine, crie um trainer e salve como rapidformer_pretrain_megatron_gpt_trainer.py.
```
engine = RapidformerEngine()
trainer = MyGPTPreTrainer(engine=engine)
trainer.train()
```

Execute via CLI com paralelismo de tensor e paralelismo de pipeline ativados.

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MASTER_ADDR=localhost
export MASTER_PORT=6010
export NNODES=1
export NODE_RANK=0

DATA_PATH=book_wiki_owtv2_small_text_sentence

rapidformer --user-script rapidformer_pretrain_megatron_gpt_trainer.py \
            --tensor-model-parallel-size 2 \
            --pipeline-model-parallel-size 2 \
            --num-layers 12 \
            --hidden-size 768 \
            --num-attention-heads 12 \
            --micro-batch-size 16 \
            --global-batch-size 128 \
            --seq-length 512 \
            --tokenizer-type GPT2BPETokenizer \
            --max-position-embeddings 512 \
            --train-iters 100 \
            --data-path $DATA_PATH \
            --vocab-file gpt2-vocab.json \
            --merge-file gpt2-merges.txt \
            --data-impl mmap \
            --split 980,20 \
            --lr 1e-3 \
            --lr-decay-style linear \
            --weight-decay 1e-2 \
            --clip-grad 1.0 \
            --lr-warmup-fraction .01 \
            --log-interval 1 \
            --zero-2-memory-optimization \
            --checkpoint-activations \
            --mixed-precision

Este exemplo ativa todas as otimizações específicas do Megatron disponíveis: paralelismo de tensor (--tensor-model-parallel-size 2), paralelismo de pipeline (--pipeline-model-parallel-size 2), acumulação de gradiente via tamanho de lote global, carregamento de dados mapeados em memória, otimização de memória ZeRO Estágio 2, checkpointing de gradiente e treinamento de precisão mista.

Próximos passos

Parameter settings guide — referência completa de parâmetros da CLI
Rapidformer API — assinaturas de métodos de Finetuner, PreTrainer e RapidformerEngine