Optimal hardware selection for AI-model learning

Sammanfattning

Artificial Intelligence (AI) is a field of study that has seen rapid advancements in recent times, particularly due to the development of larger and more complicated models. The scale of recent models, in particular Large Language Models (LLMs), has led to a massive increase in computation required for development of these state-of-the-art models. Out of necessity the workload centered around these large-scale models is therefore oftentimes delegated to purpose-built computational clusters. These clusters are usually heterogenous when it comes to the types of deployed hardware which leads to increased difficulty in scheduling incoming jobs. This thesis was carried out in collaboration with Ericsson AB with the goal of developing a recommendation system for research clusters. Specifically, the system should be able to make accurate predictions about key running time metrics of AI training jobs using only the source code as input, which could help improve cluster efficiency. To this end, the system was broken down into several sequential components which are able to operate independently. The final system consists of a Transformer-CRF model for feature extraction of the source code followed by a Set-Transformer based architecture for regression of the features into a set of key metrics. These models are combined with algorithmic and mathematical solutions to form the purpose-built system. The mixed results on practical data will be put in context of the available data at hand along with training time constraints, forming a framework to carry the system into a production-ready state.