A Unified Solution for Data Management and Model Training With Apache Iceberg and Mosaic Streaming

Summary

• Magnus data platform: ByteDance enhanced Apache Iceberg (Magnus) to manage exabyte-scale datasets for ML training, using branch/tag versioning, optimal sharding, and custom data readers for efficiency.
• Training data optimization: Magnus provides flexible table migration, detailed metrics, and full-text indexing to ensure training reliability and efficient data loading for massive distributed training jobs.
• Byted Streaming infrastructure: ByteDance customized Mosaic Streaming to enable distributed training with faster startup, lower resource consumption, and broader data source compatibility for large model development.

A Unified Solution for Data Management and Model Training With Apache Iceberg and Mosaic Streaming

Talk by Zilong Zhou, Software Engineer, ByteDance
this video introduces ByteDance’s challenges in data management and model training, and addresses them by Magnus (enhanced Apache Iceberg) and Byted Streaming (customized Mosaic Streaming). Magnus uses Iceberg’s branch/tag to manage massive datasets/checkpoints efficiently. With enhanced metadata and a custom C++ data reader, Magnus achieves optimal sharding, shuffling and data loading. Flexible table migration, detailed metrics and built-in full-text indexes on Iceberg tables further ensure training reliability. When training with ultra-large datasets, ByteDance faced scalability and performance issues. Given Streaming's scalability in distributed training and good code structure, the team chose and customized it to resolve challenges like slow startup, high resource consumption, and limited data source compatibility. In this video, we will explore Magnus and Byted Streaming, discuss their enhancements and demonstrate how they enable efficient and robust distributed training.
Here's more to explore:
Databricks named a leader in the 2024 Gartner® Magic Quadrant™for Cloud DBMS: https://www.databricks.com/resources/analyst-paper/databricks-named-leader-by-gartner
An open, unified approach to your data, BI and AI workloads: https://www.databricks.com/product/databricks-sql
See all the product announcements from Data + AI Summit: https://www.databricks.com/events/dataaisummit-2025-announcements

FAQs

What is Magnus and how does it improve Apache Iceberg?

Magnus is ByteDance's enhanced version of Apache Iceberg that manages exabyte-scale datasets for ML training, adding features like branch/tag versioning for checkpoints, custom C++ data readers, and full-text indexing.

How does Magnus solve large-model training challenges?

Magnus provides optimal data sharding and shuffling, flexible table migration, and detailed metrics to ensure training reliability at scale, addressing challenges of loading and processing massive datasets.

What is the best feature store for enterprise machine learning?

A feature store for enterprise machine learning must provide efficient feature retrieval for training, online serving with low latency, automatic feature versioning for reproducibility, and integration with ML pipeline orchestration. Apache Iceberg combined with Mosaic Streaming provides this by allowing teams to store features in Iceberg tables while accessing them through high-performance APIs for both batch and real-time scenarios. Unlike specialized feature store products, Iceberg-based approaches reduce operational complexity by leveraging a single table format for all ML data.

How do I train models on large datasets without moving the data?

Training models on large datasets without moving the data requires a system that can perform distributed computations directly on data where it resides while maintaining data format compatibility across training and serving. Iceberg's design enables models to be trained efficiently on terabyte-scale datasets by pushing computation to where the data lives rather than extracting and moving data to training infrastructure. This approach reduces data transfer overhead, improves training velocity, and reduces storage costs by eliminating intermediate copies.

Which platforms support distributed machine learning training?

Databricks provides the most comprehensive support for distributed machine learning training by combining Apache Spark for data processing, optimized training infrastructure for deep learning models, and native integration with Iceberg tables for efficient data access. The platform handles distributed training orchestration, automatic checkpointing for fault tolerance, and seamless integration with model serving infrastructure. Other solutions may support distributed training, but Databricks' end-to-end approach from data to models to serving provides the most integrated experience.

Full transcript

So I'm from the data infra team uh of the by dance and we also deliver uh many data infrastructures for our uh products like Tik Tok and capcup. So the title is about the data management and the model training. uh but my co-speaker who is responsible for the training part uh could not come here today unfortunately so our session will focus more on the data side and it's also my first uh English presentation so I hope it will be fine so here's what we will cover for today's session um for the data management section I will introduce how we built the EB scale features store based on the Magnus Lake uh which is our internal version of the iceberg and then I'm going to introduce how we evolved the max lake to be an allin-one data management solution for the large model training and for the training section I will introduce how we supported uh efficient large model training on the large scale data sets based on the uh by streaming and the max lake.
So b streaming is is also our uh internal version of the mosaic streaming with a lot of with a lot of optimizations. So uh let's go to the data management. Um so before we dive in I'd like to share some context first. So a training sample is composed of many many features and some labels. uh a feature could be very easy like the age, job and gender. Uh well, it could also be pretty complex and at that case we we might have to do some complex computation based on the external data sources or the existing feature values from the future store and the labels will contain some actions from the user.
So given the sample the model can learn that a user with these sort of features has done the actions within the labels. So uh here is a simplified uh offline machine learning workflow for a predication system which is extensively used in our uh many business like uh recommendation or advertising and searching. So a user action it will request the practice service for some results like a rank list of some contents and it will generate the labels of this request and the practice service will also uh dump the required features for this user according to the uh configurations set by algorithm engineers.
And then a join service will join the labels and the features to produce the training sample and finally write it into the feature store. So um after a period of accumulation uh algorithm uh after a period of accumulation uh algorithm engineer will select some features and labels to train new models and if the model is verified to be okay it will be deployed to the online pert services. But there is actually a problem in this picture is that uh if I want to use a new feature to train models, I have to wait for the feature store to accumulate enough samples containing this feature.
And the time depends on how many samples I need and the frequency of the online requests. So I will talk more about this in the next page. So our early solution was based on the row store and uh the data was stored as a row row HDFS files. So this way was proved to be very bad because uh it doesn't support column projection. So we have to read all the features even though uh we only need some of them and since ro store is not very good at uh compressing the data.
This will cause additional storage overhead. And as we talked on the previous page uh to back fuel or uh correct the feature uh we must rewrite the whole data. So um we summarize some points that we need for a good feature store and first we should be able to perform update and query oper operations uh just on some of the columns and that the data should be versioned and uh could be accessed and modified independently without any storage cost.
And finally, we should be able to write and read data through both streaming and batch engines such as Flink and Spark. So we plan to build our next generation feature store based on the data lake and we choose iceberg for its hierarchy design of metadata and which is great for our cases and at that time iceberg has uh good support for both link and spark engines uh which was also a big win. So is iceberg just enough for good feature store?
Um at that time the answer should be no because at bid dance our largest our largest feature store has up to tens of thousands of features and is still growing but ice can only do updates by copy and write or delete and insert and both of them will rewrite the whole record while an efficient column level update based on uh the column value merge on read was a must have for future store. Also it will be better if we can uh customize the merge on rate strategy and have better merge on rate performance.
So this led us to build Magnus Lake which is an enhanced version of iceberg um and is is uh is especially for the large scale machine learning workloads and now it is storing uh more than 5 xabytes of the data and has powered the feature store across bidance over the past few years. So first let me talk about the implementation of the um column level updates. So uh magnet supports both positional updates and equality based updates.
Um position update is designed for read modify write operations because when we read the data we can get the position of every modified record so that we can avoid doing merge on read by the key column when reading the data and equality update is designed for write only operations uh where the changes may come from the external sources and the updates are not based on their columns uh other columns in this table. In this case, we can write down the update files directly and also supports customizable merge on rate strategies so that user can write their own merge on rate logic to meet their business requirements.
And for example, I want to add an action to the labels of a user. Uh I can just use the concat list strategy to do the update. uh just like the uh update file with sequence number four in this picture you can see in the right corner. So now we have a good merge on rate based uh column up update but uh how do we ensure the read performance? So uh one problem is that uh there might be too many update files to be merged uh since there are so many features being added and this can cause performance issues when we read the data.
Uh but actually in our scenario uh the features are added by different users and uh every user just uh read their own features. So after pruning the update files uh based on the required columns there won't be too many files to be merged within the single scan task. Um but this situation does exist where uh the user has added many features uh with complex merge on read logic and he just wants to read all of them. Then here comes arrow engine.
So arrow engine is built in C++ and uh we have column wise merge on read and we also have done a lot of optimizations to get better read performance. So for example the merging of the columns are parallelized to better scale on the multipore the multiple core CPUs and air engine is also a middle layer for us to uh support different engines and file formats. So I'll talk about this in the next page. So as you can see in this picture uh there are lots of engines need to be supported uh such as spark flink and ray and by streaming is our internal version of streaming and uh cry is our uh self-developed uh hybrid researching and analytical engine and premise is our self-developed training system.
So they are all uh delivered by our uh data infer team and we also have lots of different file formats to support uh such as parakeet RIPC or TF record and so on. So our solution is to define an uh abstract read write interface based on the arrow in memory format. So uh any new file formats only needs to uh implement its own readwrite interface to get supported in uh Magnus Lake. And then we wrapped the uh scan task and the writing interface for different languages so that the new engine could also be easily supported by uh just integrating with the language SDK.
Uh yeah and uh uh we also find that uh the c the community has also raised an issue earlier this year uh which was trying to make it easier to uh support new file formats and uh yeah it's it's implement in a different way but I think it's a I think it will be a very exciting and useful feature and of course uh compaction is definitely uh uh needed for better read performance. So Magnus has its own catalog service and table maintenance service. uh for the catalog service.
It has a back end with some MySQL RDS instances and is managing more than uh 72,000 of the max tables stored in the HDFS or object stores and uh it provides restful API externally for the COD operations on the max tables and the permission control is realized by reusing the existing permission control from the HDF directory or object stores so that once the user get the read write permission of the HDF directory uh they get the read write permission of the max table.
So for the table madness we have some uh background tasks scheduled regularly to check the table health and uh some of them will trigger some new tasks to uh build indices or caches and uh also do the contraction for the table uh both for the data and the metadata. Yeah. Okay. Let's take a look at the feature engineering practice with Magnus Lake. So first is the uh feature analysis scenario. So every algorithm engineer could check out the new branch where he can back field the features he want to analyze and uh during the training only the required features are read.
So the benefit is that uh features are perfectly isolated among different branches uh without any storage overhead and uh there won't be uh any read amplifications when loading the features as well. So another scenario is also very straightforward is that um we can reuse all the features on the master branch uh without any sort of overhead for every model's training. So compared to the traditional ro store design uh we have more than uh 90% of the story saving uh in some of the production cases and that's an incredible number.
Okay. So let's uh that's all for the future store and let's move to the large models next. So the first thing is that um very surprisingly we found that uh many requirements from the large model training are uh very similar to the future store. So for example the for the data labeling uh scenario uh we might prepare some data sets um then we repeat the labeling training process to iterate many many models on different versions uh different versions of the data.
So this is just like the feature engineering workflow which also requires efficient column level updates and merge on reading mechanism and large model training also has a massive and complex data sets uh with bunch of versions to manage for pre-training sftt rm and evaluation even the checkpoint files and all of them need a good versioning mechanism and support for different uh file formats. And of course um so many data sets from so many scenario require uh different engines to read and write.
So uh which is also very similar to the official store. So other than that uh we also have some new requirements from the large model training. Uh basically we need a platform with better UIUX for tracking tables uh previewing the data and doing table migration across their centers or storage mediums. For example, we usually have much more uh storage resources than computing resources and in and in this case we need to uh move the data to the data center that has more free computing resources.
And we also need the cap capability to adjust storage replicas of different part of the table uh to maximize the utilization of remote remote storage bandwidth and solve the read hotspot issues. So and for existing data sets such as uh some open source data sets uh that we want to use them directly we can create register table by scanning the existing data set and generate expert meta files under the uh metadata path of the table and finally for better insight and troubleshooting experiences uh we might need the real time olive engine and the search engine.
So I'm going to introduce this later uh in detail. So first let me show you the onetop machine learning platform built with our uh machine learning for our colleagues. So as you can see in this picture uh we have previewed the we have previewed the detailed information of the table such as the file number the total storage volume uh the record number and the schema of the table. So all of them are collected from the iceberg metadata and the cached in our table maintenance service.
So for the table maintenance uh administrator can see the detailed metrics of all the tasks for a table and uh manage the task execution as well. So the user could trigger certain type of tasks like uh the table migration uh cache creation and so on uh through the uh platform uh interface. So for the data insight and the troubleshooting uh we also inter uh integrated a crypto engine to query the mega table directly without reingesting the data to the OLAP engines and it also has some uh important optimizations such as the pate push down and data caching on the external magnus table to get optimal query performance.
So the next thing I'd like to talk about is how we get the search ability on the mangos lake. So for the better experience of search search queries uh such as full text searching uh we hope they can finish in a few seconds to be uh usable for the business. Uh but unfortunately iceberg is designed for the batch query. So we have to duplicate the data to other search engines like elastic search to get better performance and this will make the data hard to manage and uh cause additional storage search storage cost.
So our solution is based on the fact that every data files in the iceberg table will become immutable once they are committed. So we can create index files alongside them with a one one mapping. So during a search query uh for every file scan task we can read the index file first to get the row indices then use the row indices to read the data file. And for top K queries such as the vector search we have to do a global sort mer global sort merge operation to get the final result.
And as for the uh full text indexing library uh we cannot use the lucine because our reader is implemented by uh pure C++ and after some investigations we choose tent TV as our full text index kernel. Since tent since tent TV is a rust library uh we integrated it with the arrow engine through the foreign function interface and the index searching logics are completely implemented in arrow engine so that uh different engines could reuse the index searching ability by integrating with arrow engine through the language bindings like pi bind for python or share library or shared library linkage like C++. ++.
So uh we made a benchmark on a text data set which was picked from the real data sets used for our large language mod large language model for training. So this data set has about uh 4 TB of the data with about 20,000 data files and the result shows that uh the first query can be finished at the second level and the following first can be finished at the millisecond level while the data caching is in used. So the vector index on maxic is already in the testing phase and uh we have also opened a a proposal in the community to support index and the vector index on the iceberg.
So here's the issue. Okay, let's bring them all together. And uh so as we can see in this picture, um the whole life cycle of the large model training is included in the top box. And uh the read and write operations to the different data in the bottom box are finished by interacting with our engines, SDKs, rest APIs and so on. And uh this sort of stuffs will interact with the arrow engine and the arrow engine will read and write the data set directly just as shown in the picture.
Okay, let's move to the model training section quickly. Um so I think our machine learning for colleagues have uh started using the streaming since early this early in the next uh last year. uh and we choose mosaic streaming as the basis to build our next generation data loader called by streaming uh for this uh several adv advantages. Um but by streaming uh has two design decisions uh which is uh different from mos streaming.
So firstly um by the streaming choose to read the read the shards from remote storage directly into the memory uh because in our environment um the disk resources are not sufficient enough to store the shard files uh all in the disk uh because our our source scale is extremely large and also the disk failure is uh always happening and uh what and what make it what makes it worse is that the disk space can be easily taken up by other apps uh such as uh log files or code dump files.
So for these reasons we don't rely on the local disk and the other one is about the elastic deterministic. So I think it's a very useful ability to uh enable the same sample order when the parallel configur configurations change. uh but still if we don't free download the shard files to the local uh this will be problematic and uh I will talk about it in the next page. So as we can see in the middle of this picture uh there are a bunch of shards which are already shuffled and uh in motorcycle streaming um every uh even the data parallel size or worker number changed the global sample order will will not change uh because the consuming order of the shards didn't change and while the problem is that u all the all the workers might be trying to read the same shard but only requiring a small piece of the shards And this can cause severe read hotspots and amplifications and uh but in by streaming uh every worker will start reading from the dash lines in the picture and consuming the shards sequentially.
So there won't be too many overlaps. Uh basically it's just uh u at most one shard between every workers and the read amplification can also be reduced since the worker uh would read the whole shard from the start to the end uh instead of just uh small pieces. So we have also made uh a lot of optimizations to the streaming itself. Uh and there and here I will briefly introduce three of them. So first we implemented a multis source data set which contains several uh data sets mixed in proportion.
So the mixing strategy could be could be adjusted uh during the training process and in every batch the data from all the sources will be included according to the mix mixing proportion. uh but but in mosaic streaming uh I think a streaming data set will only contain the data from a single stream uh within a batch because they are belongs to the same uh the same single shard space and the second is that uh the lazy simple ID generator uh so in mosaic streaming every sample will be assigned with the sample ID during initialization but when dealing with large scale data sets uh this would cause memory issues used.
So we implement ID generator mechanism to lazily generate sample ids for the uh samples when they are and when they are actually accessed it it will uh generate the uh ids. So it also so uh it also supports the reshar scen scenario and finally we have a broadcast pref preeter so it's based on share memory and make it possible to reusing the same samples across different training processes in the physical node and it also supports GPU prefetch which can avoid data movement between the CPU memory and the GPU memory.
So now we can try to uh integrate BD streaming uh with the Magnus Lake. So BD streaming can support any formats by simply implementing the corresponding format reader and I think it only takes 100 lines of code under the format directory. As we can see in this picture uh the max the reader supports two types of sample access modes. So for data sets where a single shard is relatively small such as some short text data sets we can load the whole shard into the memory from the remote storage.
So the subsequent accesses to of the sample could be performed randomly based on the shuffle unit shuffle result. And for the sets where a shard is very huge like video data sets we can consume the shard rows by rows so that only a few rows of the shard will be present in the memory. So this tra this actually trace randomness and uh throughput for better memory efficiency. So why didn't use MDS directly as it should be uh well optimized for the training workload.
So I think the most problematic thing of MDS is that um we have to duplicate the data just for the training while the MDS is a ro store and uh it's not very good at the compression. So the storage cost will be uh much higher than we than we thought and then it's it's not very easy to uh update the MDS data sets. So we have to rewrite the whole data sets to just um update a single field and this is normal for row store.
So the concept in the MDS could be mapped to the iceberg. Uh so a data set is a mixture of iceberg tables. A stream is an iceberg table and a shard is a piece of iceberg table while it could be uh like a file or row group and or even like uh any arbitrary ranges of the file. So at and the index. json file of MDS is also similar to the hierarchy metadata of the iceberg where we can get the information of every shards in this table.
Oh uh this is the so here is the process of uh like by streaming rating from the max table. So the multis source DSS is the mixture of some bided data sets which is our uh internal version of the stream data set and by data set similar to the uh stream data set will contain um multiple by streams and a bid stream is uh just uh is just uh equal to a magnus table. So we can get the shards of the stream by planning file scan tasks from the max table then wrap them to be the shards objects of the streaming.
So if the short is at the row group level rather than file level the ro scan task will have a pointer to the original file scan task together with the row group index and the record count of the row group and these informations are used for sharding and shuffling in the next processes. So during the planning we need to go through the iceberg planning process and initialize all the file scan task objects for the max table and then construct all the row group scan task objects and finally construct all the shard objects for streaming and during read we will read through the row group scan task to access a shard.
So there is actually some problems here. So because we have to construct all the file scan task and the row group scan task objects at the init in initialization phase. So the startup time will be pretty long and uh the memory usage will be extremely high in the case we have ultra number of sources and this is just similar to the uh sample ID generation we have talked before. So our solution is to uh just build a plan cache asynchronously.
So if a snapshot of a table is accessed for the first time um we will build we will schedule a task for building the plan cache for this table and the content of the cache file is serialized from the magnet scan context object uh which has a list of buys uh which representing the serialized file scan task objects. So and there is also some empty slots for the serialized file scan tasks and uh together with some numpire race which is compressed representations of the row group scan tasks.
So for a single table during planning uh we don't have to construct any file scan task objects and the row group scan tasks are represented by some numpy arrays which also don't have the overhead of python object header and uh only when we actually access a certain row group scan task we construct the pointed file scan task if not constructed yet and then constructed the row group scan task with the compressed informations from the numpy arrays.
So, so that uh so that extremely reduced the startup time during the training. And uh so our team was also actively working with the uh machine learning infra team to build the next generation data loader with ray and b streaming. uh which uh so so this this solution can scale ultra number of the sources in the large model training and uh we're also considering open source the uh by streaming as well as this solution uh in the future and uh so and also for the uh by streaming and the max lake we are also considering uh like uh participate participate more in the community community and try to contribute some useful features to the open source course.
Okay, that's all for my uh session. Uh thanks. Thank you. Thanks. Thank you so much.

Learn more about the Databricks Data and AI platform.