Add tvm.testing.run_with_gpu_lock backed by the existing
tvm_ffi.utils.FileLock. Migrate live local GPU tests to acquire the
machine-local lock around device execution, synchronization, host
transfer, and checks while leaving target construction and compilation
outside the critical section.
Replace the custom xdist scheduler with standard xdist_group placement
for the order-dependent test family. RPC tests retain dynamic port
allocation and per-test process isolation rather than gaining a broad
category lock.
This PR introduces Relax AnyType as the primary top/base type spelling,
replacing the previous ObjectType naming for the type that represents
any Relax value.
Changes:
- Add AnyType/AnyTypeNode with relax.AnyType registration and keep
ObjectType/R.Object compatibility aliases.
- Update Relax type analysis, type visitors, opaque function defaults,
and script printer/parser handling to use AnyType/R.Any.
- Migrate affected Python/C++ call sites, docs, and focused tests to the
new spelling.
Validation:
- cmake --build build --parallel 16
- Focused Relax/TVMScript pytest: 704 passed, 1 xfailed
- pre_commit run --files <changed files>
## Summary
- unify Relax's former StructInfo surface into the Type vocabulary and
Expr.ty storage path
- remove leftover DependentTypeNode and legacy OpNode::op_type storage
- keep base Type nullable while concrete Relax/DTensor type refs are
non-nullable
- clean stale StructInfo/TensorStructInfo/sinfo vocabulary in
Python/docs and distributed-op macros
- address Gemini follow-ups for parser annotations, BlockBuilder
docstring, and Adreno TensorType cast audit
This pr modernizes test gating. It replaces the heavy
`tvm.testing.Feature` machinery with a thin `tvm.testing.env` module of
`has_*()` capability probes, used via standard pytest.mark + skipif. And
markers move to `pyproject.toml`
The tvm_ffi Object metaclass now gives every subclass `__slots__ = ()`,
so the Disco Python wrappers can no longer store instance attributes and
every session construction fails with AttributeError. Declare the
attributes each
wrapper actually stores as named slots, fix the NVSHMEM `dist_gemm.cu`
so TVM builds with `USE_NVSHMEM = ON`, and gate the disco tests on the
disco runtime being present so they skip cleanly on builds (e.g. the pip
wheel) that report `USE_NCCL` / `USE_NVSHMEM = ON` without shipping it.
### Session attribute storage
- `DPackedFunc` / `DModule`: `__slots__ = ("session",)`.
- `Session`: `__slots__ = ("_cache", "_import_python_module")`
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 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.
* [Disco] Add loader for presharded params.
Prior to this commit, sharding of model weights was always performed
when initializing the model. This could cause slow initialization,
especially for larger numbers of GPUs, as all model weights are
initially transferred to GPU-0, before being scattered to all workers.
This commit updates the `tvm::runtime::ShardLoaderObj` to also allow
loading of pre-sharded model weights. With pre-sharded model weights,
the tensors are sharded while the model is being built, and each
worker independently loads the specific model weights that it
requires.
* Update based on review comments.
* Removed commented-out print statements
---------
Co-authored-by: Chris Sullivan <csullivan@octoml.ai>
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 `ShardLoader`, an object that allows convenient
sharding of each parameter, assuming there is a single shard dimension
and the sharding strategy is even. The sharding can be performed
efficiently on device (e.g. CUDA) and scattered to each worker with
NCCL support.
The shard loading process could be further compiled to Relax IRModule
to be executed by each worker with the function signature below:
```python
@R.function
def main(loader: R.Object) -> R.Tuple(
R.Tensor((64, 64), "float32"),
R.Tensor((16, 128), "float32"),
):
R.func_attr({"global_symbol": "main"})
with R.dataflow():
lv0 = R.call_pure_packed(
"runtime.disco.ShardLoaderLoad",
loader,
R.shape([0]),
sinfo_args=R.Tensor((64, 64), "float32"),
)
lv1 = R.call_pure_packed(
"runtime.disco.ShardLoaderLoad",
loader,
R.shape([1]),
sinfo_args=R.Tensor((16, 128), "float32"),
)
lv2 = R.tuple(lv0, lv1)
R.output(lv2)
return lv2
```
Tianqi Chen