Implementation:Microsoft DeepSpeedExamples BingBert AsyncWorker

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Data Loading, Concurrency
Last Updated	2026-02-07 12:00 GMT

Overview

A thread-based asynchronous data worker that prefetches training batches from multiple dataloaders in a background thread using producer-consumer queues.

Description

AsyncWorker extends threading.Thread to provide background data prefetching for BERT pretraining. It uses a producer-consumer pattern with two queue.Queue instances: a req_queue for incoming dataset type requests and a ret_queue for completed batch results. On initialization, the worker pre-fills the request queue with the first three entries from the dataset picker to ensure batches are ready before the training loop begins.

The run method continuously dequeues dataset type indices from the request queue, fetches the next batch from the corresponding dataloader, and places the result onto the return queue. The loop terminates gracefully when a None sentinel value is received. The prefetch method incrementally adds the next dataset type to the request queue, maintaining a rolling prefetch window.

This design decouples data loading from the training loop, allowing GPU computation and data preparation to overlap. The worker is designed to be used with BingBertDatasetProvider, which manages the lifecycle of starting and stopping the async worker.

Usage

Use this worker class when asynchronous data loading is enabled in the Bing BERT training configuration. It is instantiated and managed by BingBertDatasetProvider and should not typically be used independently. Enable it by including an 'async_worker' key in the training configuration.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/bing_bert/data_worker.py
• Lines: 1-37

Signature

class AsyncWorker(threading.Thread):
    def __init__(self, dataloaders, dataset_picker):
    def run(self):
    def get(self):
    def prefetch(self):
    def stop(self):

Import

from data_worker import AsyncWorker

I/O Contract

Inputs

Name	Type	Required	Description
dataloaders	dict[int, generator]	Yes	Dictionary mapping dataset type indices to dataloader generators
dataset_picker	list[int]	Yes	Ordered list of dataset type indices specifying the batch fetch sequence

Outputs

Name	Type	Description
batch	tuple	A single training batch retrieved from the return queue via get()

Usage Examples

from data_worker import AsyncWorker

# dataloaders: dict mapping dataset indices to generator-based dataloaders
# dataset_iterator: list of dataset type indices for batch ordering
worker = AsyncWorker(dataloaders, dataset_iterator)
worker.start()

# Fetch prefetched batches during training
for step in range(num_steps):
    batch = worker.get()
    # ... process batch ...
    worker.prefetch()

# Stop the worker thread
worker.stop()