Megatron-LM

Megatron-LM memungkinkan pelatihan model bahasa transformer berskala besar. Ia menyediakan paralelisme model berbasis tensor, pipeline, dan sequence yang efisien untuk pre-training model bahasa berbasis transformer seperti GPT (Decoder Only), BERT (Encoder Only), dan T5 (Encoder-Decoder). Untuk informasi detail dan cara kerjanya di balik layar, silakan merujuk ke repo github.

Apa yang terintegrasi?

Accelerate mengintegrasikan fitur berikut dari Megatron-LM untuk memungkinkan pre-training/finetuning skala besar dari model BERT (Encoder), GPT (Decoder) atau T5 (Encoder dan Decoder):

a. Tensor Parallelism (TP): Mengurangi jejak memori tanpa banyak komunikasi tambahan pada rank intra-node. Setiap tensor dipecah menjadi beberapa chunk, dengan setiap shard berada di GPU yang terpisah. Pada setiap langkah, mini-batch data yang sama diproses secara independen dan paralel oleh setiap shard, diikuti dengan sinkronisasi di semua GPU (operasi all-reduce). Pada lapisan transformer sederhana, hal ini menghasilkan 2 all-reduces di jalur forward dan 2 di jalur backward. Untuk detail lebih lanjut, silakan merujuk ke makalah penelitian Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism dan bagian ini dari blogpost The Technology Behind BLOOM Training.

b. Parallelisme Pipeline (PP): Mengurangi jejak memori dan memungkinkan pelatihan skala besar melalui paralelisasi antar-node. Mengurangi gelembung dari PP naif melalui jadwal PipeDream-Flush/jadwal 1F1B dan jadwal Interleaved 1F1B. Lapisan didistribusikan secara seragam di antara tahap PP. Misalnya, jika sebuah model memiliki 24 lapisan dan kami memiliki 4 GPU untuk paralelisasi pipeline, setiap GPU akan memiliki 6 lapisan (24/4). Untuk detail lebih lanjut tentang jadwal untuk mengurangi waktu idle PP, silakan merujuk ke makalah penelitian Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM dan bagian ini dari blogpost The Technology Behind BLOOM Training.

c. Sequence Parallelism (SP): Mengurangi jejak memori tanpa komunikasi tambahan. Hanya berlaku saat menggunakan TP. Hal ini mengurangi memori aktivasi yang diperlukan karena mencegah salinan yang sama berada di rank tensor parallel setelah all-reduce dengan menggantinya menggunakan reduce-scatter dan operasi no-op akan diganti dengan all-gather. Karena all-reduce = reduce-scatter + all-gather, ini menghemat banyak memori aktivasi tanpa biaya komunikasi tambahan. Secara sederhana, ini memecah output dari setiap lapisan transformer sepanjang dimensi sekuens, misalnya, jika panjang sekuens adalah 1024 dan ukuran TP adalah 4, setiap GPU akan memiliki 256 token (1024/4) untuk setiap sampel. Ini meningkatkan ukuran batch yang dapat didukung untuk pelatihan. Untuk detail lebih lanjut, silakan merujuk ke makalah penelitian Reducing Activation Recomputation in Large Transformer Models.

d. Data Parallelism (DP) via Distributed Optimizer: Mengurangi jejak memori dengan membagi (sharding) state optimizer dan gradien di antara rank DP (berbanding dengan metode tradisional yang mereplikasi state optimizer di antara rank data parallel). Sebagai contoh, ketika menggunakan optimizer Adam dengan pelatihan mixed-precision, setiap parameter menggunakan 12 byte memori. Ini didistribusikan secara merata di antara GPU, yaitu, setiap parameter akan menggunakan 3 byte (12/4) jika kita memiliki 4 GPU. Untuk detail lebih lanjut, silakan merujuk pada makalah penelitian ZeRO: Memory Optimizations Toward Training Trillion Parameter Models dan bagian berikut dari blog The Technology Behind BLOOM Training.

e. Expert Parallelism (EP) Expert parallelism pada Megatron-LM digunakan untuk lapisan Mixture-of-Experts (MoE), di mana banyak “experts” (jaringan feed-forward kecil) ada tetapi hanya beberapa yang diaktifkan untuk setiap token. Alih-alih menempatkan semua experts di setiap GPU, Megatron mendistribusikan experts yang berbeda ke GPU yang berbeda—ini adalah expert parallelism. Selama pelatihan, token di-rutekan ke GPU yang menampung experts yang dipilihnya, dihitung di sana, dan kemudian dikirim kembali, sehingga mengurangi biaya memori. Hal ini sering dikombinasikan dengan tensor/pipeline parallelism untuk model berskala besar. f. Full Activation Recomputation: Mengurangi jejak memori dari activations secara signifikan melalui smart activation checkpointing. Hal ini tidak menyimpan activations yang menempati memori besar tetapi cepat untuk dihitung ulang, sehingga mencapai tradeoff yang sangat baik antara memori dan recomputation. Sebagai contoh, untuk GPT-3, hal ini menghasilkan pengurangan 70% pada memori yang dibutuhkan untuk activations dengan biaya hanya 2.7% overhead FLOPs untuk recomputation dari activations. Untuk detail lebih lanjut, silakan merujuk ke makalah penelitian Reducing Activation Recomputation in Large Transformer Models.

g. Fused Kernels: Fused Softmax, Mixed Precision Fused Layer Norm, dan Fused gradient accumulation ke dalam perhitungan weight gradient dari linear layer. PyTorch JIT compiled Fused GeLU dan Fused Bias+Dropout+Residual addition.

h. Dukungan untuk dataset berindeks: Format biner yang efisien untuk dataset guna pelatihan skala besar. Mendukung format file indeks mmap, cached, dan loader lazy.

i. Pembentukan ulang dan interoperabilitas checkpoint: Utilitas untuk membentuk ulang checkpoint Megatron-LM dengan ukuran paralel tensor dan pipeline yang bervariasi menjadi checkpoint sharded Transformers yang populer, karena memiliki dukungan yang sangat baik dengan berbagai macam alat seperti Accelerate Big Model Inference, Megatron-DeepSpeed Inference, dan lain-lain. Dukungan juga tersedia untuk mengonversi checkpoint sharded Transformers menjadi checkpoint Megatron-LM dengan ukuran paralel tensor dan pipeline yang bervariasi untuk pelatihan skala besar.

Prasyarat

Anda perlu menginstal rilis terbaru pytorch, cuda, nccl, dan NVIDIA APEX serta pustaka nltk. Lihat dokumentasi untuk detail lebih lanjut. Cara lain untuk menyiapkan lingkungan adalah dengan menarik NVIDIA PyTorch Container yang sudah dilengkapi dengan semua instalasi yang diperlukan dari NGC.

Berikut adalah metode langkah demi langkah untuk mengatur lingkungan conda:

1. Buat lingkungan virtual

conda create --name ml

2. Dengan asumsi bahwa mesin tersebut telah menginstal CUDA 11.3, menginstal versi GPU PyTorch yang sesuai

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

3. Install 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. Menginstal 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 .

Siapkan checkpoint Megaton-LM

Jika Anda ingin melakukan fine-tuning pada model, pastikan Anda memiliki checkpoint dalam format torch dist yang siap digunakan. Jika Anda hanya memiliki akses ke model huggingface, pertimbangkan untuk mengonversinya menjadi checkpoint dalam format torch dist yang dapat diterima oleh Megatron. Satu contoh adalah dengan menggunakan skrip dari slime, dengan mengambil model GLM sebagai contoh:

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

Setelah konversi, pastikan: 1. di bawah /your/path/to/megatron/models/GLM4.5-355B-A32B_torch_dist: ubah isi latest_checkpointed_iteration.txt dari release menjadi 0 dan ganti nama direktori release menjadi iter_0000000; 2: di config, pastikan megatron_lm_no_load_optim bernilai true sehingga tidak diperlukan state optimizer.

Plugin Accelerate Megatron-LM

Fitur penting didukung secara langsung melalui perintah accelerate config. Contoh pertanyaan yang sesuai untuk menggunakan fitur Megatron-LM ditunjukkan di bawah ini:

:~$ 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

Konfigurasi yang dihasilkan ditampilkan di bawah ini:

~$ 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

Kami akan mengambil contoh pra-pelatihan GPT. Perubahan minimal yang diperlukan pada run_clm_no_trainer.py resmi untuk menggunakan Megatron-LM adalah sebagai berikut:

1. Karena Megatron-LM menggunakan implementasi Optimizer-nya sendiri, scheduler yang kompatibel dengannya perlu digunakan. Oleh karena itu, dukungan hanya tersedia untuk scheduler Megatron-LM. Pengguna perlu membuat accelerate.utils.MegatronLMDummyScheduler. Contoh diberikan di bawah ini:

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. Mendapatkan detail ukuran batch total sekarang perlu mempertimbangkan ukuran paralel tensor dan pipeline. Contoh mendapatkan ukuran batch total efektif ditunjukkan di bawah ini:

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. Saat menggunakan Megatron-LM, kerugian (losses) sudah dirata-rata di seluruh grup paralel data

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. Untuk Megatron-LM, kita perlu menyimpan model menggunakan 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
    )

Itu saja! Kita sudah siap 🚀. Silakan temukan skrip contoh di folder examples pada path accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py. Mari kita jalankan untuk arsitektur model gpt-large menggunakan 4 GPU 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"

Berikut adalah beberapa kutipan penting dari log output:

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]

Terdapat banyak opsi/fitur lain yang dapat diatur menggunakan accelerate.utils.MegatronLMPlugin.

Fitur lanjutan untuk memanfaatkan penulisan langkah pelatihan kustom dan Dataset Berindeks Megatron-LM

Untuk memanfaatkan lebih banyak fitur, silakan perhatikan detail di bawah ini.

1. Berikut adalah contoh perubahan yang diperlukan untuk menyesuaikan Train Step saat menggunakan Megatron-LM. Anda akan mengimplementasikan accelerate.utils.AbstractTrainStep atau mewarisi dari anak-anak yang sesuai accelerate.utils.GPTTrainStep, accelerate.utils.BertTrainStep atau 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. Untuk menggunakan dataset Megatron-LM, beberapa perubahan tambahan diperlukan. Dataloader untuk dataset ini hanya tersedia pada rank 0 dari setiap grup tensor parallel. Dengan demikian, ada rank di mana dataloader tidak akan tersedia, dan hal ini memerlukan penyesuaian pada training loop. Kemampuan melakukan semua ini menunjukkan betapa fleksibel dan dapat diperluasnya Accelerate. Perubahan yang diperlukan adalah sebagai berikut.

a. Untuk dataset terindeks Megatron-LM, kita perlu menggunakan MegatronLMDummyDataLoader dan meneruskan argumen dataset yang diperlukan ke dalamnya, seperti data_path, seq_length, dan sebagainya. Lihat di sini untuk daftar argumen yang tersedia.

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 diulang 3 kali untuk mendapatkan dataloader pelatihan, validasi, dan pengujian sesuai proporsi args.splits_string

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

c. Perubahan pada loop pelatihan dan evaluasi karena dataloader hanya tersedia pada rank tensor parallel 0 Jadi, kita perlu melakukan iterasi hanya jika dataloader bukan None, selain itu berikan dict kosong Dengan demikian, kita melakukan loop menggunakan loop while dan berhenti ketika completed_steps sama dengan args.max_train_steps Ini mirip dengan setup Megatron-LM, di mana pengguna harus menyediakan max_train_steps saat menggunakan dataset berindeks Megaton-LM. Hal ini menunjukkan betapa fleksibel dan dapat diperluasnya Accelerate.

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)

Utilitas untuk pembentukan ulang dan interoperabilitas checkpoint

1. Skrip untuk ini tersedia di pustaka Transformers di bawah model yang bersangkutan. Saat ini, tersedia untuk model GPT checkpoint_reshaping_and_interoperability.py

2. Di bawah ini adalah contoh konversi checkpoint dari Megatron-LM menjadi checkpoint sharded universal 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. Konversi checkpoint dari transformers ke megatron dengan tp_size=2, pp_size=2 dan 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

Model GPT Megatron-LM mendukung pengembalian logits dan fungsi megatron_generate untuk generasi teks

1. Mengembalikan logits memerlukan pengaturan require_logits=True di MegatronLMPlugin seperti yang ditunjukkan di bawah ini. Nilai-nilai ini akan tersedia di tahap terakhir pipeline.

megatron_lm_plugin = MegatronLMPlugin(return_logits=True)

2. Metode megatron_generate untuk model GPT Megatron-LM: Metode ini akan menggunakan Tensor dan Pipeline Parallelism untuk menyelesaikan generasi untuk batch input saat menggunakan greedy dengan/without top_k/top_p sampling dan untuk input prompt individual saat menggunakan beam search decoding. Hanya subset fitur dari transformers generate yang didukung. Hal ini akan membantu dalam menggunakan model besar melalui tensor dan pipeline parallelism untuk generasi (sudah melakukan key-value caching dan menggunakan fused kernels secara default). Ini memerlukan data parallel size untuk menjadi 1, sequence parallelism dan activation checkpointing untuk dinonaktifkan. Ini juga memerlukan penentuan path ke file vocab tokenizer dan file merges. Contoh di bawah ini menunjukkan cara mengonfigurasi dan menggunakan metode megatron_generate untuk model GPT 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. Contoh end-to-end penggunaan metode megatron_generate untuk model GPT Megatron-LM tersedia di megatron_gpt2_generation.py dengan file konfigurasi megatron_lm_gpt_generate_config.yaml. Skrip bash dengan perintah accelerate launch tersedia di megatron_lm_gpt_generate.sh. Log output dari skrip tersebut tersedia di megatron_lm_gpt_generate.log.

Dukungan untuk ROPE dan ALiBi Positional embeddings serta Multi-Query Attention

1. Untuk perhatian ROPE/ALiBi, berikan position_embedding_type dengan ("absolute" | "rotary" | "alibi") ke MegatronLMPlugin seperti yang ditunjukkan di bawah ini.

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

2. Untuk Multi-Query Attention, berikan attention_head_type dengan ("multihead" | "multiquery") ke MegatronLMPlugin seperti yang ditunjukkan di bawah ini.

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

Catatan

1. Mendukung model Transformers GPT2, Megatron-BERT, dan T5. Ini mencakup kelas model Decoder only, Encode only, dan Encoder-Decoder.

2. Hanya loss yang dikembalikan dari forward pass model karena terdapat interaksi yang cukup kompleks antara pipeline, tensor, dan data parallelism di balik layar. Panggilan model(**batch_data) mengembalikan loss(es) yang dirata-rata di seluruh data parallel ranks. Hal ini sudah memadai untuk sebagian besar kasus di mana tugas pre-training dijalankan menggunakan fitur Megatron-LM dan Anda dapat dengan mudah menghitung perplexity menggunakan loss. Untuk model GPT, pengembalian logits di samping loss(es) juga didukung. Logits ini tidak dikumpulkan di seluruh data parallel ranks. Gunakan accelerator.utils.gather_across_data_parallel_groups untuk mengumpulkan logits di seluruh data parallel ranks. Logits ini bersama dengan labels dapat digunakan untuk menghitung berbagai metrik kinerja.

3. Proses utama adalah rank terakhir karena losses/logits tersedia di tahap terakhir pipeline. accelerator.is_main_process dan accelerator.is_local_main_process mengembalikan True untuk rank terakhir saat menggunakan integrasi Megatron-LM.

4. Dalam panggilan accelerator.prepare, model Megatron-LM yang sesuai dengan model Transformers tertentu dibuat dengan bobot acak. Silakan gunakan accelerator.load_state untuk memuat checkpoint Megatron-LM dengan partisi TP, PP, dan DP yang sesuai.

5. Saat ini, dukungan untuk pembentukan ulang checkpoint dan interoperabilitas hanya tersedia untuk GPT. Segera akan diperluas ke BERT dan T5.

6. gradient_accumulation_steps harus bernilai 1. Saat menggunakan Megatron-LM, micro batches dalam pengaturan pipeline parallelism bersinonim dengan gradient accumulation.

7. Saat menggunakan Megatron-LM, gunakan accelerator.save_state dan accelerator.load_state untuk menyimpan dan memuat checkpoint.

8. Berikut adalah pemetaan dari arsitektur model Megatron-LM ke arsitektur model transformers yang setara. Hanya arsitektur model transformers ini yang didukung.

a. Megatron-LM BertModel : model transformers dengan megatron-bert pada tipe model di config, misalnya MegatronBERT

b. Megatron-LM GPTModel : model transformers dengan gpt2 pada tipe model di config, misalnya OpenAI GPT2

c. Megatron-LM T5Model : model transformers dengan t5 dalam tipe model config, misalnya T5 dan MT5