Эффективное обучение для больших моделей
BMTrain — это эффективный инструментарий для обучения больших моделей с десятками миллиардов параметров. Он позволяет обучать модели распределённым образом, сохраняя при этом код простым, как при автономном обучении.
Что нового
Документация Наша документация (https://bmtrain.readthedocs.io/en/latest/index.html) содержит более подробную информацию о пакете.
Установка
Установка BMTrain может занять от нескольких минут до десяти минут, так как она требует компиляции исходного кода c/cuda во время установки. Мы рекомендуем компилировать BMTrain непосредственно в среде обучения, чтобы избежать потенциальных проблем, вызванных различными средами.
Использование
Прежде чем вы сможете использовать BMTrain, вам необходимо инициализировать его в начале вашего кода. Так же, как использование распределённого модуля PyTorch требует использования init_process_group в начале кода, использование BMTrain требует использования init_distributed в начале кода.
import bmtrain as bmt
bmt.init_distributed(
seed=0,
# ...
)
Примечание: не используйте распределённый модуль PyTorch и связанные с ним функции связи при использовании BMTrain.
Чтобы включить оптимизацию ZeRO, вам нужно сделать несколько простых замен в исходном коде модели.
И оберните блоки трансформатора в bmtrain.Block.
Вот пример.
Исходный код:
import torch
class MyModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.param = torch.nn.Parameter(torch.empty(1024))
self.module_list = **Часть 1: Определение модели**
├── layers │ ├── attention.py │ ├── embedding.py │ ├── feedforward.py │ ├── init.py │ ├── layernorm.py │ └── linear.py └── models ├── gpt.py └── init.py
Выше представлена структура каталогов кода в части определения модели.
Мы определили все слои. **Часть 2: Инициализация BMTrain**
```python
bmtrain.init_distributed(seed=0)
model = GPT(
num_layers=8,
vocab_size=10240,
dim_model=2560,
dim_head=80,
num_heads=32,
dim_ff=8192,
max_distance=1024,
bias=True,
dtype=torch.half
)
bmtrain.init_parameters(model) # или загрузка контрольной точки с использованием `bmtrain.load`
# ... другая инициализация (датасет) ...
Функция bmtrain.init_distributed(seed=0)
используется для инициализации распределённой среды обучения и установки случайного начального числа для воспроизводимости результатов.
Функция bmtrain.init_parameters(model)
используется для инициализации параметров модели в распределённом режиме.
Часть 3: Инициализация оптимизатора и планировщика скорости обучения
loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)
optimizer = bmtrain.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2)
lr_scheduler = bmtrain.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)
BMTrain поддерживает все встроенные оптимизаторы PyTorch и функции потерь, а также можно использовать объединённый оптимизатор, предоставляемый BMTrain, для обучения со смешанной точностью.
Кроме того, BMTrain также предоставляет общий планировщик скорости обучения в модуле bmtrain.lr_scheduler
.
Часть 4: Цикл обучения
# создать новый экземпляр менеджера оптимизаторов
optim_manager = bmtrain.optim.OptimManager(loss_scale=1024)
# позволить optim_manager управлять всеми оптимизаторами и (необязательно) соответствующими им lr_scheduler
optim_manager.add_optimizer(optimizer, lr_scheduler)
# add_optimizer можно вызывать несколько раз для добавления других оптимизаторов.
for iteration in range(1000):
# ... загрузить данные для каждого ранга ...
# прямой проход и расчёт потери
pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
logits = model(
enc_input,
pos,
pos < enc_length[:, None]
)
batch, seq_len, vocab_out_size = logits.size()
loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
global_loss = bmtrain.sum_loss(loss).item() # суммирование потери по всем рангам. Это используется только для журнала обучения
# нулевой градиент
optim_manager.zero_grad() # вызов zero_grad для каждого оптимизатора
# потеря масштаба и обратный проход
optim_manager.backward(loss)
# ограничение нормы градиента
grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0)
# шаг оптимизатора
optim_manager.step()
# ... сохранить контрольную точку или распечатать журналы ...
Цикл обучения будет немного длиннее, но, как и обычный цикл обучения, вам не нужно сильно адаптироваться к распределённому обучению.
Вы можете следовать комментариям в коде, чтобы понять, что делает каждый раздел кода.
Единственное дополнительное замечание — это optimizer
. После использования BMTrain некоторые детали в оптимизаторах должны быть скорректированы. Мы реализовали все эти детали, необходимые в optim_manager
. Всё, что вам нужно, — это позволить optim_manager
управлять всеми оптимизаторами с помощью add_optimizer
, и позволить optim_manager
выполнять zero_grad()
, backward()
, clip_grad_norm()
и step()
вместо этого.
Если вы не используете обучение со смешанной точностью, вы можете тренироваться без loss_scale
. Просто установите loss_scale
на None
в функции __init__
OptimManager(loss_scale=None)
, что также является значением по умолчанию.
Если вы используете обучение со смешанной точностью, потеря масштаба — это широко используемый метод в обучении со смешанной точностью для предотвращения потери градиента. Используя optim_manager.backward(loss)
для масштабирования loss
перед обратным проходом и установив loss_scale
на некоторое число с плавающей запятой в функции __init__
OptimManager
, вы сможете адаптивно настроить loss_scale
во время обучения на основе градиента. 5120 hidden dimension
Длина последовательности 512
Размер пакета | 8 | 16 | 24 | 32 |
---|---|---|---|---|
BMTrain | 24,15 | 26,94 | 29,42 | 28,28 |
ZeRO3(mp=1) | 14,88 | 21,69 | 24,38 | - |
ZeRO3(mp=4) | 15,51 | - | - | - |
ZeRO3(mp=8) | 15,51 | - | - | - |
ZeRO2(mp=1) | - | - | - | - |
ZeRO2(mp=4) | 22,85 | - | - | - |
ZeRO2(mp=8) | 21,33 | - | - | - |
ZeROa(mp=b) означает DeepSpeed + Megatron ZeRO этап a и параллелизм модели = b.
Мы перенесли большинство распространённых моделей в NLP на BMTrain. Вы можете найти список поддерживаемых моделей в репозитории ModelCenter (https://github.com/OpenBMB/ModelCenter).
Приглашаем всех внести свой вклад в код, следуя нашим рекомендациям по внесению вклада (https://github.com/OpenBMB/BMTrain/blob/master/CONTRIBUTING.md).
Вы также можете найти нас на других платформах:
— QQ Group: 735930538;
— Веб-сайт: https://www.openbmb.org;
— Weibo: http://weibo.cn/OpenBMB;
— Twitter: https://twitter.com/OpenBMB.
Пакет выпущен под лицензией Apache 2.0 (https://github.com/OpenBMB/BMTrain/blob/master/LICENSE).
BMTrain вносит основные изменения в PyTorch, поэтому если ваша программа выдаёт неожиданные результаты, вы можете сообщить об этом в выпуске.
Вы можете оставить комментарий после Вход в систему
Неприемлемый контент может быть отображен здесь и не будет показан на странице. Вы можете проверить и изменить его с помощью соответствующей функции редактирования.
Если вы подтверждаете, что содержание не содержит непристойной лексики/перенаправления на рекламу/насилия/вульгарной порнографии/нарушений/пиратства/ложного/незначительного или незаконного контента, связанного с национальными законами и предписаниями, вы можете нажать «Отправить» для подачи апелляции, и мы обработаем ее как можно скорее.
Комментарии ( 0 )