apache / apache/tvm

[Bug][Relax][Distributed] DeviceMesh accepts invalid shape/device-id cardinality after 32-bit product narrowing

Open
#20,272 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

## Summary

`relax.distributed.DeviceMesh` validates the number of `device_ids` by multiplying the logical mesh shape into a C++ `int`. For a valid `ffi::Shape` whose mathematical product is larger than `INT32_MAX`, the product is narrowed before it is compared with `device_ids.size()`. On the frozen `v0.25.0.post1` snapshot, a mesh with shape `(2**31, 2)` and an empty device-id list is accepted even though the mathematical mesh contains `4,294,967,296` positions.

The reproducer observes a metadata validation failure; it does not allocate the logical mesh or launch any device work.

## Latest upstream source check

As of 2026-09-05, the same `int prod = 1; prod *= shape[i];` pattern remains in both the upstream `main` branch and release `v0.26.0`:

- [upstream main](https://github.com/apache/tvm/blob/main/src/relax/distributed/global_info.cc#L29-L55)
- [upstream v0.26.0](https://github.com/apache/tvm/blob/v0.26.0/src/relax/distributed/global_info.cc#L29-L55)

This is a source-level version check. The runtime output above was obtained from `v0.25.0.post1`; this report does not claim a v0.26.0 binary replay.

## Environment

- TVM commit: `b3e249b7d75f8f3bc7cbee48188d3c80ae323437` (`v0.25.0.post1`)
- Python: `3.11`
- Platform: Ubuntu 22.04 under WSL2, x86_64

## Affected code

[`src/relax/distributed/global_info.cc`](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/relax/distributed/global_info.cc#L29-L39)
contains:

```cpp
int prod = 1;
for (int i = 0; i < static_cast(shape.size()); i++) {
prod *= shape[i];
}
TVM_FFI_ICHECK_EQ(prod, static_cast(device_ids.size()))
<< "The number of device ids must match the product of the shape";
```

The `Range` overload repeats the same `int` accumulator at [lines 50–55](https://github.com/apache/tvm/blob/b3e249b7d75f8f3bc7cbee48188d3c80ae323437/src/relax/distributed/global_info.cc#L50-L55).

## Minimal reproduction

Run the following standalone Python program in the frozen TVM environment. It does not require any project-local file:

```python
import math

from tvm.relax.distributed import DeviceMesh

def accepted(shape, device_ids):
try:
DeviceMesh(shape, device_ids)
except Exception as exc:
print(f"shape={shape}, ids={len(device_ids)}: REJECTED ({type(exc).__name__})")
return False
print(f"shape={shape}, ids={len(device_ids)}: ACCEPTED")
return True

control_ok = accepted((2, 2), [0, 1, 2, 3])
one_id_ok = accepted((2**31, 2), [0])
overflow_ok = accepted((2**31, 2), [])
print(f"mathematical overflow mesh size: {math.prod((2**31, 2))}")

if control_ok and not one_id_ok and overflow_ok:
print("BUG REPRODUCED: invalid overflowed DeviceMesh cardinality was accepted")
else:
print("BUG NOT REPRODUCED: implementation rejected the invalid mesh")
```

Expected behavior:

```text
control shape=(2, 2), ids=4: ACCEPTED
shape=(2147483648, 2), ids=1: REJECTED
shape=(2147483648, 2), ids=0: REJECTED
```

Observed on the frozen snapshot:

```text
control shape=(2, 2), ids=4: ACCEPTED
shape=(2147483648, 2), ids=1: REJECTED
shape=(2147483648, 2), ids=0: ACCEPTED <-- invalid cardinality accepted
mathematical overflow mesh size: 4294967296
BUG REPRODUCED
```

The one-ID negative control remains rejected while the zero-ID case is accepted. Together, these controls show that the check observes the narrowed product rather than the mathematical product.

## Impact

The constructor can accept a Relax metadata state in which the number of device IDs does not match the mathematical product of the logical mesh shape. Such metadata may then be consumed by later distributed transformations without this constructor reporting the mismatch.

## Suggested fix

- Perform checked multiplication in a type wide enough for the supported shape and device-id cardinality, without narrowing before validation.
- Reject negative dimensions and report overflow as a diagnostic rather than accepting a wrapped product.
- Apply the same checked logic to both the explicit `device_ids` and `Range` overloads.
- Add regression tests for products just below `INT32_MAX`, products that narrow to zero, and a normal `(2, 2)` control mesh.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in src/relax/distributed/global_info.cc and inspect the explicit device_ids and Range overloads that compute the mesh cardinality. Run the standalone Python reproducer against the frozen snapshot, then add regression coverage for the normal control, near-limit product, narrowing-to-zero product, negative dimensions, and overflow cases. Done means invalid cardinalities are rejected without narrowing and valid meshes remain accepted.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, distributed-systems
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.