Introduce FlexShard for flexible bucketed parameter sharding

[weifengpy]

Stack from ghstack (oldest at bottom):

• Add communication-free Muon for FlexShard #3502
• [flex_shard] Plan: fp8 all-gather on GroupedRaggedShard (block-wise) #3537
• [flex_shard] Add DeepSeek V3 eager training entry point #3603
• Make FlexShard traceable by torch.compile #3317
• -> Introduce FlexShard for flexible bucketed parameter sharding #3239

Introduce experimental FlexShard, a placement-driven API for flexible bucketed parameter sharding. FlexShard lets users group parameters into explicit communication buckets and assign each parameter a pluggable placement implementation. The initial runtime is eager for now.

• flex_shard() API with explicit BucketSpecs

• Per-bucket flat byte storage for sharded parameters

• Pluggable Placement contract for local layout, unshard collectives, and gradient reduction

• Shard, Owned, and RaggedShard example placements to prove the placement contract

• RaggedShard supports uneven flattened-prefix sharding, bucket unshard/reduce-grad, and focused CPU/CUDA runtime coverage

• Eager parameter accessors backed by batched bucket unshard hooks

• Per-bucket MixedPrecisionPolicy

• reshard-after-forward path using selective activation checkpointing

• Runtime support is eager-only; will explore graph capture and torch.compile in next PR

• Each bucket currently requires a uniform parameter dtype and placement tuple.

• Mixed precision is supported per bucket.

• CPU offload is represented in the API but rejected until implemented

Test Plan:

• python -m pytest -q torchtitan/experiments/flex_shard/tests/test_flex_shard_ragged_shard.py