This PR cleans up the python API to make things more consistent
with existing python array api and torch.
Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device
API updates:
- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func
* [FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This PR establishes tvm_ffi as a standalone python module.
The ffi is structured as a minimal pip module that can be
directly install by path or url.
examples/get_started provided a minimal example.
This is a major change as we are decoupling tvm_ffi as a
separate package, users need to install tvm_ffi separately.
Thanks to its minimal dependency, tvm_ffi can be easily installed
even just from the source by pip install ./ffi
This change would enable future improvement for library plugins
to have lightweight dependencies by just working on top of
the tvm_ffi, while the main compiler toolchain and runtime
can be layered on top.
* [FFI] Improve traceback setups
This PR improves traceback related setups
This PR phases out tvm._ffi redirections in favor of new FFI
new functions are now called via tvm.ffi.
We also enabled limited API support for python 3.12+
so the compiled binary can be forward compatible to future
python versions.
Prior to this PR, module "psutil" is imported at the top level
of the disco process pool. The pool will try to kill all the processes
at the time of destruction (when `__del__` is implicitly invoked).
The `__del__` function eventually calls into a function that
uses `pstuil`. But it is possible that the top-level `psutil`
has already been released by Python, which leads to a KeyError
as follows:
```
Exception ignored in: <function DiscoPopenWorker.__del__ at 0x7f2c922bfe20>
Traceback (most recent call last):
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 67, in __del__
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 81, in kill
File "/home/ruihangl/Workspace/tvm/python/tvm/runtime/disco/process_pool.py", line 162, in _kill_child_processes
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 323, in __init__
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/__init__.py", line 353, in _init
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_pslinux.py", line 1738, in __init__
File "/home/ruihangl/Workspace/miniconda3/envs/python311/lib/python3.11/site-packages/psutil/_common.py", line 864, in get_procfs_path
KeyError: 'psutil'
```
This PR fixes the issue by lazily importing `psutil` when needed.
Removed instances of accidentally repeated words from comments. There
are cases where duplicated words appear legitimately, those cases remain
unmodified.
In our previous implementation, parameter sharding relies on pre-quantization weight processing,
meaning each set of quantized weights corresponds strictly to a hardcoded constant `num_shards`,
and re-quantization is strictly required upon each change of #GPUs, e.g. from 4-GPU to 8-GPU
setting. This PR makes it possible to move parameter sharding to post-quantization loading-time.
During loading, we iterate over all parameters and apply the sharding operation based on the
provided sharding information.
To make this happen, this PR makes an enhancement to the existing `shard_info.json` to include the
sharding function being used at loading time. Each parameter is attached to a list of loading-time
preprocessing methods that are serially applied to it to transform this parameter to the desired
shape, as shown in the example below:
```python
shard_info = {
"x_0": [ # name of the parameter
[ # a list of preprocessing functions to be applied
"tests.disco.shard_dim_1", # name of the sharding function
[(num_shards, 64, 64), "float16"], # output shape/dtype of `tests.disco.shard_dim_1`
num_shards, # extra inputs to `tests.disco.shard_dim_1`
],
],
"x_1": [...],
}
```
To parameter `x_0`, it means we will call method `tests.disco.shard_dim_1` which has the signature:
```python
def shard_dim_1(
input: NDArray,
num_shards, # extra inputs
output: NDArray, # and its shape is (num_shards, 64, 64), and dtype is "float16"
) -> None: ...
```
This approach simplifies parameter sharding for users and ensures correctness.
This PR introduces `ProcessSession`, a new session implementation based
on multi-processing.
`ProcessSession` shares exactly the same communication protocol with
`ThreadedSession`, but all workers except for worker 0 are launched in a
separate process than thread. Workers communicate with the controller
via pipe provided by the OS, rather than SPSC message queue between
threads.
In our implementation, Python's `subproces.popen` is used to create
subprocesses, and the Python executable, or more specifically,
`sys.executable` calls into `tvm.exec.disco_worker` as the entrypoint.
Besides the launching logic that is only executed once in the very
beginning, the rest of the implementation resides in a C++-only
environment, including reads/writes to pipe file descriptors,
serialization and deserialization of messages, worker interpretation of
each message, etc.
Detailed engineering elements included in this PR:
- Refactors the MinRPC-based communication protocol out to be shared by
`ProcessSession` and `ThreadedSession` as `protocol.h`;
- Refactors a controller-side worker thread into `DiscoWorkerThread`,
which is shared by both session implementation to launch worker-0;
- Added two instructions `kDebugGetFromRemote` and `kDebugSetRegister`,
which are used to communicate with workers other than worker-0 in
debug mode;
- Introduces multi-processing infra including: `tvm.exec.disco_worker`
serving as the entrypoint that launches workers, and
`tvm/runtime/disco/process_pool.py` that exposes APIs to launch worker
processes. `tvm.exec.disco_worker` calls into a global function
`runtime.disco.WorkerProcess` that executes the worker main loop in
pure C++;
- Introduces `src/support/process_id.h` that provides cross-platform pid
and tid printing utilities;
- Refactors Disco's NCCL integration that get rids of initialized-once
global NCCL context, and switches to broadcasting `ncclUniqueId` from
controller to all workers, and then create NCCL communicators in each
worker thread/process accordingly. This is a thread/process-agnostic
way of using NCCL.
Tianqi Chen