{"id":26148825,"url":"https://github.com/hpcaitech/tensornvme","last_synced_at":"2026-03-27T02:10:02.096Z","repository":{"id":47731429,"uuid":"509378185","full_name":"hpcaitech/TensorNVMe","owner":"hpcaitech","description":"A Python library transfers PyTorch tensors between CPU and NVMe","archived":false,"fork":false,"pushed_at":"2024-11-27T10:19:54.000Z","size":305,"stargazers_count":115,"open_issues_count":6,"forks_count":25,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-05-24T04:07:29.174Z","etag":null,"topics":["colossal-ai","deep-learning","nvme","pytorch"],"latest_commit_sha":null,"homepage":"","language":"C++","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/hpcaitech.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-07-01T08:28:52.000Z","updated_at":"2025-05-19T05:47:28.000Z","dependencies_parsed_at":"2024-06-19T11:18:10.870Z","dependency_job_id":"fe8f502e-df15-4205-a756-ac71d226c445","html_url":"https://github.com/hpcaitech/TensorNVMe","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/hpcaitech/TensorNVMe","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hpcaitech%2FTensorNVMe","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hpcaitech%2FTensorNVMe/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hpcaitech%2FTensorNVMe/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hpcaitech%2FTensorNVMe/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/hpcaitech","download_url":"https://codeload.github.com/hpcaitech/TensorNVMe/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/hpcaitech%2FTensorNVMe/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":268338285,"owners_count":24234540,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-08-02T02:00:12.353Z","response_time":74,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["colossal-ai","deep-learning","nvme","pytorch"],"created_at":"2025-03-11T05:21:52.466Z","updated_at":"2026-03-27T02:09:57.048Z","avatar_url":"https://github.com/hpcaitech.png","language":"C++","funding_links":[],"categories":[],"sub_categories":[],"readme":"# TensorNVME\n\nA Python Library provides APIs to move PyTorch Tensors between CPU and NVMe.\n\n## Dependencies\n\n- [liburing](https://github.com/axboe/liburing)\n- [libaio](https://pagure.io/libaio)\n\n## Install\n\nThis package is only supported on Linux. `liburing` and `libaio` can be automatically installed. `liburing` is supported on Linux \u003e= `5.10`, and it won't be installed if the version of your Linux \u003c `5.10`.\n\nIt will search `libaio` and `liburing` in `/usr/lib`, `/usr/lib64` and `$LD_LIBRARY_PATH`. If not found, backends will be installed in `~/.tensornvme`, and `~/.bashrc` will be modified to set `$LD_LIBRARY_PATH` correctly. **Please `source ~/.bashrc` after installation.** If you use other shells, please make sure `$LD_LIBRARY_PATH` is set correctly.\n\n\u003e You must install pytorch and cmake before installing tensornvme. Once you upgrade pytorch, remember to reinstall tensornvme.\n\n### From source\n\n```shell\ngit clone https://github.com/hpcaitech/TensorNVMe.git \u0026\u0026 cd TensorNVMe\n```\n\nFirst, install requirements:\n\n```shell\npip install -r requirements.txt\n```\n\nTo install `tensornvme` with `liburing` and `libaio`:\n\n```shell\npip install -v --no-cache-dir .\n```\n\nTo install `tensornvme` with only `liburing`:\n\n```shell\nDISABLE_AIO=1 pip install -v --no-cache-dir .\n```\n\nTo install `tensornvme` with only `libaio`:\n\n```shell\nDISABLE_URING=1 pip install -v --no-cache-dir .\n```\n\nIf you want to install `libaio` or `liburing` for system:\n\n```shell\nWITH_ROOT=1 sudo pip install -v --no-cache-dir .\n```\n\nThen they will be installed in `/usr` and `~/.bashrc` will not be modified. Make sure you have root access.\n\n### From PIP\n\n```shell\npip install packaging\npip install tensornvme\n```\n\nAll acceptable environment variables are the same as those when installing from source.\n\n## Use docker\n\n```shell\ngit clone https://github.com/hpcaitech/TensorNVMe.git \u0026\u0026 cd TensorNVMe/docker \u0026\u0026 docker build -t tensornvme .\n```\n\n## CLI\n\nWe provide a CLI to test whether backends work well.\n\n```shell\ntensornvme check\n```\n\n## Usage\n\nIt provide both synchronize and asynchronize I/O API.\n\n\u003e Only CPU and contiguous tensors can be offloaded.\n\nSynchronize API:\n\n```python\nimport torch\nfrom tensornvme import DiskOffloader\n\nx = torch.rand(2, 2)\ny = torch.rand(4, 4, 4)\noffloader = DiskOffloader('./offload')\noffloader.sync_write(x)\n# x is saved to a file on disk (in ./offload folder) and the memory of x is freed\noffloader.sync_read(x)\n# x is restored\noffloader.sync_writev([x, y])\n# x and y are offloaded\noffloader.sync_readv([x, y])\n# x and y are restored.\n# sync_writev() and sync_readv() are order sensitive\n# E.g. sync_writev([x, y]) and sync_writev([y, x]) are different\n```\n\nAsynchronize API:\n\n```python\nimport torch\nfrom tensornvme import DiskOffloader\n\nx = torch.rand(2, 2)\ny = torch.rand(4, 4, 4)\noffloader = DiskOffloader('./offload')\noffloader.async_write(x)\n# x is being offloaded in the background\noffloader.sync_write_events()\n# x is offloaded and the memory of x is freed\noffloader.async_read(x)\n# x is being restored in the background\noffloader.sync_read_events()\n# x is restored\noffloader.async_writev([x, y])\n# x and y are being offloaded in the background\noffloader.synchronize()\n# synchronize() will synchronize both write and read events.\noffloader.async_readv([x, y])\noffloader.synchronize()\n# x and y are restored.\n# async_writev() and async_readv() are also order sensitive\n```\n\nYou can use asynchronize API to overlap computation and data moving.\n\n```python\ntensors = []\n\nfor _ in range(10):\n    tensor = torch.rand(2, 2)\n    tensors.append(tensor)\n    offloader.sync_write(tensor)\n\noffloader.sync_read(tensors[0])\n\n# prefetch=1, writing tensor[i] and reading tensor[i+1]\nfor i, tensor in enumerate(tensors):\n    offloader.sync_read_events()\n    if i + 1 \u003c len(tensors):\n        offloader.async_read(tensors[i+1])\n    tensor.mul_(2.0) # compute\n    offloader.sync_write_events()\n    offloader.async_write(tensor)\noffloader.synchronize()\n```\n\n## How to test\n\nWe have C++ test scrpits for `AsyncIO` and `SpaceManager` class. Make sure you have installed `liburing` and `libaio`, and set environment variables correctly before testing. To run the tests:\n\n```shell\nmkdir build\ncd build\ncmake ..\nmake\n./test_asyncio\n./test_space_mgr\n```\n\nWe also have python unit tests. Make sure you have installed `pytest`. To run:\n\n```shell\npytest ./tests\n```\n\n## How to benchmark\n\nWe have benchmarks for `Adam` and `CpuAdam`  with different backend and prefetch depth to validate TensorNVME's speed. To run the benchmark:\n\n```shell\ncd benchmark\npython benchmark_adam.py\npython benchmark_cpuadam.py\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhpcaitech%2Ftensornvme","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fhpcaitech%2Ftensornvme","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhpcaitech%2Ftensornvme/lists"}