Youmu: Efficient Columnar Data Pipeline for LLM Training

Youmu: Efficient Columnar Data Pipeline for LLM Training

Speakers: Tianle Zhong, University of Virginia; Jiechen Zhao, University of Toronto; Qiang Su, Chinese University of Hong Kong; Geoffrey Fox, University of Virginia

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of large language model (LLM) training, data pipeline efficiency remains a critical bottleneck. This talk introduces Youmu, an innovative system designed to streamline the data ingestion and shuffling process for LLM training by directly leveraging Parquet, a widely adopted columnar data format. Presented by Tianle Zhong and his collaborators, Youmu tackles the pervasive problem of "costly pit stops"—the inefficient conversion of data from optimized storage formats like Parquet to less efficient ones solely for the purpose of data shuffling during model training.

The core challenge Youmu addresses stems from the inherent mismatch between Parquet's chunk-based I/O, ideal for analytical queries and compression, and the fine-grained random access required for effective data shuffling to prevent model overfitting. Youmu's significance lies in its ability to bridge this "granularity gap," enabling high-quality data shuffling and high throughput with a controlled memory footprint, all while maintaining the benefits of columnar storage. This work is crucial for practitioners and infrastructure teams seeking to optimize their LLM training pipelines, reduce operational overhead, and accelerate the development cycle by eliminating redundant data transformations and storage requirements.

Background

▶ Watch: Introduction: JSONL vs. Parquet for LLM Training (0:00)

The journey of data in LLM training typically begins with massive, often unstructured, datasets collected from the web. These raw inputs demand extensive pre-processing and cleaning before they can be fed to compute-intensive GPUs for model training. While many public datasets are released in formats like JSONL, which are often suboptimal for large-scale processing, the industry standard for data collection, cleaning, and analytical pipelines has gravitated towards Parquet. Parquet is a columnar format renowned for its efficiency in scanning operations, superior data compression, and optimized block storage—qualities that make it a cornerstone of big data analytics.

Despite Parquet's advantages, its adoption for the training phase of LLMs has been limited. The primary reason for this disconnect is Parquet's perceived inefficiency with random access, a capability that is absolutely essential for data shuffling. Shuffling data before and during training is a fundamental technique to ensure the model sees diverse examples in each batch, preventing overfitting and improving generalization performance. The traditional wisdom held that Parquet's chunk-based I/O model made fine-grained random row access prohibitively expensive, leading to what the speakers term "costly pit stops."

These pit stops involve converting data from its optimized columnar Parquet format into other formats (often row-oriented or less compressed) specifically for shuffling. This conversion process introduces several significant problems:

1. Increased Storage and DRAM Footprint: The converted data often requires substantial additional storage capacity and occupies valuable DRAM, especially problematic when large models already consume significant memory.
2. Redundant Operations on Updates: Data datasets are frequently updated—either with new data or due to refinements in pre-processing pipelines. Each update necessitates repeating the costly conversion process, leading to substantial computational and time overhead.
3. Loss of Data Integrity: Creating multiple copies in different formats breaks the "single source of truth" principle, complicating data governance and consistency.
4. Operational Complexity: Maintaining these conversion pipelines requires dedicated engineering effort, adding to operational costs.

Existing solutions attempting to directly train on Parquet data have their own limitations:

• Ray Data: While a powerful distributed data processing framework, Ray data often requires loading large portions of the dataset into DRAM for fine-grained shuffling. This can be impractical given the memory demands of LLM training (e.g., for model checkpointing). Furthermore, Ray's distributed memory management, particularly its object store interactions with RDMA, has known complexities, and data exchange between training nodes can lead to network contention with critical GPU traffic.
• Memory Mapping: For large-scale datasets, memory mapping typically results in unsatisfactory throughput due to thrashing, as it's not designed for the highly random access patterns of shuffling.
• Streaming I/O: This approach reads data sequentially from disk but relies on in-memory buffers for shuffling. While it can improve disk I/O, it often leads to a significantly degraded shuffle quality, potentially harming model accuracy.

Recognizing these challenges, the Youmu project set out with clear goals: to provide a controlled DRAM footprint, ensure sufficient throughput to keep GPUs busy, and guarantee high shuffle quality to preserve model accuracy, all while working directly with existing Parquet data.

Key Findings

▶ Watch: Youmu's Goals and Key Observations for efficiency (4:15)

Youmu's design is predicated on two critical observations that challenge conventional approaches to data loading and shuffling for LLM training:

1. Ineffectiveness of Caching for Random Access: For globally random access patterns, caching data in memory provides limited benefit. Since a significant portion of data will inevitably come from disk, the overall training pipeline frequently becomes bottlenecked by disk I/O regardless of caching efforts. Youmu therefore consciously abandons the use of a traditional cache, opting instead for a controlled in-memory buffer.

1. The Granularity Gap: This is the central insight. Data shuffling requires fine-grained access to individual rows or small batches to achieve high quality and prevent overfitting. However, standard columnar I/O, as implemented in Parquet, operates on chunk-based granularity. If only a single random row is needed from a chunk, the entire chunk must be read, leading to substantial wasted I/O bandwidth. Experiments showed this could reduce effective bandwidth to less than 0.1%, severely bottlenecking the training system.

Based on these observations, Youmu proposes a core idea: page-level granularity for I/O. Within Parquet's columnar chunks, data is further organized into smaller data pages, typically ranging from 10 to 100 kilobytes. Historically, these pages were primarily units for data compression, not for I/O. Youmu extends their utility, making them the fundamental unit for I/O operations. This finer granularity allows Youmu to read only the necessary pages when performing random access, drastically reducing wasted I/O and improving effective bandwidth.

The key findings and contributions of Youmu are:

• Direct Parquet Integration: Youmu works directly with existing Parquet datasets, eliminating the need for costly data format conversions.
• Page-Level I/O: By adopting page-level granularity (10-100KB pages), Youmu overcomes the granularity gap, enabling efficient fine-grained random access for shuffling without reading entire chunks.
• Global Page Index: A novel indexing mechanism is built by reading Parquet file headers and metadata to compute offsets, allowing for rapid calculation of any page's physical location given a random ID.
• Aggressive Buffer Shuffle: A specialized buffering and shuffling strategy that shuffles the in-memory buffer every time a new page is loaded, maximizing shuffle opportunities within a controlled memory footprint.
• Superior Performance: Experimental validation shows Youmu achieves model accuracy comparable to traditional row-based shuffling, maintains a significantly lower memory footprint compared to solutions like Ray Data, and provides sufficient throughput to prevent GPU idleness.
• Practicality: The system is implemented in Rust with Python APIs, making it easy to integrate into existing ML workflows via standard Python dataset interfaces.

In essence, Youmu preserves I/O and memory efficiency for LLM training on columnar data storage while simultaneously ensuring high shuffle quality and throughput, fundamentally transforming how large-scale data is prepared for deep learning.

Technical Deep Dive

▶ Watch: Core Idea: Page-Level Granularity for IO (6:00)

Youmu’s technical innovation centers on its ability to transform Parquet's internal data organization into an efficient mechanism for random access, specifically tailored for LLM training workloads.

At the heart of Parquet's structure are row groups, which contain column chunks for each column. Within these column chunks, data is further segmented into data pages. These pages are crucial for Parquet’s compression schemes, allowing independent compression and encoding of small data blocks. Youmu's breakthrough is to elevate these data pages (typically 10-100KB in size) from mere compression units to the primary units of I/O. This is a departure from traditional big data systems that operate on chunk-level I/O, which is efficient for sequential scans but highly inefficient for random access. By making pages the I/O granularity, Youmu ensures that when a specific row is needed for shuffling, only the small page containing that row is fetched, rather than the entire, much larger column chunk. This dramatically reduces the amount of "wasted" I/O, improving effective disk bandwidth.

To enable efficient random access to these granular data pages, Youmu constructs a Global Page Index. This index is built by parsing the Parquet file headers and metadata. These metadata sections contain crucial information about the offsets and lengths of row groups, column chunks, and critically, the data pages within them. By accumulating this hierarchical offset information, Youmu can, given a random logical row ID (or a page ID derived from it), quickly compute the exact physical location (file offset and length) of the corresponding data page. This lookup is designed to be highly efficient, likely employing techniques such as binary search over the indexed offsets. This index allows Youmu to perform targeted, fine-grained reads without scanning large portions of the file.

The next critical component is the aggressive buffer shuffle mechanism. Given that data pages are small (10-100KB), a large number of pages would need to be loaded to fill a traditional buffer for shuffling. Youmu optimizes this by continuously shuffling the in-memory buffer as new pages arrive. Instead of waiting for the buffer to be completely filled before shuffling, Youmu performs shuffling operations incrementally. Each time a new data page is read and added to the buffer, the shuffling algorithm re-evaluates and potentially reorders the elements within the buffer. This "aggressive" approach maximizes the opportunities for shuffling, ensuring a high degree of randomness even with a relatively small, controlled in-memory buffer. This strategy is essential for achieving high shuffle quality while adhering to strict DRAM footprint constraints.

Youmu is implemented in Rust for its performance benefits, particularly in memory management and I/O-intensive operations. To ensure ease of adoption and integration into existing ML workflows, it provides Python APIs and adheres to standard Python dataset interfaces. This allows data scientists and ML engineers to use Youmu as a drop-in replacement for their existing data loaders, requiring minimal code changes.

The system also offers flexibility regarding shuffle quality and throughput. While its default mode emphasizes random page consumption (picking random pages to form batches) for optimal throughput and memory efficiency, it can also support full shuffle at the row level. The speaker notes that full row-level shuffling would sacrifice some throughput but is still "good enough" in most experimental cases. Importantly, Youmu's architecture decouples the control plane from the data plane. This means users can configure the desired granularity of I/O (e.g., opting for larger chunk reads if fine-grained shuffling isn't strictly necessary for a particular task, or adjusting buffer sizes) without needing to modify the underlying data format or the core data processing logic. This configurability allows Youmu to adapt to diverse training requirements and resource constraints.

Experimental Setup & Results

▶ Watch: Youmu Architecture: Global Page Index and Shuffle (7:00)

The evaluation of Youmu focused on validating its core claims: preserving model accuracy through high-quality shuffling, achieving a low memory footprint, and providing sufficient throughput to prevent GPU idleness. While specific hardware configurations, dataset names, or detailed baseline implementations beyond "Ray Data" are not explicitly detailed in the transcript, the results presented offer clear comparative insights.

Model Accuracy:

A primary concern with any novel shuffling mechanism is its impact on the downstream model's accuracy. Youmu was tested by consuming random pages from the dataset, a strategy designed to maximize I/O efficiency. The experiments demonstrated that this page-level random consumption provides sufficient model accuracy, performing "just as good as row-based shuffling." This is a critical finding, as it validates Youmu's ability to achieve high shuffle quality through page-level granularity without compromising the model's learning capabilities. The KB-level page size was found to be granular enough to ensure this high shuffle quality.

Memory Footprint:

Youmu's memory efficiency was directly compared to Ray Data, a prominent existing solution for distributed data processing. The results indicated that Youmu achieves a "really low" memory footprint. This is a direct consequence of its design philosophy: abandoning a global cache and instead relying on a controlled, aggressively shuffled in-memory buffer. By minimizing the amount of data held in DRAM at any given time, Youmu frees up valuable CPU memory for other critical training operations, such as model checkpointing, which can be a significant constraint in large-scale LLM training.

Throughput:

The system also demonstrated its ability to provide sufficient throughput to prevent GPUs from becoming idle. GPU idleness is a common and costly problem in LLM training, as the compute units sit waiting for data. Youmu's optimized page-level I/O, coupled with its efficient indexing and aggressive buffering, ensures that data is fed to the GPUs at a rate that keeps them consistently utilized. While specific throughput numbers (e.g., GB/s) are not provided, the qualitative statement "sufficient throughput" indicates that Youmu effectively removes data loading as a bottleneck in the training pipeline. The ability to avoid thrashing (a common issue with memory mapping for large datasets) further contributes to its stable and high throughput.

Trade-offs (as discussed in Q&A):

The speaker acknowledged a theoretical trade-off where random page consumption might sacrifice some shuffle quality compared to true row-level shuffling. However, the experimental results effectively mitigated this concern by showing no loss in model accuracy. Furthermore, Youmu provides the option for full row-level shuffle, albeit at a potential throughput cost, demonstrating flexibility. Another potential trade-off, breaking down full chunk reads into many small page reads, can also be configured by the user through Youmu's decoupled control plane, allowing for customization based on specific workload needs.

In summary, Youmu successfully validated its design principles through experimental evaluation, confirming its ability to deliver high-quality data shuffling, low memory overhead, and robust throughput for LLM training directly on Parquet data.

Practical Implications

▶ Watch: Experimental Results: Model accuracy and memory footprint (8:00)

Youmu presents a paradigm shift for data management in large language model training, offering profound practical implications for various stakeholders within the ML ecosystem.

For practitioners and model builders, Youmu eliminates the cumbersome and often error-prone step of converting data from Parquet to another format solely for shuffling. This means they can directly use their cleaned and pre-processed data, which is already stored efficiently in Parquet, without incurring "costly pit stops." This direct integration simplifies the data pipeline, accelerates experimentation, and reduces the time-to-model deployment. The guarantee of high shuffle quality without sacrificing model accuracy is a crucial benefit, as it ensures that the model can still learn effectively and generalize well.

Infrastructure teams stand to gain significantly from Youmu's efficiencies. By avoiding redundant data copies, storage requirements are dramatically reduced. The lower DRAM footprint frees up valuable CPU memory, which is often a contended resource in GPU-accelerated training systems. The elimination of complex data conversion pipelines also means less engineering effort is required for pipeline maintenance and fewer specialized engineers are needed to manage disparate data formats. This translates directly into cost savings on storage, compute resources, and personnel. Furthermore, Youmu's ability to maintain a single source of truth for data (i.e., the Parquet files) simplifies data governance, versioning, and debugging, enhancing data consistency across analysis and training stages.

For deployers and MLOps teams, Youmu offers a more robust and scalable data pipeline. The system's ability to handle frequently updated datasets without requiring repeated, full-scale data conversions ensures agility. As new data becomes available or pre-processing techniques improve, Youmu can quickly adapt, feeding the latest data to models without introducing significant delays or operational overhead. The decoupled control and data planes provide flexibility, allowing teams to fine-tune I/O granularity and buffering strategies based on the specific characteristics of their hardware, datasets, and training objectives without modifying core data processing logic.

Trade-offs and Limitations:

While Youmu offers substantial benefits, it's important to consider its trade-offs. The primary theoretical trade-off is that random page consumption, while experimentally shown to maintain model accuracy, is not strictly equivalent to a full row-level shuffle. However, the empirical evidence suggests this is not a practical limitation for LLM training. For scenarios demanding absolute, perfectly uniform row-level shuffling, Youmu does offer a "full shuffle" mode, though it comes with a potential reduction in throughput.

The decision to stick with Parquet, rather than inventing a new format optimized specifically for random access, is a deliberate one. The speaker highlighted that Parquet is already widely adopted by companies and for data storage, and it is continuously evolving. This choice minimizes user adoption efforts, as users don't need to convert their existing data or learn a new data format. It also allows for a unified ecosystem where the same data can be used for both big data analysis and LLM training, fostering greater integration. While a purpose-built format might theoretically offer marginal gains, the practical benefits of leveraging an established, feature-rich format like Parquet, especially when it can be made highly efficient as Youmu demonstrates, often outweigh the advantages of creating a new standard.

In essence, Youmu empowers the ML community to build more efficient, less costly, and more agile LLM training pipelines by intelligently leveraging existing, powerful data formats.

Key Takeaways

• Eliminates "Costly Pit Stops": Youmu enables direct, efficient LLM training on existing Parquet columnar data, removing the need for expensive and redundant data format conversions solely for shuffling.
• Page-Level I/O Revolutionizes Random Access: By extending the role of Parquet's internal data pages (10-100KB) from compression units to I/O units, Youmu effectively bridges the "granularity gap," allowing fine-grained random access for shuffling without wasting I/O bandwidth.
• High Shuffle Quality with Low Memory Footprint: Youmu's Global Page Index and aggressive buffer shuffle strategy ensure high shuffle quality comparable to row-based shuffling, while significantly reducing DRAM consumption compared to traditional solutions like Ray Data.
• Maintains Model Accuracy and Throughput: Experiments confirm that Youmu's page-level shuffling strategy does not degrade model accuracy and provides sufficient throughput to keep GPUs fully utilized, preventing costly idleness.
• Practical and Flexible: Implemented in Rust with Python APIs, Youmu easily integrates into existing ML pipelines. Its decoupled control and data planes offer configuration flexibility, adapting to diverse training needs and resource constraints.
• Leverages Existing Ecosystem: By optimizing for Parquet, Youmu avoids the burden of new data format adoption, allowing organizations to unify their data analysis and LLM training pipelines on a single, widely-used, and evolving standard.

About the Speaker(s)

The work on Youmu was a collaborative effort presented by Tianle Zhong from the University of Virginia. He was supported by his collaborators, Jiechen Zhao from the University of Toronto, Qiang Su from the Chinese University of Hong Kong, and his advisor, Professor Geoffrey Fox, also from the University of Virginia. Their combined expertise in systems, data management, and machine learning research contributed to the development of this innovative data pipeline solution.