All Products
Search

This Product

    Platform For AI:Use case: Accelerate Transformer model training

    Document Center

    Platform For AI:Use case: Accelerate Transformer model training

    Last Updated:Apr 01, 2026

    PAI-Rapidformer accelerates PyTorch Transformer model training through two methods: black box (CLI-only, no code changes) and white box (code templates for custom training loops).

    Prerequisites

    Before you begin, ensure that you have:

    Choose an acceleration method

    Use this table to select the right method before reading further.

    Black boxWhite box
    How it worksPass your model and dataset names to the Rapidformer CLI. No code changes required.Subclass a Rapidformer template (Finetuner or PreTrainer), implement required methods, then invoke via CLI with --user-script.
    Use whenYour model and dataset are registered with Hugging Face or Megatron and you do not need to customize the training loop.You need to control data loading, model architecture, or forward pass logic — or you are using a fully custom Megatron model.
    CustomizationNoneFull control over data, model, optimizer, and forward pass
    Acceleration featuresMixed-precision training, ONNX Runtime graph optimization, ZeRO memory optimization, FSDPSame as black box, plus tensor parallelism and pipeline parallelism (Megatron only)
    For programs with existing custom Trainers, Rapidformer provides limited intrusive acceleration — Apex optimizer, model state partitioning, and computational graph optimization only. Mixed-precision training requires significant modifications. Use the white box template method instead for full acceleration coverage.

    Acceleration flags reference

    All CLI examples in this document use combinations of these flags. Refer to this table to understand what each flag does and when to enable it.

    FlagAcceleration techniqueWhat it doesCombine with
    --mixed-precisionMixed-precision trainingTrains in FP16 to reduce memory use and speed up computationAny other flag
    --onnx-runtime-trainingComputational graph optimizationUses ONNX Runtime to optimize the computation graphAny other flag
    --zero-1-memory-optimizationZeRO Stage 1 — optimizer state partitioningPartitions optimizer states across GPUs--mixed-precision
    --zero-2-memory-optimizationZeRO Stage 2 — optimizer + gradient partitioningPartitions optimizer states and gradients across GPUs--mixed-precision
    --zero-3-memory-optimizationZeRO Stage 3 — full model state partitioningPartitions optimizer states, gradients, and model parameters--mixed-precision
    --fsdp-memory-optimizationFSDP — Fully Sharded Data ParallelShards model states using PyTorch FSDP--mixed-precision
    --tensor-model-parallel-size <N>Tensor parallelismSplits operators across N GPUsMegatron pre-training only
    --pipeline-model-parallel-size <N>Pipeline parallelismSplits model layers across N GPUsMegatron pre-training only
    --checkpoint-activationsGradient checkpointingRecomputes activations during backward pass to save VRAMAny other flag
    --data-impl mmapMemory-mapped dataReads dataset from disk using memory-mapped files, reducing data loading overheadPre-training only

    For a complete parameter reference, see the Parameter settings guide.

    Black box: Hugging Face fine-tuning

    No code required. Register your dataset and model with Hugging Face, then invoke the CLI. This example enables mixed-precision training, ONNX Runtime graph optimization, and ZeRO Stage 1 optimizer state partitioning.

    1. Register your dataset with Hugging Face, or use an existing one. Pass the dataset to Rapidformer using --dataset-name.

    2. Register your model with Hugging Face, or use an existing one. Pass the model to Rapidformer using --pretrained-model-name-or-path.

    3. Run the Rapidformer CLI to start training.

      #!/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

    Black box: Hugging Face pre-training

    No code required. Create an mmap dataset, register a model, then invoke the CLI. This example enables memory-mapped data loading, mixed-precision training, ONNX Runtime graph optimization, and FSDP model state partitioning.

    1. Create an mmap-format dataset for pre-training using the 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

    2. Register your model with Hugging Face, or use an existing one. Pass the model to Rapidformer using --pretrained-model-name-or-path.

    3. Run the Rapidformer CLI to start training.

      #!/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

    White box: Hugging Face fine-tuning with the Finetuner template

    Subclass Finetuner and implement four methods. Rapidformer handles distributed training, mixed-precision, ZeRO, and FSDP automatically — your code defines only data loading, model construction, forward logic, and evaluation metrics.

    Template structure

    The Finetuner base class requires four methods:

    MethodInputOutputPurpose
    train_valid_test_datasets_providerNonetrain_dataset, valid_dataset, test_dataset, collate_fnCreate train, validation, and test datasets
    model_optimizer_lr_scheduler_providerNonemodel, optimizer, lr_schedulerConstruct the model, optimizer, and learning rate scheduler
    run_forward_stepbatch_or_iterator, modellossDefine forward pass logic
    run_compute_metricsmodel, eval_dataloadermetricCompute evaluation metrics (fine-tuning only)

    See the Rapidformer API for full method signatures and return type details.

    Implementation steps

    Prepare the dataset and model following the same approach as in Black box: Hugging Face fine-tuning. Then:

    1. Import the Rapidformer and Hugging Face interfaces.

      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

    2. Implement the four methods in your Finetuner subclass.

      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

    3. Initialize the engine, create a trainer, and save as rapidformer_finetune_huggingface_bert_trainer.py.

      engine = RapidformerEngine()
trainer = MyFintuner(engine=engine)
trainer.train()

    4. Run via the CLI, pointing --user-script at the file from step 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

    White box: Hugging Face pre-training with the PreTrainer template

    Subclass PreTrainer and implement three methods. Unlike the Finetuner, pre-training reads data via an iterator, so import mpu for data parallelism broadcast.

    Template structure

    The PreTrainer base class requires three methods:

    MethodInputOutputPurpose
    train_valid_test_datasets_providertrain_val_test_num_samplestrain_ds, valid_ds, test_dsCreate train, validation, and test datasets
    model_optimizer_lr_scheduler_providerNonemodel, optimizer, lr_schedulerConstruct the model, optimizer, and learning rate scheduler
    run_forward_stepdata_iterator, modellossDefine forward pass logic (iterator-based)

    See White box: Hugging Face fine-tuning with the Finetuner template for method return type details.

    Implementation steps

    Prepare the dataset and model as in Black box: Hugging Face fine-tuning. Then:

    1. Import the Rapidformer and Hugging Face interfaces.

      Pre-training reads data via an iterator. Import mpu to broadcast data across parallel workers.
      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

    2. Implement the three methods in your PreTrainer subclass.

      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']

    3. Initialize the engine, create a trainer, and save as rapidformer_pretrain_huggingface_bert_trainer.py.

      engine = RapidformerEngine()
trainer = MyBertPreTrainer(engine=engine)
trainer.train()

    4. Run via the CLI, pointing --user-script at the file from step 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

    White box: Hugging Face fine-tuning with a custom Trainer

    For existing Hugging Face training code that uses a custom Trainer, add Rapidformer APIs incrementally. This section shows the two most impactful changes: data parallelism and a faster optimizer.

    This method provides limited acceleration — Apex optimizer, model state partitioning, and computational graph optimization only. Mixed-precision training requires significant modifications to your training loop. For full acceleration, use the Finetuner template method instead.

    Add data parallelism

    Replace the standard DataLoader with finetuner.build_data_loader. This loader supports data parallelism and automatically moves batches to GPU, so remove batch.to(device) from your training loop.

    + 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"])

    Use the Apex Fused Adam optimizer

    Replace AdamW with Rapidformer's Apex Fused Adam optimizer. Call engine.compose to wrap the model, optimizer, and learning rate scheduler together.

    - 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)
    Combining Apex optimizer with mixed-precision training in a custom Trainer loop requires switching the model to FP16 and adding loss scaling — a significant modification. The Finetuner template integrates data parallelism, Apex, PyTorch mixed-precision training, Megatron optimizer mixed-precision, and VRAM optimization from FairScale and DeepSpeed automatically.

    White box: Megatron pre-training with the PreTrainer template

    For fully custom Megatron models — without Hugging Face or the Data/Model Hub — implement all three PreTrainer methods with custom logic. Datasets must inherit from torch.utils.data.Dataset; models must inherit from torch.nn.Module.

    This method also enables tensor parallelism and pipeline parallelism, which are available only with Megatron models.

    1. Create an mmap-format dataset using the 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

    2. Implement train_valid_test_datasets_provider with custom dataset creation logic.

      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

    3. Implement model_optimizer_lr_scheduler_provider with your custom model.

      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

    4. Implement run_forward_step with your custom forward pass logic.

      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

    5. Initialize the engine, create a trainer, and save as rapidformer_pretrain_megatron_gpt_trainer.py.

      engine = RapidformerEngine()
trainer = MyGPTPreTrainer(engine=engine)
trainer.train()

    6. Run via the CLI with tensor parallelism and pipeline parallelism enabled.

      #!/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

      This example enables all available Megatron-specific optimizations: tensor parallelism (--tensor-model-parallel-size 2), pipeline parallelism (--pipeline-model-parallel-size 2), gradient accumulation via global batch size, memory-mapped data loading, ZeRO Stage 2 memory optimization, gradient checkpointing, and mixed-precision training.

    What's next