Will Zero-Infinity work on Windows?

[lostmsu]

I also can't find any tutorial to use it.

[tjruwase]

@lostmsu, thanks for your question.

Sorry, ZeRO-Infinity does not currently work on Windows.

Are you currently running DeepSpeed on Windows?

[lostmsu]

@tjruwase I am. I also just improved Windows support in OpenAI's Triton, so their next public release might have it.

[tjruwase]

Can you please explain a bit more about your Windows support for DeepSpeed? Also, how does it relate to the next OpenAI's Triton public release? We are very curious indeed. Thanks.

[lostmsu]

@tjruwase not sure what the question is, but here's a summary:

I am personally interested in DeepSpeed and specifically Zero-Infinity running on Windows. I am a deep learning enthusiast interested in research around reducing the hardware requirements for training large models, and I run Windows because it has GPU scheduler, meaning I can still use my PC when something is training at 100% GPU load. In fact I was playing with bypassing random transformer layers during training with offloading them to CPU RAM until I found out about DeepSpeed project.

Now to DeepSpeed and Triton: I am just starting with DeepSpeed, and it seems to have numerous issues when installing on Windows. Triton is listed in the dependencies for sparse attention which I was considering a nice feature and their official builds do not support Windows, so I contributed that back to them. Locally, Triton runs on my Windows machine, so I should be able to get sparse attention on Windows once DeepSpeed installs correctly.

I also noticed cpuadam and cpuadagrad are only supported on Linux, so I have a work-in-progress fix for that.

Once I get that running and upstreamed, I plan to start working on the offload support.

[tjruwase]

Thanks for sharing. This very helpful information. In particular, your goal of reducing hardware costs of training large model is well aligned with ours. I look forward to close collaboration with you.

[lostmsu]

@tjruwase can you point me to the part of code, that is responsible for Infinity offloading? I might be able to port it. Does it depend on libaio?

[tjruwase]

Thanks for your interest. Yes, the nvme offloading uses libaio. It will be exciting if you could help out with this. You can get started with the following references to libaio.h and other linux headers.

1. deepspeed_aio_utils.h
2. deepspeed_aio_types.h
3. deepspeed_aio_common.cpp

Please let me know how I can be of help. Thanks once again for your interest.

[lostmsu]

@tjruwase do you know if other, simpler options were tried before deciding on aio? For example, memory-mapped files. I am tempted to try the standard Python mmap module out. Are there benchmarks for IO performance?

Does offload have C++ interface, or is all high-level stuff (e.g. splitting tensors, deciding which ones to offload) implemented in Python?

[lostmsu]

I just made a simple test script, and it gives promising results for mmap (temp-12G 12GB file must preexist):

import mmap
import random
import time

file = open("temp-12G", "r+b")

map = mmap.mmap(file.fileno(), length=0)

MEGABYTES = 12 * 1024

assert map.size() == MEGABYTES * 1024 * 1024

def run_test_once():
    random_1MB = bytearray(random.getrandbits(8) for _ in range(1024 * 1024))

    start = time.time()

    for mb_index in range(MEGABYTES):
        map[mb_index * 1024 * 1024 : (mb_index + 1) * 1024 * 1024] = random_1MB
    
    map.flush()

    end = time.time()

    return end - start

buf = bytearray(1024 * 1024)

def run_read_test_once():
    start = time.time()

    for mb_index in range(MEGABYTES):
        buf = map[mb_index * 1024 * 1024 : (mb_index + 1) * 1024 * 1024]

    end = time.time()

    return end - start

for i in range(4):
    took = run_test_once()
    print(f"run {i} time: {took}s write speed: {MEGABYTES / took}MB/s")

for i in range(4):
    took = run_read_test_once()
    print(f"run {i} time: {took}s read speed: {MEGABYTES / took}MB/s")

Gives back write 2140MB/s, read 2930MB/s on my WD SSD, which is only about 1/2 of the optimal, but seems to be sufficient for the offloading requirements.

Might be even closer to optimal on Linux, where Python also exposes madvice (I will try to test tomorrow).

[tjruwase]

@tjruwase do you know if other, simpler options were tried before deciding on aio? For example, memory-mapped files. I am tempted to try the standard Python mmap module out. Are there benchmarks for IO performance?

We went with aio because we wanted to sustain >20GB/sec I/O speeds in a DGX-2 box. At that time aio was the best way to achieve ~90% of the peak nvme device speeds in each box. In particular aio enables us to leverage the parallelism (e.g., block size and queue depth) of the nvme device for faster IO. We also leverage asynchrony so that the I/O operations can be overlapped with GPU computation. Fast I/O is critical to prevent offloading becoming the bottleneck of large model training as discussed in the paper.

Does offload have C++ interface, or is all high-level stuff (e.g. splitting tensors, deciding which ones to offload) implemented in Python?
Those high-level stuff are implemented in PyTorch since they are application dependent. The C++ backend implemented here exposes a Python aio_handle object with read/write methods that take a file and a tensor for I/O operation. Providing a C++ interface to aio_handle is trivial.

[tjruwase]

I just made a simple test script, and it gives promising results for mmap (temp-12G 12GB file must preexist):

Gives back write 2140MB/s, read 2930MB/s on my WD SSD, which is only about 1/2 of the optimal, but seems to be sufficient for the offloading requirements.

Might be even closer to optimal on Linux, where Python also exposes madvice (I will try to test tomorrow).

Thanks for sharing this interesting result. If I recall correctly, we also noticed that 50% of peak device speeds was achievable using conventional I/O stack, which was insufficient for us. Thanks to fio we discovered that libaio could get us to the >90% of peak that we wanted by exploiting the nvme stack. Recently, it seems that io_uring is even faster than libaio, and so we need to explore that as well.

I will try out your script in my environment. Also, can you check the speeds achievable using fio and libaio engine?

[tjruwase]

Below is the output of your script on a dgx-2 box. temp-12G file exists on NVMe device with peak GB/sec: read=14 and write=12.4

run 0 time: 11.538774251937866s write speed: 1064.9311384124103MB/s
run 1 time: 12.798313856124878s write speed: 960.1264774515076MB/s
run 2 time: 12.075251340866089s write speed: 1017.6185698440834MB/s
run 3 time: 11.411003351211548s write speed: 1076.855349332216MB/s
run 0 time: 1.5713260173797607s read speed: 7820.146719450783MB/s
run 1 time: 1.664602518081665s read speed: 7381.942455644629MB/s
run 2 time: 1.7217073440551758s read speed: 7137.101460610488MB/s
run 3 time: 1.7196314334869385s read speed: 7145.7172512154675MB/s

As you can see, the reads are 7X faster than writes which suggests to me that the reads are actually hitting the buffer cache rather than the disk. Thus, it seems to me that the writes are more reflective of the sustainable I/O performance ~1GB/sec.
I am really curious as to what you find out with fio on your hardware. Since fio supports libaio, io_uring, and mmap as io_engine options you would be able to compare all three, which could be interesting.

[tjruwase]

You might also be interested in the following discussion.