How GPT‑MoE Cuts Training Costs: Sparse Transformer Techniques and Performance Insights

Overview

GPT models excel at text generation tasks such as completion, QA, summarization, and creative writing, but their training costs are extremely high. For example, training a 175‑billion‑parameter GPT‑3 on 1024 A100 GPUs would take about 34 days.

MoE Sparse Training

Mixture‑of‑Experts (MoE) introduces sparsity by selecting one expert (an MLP layer) per token during training, allowing the model to increase parameter count without raising FLOPs. This yields up to 1.2× training throughput and 1.3× inference throughput compared to dense models of similar quality.

We combine MoE with a GPT decoder architecture because decoder‑based routing performs better than encoder‑based routing for language modeling.

Expert Routing

We adopt the top‑1 routing mechanism from Switch Transformer, where each expert receives a probability from a softmax function and the highest‑probability expert processes the token.

W_r = ...  // routing weight matrix learned during training

Performance Analysis

Eight GPT model configurations (dense and MoE) were evaluated. The 1.3B+MoE‑32/64 models achieved lower validation loss than the dense 1.3B model, and the 0.35B+MoE‑64 model showed the fastest training throughput (≈2× other models).

Inference throughput favored the 1.3B dense model for memory usage, while the 0.35B+MoE‑64 model achieved the lowest latency.

Training on a single A100 node for 200 hours demonstrated that the 1.3B+MoE‑64 model converged 1.17× faster than a 2.7B dense model, while the 1.3B+MoE‑32 model lagged by 15%.

Zero‑Shot NLU and Text Generation Evaluation

We benchmarked Chinese zero‑shot NLU and various text‑generation tasks (completion, poetry, advertising copy, essay) using the MoE models, showing competitive or superior quality compared to dense baselines.

Algorithm Innovations

Top‑1 Gating Limitations

Top‑1 gating can cause load imbalance among experts. To address this, we designed a new routing algorithm that allows each token to be processed by multiple experts, fixing a capacity per expert and using weighted sums of expert outputs.

Training Techniques

Mixed‑precision training to halve memory usage and speed up computation.

Selective activation recomputation to checkpoint intermediate activations and reduce memory.

Zero Redundancy Optimizer (ZeRO‑1) for distributed memory savings.

Sequence parallelism to split long sequences across devices, reducing per‑device workload.

Practical Deployment on PAI

Pre‑training with PAI DLC

Using Alibaba Cloud PAI DLC, we launch a training job with a container image, dataset mounted from OSS, and command line arguments such as model size, MoE expert count, batch size, and ZeRO configuration.

cd /workspace/RapidformerPro/examples/megatron && \
bash dlc_run_pretrain_megatron_gpt.sh run 1 jiebabpe 0.125B 8 0 1 1 sel none 10000

Fine‑tuning with PAI DSW

After pre‑training, we fine‑tune the model on downstream tasks (e.g., poetry generation) using PAI DSW, again mounting checkpoints from OSS and running a similar script.

!wget https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/text_generation_datasets/poetry/train.tsv
!wget https://atp-modelzoo-sh.oss-cn-shanghai.aliyuncs.com/release/tutorials/text_generation_datasets/poetry/dev.tsv

Training scripts import necessary libraries and initialize the model, tokenizer, and engine before calling finetuner.train().

Online Inference Deployment

We convert the fine‑tuned model to FasterTransformer format and develop a processor for PAI‑EAS that handles input JSON, tokenization, generation parameters (max_length, top_k, temperature, etc.), and returns generated text.

def process(self, data):
    data_str = data.decode('utf-8')
    data_json = json.loads(data_str)
    # tokenization and generation logic using the GPT model
    result_dict = {'text': outputs[0]}
    return get_result_str(result_dict=result_dict)

The service is then deployed, providing online text‑generation APIs for poetry, advertising copy, and essay generation.

References

[1] Outrageously Large Neural Networks: The Sparsely‑Gated Mixture‑of‑Experts Layer [2] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding [3] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity [4] BASE Layers: Simplifying Training of Large, Sparse Models [5] Hash Layers For Large Sparse Models [6] Taming Sparsely Activated Transformer with Stochastic Experts [7] GLaM: Efficient Scaling of Language Models with Mixture‑of‑Experts [8] Unified Scaling Laws for Routed Language Models [9] Designing Effective Sparse Expert Models [10] Large Margin Deep Networks for Classification