[DO NOT REVIEW] Handle unevenly sharded expert weights in FSDP compute weights

[anijain2305]

An expert count that does not divide the FSDP degree leaves the last rank
with a short shard. FSDP's contract is that fsdp_pre_all_gather returns the
padded sharded size and carries the logical size in metadata, so that
fsdp_post_all_gather can drop the padding. The compute-weight hook did
neither: it rejected the layout outright with a message about dimension 0
divisibility.

Pad the communication tensor to ceil(outer_size[0] / mesh.size()) and return
outer_size as metadata; narrow back to it before quantizing. Without the
narrow, padding rows would occupy real 32x32 scale tiles and, for a grouped
weight, appear to the grouped GEMM as extra experts.

Reject Shard(1) explicitly rather than by implication. TorchTitan selects it
for grouped experts when the FSDP degree exceeds the expert count, which any
job with more ranks than experts hits -- gpt_oss has 8. The all-gather then
concatenates dim-1 shards along dim 0, so the hook would have to rebuild the
logical weight before quantizing; that is left for later, and the error now
names the condition and the remedies. The shard dimension is recovered by
comparing the local shard against the logical size, since FSDP does not pass
it.

Add coverage for both: uneven dim 0 asserts finite gradients of the right
shape rather than mere completion, and Shard(1) is a strict xfail so the
marker has to be removed when support lands.

Note the expert count must stay a power of two independently of any of this:
TorchAO's scale-rearrange kernel does tl.arange over the token groups, so
E=3, 5, 6 and 12 fail to compile the kernel with or without FSDP.

---

_{Stack created with GitHub Stacks CLI • Give Feedback 💬}