feifeibear/long-context-attention

Что это за инструмент

YunChang — это библиотека для оптимизации обучения и инференса длинных контекстов в LLM, объединяющая преимущества Ulysses и Ring Attention в единый гибридный подход (USP). Она позволяет эффективно масштабировать модели на множество GPU, преодолевая ограничения отдельных методов параллелизма.

Зачем нужен

Инструмент решает проблему неэффективности или неприменимости стандартных методов распределенного внимания (Ulysses и Ring) при работе с длинными контекстами и специфическими архитектурами (GQA/MQA). Он полезен тем, кто хочет максимизировать пропускную способность и стабильность распределенных вычислений, используя сильные стороны обоих подходов в одной реализации.

Что можно реализовать

• Обучение и инференс LLM с очень длинным контекстом (Long Context) на кластерах NVIDIA GPU (A100, H100).
• Работа с моделями, использующими GQA (Grouped Query Attention) или MQA (Multi-Query Attention), где классический Ulysses неэффективен.
• Интеграция оптимизированного внимания в существующие пайплайны через замену стандартного Attention на LongContextAttention.
• Запуск вычислений на аппаратных платформах без поддержки FlashAttention (например, NPU) через бэкенд torch.

Ключевые возможности

• Гибридный подход (USP): синергия DeepSpeed-Ulysses и Ring Attention для лучшей универсальности и производительности.
• Поддержка FlashAttention v2 (для A100/L40) и v3 (для H100/B100), а также fallback-режим на чистом PyTorch.
• Готовая интеграция с NVIDIA TransformerEngine через API attn_forward_func_with_cp.
• Гибкая настройка степеней параллелизма (sp_ulysses_degree и sp_ring_degree) для балансировки нагрузки.
• Поддержка стратегий load balancing, таких как zigzag и stripe, для каузальных (causal) моделей.

👤 Кому подойдёт: ML-инженеры, разработчики инфраструктурных решений для LLM, исследователи, работающие с распределенным обучением больших языковых моделей.