fix shardedtensor init (#2586)

Summary:
Pull Request resolved: #2586

we have global metadata for shardedtensor to use for init - we can avoid allgather altogether.

Reviewed By: kausv

Differential Revision: D66394789

fbshipit-source-id: 6edde29a711110f0588c5040db0e330d1b796d1b

torchrec/distributed/embeddingbag.py

@@ -953,9 +953,9 @@ def _initialize_torch_state(self) -> None:  # noqa
                 )
 
                 self._model_parallel_name_to_sharded_tensor[table_name] = (
-                    ShardedTensor._init_from_local_shards(
-                        local_shards,
-                        self._name_to_table_size[table_name],
+                    ShardedTensor._init_from_local_shards_and_global_metadata(
+                        local_shards=local_shards,
+                        sharded_tensor_metadata=metadata,
                         process_group=(
                             self._env.sharding_pg
                             if isinstance(self._env, ShardingEnv2D)