Uni-Core Distributed PyTorch Framework

Uni-Core is an efficient distributed PyTorch framework optimized for high-performance Transformer-based models. Key features include:

• Distributed Training: Support for multi-GPU and multi-node setups.
• Mixed-Precision: Support for fp16 and bf16 training.
• Performance Optimizations: High-performance fused CUDA kernels and a buffered data loader (GPU-CPU overlapping).
• Training Utilities: Model checkpoint management, gradient accumulation, friendly logging, and common optimizers/LR schedulers.
• Model Development: Tools to easily create new models.

You can build Uni-Core from source using pip or setup.py.

Important Requirements:

• The CUDA version in your build environment must match the CUDA version used by your PyTorch installation.

If CUDA is not available in your environment, you can build the package without the CUDA extension operator by using the --disable-cuda-ext flag.

# Standard installation
pip install .
# OR
python setup.py install

# Installation without CUDA extensions
python setup.py install --disable-cuda-ext

Pre-compiled Python wheels are available on the GitHub Releases page.

When selecting a wheel, ensure there is a match for your:

1. Python version
2. PyTorch version
3. CUDA version

Example: For PyTorch 1.12.1, Python 3.7, and CUDA 11.3, use a wheel named like unicore-0.0.1+cu113torch1.12.1-cp37-cp37m-linux_x86_64.whl.

A Docker image is provided for a pre-configured environment.

To use GPUs within the container, you must have nvidia-docker-2 installed on your host machine.

You can pull the image using the following command (example for version 0.0.1 with PyTorch 1.11.0 and CUDA 11.3):

docker pull dptechnology/unicore:0.0.1-pytorch1.11.0-cuda11.3

To run the BERT example, you must download the Wikitext-2 dataset, unzip it, and then preprocess the token files into LMDB format using the provided preprocess.py script. Follow these steps in order:

1. Download and unzip data: Use wget to fetch the dataset and unzip to extract it.
2. Preprocess training data: Run preprocess.py with the training tokens file as the source and a target .lmdb filename.
3. Preprocess validation data: Run preprocess.py with the validation tokens file as the source and a target .lmdb filename.

# 1. Download and unzip data
wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip && unzip wikitext-2-v1.zip

# 2. Preprocess training data
python preprocess.py ./wikitext-2/wiki.train.tokens ./train.lmdb

# 3. Preprocess validation data
python preprocess.py ./wikitext-2/wiki.valid.tokens ./valid.lmdb