nlp_llm_001

Отчет LLM (Сидоров)

Архитектура модели

1. Всю конфигурацию я вынес в LLaMAConfig, в который входит размер словаря vocab_size, размер скрытых слоев hidden_size, количество слоев внимания num_attention_heads, количество скрытых слоев num_hidden_layers, максимальное количество позиций max_position_embeddings и размер промежуточного слоя intermediate_size. В качестве значений я в первую очередь исходил из времени обучения (например, при большом vocab_size время обучения может составлять >20ч). В итоге я остановился на конфигурации LLaMAConfig(vocab_size=100000, hidden_size=1000, num_attention_heads=8, num_hidden_layers=10, max_position_embeddings=1024, intermediate_size=512). Я еще пробовал брать vocab_size in (1e3, 1e5, 1e8), качество немного менялось, но по-моему такое небольшое уменьшение лосса не стоило заметного прироста во времени. Еще (я не понял, почему) иногда при проставлении слишком большого vocab_size cuda выдавала непонятную ошибку

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1

которая фиксилась только рестартом ядра

1. Далее в архитектуре я использовал блоки нормализации RMS RMSNorm, внимания RoPE RoPEAttention, Feed-Forward слой с SwiGLU SwiGLU и блоки LLaMA LLaMABlock, здесь просто делал по статье https://arxiv.org/abs/2302.13971
2. В основную модель LLaMAModel я включил эмбеддинги, несколько блоков LLaMABlock, финальную нормализацию и выходной слой для предсказания токенов. Здесь я столкнулся с проблемой, что почему-то при обучении не все данные модели отправлялись на cuda, из-за чего возникали ошибки. Поэтому такое перемещение я дополнительно вынес в класс модели

Цикл обучения

Отдельной болью тут были ошибки во время установки библиотек, связанные с несовместимостью версий. В итоговом ноутбуке часть строк закоменчена, если вдруг кто-нибудь захочет позапускать ячейки, возможно, все сломается :)

Для ускорения обучения для теста работы модели я реализовал функцию downsample_dataset, которая рандомно сокращает датасет. В итоге я ее использовал и для финального обучения, чтобы не ждать вечность, но в итоге лосс оказался приемлимым (если читать условие задания6 там лосс = 4 считался нормальным). И в том числе для ускорения я использовал Accelerator и mixed_precision="fp16", на качестве этот никак не сказалось

В основном цикле обучения для каждой эпохи просто происходит итерация по батчам, внутри каждой итерации происходит вычисление инпута и меток, вычисление потерь и обратное распространение ошибки и обновление градиентов. В том числе я через writer.add_scalar пытался сохранять все, что можно, но в какой-то момент у меня случилось переполнение памяти в датасфера, из-за чего я иставил только сохранение модели после каждой эпохи и массив с лоссами. Для предотвращения проблемы взрывных градиентов я использовал клиппинг градиентов с помощью функции torch.nn.utils.clip_grad_norm_, но и без него модель тоже работала и выдавала похожий лосс.

Как я уже писал выше, здесь у меня тоже возникли проблемы с неправильными размерами входных данных. Я использовал torch.tensor(batch["input_ids"]).to(accelerator.device), то есть заработало именно с accelerator.device. Также здесь были проблемы с переполнением памяти, но это решилось перебором батчей (в итоге использовал batch_size=1024). Кроме того, я использовал gradient_accumulation_steps, чтобы более эффективно распараллелить обучение и оптимизировать использование ресурсов.