NVIDIA NeMo Çerçevesi ile LLM Modeli İnceltme ve Bilgi Damıtma Yöntemleri

Model budama ve bilgi damıtma teknikleri, daha büyük bir dil modelinden daha küçük, daha etkili modeller elde etmenin güçlü ve maliyet etkin yollarıdır.

Budama: Katmanları çıkarmak (derinlik-budama) veya nöronları, dikkat başlıklarını ve gömme kanallarını çıkarmak (genişlik-budama).
Bilgi Damıtma: Büyük bir öğretici modelden küçük bir öğrenci modele bilgi aktarımı sağlamak, böylece daha hızlı ve daha az kaynak tüketen bir model oluşturmak hedeflenir.

NVIDIA‘nın Llama-3.1 8B modelini NVIDIA Llama-3.1-Minitron 4B modeline budama ve damıtma süreci hakkında detaylar sunan nasıl yapılır yazısında, derinlik, genişlik, dikkat ve MLP budama ile bilgi damıtması tabanlı yeniden eğitim uygulamalarının en iyi yöntemleri ele alınmıştır.

Genel Bilgi

Bu yazı, veri setini hazırlama, WikiText-103-v1 veri setinde öğretici modelin ince ayarını yapma, ardından modeli budama ve damıtma sürecini oluşturan basit bir hattı incelemeye odaklanmaktadır. WikiText-103-v1 veri seti, Wikipedia’da “İyi” ve “Özellikli” makalelerden çıkarılan 100 milyondan fazla token içeren bir koleksiyondur ve Hugging Face üzerinden kamuya açıktır.

Bu yazıda, aşağıdaki yüksek düzeydeki adımları içeren budama ve damıtma hattını tanımlayacağız:

A workflow diagram shows downloading the dataset, tokenizing, fine-tuning the 8B teacher dataset, pruning the teacher model, and distilling knowledge from teacher to student. — *Şekil 1. Veri setinin elde edilmesinden damıtılmış 4B modele ulaşma adımları.*

Hazırlık:
- Veri setini indirin ve JSONL formatına dönüştürün.
- Veri setini özelliklerine göre tokenize edin.
- Öğretici modeli veri seti üzerinde ince ayar yapın.
- İnce ayar yapılmış öğretici modelin derinliğini budayın. Derinliği budanmış model, öğrenci ağı için başlangıç noktasıdır.
- İnce ayar yapılmış öğretici modelin genişliğini budayın. Genişliği budanmış model, öğrenci ağı için başlangıç noktasıdır.
Öğreticiden öğrenciye bilgi damıtma işlemi, 8B modelini öğretici ve 4B budanmış modelini öğrenci olarak kullanarak gerçekleştirilir.

Bu eğitime ait Jupyter not defterlerine erişmek için NVIDIA/NeMo GitHub reposuna göz atabilirsiniz.

Gereksinimler

En az sekiz adet 80 GB bellek kapasitesine sahip NVIDIA GPU’ya (örneğin H100-80GB veya A100-80GB) ve bir Docker destekli ortama erişiminiz olmalıdır.

NeMo çerçevesini kurmak, Meta-Llama-3.1-8B öğretici modelini indirmek ve Hugging Face erişim jetonunu almak için projenin README dosyasını takip edebilirsiniz.

Veri setini indirin

WikiText-103-v1 veri setini indirin ve eğitim, test ve doğrulama parçalarını JSONL dosyalarına dönüştürmek için aşağıdaki kodu çalıştırın ya da tanıtım not defterini kullanın:

# Eğitim, test ve doğrulama dosyalarına ayırma

import json
import os
from datasets import load_dataset

# WikiText-103 veri setini yükleyin
dataset = load_dataset("wikitext", "wikitext-103-v1")

# Hedef klasörü tanımlayın
data_folder = 'wikitext-data'
os.makedirs(data_folder, exist_ok=True)

# Dosya yollarını ve hedef yolları tanımlayın
file_paths = {
    'train': os.path.join(data_folder, 'wikitext-train.jsonl'),
    'validation': os.path.join(data_folder, 'wikitext-val.jsonl'),
    'test': os.path.join(data_folder, 'wikitext-test.jsonl')
}

# JSONL dosyasına veri kümesi parçasını kaydetmek için fonksiyon
def save_to_jsonl(file_path, data):
    with open(file_path, 'w') as file:
        for item in data:
            file.write(json.dumps(item) + 'n')

# Tanımlı parçalar
splits = ["train", "validation", "test"]

# JSONL dosyalarına parçaları kaydedin ve boyutlarını hesaplayın
for split in splits:
    if split in dataset:
        save_to_jsonl(file_paths[split], dataset[split])
    else:
        print(f"{split} parçası veri setinde bulunamadı.")

Veri setini hazırlayın

Budama ve damıtma betikleri, veri dosyalarının Meta-Llama-3.1-8B tokenizasyon modelini kullanarak ön işlenmesi gerektiğini gerektirir. Bu işlem, NeMo çerçevesindeki preprocess_data_for_megatron.py ön işleme betiği ile yapılabilir.

Aşağıdaki betiği eğitim parçası üzerinde çalıştırarak veri setini budama ve damıtma için hazırlayın:

!python /opt/NeMo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py 
--input="./wikitext-data/wikitext-train.jsonl" 
--tokenizer-library='huggingface' 
--tokenizer-type='meta-llama/Meta-Llama-3.1-8B' 
--output-prefix=wikitext_tokenized_train 
--append-eod 
--workers=32

Aynı işlemi test ve doğrulama parçaları için de çalıştırın. Veri hazırlama not defteri, öğretici modeli ince ayarlamak için kullanılabilecek tokenleştirilmiş wikitext_tokenized_{train/val/test}_text_document.{idx/bin} dosyalarını oluşturmak için gereken tüm betikleri içerir.

Öğretici modeli veri setinde ince ayar yapın

Hazırlanan veri seti ile budanmamış öğretici model üzerinde hafif bir ince ayar yapın. Bu bölüm, betiklerin kullanımını göstermektedir; bu nedenle, ince ayar ayarları GLOBAL_BATCH_SIZE değeri 128 ve STEPS değeri 30 olarak ayarlanmıştır.

A workflow diagram shows multiple steps: input token, embedding, transformer layers, LM head, Softmax, Logits, cross-entropy loss, and next token. Steps are marked as trainable or loss. — *Şekil 2. Öğretici ince ayarı*

Distribütör kaybını düzeltmek için megatron_gpt_pretraining.py betiğini çalıştırın. Distribütör kaybı düzeltmesi yapılmazsa, öğretici modelin veriler üzerinde sağladığı rehberlik suboptimal olacaktır.

%%bash 

export CUDA_DEVICE_MAX_CONNECTIONS=1

# Farklı path'ler tanımlayın, eğer gerekirse:

MODEL="/workspace/llama-3_1-8b-nemo_v1.0/llama3_1_8b.nemo"

# Kaynaklara uygun olarak bunları değiştirebilirsiniz:

TENSOR_PARALLEL_SIZE=8
NODES=1
MICRO_BATCH_SIZE=4

# Aşağıdakileri değiştirmeyin:

EXPERIMENT_DIR="distill_trainings"
EXPERIMENT_NAME="megatron_llama_ft"

DATA_TRAIN='wikitext_tokenized_train_text_document'
DATA_VAL='wikitext_tokenized_test_text_document'
DATA_TEST='wikitext_tokenized_val_text_document'

STEPS=30
GLOBAL_BATCH_SIZE=128

LOG_INTERVAL=1
VAL_INTERVAL=10
NUM_VAL_BATCHES=5

LR=1e-4
MIN_LR=1e-5
WARMUP_STEPS=2

cmd="torchrun --nproc-per-node=${TENSOR_PARALLEL_SIZE}"

${cmd} /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py 
    --config-path /opt/NeMo/examples/nlp/language_modeling/conf/ 
    --config-name megatron_llama_distill.yaml 
    
    name=${EXPERIMENT_NAME} 
    
    exp_manager.exp_dir=${EXPERIMENT_DIR} 
    exp_manager.checkpoint_callback_params.save_top_k=1 
    exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True 
    
    trainer.max_steps=${STEPS} 
    trainer.log_every_n_steps=${LOG_INTERVAL}

Bu betiği çalıştırmak veya öğretici ince ayar not defterini çalıştırmak, ince ayarlı öğretici modeli oluşturur.

İnce Ayarlanmış Öğretici Modeli Budama

İnce ayarlı öğretici modelinizi budamak için iki yöntem kullanabilirsiniz: derinlik budama ve genişlik budama.

Derinlik budama yöntemi, bir 8B modelden 4B modele geçmek için son 16 katmanı (katman 16-31) budamayı içerir. Derinliği budamak için, megatron_gpt_drop_layers.py betiğini çalıştırın:

!python -m torch.distributed.launch --nproc_per_node=8 
     /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_drop_layers.py 
     --path_to_nemo "./distill_trainings/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo" 
     --path_to_save "/workspace/4b_depth_pruned_model.nemo" 
     --tensor_model_parallel_size 8 
     --pipeline_model_parallel_size 1 
     --gpus_per_node 8 
     --drop_layers 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

Bu betiği çalıştırmak veya derinlik budama not defterini uygulamak, son 16 katmanı kaldırarak daha küçük bir kontrol noktası oluşturur: 4b_depth_pruned_model.nemo.

Genişlik Budama

İkinci yöntem genişlik budama yöntemidir. 8B modelden 4B modele geçmek için MLP ara boyutunu ve gizli boyutu azaltarak dikkat başlıklarını ve katman sayılarını yeniden eğiterek budama işlemini gerçekleştirin.

Genişlik budamak için megatron_gpt_prune.py betiğini çalıştırın:

!torchrun --nproc-per-node=8 /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_prune.py 
     model.restore_from_path="./distill_trainings/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo" 
     model.tensor_model_parallel_size=1 
     model.pipeline_model_parallel_size=8 
     +model.dist_ckpt_load_strictness=log_all 
     inference.batch_size=64 
     trainer.num_nodes=1 
     trainer.precision=bf16 
     trainer.devices=8 
     prune.ffn_hidden_size=9216 
     prune.num_attention_heads=null 
     prune.num_query_groups=null 
     prune.hidden_size=3072 
     export.save_path="/workspace/4b_width_pruned_model.nemo"

Bu betiği çalıştırmak veya genişlik budama not defterini çalıştırmak, daha küçük bir genişlik budanmış kontrol noktası oluşturur: 4b_width_pruned_model.nemo.

Öğreticiden Öğrenciye Bilgi Damıtma

Damıtma süreci, ince ayar yapılmış modeli (8B) öğretici model olarak ve budanmış modeli (4B) öğrenci modeli olarak kullanarak damıtma işlemini gerçekleştirir. NeMo’da yalnızca logit kaybı şu an kullanılmaktadır.

A workflow diagram shows classical knowledge distillation from teacher to student, with loss function from several layers of the transformer architecture. A student model with N layers is distilled from a teacher model with M layers. The student learns by minimizing a combination of embedding output loss, logit loss and transformer encoder specific losses mapped across student block S and teacher block T. — *Şekil 4. Damıtma iş akışı*

Bu bölümde, öğretici modelden her iki öğrenci modeline bilgi damıtacağız ve bunları karşılaştıracağız:

İnce ayar yapılmış öğreticiden derinlik budaması yapılmış öğrenciye bilgi damıtma.
İnce ayar yapılmış öğreticiden genişlik budaması yapılmış öğrenciye bilgi damıtma.

Derinlik Budaması Yapılmış Öğrenciye Bilgi Damıtma

Öğretici modelden derinlik budamalı öğrenci modeline bilgi damıtmak için megatron_gpt_distillation.py betiğini çalıştırın.

%%bash 

export CUDA_DEVICE_MAX_CONNECTIONS=1

# Kaynaklara uygun olarak bunları değiştirebilirsiniz:

TENSOR_PARALLEL_SIZE=8
NODES=1
MICRO_BATCH_SIZE=4

# Aşağıdakileri değiştirmeyin:

EXPERIMENT_DIR="distill_trainings"
EXPERIMENT_NAME="megatron_llama_distill_depth_pruned_student"

TEACHER="${EXPERIMENT_DIR}/megatron_llama_ft/checkpoints/megatron_llama_ft.nemo"
STUDENT="/workspace/4b_depth_pruned_model.nemo"

FINAL_MODEL_PATH="${EXPERIMENT_DIR}/${EXPERIMENT_NAME}/checkpoints/depth_pruned_distilled_4b_model.nemo"

DATA_TRAIN='wikitext_tokenized_train_text_document'
DATA_VAL='wikitext_tokenized_test_text_document'
DATA_TEST='wikitext_tokenized_val_text_document'

STEPS=30
GLOBAL_BATCH_SIZE=128

LOG_INTERVAL=1
VAL_INTERVAL=10
NUM_VAL_BATCHES=5

LR=1e-4
MIN_LR=1e-5
WARMUP_STEPS=2

cmd="torchrun --nproc-per-node=${TENSOR_PARALLEL_SIZE}"

${cmd} /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_distillation.py 
    name=${EXPERIMENT_NAME} 
    
    exp_manager.exp_dir=${EXPERIMENT_DIR} 
    exp_manager.checkpoint_callback_params.save_top_k=1 
    
    trainer.max_steps=${STEPS} 
    trainer.log_every_n_steps=${LOG_INTERVAL} 
    trainer.val_check_interval=${VAL_INTERVAL} 
    trainer.limit_val_batches=${NUM_VAL_BATCHES} 
    +trainer.num_sanity_val_steps=0 
    
    trainer.precision=bf16 
    trainer.devices=${TENSOR_PARALLEL_SIZE} 
    trainer.num_nodes=${NODES} 
    
    "model.data.data_prefix={train:[1.0,$DATA_TRAIN],validation:[$DATA_VAL],test:[$DATA_TEST]}" 
    
    model.restore_from_path=${STUDENT} 
    model.kd_teacher_restore_from_path=${TEACHER} 
    model.nemo_path=${FINAL_MODEL_PATH} 
    
    model.tensor_model_parallel_size=${TENSOR_PARALLEL_SIZE} 
    model.sequence_parallel=True 
    model.micro_batch_size=${MICRO_BATCH_SIZE} 
    model.global_batch_size=${GLOBAL_BATCH_SIZE} 
    
    model.optim.name=distributed_fused_adam 
    model.optim.lr=${LR} 
    model.optim.sched.min_lr=${MIN_LR} 
    model.optim.sched.warmup_steps=${WARMUP_STEPS}

Bu betiği çalıştırmak veya derinlik budamalı öğrenci damıtma not defterini çalıştırmak, damıtılmış model oluşturur: depth_pruned_distilled_4b_model.nemo.

Genişlik Budaması Yapılmış Öğrenciye Bilgi Damıtma

Öğretici model ile genişlik budamalı öğrenci modeline bilgi damıtmak için, megatron_gpt_distillation.py betiğini yine çalıştırın. Burada, öğrenci modelini (STUDENT) ve damıtılmış modelin kaydedileceği dizini (FINAL_MODEL_PATH) değiştirerek çalıştırmalısınız.

Bu işlemi çalıştırarak, width_pruned_distilled_4b_model.nemo adında damıtılmış bir model oluşturabilirsiniz.

Doğrulama Kaybını Sergile

Doğrulama kaybını görselleştirmek için aşağıdaki kodu çalıştırın veya sonuçlar not defterini kullanın. Çalıştırmadan önce kontrol noktasının yolunu değiştirin:

%load_ext tensorboard
%tensorboard --logdir "distill_trainings/megatron_llama_distill/" --port=6007

Damıtma betiğini 30 STEPS değeri ile çalıştırdıktan sonra, derinlik ve genişlik budamalı öğrencilerin doğrulama kaybını Şekil 5 ve Şekil 6’da görebilirsiniz.

A plot shows the validation loss under 8 after running the training step in the distillation script for 30 steps with the depth-pruned student. — *Şekil 5. Derinlik budamalı doğrulama kaybı, 30 adımda*

A plot shows the validation loss under 8 after running the training step in the distillation script for 30 steps with the width-pruned student. — *Şekil 6. Genişlik budamalı doğrulama kaybı, 30 adımda*

Pipelinesinizi kendi kullanım durumunuza göre ayarlamak için, betikleri çoklu düğümlü bir küme üzerinde yüksek GLOBAL_BATCH_SIZE, STEPS ve VAL_INTERVAL değerleri ile çalıştırarak doğrulama kaybını geliştirebilirsiniz.

Şekil 7 ve Şekil 8, 880 STEPS ve 2048 GLOBAL_BATCH_SIZE değerinde çalıştırıldığında doğrulama kaybının azaldığını göstermektedir.

A plot shows the validation loss under 2.5 after running the training step in the distillation script with the depth-pruned model as the student. — *Şekil 7. Derinlik budamalı doğrulama kaybı, 880 adımda (GBS=2048 ile)*

A plot shows the validation loss drop to under 2.5 after running the training step in the distillation script with the width-pruned model as the student. — *Şekil 8. Genişlik budamalı doğrulama kaybı, 880 adımda (GBS=2048 ile)*

Sonuç

Budama ve damıtma, dil model optimizasyonu alanında büyük bir ilerleme temsil etmektedir. Llama-3.1-Minitron-4B gibi daha küçük ve daha verimli modellerin, kaynak kısıtlı ortamlarda performansı koruyarak kullanıma sunulabilmesi, AI endüstrisi için dönüştürücü bir etkendir.

Mistral-NeMo-Minitron-8B modelinin bu yaklaşım kullanılarak geliştirildiği ve Llama-3.1-8B modelinin birçok testte daha iyi performans gösterdiği belirtilmektedir.

Bu yöntem, hesaplama maliyetlerini ve enerji tüketimini azaltarak verimlilik sağlamanın yanı sıra, mobil cihazlar, kenar bilişimi ve kısıtlı kaynak ayarlarında ileri düzey NLP yeteneklerinin kullanımını demokratikleştiriyor. Bu teknikler gelişmeye devam ettikçe, daha kompakt ancak güçlü dil modelleri beklenmektedir ve bu da teknolojinin çeşitli endüstrilere yayılma kapılarını açmaktadır.

Daha fazla bilgi için şu kaynaklara göz atabilirsiniz: