Machinelearning

🌟Adam-mini: облегченная версия оптимизатора Adam.

Основная идея Adam-mini заключается в том, что матрица Гессиана нейронных сетей, особенно трансформеров, имеет почти блочно-диагональную структуру. Такая структура подразумевает, что для оптимальной работы различных блоков может потребоваться разная скорость обучения.

Adam-mini решает эту проблему, разбивая параметры модели на блоки по наименьшим плотным подблокам в матрице Гессиана. Каждому блоку присваивается одна скорость обучения. Скорость обучения для каждого блока в Adam-mini определяется путем усреднения значений вектора импульса второго порядка Adam (v) в пределах этого блока.

Эта методика сокращает количество необходимых LR, что приводит к значительной экономии памяти. Например, на LLM Adam-mini может сократить до 90% LR, по сравнению с Adam, что в итоге экономит использования памяти на 45-50 %.

Эффективность Adam-mini была проверена сравнением с показателями AdamW в различных сценариях:

Pre-training: на Llama2-7B Adam-mini сокращает использование памяти на 48,04 %, сохраняя при этом сопоставимые с AdamW потери при проверке.

SFТ и RLHF: превосходит AdamW в задачах на основе LoRA и RLHF, удерживая низкое значение perplexity.

Non-LLM Tasks: в задачах, не связанных с LLM - модели СV, ResNet, диффузионные модели, GCN и GAT демонстрирует сравнимую или лучшую производительность, чем AdamW, при этом используя меньше памяти.

Пропускная способность: при предварительном обучении Llama2-7B на 2×A800-80GB Adam-mini показывает производительность на 49,6 % выше, чем AdamW, экономя при этом 33,1 % времени.

▶️Текущая реализация Adam-mini поддерживает популярные фреймворки:

🟢DDP distributed framework;
🟢FSDP distributed framework;
🟢DeepSpeed;
🟢Hugginface Trainer;
🟢Torchtitan.

В репозитории проекта представлены примеры кода для SFT и RLHF претрейна LLM:

🟠

GPT2 (125M-1.5B), NanoGPT codebase на фреймворке DDP

🟠

Llama3-8B, Torchtitan code base на фреймворке FSDP

🟠

SFT и RLHF Llama2-7B, ReMax codebase на фреймворке DeepSpeed

▶️Локальный запуск :

# # import from source
git clone https://github.com/zyushun/Adam-mini
cd Adam-mini
pip install -e .

# Then use Adam-mini optimizer as follows
from adam_mini import Adam_mini

optimizer = Adam_mini(
            named_parameters = model.named_parameters(), 
            lr = lr, 
            betas = (beta1,beta2), 
            eps = eps,
            weight_decay = weight_decay, 
            model_sharding = True,
            dim = model_config.dim,
            n_heads = model_config.n_heads,
            n_kv_heads = model_config.n_kv_heads,
            )
# all the hyperparameters, including learning rate (lr), weight_decay, beta1, beta2, eps, its recommend using the same values as for AdamW

🟡

Arxiv

🖥

Github [ Stars: 226 | Issues: 8 | Forks: 9]

@ai_machinelearning_big_data

#AI #ML #Adam #Pytorch #Train

Please open Telegram to view this post