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

[LittlehamsterXu]

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
• upstream v0.26.0

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
contains:

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.

Minimal reproduction

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

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:

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

Observed on the frozen snapshot:

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.