Тайният сос зад 100K контекстен прозорец в LLMs: всички трикове на едно място

Тайният сос зад 100K контекстен прозорец в LLMs: всички трикове на едно място

tldr; техники за ускоряване на обучението и извода на LLMs за използване на голям контекстен прозорец до 100K входни токени по време на обучение и извод: позиционно вграждане на ALiBi, Оскъдно внимание, FlashAttention, внимание с множество заявки, Условно изчисление, и 80GB A100 графични процесори.

Наскоро имаше няколко съобщения за нови големи езикови модели (LLM), които могат да използват изключително голям контекстен прозорец, като 65K токени(MPT-7B-StoryWriter-65k+ от MosaicML) или дори 100K токени(Представяне на 100K контекстни прозорци от Antropic). В „техническия доклад“ на Palm-2 Google не разкрива размера на контекста, но споменава, че те „увеличават значително дължината на контекста на модела.“

За сравнение, настоящият модел GPT-4 може да работи с дължина на контекста от 32K входни токена. И повечето от LLM с отворен код имат дължина на контекста 2K токена.

Това е впечатляващо, тъй като наличието на толкова голяма дължина на контекста означава, че подканата може да бъде буквално с размер на книга. Великият Гетсби е 72K жетона, 210 страници и 6 часа четене при скорост от 1,7 минути/страница. Така моделът може да сканира и съхранява това количество „персонализирана“ информация за обработка на заявки!

Опитвах се да си обясня как това е технически възможно, така че в тази публикация в блога събирам разпръснати части от информация (тази „нишка“ беше първата следа) и обхващам следното:

• Защо дължината на контекста има значение и защо може да промени играта
• Какви са основните ограниченияв оригиналната архитектура на Transformer при работа с големи контекстни дължини
• Изчислителната сложност на трансформаторната архитектура
• Какви техники за оптимизация съществуват в момента за ускоряване на трансформатора и увеличаване на дължината на контекста до 100K

"Кратко обобщение

Тук и по-късно използваме „дължина на контекста“, „прозорец на контекста“ и „броя входни токени“ взаимозаменяемо, обозначавайки ги като n.

Публикацията в блога е малко дълга, така че има обобщение с основните точки и трикове:

• Първият проблем е квадратичната времева и пространствена сложност на изчисленията на слоя внимание w.r.t. броя на въведените токени n.
• Когато размерът на вграждане d › n, вторият проблем е квадратната времева сложност на линейните слоеве w.r.t. размер за вгражданеd.
• 3-ти проблем е Позиционно синусоидално вграждане, използвано в оригиналната архитектура.
• В архитектурата на Transformer формите на обучаемите тегла на матрицата са агностик спрямо броя на входните токени n.
• И така, обучен Трансформатор в 2K контекстни дължини може да консумира токени с всякаква дължина, дори 100K. Но моделът няма да даде значими резултати за 100K токени по време на извод, ако не е обученна 100K.
• Обучението на ванилия трансформатор върху гигантски корпус и само върху голяма дължина на контекста е неосъществимо скъпо поради квадратичната сложност спрямо n и d. LLaMA на дължина на контекста 2K беше „оценено“ да бъде обучен за ~$3 милиона. По този начин LLaMA на 100K ще струва ~$150 милиона.
• Една от възможностите е да обучите модела в контекст на 2K токени и след това да го настроите фино в по-дълги контексти (например 65K). Но няма да работи с оригиналния трансформатор поради позиционното синусоидално кодиране.
• [Трик №1]За да разрешите това, премахнете позиционното синусоидално кодиране и използвайте ALiBi, просто и елегантно позиционно вграждане, което не вреди на точността. След това можете да тренирате на 2K и да настройвате фино на 100K.
• [Трик №2] Не е необходимо да изчислявате оценки за внимание между всички токени. Някои токени са по-важни от други, така че може да се използва Sparse Attention. Ще ускори както обучението, така и извода.
• [Трик #3] Flash Attention ефективно прилага слоя за внимание за GPU. Той използва подреждане и избягва материализирането на големи междинни матрици (n, n), които не се вписват в GPU SRAM. Ще ускори както обучението, така и изводите.
• [Трик #4]Внимание с няколко заявки вместо внимание с няколко глави. Това означава, че споделяте тежести във всички глави, когато проектирате линейно K и V. Това драстично ускорява инкременталното заключение.
• [Трик №5] Условното изчисление избягва прилагането на всички параметри на модела към всички токени от входната последователност. CoLT5 прилага тежки изчисления само към най-важните токени и обработва останалите токени с по-лека версия на слоевете. Ще ускори както обучението, така и извода.
• [Трик №6]За да поберете голям контекст, имате нужда от много RAM в GPU, така че хората използват 80GB A100 GPU.

За да обобщим, колкото повече ускорявате обучението и изводите, толкова по-голяма дължина на контекста можете да използвате.

Нека сега обсъдим всички тези точки по-подробно.

Защо дължината на контекста има значение

Дължината на контекста е едно от критичните ограничения на LLM. И увеличаването му до вече 100K е невероятно постижение (чудя се как ще изглежда това твърдение след година).

Един от важните случаи на употреба, при който хората искат да прилагат LLM, е „пускането на голяма купчина персонализирани данни в LLM“ (документи, свързани с компанията или конкретен проблем, различни разнородни текстове и т.н.) и задаването на въпроси относно тези конкретни данни, не някакви абстрактни данни от интернет, които LLM видя по време на обучение.

За да преодолеят това ограничение сега, хората правят различни неща:

• Изпробване на техники за обобщаване и сложни верижни подкани
• Поддържане на векторни бази данни за запазване на вграждания за персонализирани документи и след това „търсене“ в тях по някакъв показател за сходство
• Фина настройка на LLM с персонализирани данни, когато е възможно (не всички търговски LLM позволяват това и това не е очевидна задача за LLM с отворен код)
• Разработване на персонализирани по-малки LLM за тези конкретни данни (отново не е очевидна задача)

Наличието на голяма дължина на контекста позволява на вече мощен LLM (който видя целия интернет) да разгледа вашия контекст и данни и да взаимодейства с вас на напълно различно ниво с по-висока персонализация. И всичко това без да променяте теглото на модела и да правите „обучението“ си в движение, „на памет“. И като цяло големият контекстен прозорец внася повече точност, плавност и креативност на модела.

Една аналогия тук може да бъде компютърната RAM памет, където операционната система поддържа контекста в реално време на всички ваши приложения. Със значителна дължина на контекста, LLM може да бъде като „компютър за разсъждение“, запазвайки много потребителски контекст.

Оригинален трансформатор и контекстна дължина

Важно е да се отбележи, че в архитектурата на Transformer формите на всички тегла на матрицата, които могат да се научат, не зависят от броя на входните токени n. Всички обучаеми параметри (вграждане на търсене, проекционни слоеве, softmax слой и слоеве за внимание) не зависят от дължината на входа и трябва да обработват входове с променлива дължина. Чудесно е,чеимаме това готово свойство на архитектурата.

Това означава, че ако сте обучили модел на Transformer с дължина на контекста 2K, можете да направите извод за последователности от токени с всякакъв размер. Единственият проблем е, че моделът няма да произведе смислени резултати за 100K токени по време на извод, ако не е обученна 100K контекстна дължина. В този случай разпределението на данните за обучение ще бъде далеч от това по време на извода, така че моделът ще се провали като всеки модел за машинно обучение в тази настройка.

Едно решение за обучение на трансформатор с голяма дължина на контекста е да го обучите на два етапа: обучете базовия модел на дължина на контекста на 2K токени и след това продължете обучението (фина настройка) на по-дълги контексти (например 65K или 100K). Точно това „направи“ MosaicML. Но проблемът е, че няма да работи с оригиналната архитектура на Transformer, така че трябва да използвате някои трикове (вижте Трик №1 по-късно в публикацията).

Обобщение на вниманието на няколко глави

Предизвикателствата на голямата дължина на контекста са свързани с изчислителната сложност на трансформаторната архитектура. За да обсъдим сложността, първо, нека обобщим как работи слоят внимание.

Q —заявки, K — ключове и V — стойности, бележки от хартията, свързани с извличането на информация, където вмъквате „заявка“ към системата и потърсете най-близкия „ключ“
n —въведеният брой токени
d — измерение за вграждане на текст
h — броят на главите на вниманието
k— размер на линейната проекция за Q и K
v — размер на линейната проекция за V

Внимание от няколко глави:

1. Имаме слой за вграждане за търсене, който за даден токен връща вектор с размер (1, d). Така за поредица от n токени получаваме матрицата за вграждане на текст X с размер (n, d). След това го обобщаваме с позиционното синусоидално вграждане.
2. Слоят Multi-Head Attention има за цел да изчисли новото вгражданеза тази последователност от токени, които могат да се считат за оригинален текст, кодиращ X но претеглени (1) по относителна важност между лексеми по отношение на контекста и (2) чрез относителни позиции на лексеми.
3. Ние обработваме тази матрица за вграждане X (n, d) паралелно с h слоеве за внимание (глави). За да получите Q, K, и Vза всички глави за внимание, проектирате линейно X към k, k,и vразмери съответно. Правите го, като умножите X по h матрици с форма (d, k), (d, k), и (d , v). Можете да мислите за това като за умножаване на (n, d) по (h, d, k), (h, d, k), и (h, d, v).
4. Attention Heads връщат h матрици с резултати за внимание с размер (n, v). След това свързваме парчета от всички глави (n, h*v) и ги проектираме линейно за следващите стъпки.

Мащабирано внимание към продукта:

Сега нека увеличим мащаба на едно внимание главата.

1. Q,K,V са 3 линейни проекции на Xс размер (n, k), (n, k) и (n, v)получени чрез умножаване на научаваеми тегла отделно за всяка глава.
2. Получаваме резултати за внимание, като изчисляваме разстоянието (точков продукт) между Q и K(транспонирано). Умножавате матрицата (n, k) по (k, n) и получавате матрицата (n, n). След това го умножаваме по матрицата на маската, за да нулираме някои от токените (необходими в декодера). След това го мащабираме и прилагаме softmax, за да бъде от 0 до 1. По този начин получаваме матрицата на формата (n, n) с n_ij - относителен резултат за внимание от 0 до 1 между i-тия и j-ия токен, който показва колко „близки“ са тези токени в този конкретен контекст с дължина n .
3. След това умножаваме тази матрица за оценка на вниманието (n, n)по „стойности“ V с размер (n, d), за да получим вграждането на текст, претеглено от тези относителни резултати за внимание.

Нека да разгледаме тази част от кода от статията „Внимание за множество заявки“. Показва как вниманието на няколко глави се изчислява с групиране и формите са ясни на всяка стъпка. Те също така включват маскиращо умножение, използвано по време на декодиране.

Сложността на трансформатора и дължината на контекста

Сложността на 2 матрично умножение (a,b)*(b,c) е O(a*b*c).
Приемаме, че k*h = O(d) за простота,и ще използваме това, за да извлечем сложността на вниманието.

Сложността на слоя внимание се състои от две части:

1. Линейни проекции за получаване на Q, K, V: умножение на матрица за вграждане с размер (n, d) по h обучаеми матрици (d, k), (d, k), и (d, v). Така сложността ~ O(nd²)
2. Умножения на Q по K, трансформирани и след това умножение по V: (n,k) * (k,n) = (n,n) и(n,n)*(n,v) = (n,v). Сложността ~ O(n²d)

И така, сложността на слоя за внимание е O(n²d + nd²), където n — е дължината на контекста (брой входни токени) и d — размер на вграждане. И така, от тук виждаме, че сложността на изчислението на слоя внимание е квадратична спрямо броя на входните токени n и квадратична спрямо размера на вграждане d.

Терминът O(nd²)е важен, когато d › n (например в LLaMa, n=2K и d=4K).
Терминът O(n²d) е важна, когато n › d (например обучение на MosaicML с n=65K и d=4K).

Само да ви напомня колко лошо е квадратичното нарастване:
2 000² = 4 000 000, 100 000² = 10 000 000 000.

Позволете ми да ви дам пример за това как тази квадратична сложност влияе върху цената на обучението за модели. Прогнозната цена на обучението LLaMa беше ~$3 милионаи има 65B параметри, 2K дължина на контекста и 4K размер на вграждане. Очакваното време е предимно време за обучение на GPU. Ако увеличим дължината на контекста от 2K на 100K (50x), времето за обучение ще се увеличи също ~50x (имаме нужда от по-малко итерации, защото контекстът е по-голям, но отнема повече време за всяка). Така че обучението на LLaMA на 100K контекст ще струва около ~$150 милиона.

Малко подробности за това изчисление:

За броя на токените, равен на n, сложността на вниманието е O(n²d + nd²)и са необходими M итерации за обучение. Ако увеличим дължината на контекста от n → p*n,ще са необходими M/p итерации, тъй като дължината на контекста стана по-голяма (да приемем за простота е линейно, може да е надценяване или подценяване в зависимост от задачата). Сега имаме 2 уравнения:
(1) Сложност за n ~M * (n²d + nd²)
(2)Сложност за p*n ~ M/p * ((p*n)²d + (p*n)d²)
След поредица от опростявания и деления, съотношението (2)/(1) ~(d + p*n)/(d + n)

Ако d ‹‹ n, увеличаването на n с фактор p ще доведе до ~ p пъти повече итерации.
Ако d ~ n, увеличаването на n с фактор p ще доведе до ~ p/2 пъти повече итерации .

Разлика между етапите на обучение и извод в Transformer

Последното нещо, което трябва да обсъдим, преди да се заровим в техниките за оптимизация, е разликата в изчисленията по време на обучението и изводите.

По време на обучението изпълнявате нещата паралелно, докато за генериране на текст по време на извод трябва да го правите последователно, защото следващият токен зависи от предишните. Простият начин за прилагане на извода е да се изчислят оценките за внимание постепенно и кеширане на предишни резултати за бъдещи токени.

Това разграничение носи различни подходи за ускоряване на обучението и изводите. Ето защо някои трикове по-долу ще оптимизират и двата етапа, но някои ще оптимизират само извода.

Техники за оптимизация за увеличаване на дължината на контекста

Сега нека поговорим за това как изследователите са преодолели всички тези предизвикателства и са успели да обучат LLM с голяма дължина на контекста.

[Трик#1]По-добро позиционно кодиране — ALiBi

Едно решение за обучение на трансформатор с голяма дължина на контекста е да се обучи на два етапа: да се обучи базовият модел на дължина на контекста на 2K токени и след това да се прецизира върху по-дълги контексти (например 65K). Но по-рано казахме, че няма да работи с оригиналната архитектура на Transformer. Защо?

Заради Позиционното синусоидално кодиране, което няма възможност за „екстраполация“. В статията ALiBI [4] авторите показаха, че позиционното синусоидално кодиране не е устойчиво на разширението на контекстния прозорец по време на извод. След още няколко символа производителността започва да се влошава. Така че липсата на способност за „екстраполация“ основно означава, че не можете да използвате по-големи дължини на контекста по време на извод/фина настройка, отколкото по време на обучение. Терминът "екстраполация" и сравнението на различни позиционни кодирания са описани в [4].

В оригиналната трансформаторна хартия Positional Sinusoidal Embedding е сумирано с токените Embeddings в долната част на архитектурата, за да добави информация за реда на думите. Ако искате да научите как се изчислява позиционното синусоидално вграждане, препоръчвам това забавно видео, където е обяснено интуитивно и с много подробности.

И така, първият трик е да премахнете позиционното синусоидално вграждане и да го замените с друго вграждане на позиция — Внимание с линейни отклонения (ALiBI).

То се прилага в главата за внимание (не в долната част на мрежата) и променя резултатите за внимание на ключове за заявки с наказание, което е пропорционално на тяхното разстояние (преди softmax).

Този трик ускорява обучението.

[Трик №2] Оскъдно внимание

Не всички токени в контекста на размер 100K са подходящи един за друг. Един от начините за намаляване на броя на изчисленията е да разглеждате само някои жетони, когато изчислявате оценките за внимание. Целта на добавянето на рядкост е да направи изчислението линейно спрямо n, а не квадратично. Има няколко подхода как да изберете връзката между токените и има отлична илюстрация за това в публикацията в блога на Google:

Например, „Вниманието на плъзгащия се прозорец“ (наричано още локално) използва внимание на прозорец с фиксиран размер около всеки токен. В този модел на внимание, при фиксиран размер на прозореца w, всеки токен се грижи за w/2 токена от всяка страна. Изчислителната сложност на този модел е O(n*w), която се мащабира линейно с дължина на входната последователност n. За да бъде ефективен, w трябва да е малък в сравнение с n. Номерът е, че информацията за вниманието „преминава“ през целия контекстен прозорец в рамките на близки токени, приближавайки се до пълната графика.

Методът за оценяване на вниманието BigBird съчетава глобални, локални и произволни механизми. В статията авторите показаха изключително важно наблюдение, че има присъщо напрежение между това колко малко оценки за сходство се изчисляват и потока на информация между различни възли (т.е. способността на един токен да си влияе един на друг ).

Този трик ускорява както обучението, така и извода.

[Трик #3] FlashAttention — ефективно внедряване на слоя за внимание за GPU

Има няколко изчислителни операции в слоя на вниманието, които се повтарят отново и отново:

1. S = Q*K
2. P= softmax(S)
3. O = P*V

Запомнете понятието за P, S и O резултати; ще го използваме по-късно. Авторите на FlashAttention са „слели“ тези операции: те са внедрили алгоритъм за слой внимание, който използва ефективно паметта на GPU и изчислява точното внимание.

За да може GPU да извърши операция, входните данни трябва да присъстват в „бързата“ памет, наречена SRAM. Данните се копират от „бавната“ HBM памет в SRAM и се връщат обратно в HBM, след като изчислението приключи. SRAM паметта е много по-бърза от HBM, но много по-малка по размер (20MB срещу 40GB в A100 40GB GPU).

Така че достъпът до HBM е скъпа операция.

Основният проблем в слоя за внимание по отношение на използването на паметта на графичния процесор са „междинните“ резултати от умножение, P, S,иO,които са големи по размер (n , n). Трябва да ги запазим в HBM и да ги прочетем отново между операциите за внимание. Преместването на P, S и O от HBM към SRAM обратно и принудително е тясното място, което авторите разрешиха в документа.

Основната идея зад алгоритъма FlashAttention е да се разделят входните Q, K и V матрици на блокове, като се зареждат тези блокове от HBM в SRAM и след това се изчислява изходът на вниманието спрямо тези блокове. Тази процедура се нарича плочка.

Операцията „умножение на матрици“ вече е оптимизирана за GPU. Може да мислите за този алгоритъм FlashAttention като за прилагане на операцията „слой на вниманието“, оптимизирана за GPU. Авторите „сляха“ операции на няколко умножения и softmax с подреждане и оптимизиран HBM достъп.

Има добър „общ преглед“ на „хартията“ на FlashAttention.

От „наскоро“ PyTorch 2.0 има „вградено“ флаш внимание. Това е FlashAttention имплементация в Triton език от авторите.

Този трик ускорява както обучението, така и извода.

[Трик #4] Внимание за множество заявки (MQA)

Оригиналното Multi-Head Attention (MHA) има отделен линеен слой за K и V матрици във всяка глава.

По време на извода ключовете и стойностите на предишните токени в декодера се кешират, за да се предотврати повторното им изчисляване, така че използването на паметта на GPU нараства с всеки генериран токен.

Внимание с множество заявки (MQA) е оптимизацията, която предлага споделяне на тегла във всички глави на вниманието при линейно проектиране на K и V, така че ще трябва да запазим само 2 матрици с размер (n, k) и (n, v). Един голям модел може да има до 96 глави (като GPT-3), което означава, че използването на MQA може да спести 96 пъти потреблението на памет от кеша на декодера ключ/стойност.

Тази оптимизация е особено полезна при генериране на дълги текстове. Например, имайки голяма дължина на контекста и изисквайки дълъг, смислен анализ или обобщение.

Основното предимство на този подход е значителнотоускоряване на изчисляването на инкременталните резултати за внимание по време на извод. Скоростта на тренировка остава почти същата. Например PaLM го „използва“.

[Трик #5] Условно изчисление

Когато d › n, тясното място в скоростта не е нивото на вниманието, а слоевете за предаване и проекция. Често срещан подход за намаляване на FLOPs е използването на някаква форма на условно изчисление, което избягва прилагането на всички параметри на модела към всички токени от входната последователност.

В раздела Sparse Attention обсъдихме, че някои токени са по-важни от други. Следвайки същата интуиция, в документа на CoLT5 авторите разделиха всички изчисления за предварителна информация и внимание на два клона: тежки и леки. Леките слоеве се прилагат към всички токени, а тежките само към важните.

„Леките и тежките клонове за предварителна връзка се различават само по своите скрити измерения, като лекият клон има по-малък скрит размер от стандартния слой за предварителна връзка T5, а тежкият клон е по-голям“.

Доказано е, че този подход превъзхожда както скоростта, така и точността на съществуващия модел LongT5 за изключително дълги последователности до 64K входни токена.

[Трик №6] Графични процесори с голяма RAM памет

Това не е трик, а необходимост. За да паснете на голям контекст, имате нужда от голяма RAM в GPU, така че хората използват 80GB A100 GPU.

Заключение

Леле, това е много. Не очаквах да завърша с толкова дълга публикация в блога :D

Надявам се да е било полезно! Научих много и се надявам, че вие ​​също, и сега можем да познаем как тези големи езикови модели с милиарди параметри са били обучени в безпрецедентни контекстни прозорци от 65-100K токени.

Вдъхновяващо е да видите как различни умни хора се справят с един и същи проблем от различни страни, оптимизират тук и там и измислят страхотни идеи. Всичко това води до смислено и елегантно решение.

Харесва ми това, което един изследовател каза за обучението на LLM с широк контекст: „Без таен сос, само добре проверени изследвания.“

Препратки

[1] Представяне на 100K контекстни прозорци от Antropic
[2] MPT-7B от MosaicML
[3] Технически доклад на Palm-2 от Google
[4] ALiBI : Train Short, Test Long: Attention with Linear Biases Enable Input Length Extrapolation
[5] FlashAttention: Бързо и ефективно за паметта точно внимание с IO-Awareness
[6] Multi-Query внимание : Декодиране с бърз трансформатор: Една глава за запис е всичко, от което се нуждаете
[8] Вниманието е всичко, от което се нуждаете
[9] Видео за позиционно синусоидално вграждане
[10] Преглед на документа за FlashAttention
[11] Внимание при плъзгащ се прозорец
[12] Изграждане на трансформатори за по-дълги последователности с методи за оскъдно внимание
[13] Внедряване на FlashAttention в езика Triton
[14] Как да ускорим пропускателната способност на HuggingFace със 193% с Triton и ClearML
[15] Обслужване на ClearML
[16] Анализиране на плюсовете и минусите на NVIDIA Triton Inference Server срещу други машини за изводи
[17] COLT5: По-бързи трансформатори за голям обсег с условно изчисление
[18] LongT5: Ефикасен трансформатор от текст към текст за дълги последователности
[19] PaLM
[20] Механизъм за внимание BigBird.