Research Engineer (LLM Training and Performance)
Амстердам, Нидерланды; Белград, Сербия; Берлин, Германия; Лимасол, Кипр; Лондон, Великобритания; Мадрид, Испания; Мюнхен, Германия; Пафос, Кипр; Прага, Чехия; Варшава, Польша; Ереван, Армения
В JetBrains код — наша страсть. С самого начала в 2000 году мы стремимся создавать самые мощные и эффективные инструменты для разработчиков в мире. Автоматизируя рутинные проверки и исправления, наши инструменты ускоряют производство, освобождая разработчиков для роста, открытий и творчества.
Мы ищем Research Engineer, который возьмёт на себя стек обучения и архитектуру моделей для нашего семейства LLM Mellum. Ваша задача проще на словах, чем на деле: сделать обучение быстрее, дешевле и стабильнее в крупных масштабах. Вы будете профилировать, проектировать и внедрять изменения в конвейер обучения — от архитектуры до пользовательских GPU ядер, по мере необходимости.
В составе нашей команды вы будете:
- Отвечать за повышение сквозной производительности многопоточных конвейеров предварительного и последующего обучения LLM.
- Профилировать узкие места (Nsight Systems/Compute, NVTX) и исправлять их, используя перекрытие вычислений и коммуникаций, слияние ядер, планирование и т.д.
- Проектировать и оценивать варианты архитектуры (глубина/ширина, варианты механизма внимания, включая GQA/MQA/MLA/Flash-style, масштабирование RoPE/NTK, а также маршрутизацию и балансировку нагрузки в MoE).
- Реализовывать пользовательские операции (Triton и/или CUDA C++), интегрировать через расширения PyTorch и по возможности отправлять изменения вверх по потоку.
- Использовать рычаги управления памятью и производительностью: FSDP/ZeRO, контрольные точки активации, FP8/TE, тензорный/конвейерный/последовательный/экспертный параллелизм, настройка NCCL.
- Укреплять крупные прогоны, создавая гибкие и отказоустойчивые настройки обучения, обеспечивая надежное создание контрольных точек, повышая воспроизводимость и улучшая устойчивость к прерываниям.
- Поддерживать быстрый путь данных с использованием потоковых и сегментированных загрузчиков данных и конвейеров токенизаторов, а также повышать общую пропускную способность и эффективность кэширования.
- Определять правильные метрики, создавать информационные панели и обеспечивать стабильные улучшения.
- Эффективно запускать как предварительное, так и последующее обучение (включая методы SFT, RLHF и GRPO) на значительных кластерах.
Мы будем рады принять вас в команду, если у вас есть:
- Большой опыт работы с PyTorch и PyTorch Distributed, опыт запуска многопоточных задач с десятками и сотнями GPU.
- Практический опыт работы с Megatron-LM/Megatron-Core/NeMo, DeepSpeed или серьёзный опыт работы с FSDP/ZeRO.
- Реальный опыт профилирования (Nsight Systems/Compute, nvprof) и опыт работы с инструментированными NVTX рабочими процессами.
- Навыки программирования для GPU с использованием Triton и/или CUDA, а также умение писать, тестировать и отлаживать ядра.
- Глубокое понимание коллективных операций NCCL, а также влияния топологии и среды передачи данных (IB/RoCE) и того, как они отображаются в трассировках.
Наш идеальный кандидат имел бы опыт работы с:
- FlashAttention-2 и 3, CUTLASS и CuTe, TransformerEngine и FP8, Inductor, AOTAutograd и torch.compile.
- MoE в крупном масштабе (экспертный параллелизм, потери маршрутизатора, управление пропускной способностью) и трюками для длинного контекста (ALiBi/YaRN/NTK scaling).
- Kubernetes или SLURM в крупном масштабе, настройка размещения и сродства, а также парки GPU в AWS, GCP и Azure.
- Крупномасштабными системами обработки данных (потоковые наборы данных, Parquet и TFRecord, производительность токенизатора), инструментами оценки и бенчмаркинга.
- Методами обеспечения безопасности и последующего обучения, такими как DPO, ORPO, GRPO и моделями вознаграждения.
- Экосистемами для вывода, такими как vLLM и paged KV.
Мы предоставляем равные возможности
Мы знаем, что отличные идеи могут прийти от кого угодно и где угодно. Поэтому мы делаем всё возможное, чтобы создать открытое и инклюзивное рабочее место — такое, которое приветствует всех, независимо от их происхождения, идентичности, религии, возраста, потребностей в доступности или ориентации.
