pytorch - 💡(How to fix) Fix [DTensor] Hinted even unbacked shard sizes still take symbolic padding path

Summary

DTensor placement logic conservatively takes the padding path for unbacked shard dimensions even when an explicit optimization hint says the dimension is evenly divisible by the mesh size.

For arbitrary unbacked SymInts that is the safe default, because guard_or_false(size % num_chunks == 0) cannot prove divisibility. However, when the frontend has explicitly provided an optimization hint for an unbacked shard dimension and the hint is divisible by the mesh size, DTensor still traces symbolic padding/unpadding. That path constructs symbolic pad sizes, and those eventually flow into padding helpers that require concrete Python padding values.

Reproducer

import torch
from torch._subclasses import FakeTensorMode
from torch.distributed.tensor.placement_types import Shard
from torch.fx.experimental.symbolic_shapes import ShapeEnv


fake_mode = FakeTensorMode(allow_non_fake_inputs=True, shape_env=ShapeEnv())
shard = Shard(0)

with fake_mode:
    dim_size = fake_mode.shape_env.create_unbacked_symint()
    torch._dynamo.override_optimization_hint(dim_size, 8)

    local_tensor = torch.empty(dim_size // 4, 4)
    full_tensor = torch.empty(dim_size, 4)

    padded = shard._maybe_pad_tensor(local_tensor, dim_size, 4)
    unpadded = shard._maybe_unpad_tensor(full_tensor, dim_size, 4)

print("padded shape:", padded.shape)
print("unpadded shape:", unpadded.shape)

Expected behavior

DTensor should be able to trace the no-padding policy for explicitly hinted, evenly sharded unbacked sizes, while preserving runtime correctness with a symbolic precondition such as:

torch._check(dim_size % 4 == 0)

Runtime sizes that satisfy the divisibility condition should continue to work, and sizes that do not should fail closed.

Actual behavior

The unbacked divisibility predicate is not guardable at trace time, so DTensor takes the conservative padding path. Even if the runtime size would need zero padding, the traced path still computes symbolic pad sizes from the unbacked SymInt and can fail when those symbolic values reach padding/unpadding helpers that need concrete Python padding values.

Notes

The optimization hint should not specialize tensor shapes or replace the symbolic dimension. It is only useful as a trace-time policy selector for the no-padding branch; the generated graph still needs to carry the real symbolic divisibility check.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx