Add logging when merging VBE embeddings from multiple TBEs (#3304)

Summary:
Pull Request resolved: #3304

Add a logging in `_merge_variable_batch_embeddings` so that we can tell if a model has multiple TBEs for lookup with VBE enabled.

This function can take a significantly long time when the world size and number of features are large. It can sometimes lead to GPU inefficiencies. See this [doc](https://docs.google.com/document/d/1h5YyeCjYmmN-CIFB98CrBf1uMksidPbNvM1rl8yZeds/edit?addon_store&fbclid=IwY2xjawL3TWBleHRuA2FlbQIxMQBicmlkETEzSzN5RFhBN0hyd0RZa2ZnAR6U1Lg7P5B-BgHUp_-eFqgx-__zTPYiXWSS0eZc9SmJdgiJeZv-fwnrCaSSLA_aem_YhBf3lxMTQ_xrKbE-GXpnQ&tab=t.0) for more context.

Reviewed By: TroyGarden

Differential Revision: D80686878

fbshipit-source-id: bf63fc5ec5073c5e5a9dd6d1c4cb4bc3306d4c86

Author: Shuangping Liu

torchrec/distributed/embedding_lookup.py

@@ -594,6 +594,18 @@ def _merge_variable_batch_embeddings(
         self, embeddings: List[torch.Tensor], splits: List[List[int]]
     ) -> List[torch.Tensor]:
         assert len(embeddings) > 1 and len(splits) > 1
+
+        logger.info(
+            "Merge VBE embeddings from the following TBEs "
+            f"(world size: {self._world_size}):\n"
+            + "\n".join(
+                [
+                    f"\t{module.__class__.__name__}:{len(split)} splits"
+                    for module, split in zip(self._emb_modules, splits)
+                ]
+            )
+        )
+
         split_embs = [e.split(s) for e, s in zip(embeddings, splits)]
         combined_embs = [
             emb