Proportional Uneven Inference Sharding

Summary:
Support bucketization aware inference sharding in TGIF for ZCH bucket boundaries from training.
A "best effort" sharding is performed across bucket boundaries proportional to memory list.

Differential Revision: D69057627

torchrec/distributed/quant_state.py

@@ -441,6 +441,7 @@ def sharded_tbes_weights_spec(
                     shard_sizes: List[int] = [table.local_rows, table.local_cols]
                     shard_offsets: List[int] = table_metadata.shard_offsets
                     s: str = "embedding_bags" if is_sqebc else "embeddings"
+                    s = ("_embedding_module." if is_sqmcec else "") + s
                     unsharded_fqn_weight: str = f"{module_fqn}.{s}.{table_name}.weight"
 
                     sharded_fqn_weight: str = (