Effizientes LLM-Training und -Serving mit heterogenem Kontext-Sharding unter Attention-Heads

Bestehende LLM-Trainings- und Inferenz-Frameworks haben Schwierigkeiten, die Effizienz durch Sparsität zu steigern und dabei die Integrität von Kontext und Modellarchitektur zu erhalten. Inspiriert vom Sharding-Konzept in Datenbanken und der Tatsache, dass Attention auf Beschleunigern über Heads parallelisiert, schlagen wir Sparsely-Sharded (S2) Attention vor, einen Attention-Algorithmus, der unterschiedliche Kontextpartitionen für verschiedene Attention-Heads zuweist, um Aufteilen und Beherrschen zu ermöglichen. S2-Attention erzwingt, dass jeder Attention-Head nur eine Teilmenge der Kontexte im strided Sparsity-Muster betrachtet, während der vollständige Kontext als Vereinigung aller Shards erhalten bleibt. Da die Attention-Heads in separaten Thread-Blöcken verarbeitet werden, kann die Kontextreduktion für jeden Head eine End-to-End-Beschleunigung und Speicherreduzierung bewirken. Bei der Inferenz können LLMs, die mit S2-Attention trainiert wurden, die KV-Cache-Reduktion als kostenlosen Vorteil mit garantierter Modellqualität nutzen. In Experimenten zeigen wir, dass S2-Attention (1) eine bis zu 25,3-fache Beschleunigung der Attention im Vergleich zu FlashAttention-2 bietet, was zu einer 6-fachen Reduktion der End-to-End-Trainingszeit und einer 10-fachen Verringerung der Inferenzlatenz führt, (2) eine modelltrainingqualität auf gleichem Niveau wie bei der standardmäßigen Attention erreicht und (3) perfekte Needle-Retrieval-Genauigkeit über ein 32K Kontextfenster ermöglicht. Aufbauend auf dem Algorithmus haben wir DKernel entwickelt, eine LLM-Trainings- und Inferenz-Kernbibliothek, die es Nutzern ermöglicht, Sparsity-Muster für ihre eigenen Modelle zu individualisieren. Wir haben DKernel als Open Source veröffentlicht und machen es kompatibel mit Megatron, Pytorch und vLLM.