تحسين كفاءة تدريب الوجه على المعانقة من خلال التعبئة مع انتباه الفلاش 2

أصبح التدريب باستخدام أمثلة ضبط التعليمات المجمعة (بدون الحشو) متوافقًا الآن مع Flash Attention 2 في Hugging Face، وذلك بفضل العلاقات العامة الحديثة وDataCollatorWithFlattening الجديد.

يمكن أن يوفر تحسينًا يصل إلى 2x في إنتاجية التدريب مع الحفاظ على جودة التقارب. واصل القراءة للحصول على التفاصيل!

مقدمة

يعد حشو تسلسلات الإدخال على دفعات صغيرة طريقة معتادة لجمع المدخلات أثناء التدريب. ومع ذلك، يؤدي هذا إلى عدم الكفاءة بسبب رموز الحشو غير ذات الصلة. يعد تجميع الأمثلة بدون حشوة، واستخدام معلومات موضع الرمز المميز، بديلاً أكثر كفاءة. ومع ذلك، فإن التطبيقات السابقة للتعبئة لم تأخذ في الاعتبار حدود المثال عند استخدام Flash Attention 2، مما أدى إلى انتباه غير مرغوب فيه للأمثلة التبادلية مما يقلل من الجودة والتقارب.

تعالج Hugging Face Transformers الآن هذه المشكلة من خلال ميزة جديدة تحافظ على الوعي بالحدود أثناء التعبئة، إلى جانب تقديم أداة تجميع البيانات الجديدة، DataCollatorWithFlattening.

عن طريق الاختيار DataCollatorWithFlattening، معانقة الوجه Trainer يمكن للمستخدمين الآن ربط التسلسلات بسلاسة في موتر واحد مع مراعاة حدود التسلسل أثناء حسابات Flash Attention 2. ويتم تحقيق ذلك من خلال flash_attn_varlen_func، الذي يحسب أطوال التسلسل التراكمي في كل دفعة صغيرة (cu_seqlens).

نفس الميزة متاحة لـ Hugging Face SFTTrainer المستخدمين في TRL المكتبة عن طريق وضع علامة جديدة، padding_free=True، عند استدعاء مجمّع البيانات DataCollatorForCompletionOnlyLM.

زيادة في الإنتاجية تصل إلى 2x

نحن نرى تحسنًا كبيرًا في إنتاجية التدريب باستخدام هذه الميزة مع الجديد DataCollatorWithFlattening. يوضح الشكل أدناه الإنتاجية المقاسة بالرموز المميزة/الثانية أثناء التدريب. في هذا المثال، يكون معدل النقل هو متوسط ​​كل وحدة معالجة رسومات (GPU) على مدى 8 معالجات رسومية A100-80 خلال فترة واحدة من عينة مختارة عشوائيًا تبلغ 20 ألفًا من مجموعتي بيانات مختلفتين لضبط التعليمات، FLAN وOrcaMath.

تحتوي FLAN على تسلسلات قصيرة في المتوسط ​​ولكن يوجد تباين كبير في طول التسلسل، لذلك قد تختلف أطوال المثال في كل دفعة بشكل كبير. وهذا يعني أن دفعات FLAN المبطنة قد تتحمل حملًا كبيرًا في رموز الحشو غير المستخدمة. يُظهر التدريب على مجموعة بيانات FLAN فائدة كبيرة باستخدام الجديد DataCollatorWithFlattening من حيث زيادة الإنتاجية. نرى زيادة في الإنتاجية بمقدار 2x في النماذج الموضحة هنا: llama2-7B، وmistral-7B، و Granite-8B-code.

يحتوي OrcaMath على أمثلة أطول وتباين أقل في طول المثال. على هذا النحو، فإن التحسن من التعبئة أقل. تُظهر تجاربنا زيادة بمقدار 1.4 مرة في الإنتاجية عند التدريب باستخدام هذا النوع من التعبئة على مجموعة بيانات OrcaMath عبر هذه النماذج الثلاثة.

يتحسن استخدام الذاكرة أيضًا من خلال التعبئة بالجهاز الجديد DataCollatorWithFlattening. يوضح الشكل التالي ذروة استخدام الذاكرة لنفس النماذج الثلاثة التي يتم تدريبها على نفس مجموعتي البيانات. يتم تقليل الذاكرة القصوى بنسبة 20% في مجموعة بيانات FLAN، والتي تستفيد بشكل كبير من التعبئة.

يبلغ الحد الأقصى لخفض الذاكرة 6% في مجموعة بيانات OrcaMath بأطوال الأمثلة الأكثر تجانسًا.

أمثلة التعبئة، عندما تقلل من عدد خطوات التحسين، قد تضر بتقارب التدريب. ومع ذلك، تحتفظ الميزة الجديدة بالدفعات الصغيرة، وبالتالي، بنفس عدد خطوات التحسين التي سيتم استخدامها مع الأمثلة المبطنة. وبالتالي، لا يوجد أي تأثير على تقارب القطار، كما نرى في الشكل التالي، والذي يوضح خسارة التحقق المتطابقة لنفس النماذج الثلاثة التي يتم تدريبها على نفس مجموعتي البيانات، سواء تم تدريب النماذج على التعبئة باستخدام الطريقة الجديدة DataCollatorWithFlattening أو مع الحشو.

كيف يعمل

ضع في اعتبارك مجموعة من البيانات بحجم الدُفعة = 4 حيث تكون التسلسلات الأربعة كما يلي:

بعد تسلسل الأمثلة، يقوم المجمّع الخالي من الحشو بإرجاع الملف input_ids, labels، و position_ids من كل مثال. وبالتالي، يوفر المجمع، لهذه الدفعة من البيانات،

التعديلات المطلوبة خفيفة الوزن وتقتصر على توفير position_ids إلى فلاش الانتباه 2.

لكن هذا يعتمد على تعريض النموذج position_ids. وحتى وقت كتابة هذا المقال، هناك 14 نموذجًا يعرضونها ويدعمها الحل. على وجه التحديد، يتم دعم Llama 2 و3 وMistral وMixtral وGranite وDBRX وFalcon وGemma وOLMo وPhi 1 و2 و3 وphi3 وQwen 2 و2 MoE وStableLM وStarCoder 2 بواسطة الحل.

ابدء

جني فوائد التعبئة مع position_ids سهل.

إذا كنت تستخدم معانقة الوجه Trainer من Transformers، مطلوب خطوتين فقط:

1. إنشاء مثيل للنموذج باستخدام Flash Attention 2
2. استخدم الجديد DataCollatorWithFlattening

إذا كنت تستخدم معانقة الوجه SFTTrainer من TRL مع DataCollatorForCompletionOnlyLM، فالخطوتان المطلوبتان هما:

1. إنشاء مثيل للنموذج باستخدام Flash Attention 2
2. تعيين padding_free=True عند الاتصال DataCollatorForCompletionOnlyLM على النحو التالي:
   collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer, padding_free=True)

كيفية استخدامه

ل Trainer للمستخدمين، يوضح المثال أدناه كيفية استخدام الميزة الجديدة.

 
import torch


from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "instructlab/merlinite-7b-lab",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)


from datasets import load_dataset
train_dataset = load_dataset("json", data_files="path/to/my/dataset")["train"]


from transformers import DataCollatorWithFlattening
data_collator = DataCollatorWithFlattening()


from transformers import TrainingArguments, Trainer
train_args = TrainingArguments(output_dir="/save/path")
trainer = Trainer(
    args=train_args,
    model=model,
    train_dataset=train_dataset,
    data_collator=data_collator
)
trainer.train()

ل TRL للمستخدمين، يوضح المثال أدناه كيفية استخدام الميزة الجديدة مع SFTTrainer.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("lucasmccabe-lmi/CodeAlpaca-20k", split="train")

model = AutoModelForCausalLM.from_pretrained(
    "instructlab/merlinite-7b-lab",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained("instructlab/merlinite-7b-lab")
tokenizer.pad_token = tokenizer.eos_token 

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = " ### Answer:"
response_template_ids = tokenizer.encode(response_template, add_special_tokens=False)[2:]
collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer, padding_free=True)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(
        output_dir="./tmp",
        gradient_checkpointing=True,
        per_device_train_batch_size=8
    ),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

trainer.train()

الاستنتاجات

أصبحت الآن أمثلة ضبط تعليمات التعبئة، بدلاً من الحشو، متوافقة تمامًا مع Flash Attention 2، وذلك بفضل العلاقات العامة الأخيرة والجديد DataCollatorWithFlattening. الطريقة متوافقة مع النماذج التي تستخدم position_ids. يمكن رؤية الفوائد في الإنتاجية وذروة استخدام الذاكرة أثناء التدريب، دون أي تدهور في تقارب التدريب. يعتمد تحسين الإنتاجية والذاكرة الفعلية على النموذج وتوزيع أطوال الأمثلة في بيانات التدريب. إن التدريب على البيانات التي تحتوي على تباين كبير في أطوال الأمثلة سيحقق أكبر فائدة، فيما يتعلق بالحشو، باستخدام DataCollatorWithFlattening. نفس الميزة متاحة ل SFTTrainer المستخدمين في TRL المكتبة عن طريق وضع علامة جديدة، padding_free=True، عند الاتصال DataCollatorForCompletionOnlyLM.

للحصول على تحليل أكثر تفصيلاً، قم بإلقاء نظرة على الورقة على https://huggingface.co/papers/2407.09105