A request to retrieve data from a memory device of a memory sub-system can be received from a machine learning (ML) framework executing on a host system, where the data comprises a plurality of logical partitions. A set of parallel I/O threads can be initiated to retrieve the data from the memory device, where each I/O thread of the set of parallel I/O threads retrieves a different portion of the data from a different corresponding logical partition and stores the different portion of the data in a I/O buffer of a set of I/O buffers corresponding to the set of I/O threads in parallel. The different portion of the data can be successively provided from each I/O buffer to the ML framework, where the set of parallel I/O threads is to continually retrieve the data from the memory device until all of the data from the logical partitions has been provided to the ML framework.