[DeepMind Chinchilla] Training Compute-Optimal Large Language Models
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, Laurent Sifre
Статья: https://arxiv.org/abs/2203.15556
Пост в блоге: https://www.deepmind.com/blog/an-empirical-analysis-of-compute-optimal-large-language-model-training
Старый долг про Шиншиллу, про которую давно надо было написать. Это та самая модель от DeepMind на 70B параметров, которая бьёт GPT-3 на 175B и другие ещё более тяжёлые модели.
Собственно, тема большая и важная — обучение огромных моделей, которое, судя по всему, сильно субоптимально, и важно разобраться, что можно сделать лучше. Тогда можно будет обучать более скромные по размерам модели до тех же результатов, что и более тяжёлые.
А по сути, вопрос в том, что имея фиксированный вычислительный бюджет (во FLOPs, параметр C), как лучше его потратить на размер модели (параметр N) и количество токенов в обучении (параметр D).
Авторы взяли архитектуру Gopher (https://t.me/gonzo_ML/742), и обучили 400 моделей с N=70М-16B параметров и D=5B-400B токенов, получив кривую лосса L(N,D) и оценив по ней фронтир.
Пробовали три разных подхода:
1. Зафиксировать размер модели и варьировать число токенов. Взяли много моделей с N=70М-16B параметров, каждый размер обучали с четырьмя разными косинусными циклами, получив непрерывное отображение FLOPs в лосс, а затем в итоге и в N и D так, что FLOPs(N,D) = C.
2. IsoFLOP, когда варьируется размер модели для набора заданных количеств FLOPs (9 разных вариантов от 6×10^18 to 3×10^21 FLOPs). Это напрямую отвечало на вопрос, каков оптимальный размер модели для заданных FLOPs.
3. Для первых двух подходов зафитить функцию определённого вида, описывающую лосс от N и D.
Везде в итоге зафитили степенные функции, предсказывающие оптимальные N_opt и D_opt, и по всем трём подходам оказалось, что при увеличении вычислительного бюджета, и размер модели, и число токенов в обучении надо скейлить примерно в одинаковой пропорции, что заметно отличалось от более ранней классической работы про скейлинг моделей (https://arxiv.org/abs/2001.08361), где получалось, что надо больше вкладываться в размер модели.
Также посчитали интересную таблицу, где для заданного размера модели определены бюджет и число токенов, чтобы результат лежал на compute-optimal фронтире. По этой логике получилось, что многие современные модели чрезмерно раздуты по размеру для имеющегося бюджета.
Если спроецировать найденные закономерности на вычислительный бюджет Суслика (Gopher, 280B модель), то оптимальный размер этой модели будет по разным из этих трёх оценок от 40 до 70B. Выбрали верхнюю границу в 70B и соответствующие им 1.4T токенов. Полученную модель назвали Шиншиллой (Chinchilla).
Итого, для сравнения:
- Gopher, N=280B, D=300B
- Chinchilla, N=70B, D=1.4T
Вычислений для обучения требуется одинаковое количество, но модель в 4 раза легче, что профит для инференса, где и памяти требоваться будет в 4 раза меньше, и считаться будет быстрее. Слоёв так же 80, размер уменьшили за счёт меньшего числа голов (64 вместо 128) и меньшей размерности модели (которая d_model).
Оптимальная GPT-3 со 175B параметров должна была бы обучаться с более солидным бюджетом на 4.2T токенов (вместо 300B), а оригинальный Gopher с 280B параметров должен был обучаться на 6.8T токенов (вместо тех же 300B). К слову, пресловутая LaMDA (https://t.me/gonzo_ML/1017) с N=137B/D=168B, получается, тоже могла бы быть сильно более sentient!
В обучении Шиншиллы были некоторые отличия от Суслика: чуть другое распределение в датасете, AdamW вместо Adam, слегка модифицированный SentencePiece токенизатор (94%+ токенов совпадает, но новый вариант лучше, например, для химии и математики), forward и backward pass считались в bfloat16, но для распределённого оптимизатора использовался float32.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, Laurent Sifre
Статья: https://arxiv.org/abs/2203.15556
Пост в блоге: https://www.deepmind.com/blog/an-empirical-analysis-of-compute-optimal-large-language-model-training
Старый долг про Шиншиллу, про которую давно надо было написать. Это та самая модель от DeepMind на 70B параметров, которая бьёт GPT-3 на 175B и другие ещё более тяжёлые модели.
Собственно, тема большая и важная — обучение огромных моделей, которое, судя по всему, сильно субоптимально, и важно разобраться, что можно сделать лучше. Тогда можно будет обучать более скромные по размерам модели до тех же результатов, что и более тяжёлые.
А по сути, вопрос в том, что имея фиксированный вычислительный бюджет (во FLOPs, параметр C), как лучше его потратить на размер модели (параметр N) и количество токенов в обучении (параметр D).
Авторы взяли архитектуру Gopher (https://t.me/gonzo_ML/742), и обучили 400 моделей с N=70М-16B параметров и D=5B-400B токенов, получив кривую лосса L(N,D) и оценив по ней фронтир.
Пробовали три разных подхода:
1. Зафиксировать размер модели и варьировать число токенов. Взяли много моделей с N=70М-16B параметров, каждый размер обучали с четырьмя разными косинусными циклами, получив непрерывное отображение FLOPs в лосс, а затем в итоге и в N и D так, что FLOPs(N,D) = C.
2. IsoFLOP, когда варьируется размер модели для набора заданных количеств FLOPs (9 разных вариантов от 6×10^18 to 3×10^21 FLOPs). Это напрямую отвечало на вопрос, каков оптимальный размер модели для заданных FLOPs.
3. Для первых двух подходов зафитить функцию определённого вида, описывающую лосс от N и D.
Везде в итоге зафитили степенные функции, предсказывающие оптимальные N_opt и D_opt, и по всем трём подходам оказалось, что при увеличении вычислительного бюджета, и размер модели, и число токенов в обучении надо скейлить примерно в одинаковой пропорции, что заметно отличалось от более ранней классической работы про скейлинг моделей (https://arxiv.org/abs/2001.08361), где получалось, что надо больше вкладываться в размер модели.
Также посчитали интересную таблицу, где для заданного размера модели определены бюджет и число токенов, чтобы результат лежал на compute-optimal фронтире. По этой логике получилось, что многие современные модели чрезмерно раздуты по размеру для имеющегося бюджета.
Если спроецировать найденные закономерности на вычислительный бюджет Суслика (Gopher, 280B модель), то оптимальный размер этой модели будет по разным из этих трёх оценок от 40 до 70B. Выбрали верхнюю границу в 70B и соответствующие им 1.4T токенов. Полученную модель назвали Шиншиллой (Chinchilla).
Итого, для сравнения:
- Gopher, N=280B, D=300B
- Chinchilla, N=70B, D=1.4T
Вычислений для обучения требуется одинаковое количество, но модель в 4 раза легче, что профит для инференса, где и памяти требоваться будет в 4 раза меньше, и считаться будет быстрее. Слоёв так же 80, размер уменьшили за счёт меньшего числа голов (64 вместо 128) и меньшей размерности модели (которая d_model).
Оптимальная GPT-3 со 175B параметров должна была бы обучаться с более солидным бюджетом на 4.2T токенов (вместо 300B), а оригинальный Gopher с 280B параметров должен был обучаться на 6.8T токенов (вместо тех же 300B). К слову, пресловутая LaMDA (https://t.me/gonzo_ML/1017) с N=137B/D=168B, получается, тоже могла бы быть сильно более sentient!
В обучении Шиншиллы были некоторые отличия от Суслика: чуть другое распределение в датасете, AdamW вместо Adam, слегка модифицированный SentencePiece токенизатор (94%+ токенов совпадает, но новый вариант лучше, например, для химии и математики), forward и backward pass считались в bfloat16, но для распределённого оптимизатора использовался float32.