[Fix] Infer hybrid FSDP replicas from world size

[primorLee]

Summary

• infer the HYBRID_SHARD replica dimension from the distributed world size when num_replicate is omitted
• validate that explicit mesh dimensions cover every distributed rank before constructing the DeviceMesh
• add CPU regression coverage for multi-node inference, explicit configurations, and invalid shapes

Problem

The training CLI currently defaults to:

num_replicate = 1
num_shard = 8

That creates a (1, 8) DeviceMesh even when a two-node launch has world_size=16. Only ranks 0-7 belong to the mesh; ranks on the second node have no mesh process-group metadata and fail inside FSDP with:

IndexError: list index out of range

This matches the rank-15 failure reported in #295.

Fix

For HYBRID_SHARD, the replica dimension is determined by:

num_replicate = world_size // num_shard

The CLI therefore uses None as the default and resolves the value at FSDP setup time. Existing single-node 8-GPU behavior remains (1, 8); a 16-GPU launch now becomes (2, 8).

If a user supplies an explicit incompatible replica count, or if world_size is not divisible by num_shard, the code now raises an actionable ValueError before PyTorch reaches the internal _dim_group_infos failure.

Validation

• python -B -m unittest discover -s tests -p test_fsdp_mesh.py -v — 5 passed
• wrapper integration test confirms a 16-rank world passes mesh_shape=(2, 8) to init_device_mesh
• Transformers HfArgumentParser accepts both the omitted Optional[int] value and --num_replicate 2
• Python compilation passes for all three changed files
• Ruff passes for all changed files with the repository's four unrelated pre-existing F541 findings in train/fsdp_utils.py excluded
• git diff --check passes

I did not run a real two-node GPU training job; the regression covers the exact mesh-shape construction that fails before model execution.

Closes #295