LittlehamsterXu opened a new issue, #20272: URL: https://github.com/apache/tvm/issues/20272
## 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<int>(shape.size()); i++) { prod *= shape[i]; } TVM_FFI_ICHECK_EQ(prod, static_cast<int>(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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
