Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

Technisch Overzicht: Beyond Task-Agnostic: Task-Aware Grouping voor Communicatie-efficiënte Multi-Task MoE Inferentie

1. Probleemstelling

Sparsely geactiveerde Mixture-of-Experts (MoE) modellen schalen de capaciteit efficiënt via conditionele berekening, maar kampen met aanzienlijke systeemuitdagingen tijdens gedistribueerde inferentie, specifiek excessieve cross-GPU communicatie-overhead en ernstige lastonbalans. Bestaande strategieën voor expert-plaatsing proberen deze kosten te mitigeren door co-geactiveerde experts samen te brengen, maar zij afleiden doorgaans een enkele, globale implementatieplanning vanuit taak-agnostische (globaal geaggregeerde) routing-traces.

De auteurs identificeren een kritiek gebrek in deze aanpak: expert co-activatie is sterk taak-geconditioneerd. Zoals aangetoond in de analyse van het paper (Figuur 1), zijn expert-paren die nauw gekoppeld zijn in één taakfamilie (bijv. codegeneratie) vaak ongecorreleerd of zelfs concurrerend in een andere (bijv. wiskundige redenering). Het middelen van deze heterogene patronen in een enkele globale affiniteitsgrafiek laat taakspecifieke expertgroepen instorten, waardoor communicatielocaliteit onbenut blijft en het niet kan adapteren aan distributieverschuivingen in online workloads.

2. Methodologie

Het voorgestelde framework, Task-Aware Coactivation Grouping (TACG) in combinatie met Generic Expert Shared Replication (GESR), pakt deze problemen aan via een tweestaps proces tijdens de implementatie:

A. Task-Aware Coactivation Grouping (TACG)
TACG zet taak-geconditioneerde co-activatie statistieken om in een capaciteits-haalbare plaatsingsplanning:

1. Scoren van Taakfamilie-voorkeur: Gebruikmakend van kalibratie-traces, berekent het systeem gebruiksvectoren en co-activatiematrices voor elke taakfamilie. Het berekent een "soft preference score" ($s^{(f)}_{l,e}$) voor elke expert ten opzichte van specifieke taakfamilies, waarbij deze scores worden genormaliseerd om een waarschijnlijkheidsverdeling ($p^{(f)}_{l,e}$) te creëren die de affiniteit van een expert voor een familie representeert.
2. Taak-gemoduleerde Coactivatiegrafiek: In plaats van een gepoolde co-activatiegrafiek te gebruiken, construeert TACG een taak-gemoduleerde grafiek ($\tilde{A}_l$). Deze grafiek herweegt de standaard co-activatie-edges door het product van de familie-voorkeuren van de experts ($p^{(f)}_{l,e} p^{(f)}_{l,e'}$). Dit zorgt ervoor dat edges tussen experts die co-activeren en een sterke familie-voorkeur delen, worden versterkt, terwijl cross-familie edges worden afgezwakt.
3. Capaciteits-haalbare Groepering: Het framework partitioneert experts in $M$ groepen (één per GPU) door de binnen-groep affiniteit op de taak-gemoduleerde grafiek te maximaliseren onder strikte capaciteitsrestricties. Dit wordt opgelost via een tweestaps spectrale clustering-aanpak gevolgd door een lokale "capaciteitsherstel"-stap om ervoor te zorgen dat exacte GPU-lastlimieten worden gehaald.

B. Generic Expert Shared Replication (GESR)
Om de lastscheefheid (load skew) aan te pakken die wordt veroorzaakt door "generieke" experts (die actief zijn over alle families en gevoelig zijn voor hotspots), introduceert GESR een lichtgewicht begeleidend mechanisme:

1. Identificatie: Het identificeert generieke experts met een score die globale centraliteit en cross-familie consistentie beloont, terwijl sterke taakspecifieke afwijking wordt bestraft.
2. Replicatie: Deze experts worden gerepliceerd naar een kleine set secundaire GPU's, gekozen op basis van structurele compatibiliteit (co-activatie affiniteit).
3. Online Guard: Tijdens het serveren kiest een localiteit- en lastbewuste selectieregel de beste kandidaat-GPU voor een generieke expert, waarbij exponentieel vervallende lastsignalen worden gebruikt om hotspots te voorkomen zonder de MoE-router te wijzigen.

3. Belangrijkste Bijdragen

• Inzicht: Het paper stelt vast dat expert co-activatie in multi-task MoE serving sterk taak-geconditioneerd is, en betoogt dat effectieve plaatsing gedreven moet worden door taak-bewuste co-activatie in plaats van een enkele globaal gemiddelde affiniteitsgrafiek.
• TACG: Een taak-bewuste strategie die de co-activatiegrafiek herweegt met behulp van expert-naar-familie voorkeuren, wat een plaatsing oplevert die strikt generaliseert ten opzichte van taak-agnostische groepering.
• GESR: Een lichtgewicht replicatiemodule die selectief cross-familie generieke experts repliceert en lastbewuste online selectie uitvoert om balans te behouden onder workloadverschuivingen.
• Empirische Validatie: Uitgebreide evaluatie op drie representatieve open-source MoE-modellen (DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B) die consistente verbeteringen laten zien ten opzichte van sterke baselines.

4. Experimentele Resultaten

Het framework werd geëvalueerd op een 16-GPU expert-parallel setup over drie modellen met vier taakfamilies (Code, Query, Math, Complex Reasoning).

• Reductie van Communicatie: Het volledige framework (TACG + GESR) verminderde de gemiddelde per-token communicatiekosten met 31,39% vergeleken met de baseline, waarmee het de sterkste eerdere plaatsingsstrategie (ExFlow) met 12,58% versloeg.
• Lastbalans: De methode handhaafde een hoge gemiddelde Jain fairness index van 0,9975 en een lage gemiddelde maximale lastschending (MaxVio) van 0,0736, wat de lastbalans van baselines evenaart of overtreft terwijl de communicatie aanzienlijk werd verminderd.
• Robuustheid: Onder zware taak-distributieverschuivingen (Out-of-Distribution single-task workloads), behield de taak-bewuste benadering stabiele communicatiereducties, terwijl taak-agnostische varianten leden onder grote prestatieschommelingen en hotspot-vorming (waarbij MaxVio met ~7x steeg bij Code-taken).
• Ablatie: Het verwijderen van het taak-gemoduleerde co-activatiesignaal verminderde de communicatiewinst van 31,39% naar 16,51%, wat bevestigt dat taak-conditionering de dominante factor is. GESR bleek essentieel voor het beheersen van lastschendingen (MaxVio) zonder de communicatie-efficiëntie op te offeren.

5. Betekenis en Claims

Het paper claimt dat de betekenis ligt in het verschuiven van het paradigma van MoE expert-plaatsing van een statisch, taak-agnostisch perspectief naar een dynamisch, taak-bewust perspectief. Door expliciet de heterogeniteit van expert co-activatie over verschillende taakfamilies te modelleren, ontsluit het framework een communicatielocaliteit die onzichtbaar is voor globale aggregatiemethoden.

De auteurs stellen dat deze aanpak een strikte verbetering biedt op de communicatie-lastbalans front. In tegen tegenstelling tot eerdere methoden die vaak lastbalans opofferen voor communicatiewinst (of vice versa), bereikt TACG + GESR een substantiële reductie in communicatie terwijl een bijna perfecte lastbalans behouden blijft. Het werk benadrukt dat effectieve multi-task serving strategieën vereist die de intrinsieke structuur van taakspecifieke expert-samenwerking respecteren, in plaats van deze te middelen.

Het paper blijft bescheiden over de reikwijdte en merkt op dat het zich richt op statische offline plaatsing en toekomstig werk openlaat voor dynamische hergroepering binnen request-windows, geheugen co-design (prefetching/caching), en gespecialiseerde experts in domein-gepartitioneerde modellen.