April 9, 2025
hadia-ahmed
Training deep learning models efficiently isn’t just about powerful GPUs—it starts with how you move and prepare your data. In this post, we’ll explore how to create a flexible system capable of scaling across heterogeneous CPU and GPU environments. Using Bodo and its autoparallelizing inferential compiler, we’ll preprocess data on CPUs at scale and send it in batches to GPUs for distributed training in PyTorch.
Bodo is designed to enable scalable and efficient parallel computing using Python. By leveraging Bodo’s capabilities to preprocess and batch distributed data on CPUs before passing it to GPUs, we can maximize computational efficiency and training speed. This approach is especially beneficial for handling large datasets that don’t fit entirely on a single GPU.
NOTE: Bodo will install MPI as part of its dependencies.
PyTorch has to be installed from source to enable using MPI as the backend
Our architecture leverages three powerful technologies in combination:
Bodo is a parallel computing framework that enables significant speedups for Python data science workloads. Key benefits include:
PyTorch provides the deep learning foundation with:
torch.distributedMPI serves as the communication backbone, offering:
We begin by initializing the process group using PyTorch’s distributed package. The script uses MPI as the backend for process group initialization, which works seamlessly with Bodo’s parallel capabilities.
Next, we determine the GPU ranks for each process to ensure that each GPU gets assigned the appropriate data batches. The function get_gpu_ranks uses MPI to collect GPU availability information and distribute GPU resources efficiently.
Efficient data loading and preprocessing are crucial for large datasets. Bodo’s JIT compilation significantly speeds up this step. In our example, we load the NYC taxi dataset, preprocess it using Scikit-learn’s StandardScaler, and transform the features:
A key aspect of this setup is sending data from CPUs to GPUs in batches to optimize memory usage. The batch_generator function uses Bodo’s rebalance to efficiently transfer batches to designated GPU ranks.
One of the key advantages of using the yield keyword in the batch_generator function is that it allows for overlapping batch preparation with model training. Since yield makes the generator function produce one batch at a time, the model can start training on the current batch while the next batch is being prepared. This significantly reduces idle time and increases the efficiency of the data pipeline, especially in distributed settings where data loading can become a bottleneck.
The model architecture consists of a simple feedforward neural network. We distribute the model across available GPUs using DistributedDataParallel, allowing for synchronized training across multiple devices:
The TaxiDemandNN model is a simple feedforward neural network that aims to train a neural network model to predict the trip_miles (how far a taxi will travel on future trips) based on features like pickup/dropoff locations, trip time, and fare details using NYC taxi data.
The training loop efficiently batches data using the generator function and processes it on the assigned GPU. By using Bodo to preprocess and rebalance the data across multiple GPUs, we reduce data transfer bottlenecks and increase training throughput.
By combining Bodo’s high-performance CPU-based preprocessing with PyTorch’s flexible GPU-based training, we’ve demonstrated a scalable and efficient architecture for large-scale machine learning. No more bottlenecks, memory meltdowns, or waiting around for data to catch up. This setup is flexible, scalable, and surprisingly elegant. Whether you’re scaling up an existing project or building something brand new, this approach gives you the flexibility and efficiency to move quickly and build smarter.
Want to try it yourself?
Join our Slack community to connect with other users, share use cases, and stay up-to-date on the latest from Bodo.
