ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

Technische Samenvatting: ParaBlock: Communicatie-Computatie Parallel Blok Coördinaat Federated Learning voor Large Language Models

Probleemstelling
Federated Learning (FL) is opgekomen als een privacy-bewarend paradigma voor het trainen van modellen over gedecentraliseerde data. Recentelijk hebben federated block coordinate descent (BCD) schema's tractie gekregen voor het trainen van grootschalige modellen, waarbij cliënten slechts een subset (blok) van de modelparameters lokaal kunnen bijwerken in plaats van het gehele model. Echter, in het tijdperk van Large Language Models (LLMs), kan zelfs een enkel blok tientallen tot honderden miljoenen parameters bevatten. Deze schaal creëert een kritieke bottleneck: de communicatielatentie die vereist is om deze grote blokken tussen cliënten en de server te verzenden wordt aanzienlijk, met name voor middelen-beperkte edge-cliënten.

Traditionele federated BCD-methoden werken met een single-threaded ontwerp (communicatie-computatie-communicatie), waarbij cliënten moeten wachten tot de modeltransmissie voltooid is voordat ze de volgende computatiestap kunnen initiëren. Hoewel deze communicatietijd verwaarloosbaar was voor kleinere modellen, belemmert het nu de schaalbaarheid en efficiëntie van federated fine-tuning voor LLMs, wat zorgt voor substantiële inactiviteit bij cliënten.

Methodologie: ParaBlock
Om deze latentie aan te pakken, stellen de auteurs ParaBlock voor, een nieuwe federated BCD-methode die een communicatie-computatie parallel architectuur introduceert. In tegenstelling tot standaard single-threaded benaderingen, vestigt ParaBlock twee parallelle threads op lokale cliënten:

1. Computation Thread: Voert lokale block coördinaat updates uit (het fine-tunen van het toegewezen blok) met behulp van lokale data.
2. Communication Thread: Synchroniseert met de server om de lokale updates van de vorige ronde te verzenden en de geaggregeerde globale updates te ontvangen.

Kernmechanismen:

• Overlapping: Cliënten voeren lokale computatie uit terwijl ze gelijktijdig met de server communiceren. Dit elimineert de inactiviteit die normaal gesproken wordt veroorzaakt door het wachten op netwerktransmissie.
• Staleness en Correctie: Omdat de communicatie-thread één ronde achterloopt (cliënten ontvangen globale updates van ronde $t-1$ terwijl ze berekenen voor ronde $t$), kunnen de lokale modelparameters inconsistent worden met het globale model. ParaBlock pakt dit aan via een correctiemechanisme. Wanneer een cliënt zijn lokale model voor de volgende ronde bijwerkt, past hij niet alleen de nieuwe lokale update toe op het huidige blok, maar corrigeert hij ook de parameters van het vorige blok met behulp van de ontvangen globale update ($\Delta_{t-1}$) om de verouderde lokale update ($\Delta^i_{t-1}$) te vervangen.
• Recursieve Consistentie: Theoretische analyse toont aan dat deze correctie ervoor zorgt dat na één globale ronde de lokale modelparameters voor het eerder bijgewerkte blok exact overeenkomen met de globale modelparameters, waardoor de impact van de één-ronde vertraging effectief wordt gemitigeerd.

Belangrijkste Bijdragen

1. Nieuwe Architectuur: Het voorstel van ParaBlock, dat het traditionele single-threaded FL-ontwerp vervangt door een twee-thread parallel schema dat specifiek is afgestemd op het verminderen van communicatielatentie bij federated LLM fine-tuning.
2. Theoretische Garanties: De auteurs bieden een rigoureuze convergentieanalyse voor niet-convexe doelen. Ze bewijzen dat ParaBlock een convergentiesnelheid van $O(1/\sqrt{T})$ bereikt, wat consistent is met standaard federated block coordinate descent methoden. Dit bevestigt dat de efficiëntiewinst door parallellisme de theoretische convergentieprestaties niet in gevaar brengt.
3. Empirische Validatie: Uitgebreide experimenten zijn uitgevoerd op het fine-tunen van Llama 3-8B en Llama 3.2-3B modellen op algemene instructie-opvolging (Alpaca-GPT4) en wiskundige redenering (MathInstruct/GSM8K) taken.

Resultaten

• Prestaties: ParaBlock behoudt sterke downstream prestaties en behaalt scores die vergelijkbaar zijn met of beter zijn dan standaard federated BCD baselines (bijv. FedBCD) en presteert significant beter dan verschillende Parameter-Efficient Fine-Tuning (PEFT) gebaseerde FL-methoden (bijv. FedIT, FLoRA) op instructie-opvolging en wiskundige redeneringstaken.
• Efficiëntie: De methode vermindert de wall-clock runtime aanzienlijk. In scenario's met een lage netwerkbandbreedte (50M/s) en kleine batchgroottes, verminderde ParaBlock de runtime met meer dan 30% vergeleken met block-coordinate baselines zoals FedCyBGD. Vergeleken met LoRA-gebaseerde methoden zoals FedIT, bereikte ParaBlock in de meeste settings een runtime reductie van ongeveer 40%.
• Robuustheid: Ablatie-studies toonden aan dat de methode effectief blijft onder heterogene datadistributies (non-IID) en over verschillende modelarchitecturen (inclusief Qwen-2.5). De methode demonstreerde ook compatibiliteit met gedeeltelijke participatie-instellingen in cross-silo scenario's.

Betekenis
Het artikel stelt dat ParaBlock een eenvoudige maar effectieve oplossing biedt voor de communicatie-bottleneck die inherent is aan federated BCD voor LLMs. Door communicatie en computatie te overlappen, maakt het de schaalbare inzet van federated fine-tuning voor grote modellen op middelen-beperkte cliënten mogelijk zonder concessies te doen aan convergentiegaranties of modelutiliteit. Het werk benadrukt dat het verminderen van communicatielatentie door architecturale parallelliteit een levensvatbare en noodzakelijke strategie is voor de volgende generatie federated LLM training, die onderscheidend is van en orthogonaal aan bestaande technieken zoals kwantisatie of pruning.