Package: tensornvme
TensorNVMEA Python Library provides APIs to move PyTorch Tensors between CPU and NVMe.DependenciesliburinglibaioInstallThis package is only supported on Linux.liburingandlibaiocan be automatically installed.liburingis supported on Linux >=5.10, and it won't be installed if the version of your Linux <5.10.It will searchlibaioandliburingin/usr/lib,/usr/lib64and$LD_LIBRARY_PATH. If not found, backends will be installed in~/.tensornvme, and~/.bashrcwill be modified to set$LD_LIBRARY_PATHcorrectly.Pleasesource ~/.bashrcafter installation.If you use other shells, please make sure$LD_LIBRARY_PATHis set correctly.You must install pytorch and cmake before installing tensornvme. Once you upgrade pytorch, remember to reinstall tensornvme.From sourcegitclonehttps://github.com/hpcaitech/TensorNVMe.git&&cdTensorNVMeFirst, install requirements:pipinstall-rrequirements.txtTo installtensornvmewithliburingandlibaio:pipinstall-v--no-cache-dir.To installtensornvmewith onlyliburing:DISABLE_AIO=1pipinstall-v--no-cache-dir.To installtensornvmewith onlylibaio:DISABLE_URING=1pipinstall-v--no-cache-dir.If you want to installlibaioorliburingfor system:WITH_ROOT=1sudopipinstall-v--no-cache-dir.Then they will be installed in/usrand~/.bashrcwill not be modified. Make sure you have root access.From PIPpipinstallpackaging
pipinstalltensornvmeAll acceptable environment variables are the same as those when installing from source.Use dockergitclonehttps://github.com/hpcaitech/TensorNVMe.git&&cdTensorNVMe/docker&&dockerbuild-ttensornvme.CLIWe provide a CLI to test whether backends work well.tensornvmecheckUsageIt provide both synchronize and asynchronize I/O API.Only CPU and contiguous tensors can be offloaded.Synchronize API:importtorchfromtensornvmeimportDiskOffloaderx=torch.rand(2,2)y=torch.rand(4,4,4)offloader=DiskOffloader('./offload')offloader.sync_write(x)# x is saved to a file on disk (in ./offload folder) and the memory of x is freedoffloader.sync_read(x)# x is restoredoffloader.sync_writev([x,y])# x and y are offloadedoffloader.sync_readv([x,y])# x and y are restored.# sync_writev() and sync_readv() are order sensitive# E.g. sync_writev([x, y]) and sync_writev([y, x]) are differentAsynchronize API:importtorchfromtensornvmeimportDiskOffloaderx=torch.rand(2,2)y=torch.rand(4,4,4)offloader=DiskOffloader('./offload')offloader.async_write(x)# x is being offloaded in the backgroundoffloader.sync_write_events()# x is offloaded and the memory of x is freedoffloader.async_read(x)# x is being restored in the backgroundoffloader.sync_read_events()# x is restoredoffloader.async_writev([x,y])# x and y are being offloaded in the backgroundoffloader.synchronize()# synchronize() will synchronize both write and read events.offloader.async_readv([x,y])offloader.synchronize()# x and y are restored.# async_writev() and async_readv() are also order sensitiveYou can use asynchronize API to overlap computation and data moving.tensors=[]for_inrange(10):tensor=torch.rand(2,2)tensors.append(tensor)offloader.sync_write(tensor)offloader.sync_read(tensors[0])fori,tensorinenumerate(tensors):offloader.sync_read_events()ifi+1<len(tensors):offloader.async_read(tensors[i+1])tensor.mul_(2.0)# compute with tensoroffloader.sync_write_events()offloader.async_write(tensor)offloader.synchronize()How to testWe have C++ test scrpits forAsyncIOandSpaceManagerclass. Make sure you have installedliburingandlibaio, and set environment variables correctly before testing. To run the tests:mkdirbuildcdbuild
cmake..
make
./test_asyncio
./test_space_mgrWe also have python unit tests. Make sure you have installedpytest. To run:pytest./tests