Model training forms the core of ML platforms, and given the diversity of use cases, we've designed our training platform to be highly adaptable.
A standard model training job generally involves four tasks: (1) training data generation, (2) model training and tuning, (3) model evaluation, and (4) model export. These tasks create a workflow graph. Depending on the ML application, there are different variations of this workflow, such as composite models—where a model combines multiple separately-trained sub-models—and model distillation, where a sophisticated, resource-intensive “teacher” model is used to train a simpler, more cost-efficient “student” model. As flexibility is crucial, we've adopted Kubeflow which is an open-source workflow standard for ML applications. Through Kubeflow, we have a variety of workflow templates supporting a wide range of scenarios. Depending on the configuration, we orchestrate the workflows on Google Kubernetes Engine or Google Cloud's Vertex AI Pipelines. The following screenshot is a training workflow example displayed in the Bento UI.
We have already discussed the training data generation step in the previous section. The model training job is a Python application running on GPU or TPU hosts. While the training application could involve any user-provided code, to enhance ML development speed—particularly for ranking and recommendation applications—we've structured the application into three composable layers: (1) Core framework, (2) User model code, and (3) Training configuration.
The Core framework is an internal library that encapsulates Tensorflow and Keras, standardizing many aspects of modeling techniques and components for recommendation use cases, such as deep cross network and transformers. This structure enables ML engineers to experiment more efficiently with various feature ideas, and promotes code reuse and sharing. ML engineers then express their model by authoring user model code. Finally, the training configuration is a YAML file that outlines how the training job should be executed. For example, it specifies the type of hardware, the location of input and output data, and other runtime options that determine the behavior of the training code. This layered structure of our training code facilitates rapid experimentation by merely modifying the job configuration. Finally, we support model training on GPUs and TPUs. Our Core framework provides sufficient abstractions to handle flexibility of running training code on different hardware types.
Model training has many hyperparameters such as learning rate, batch size, and other model parameters such as layer dimensions. Part of the ML development process is to experiment with different parameter configurations. To facilitate this, we integrate Bayesian optimization to the workflow. Bayesian optimization suggests better parameter values based on observed metrics, reducing the amount of hand tuning ML engineers need to do.
Once a model is generated, we run an evaluation task to yield model metrics. We upload these metrics, along with those from the training phase, to Tensorboard for visualization. The evaluation results for each model are stored in a database and displayed in the Bento UI.
Following the evaluation task, the Model Export task takes the trained model and prepares it for inference. At Snap, we use different hardware for inference, depending on factors like application, latency requirements, performance characteristics, and resource usage. The Model Export task creates optimized models for different types of hardware. For instance, for models destined for inference on GPUs, we examine the compute graph and determine which operations should be performed on GPUs—such as dense matrix multiplications—and which should be performed on CPUs, like feature parsing and embedding lookup. This optimization is crucial for reducing the serving costs of ranking and recommendation models, because these models exhibit a distinct pattern of dense computation that is compute-bound, along with large embedding lookup tables that are memory-size-bound.
Last but not least, to maintain prediction accuracy, models must be incrementally updated as new events are recorded. Bento 100% automates incremental training. Our orchestration handles the incremental appending, joining, and transformation of new training data, preparing it for incremental training. The incremental training itself is scheduled on Kubeflow. After passing through export and various validations, models are automatically deployed to production. We have online ML monitoring to ensure the quality of model production, which will be introduced later in this article.