Federated learning in multi-vendor environments, lessons from the NextGen project

Handling heterogenous hardware

While traditional federated learning frameworks can support heterogeneous hardware, they often require complex, stable, always-open network connections. NextGen instead puts hardware drivers in containers and uses an asynchronous "blackboard" communication pattern built on standard APIs. The approach lets the framework successfully connect NVIDIA GPUs, Intel Gaudi HPUs, and AMD graphics cards, stored locally or behind a cloud service, all within the same training loop.

Led by AI Sweden and co-funded by Vinnova, the project's goal was to build a "plug n' play" infrastructure wrapper that allows companies to collaborate securely even though their hardware setups vary. The issue is brought to light by a growing interest in data sovereignty and recent regulations like the EU AI Act. 

There is also strong financial incentive as recent studies show that the total cost of ownership (TCO) for on-premise AI compute is increasingly outcompeting third-party cloud services.

A pharmaceutical proving ground

The project participants included Intel, AstraZeneca, Zenseact, Aixia, NetApp, and Proact. To test the framework against real-world enterprise demands, focus was set on a complex primary use case: large language model (LLM) capabilities for information retrieval in a pharmaceutical RAG system. 

The project focused on fine-tuning the embedding model component of the RAG, and using them for retrieving relevant document passages from dense, multi-100-page pharmaceutical documents. A separate conversational LLM, a chatbot, acted as the generator and assistant to answer user queries based on the retrieved context, as in a typical RAG setup.

“The blackboard pattern we use is completely asynchronous and is flexible when clients drop out for whatever reason. Likewise, we can connect new clients on-the-fly. The blackboard pattern we use is completely asynchronous and is flexible when clients drop out for whatever reason. Likewise, we can connect new clients on-the-fly. Model updates are never pushed to the clients, only pulled. This puts the clients themselves in control of the information flow”, says Mauricio Muñoz.  

Mixing different hardware during training did create a challenge though, machines with less memory and processing power needed to process smaller batches of data, and this matters for training models with the special type of learning algorithm known as “contrastive learning” used in this case.

”Metaphorically, high-memory nodes can use a ’high-definition map’ of the data, while lower-memory nodes have to navigate with a ’blurry outline,’” explains Mauricio Muñoz. “Merging these learning signals into a single model update leads to interesting dynamics.”

Because contrastive learning, the training method used, relies heavily on batch size to learn effectively, averaging these "high" and "low-definition" results together slightly lowered raw data retrieval scores. However, the final AI system remained highly accurate, proving the mixed-hardware framework works successfully without breaking the end product. 

An open-source foundation to build upon

The main output artifact from the NextGen project is a functional, open-source, plug-and-play engine, released as core software repositories (nextgen-framework, nextgen-train, and nextgen-rag-eval) alongside demos and extensive technical documentation. 

By combining efficient training techniques with smart data aggregation, the framework provides a flexible, secure foundation for the future of decentralized enterprise AI.