This PR enables ruff pyupgrade (UP) rules with py310 target, auto-fixing
~5600 annotation modernizations (PEP 585 generics, PEP 604 unions,
deprecated typing imports).
Also removes from __future__ import annotations from ir/module.py and
rmsnorm.py, bumps requires-python to >=3.10, and removes absolute_import
aliases from topi/contrib files.
This PR fixes a few places where the python import of runtime tensor
is incorrect. The error wasn't revealed in the previous
NDArray->Tensor rename PR since these imports are not at the top
level.
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
This PR Updates the NDArray => Tensor.
Both tensor and ndarray are commonly used terms.
Because the term Tensor is getting more common in the context of ML,
we do the rename to stay more aligned with torch.Tensor and DLTensor.
* [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 hides StringObj/BytesObj into details and bring
implementations to directly focus on the String/Bytes.
This change will prepare us for future changes such as SmallStr support.
Also moves more ObjectRef into Any in RPC.
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.
In addition to the PackedFunc `"runtime.disco.worker_id"`, which
returns the worker ID wrapped in a `ShapeTuple`, this commit adds
`"runtime.disco.worker_rank"`, which returns the worker ID without
wrapping, and `"runtime.disco.device"`, which returns the device for
each worker.
The unit test added in this commit simulates loading of model weights
through a parameter transformation function.
Disco worker originally automatically import `tvm.testing.disco` for
convenient unittesting. However, `tvm.testing` is a special subpackage
that introduces many unnecessary dependencies, for example, pytest. This
PR removes such dependencies by directly moving the testing function
registration logic to the entry file.
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