Research Engineer (LLM Training and Performance)

ПЕРЕЕЗД: Амстердам, Нидерланды; Белград, Сербия; Берлин, Германия; Лимасол, Кипр; Лондон, Великобритания; Мадрид, Испания; Мюнхен, Германия; Пафос, Кипр; Прага, Чехия; Варшава, Польша; Ереван, Армения

В JetBrains мы увлечены кодом. С момента основания в 2000 году мы стремимся создавать самые мощные и эффективные инструменты для разработчиков в мире. Автоматизируя рутинные проверки и исправления, наши инструменты ускоряют производство, освобождая разработчиков для роста, открытий и творчества.

Мы ищем инженера-исследователя, который будет отвечать за тренировочный стек и архитектуру моделей нашего семейства 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).
• Kubernetes или SLURM в больших масштабах, настройка размещения и привязки, а также GPU-кластеры AWS, GCP и Azure.
• Масштабная обработка данных (потоковые наборы данных, Parquet и TFRecord, производительность токенизатора), оценочные среды и бенчмаркинг.
• Методы обеспечения безопасности и пост-обучения, такие как DPO, ORPO, GRPO и модели вознаграждения.
• Системы инференса, такие как vLLM и paged KV.