update mtia info in torchrec (#3391)

Summary: Pull Request resolved: #3391

Reviewed By: egienvalue

Differential Revision: D83037454

fbshipit-source-id: 824b101749fe6e9022b54323e072828f4675cdb2

Author: optimisea

torchrec/distributed/planner/constants.py

@@ -77,9 +77,9 @@ def kernel_bw_lookup(
         ("cpu", EmbeddingComputeKernel.QUANT.value): 1 * ddr_mem_bw,
         # TODO: Determine the correct value later. MTIA uses values same as CPU's.
         # MTIA
-        ("mtia", EmbeddingComputeKernel.DENSE.value): 0.5 * ddr_mem_bw,
-        ("mtia", EmbeddingComputeKernel.FUSED.value): 1 * ddr_mem_bw,
-        ("mtia", EmbeddingComputeKernel.QUANT.value): 1 * ddr_mem_bw,
+        ("mtia", EmbeddingComputeKernel.DENSE.value): 0.5 * hbm_mem_bw,
+        ("mtia", EmbeddingComputeKernel.FUSED.value): 1 * hbm_mem_bw,
+        ("mtia", EmbeddingComputeKernel.QUANT.value): 1 * hbm_mem_bw,
         # CUDA
         ("cuda", EmbeddingComputeKernel.DENSE.value): 0.5 * hbm_mem_bw,
         ("cuda", EmbeddingComputeKernel.FUSED.value): 1 * hbm_mem_bw,

torchrec/distributed/planner/stats.py

@@ -705,7 +705,10 @@ def _log_rank_mem_usage_and_perf(
             used_hbm_gb = bytes_to_gb(used_hbm[rank])
             used_hbm_ratio = (
                 used_hbm[rank] / ((1 - reserved_hbm_percent) * device.storage.hbm)
-                if topology.compute_device == "cuda"
+                if (
+                    topology.compute_device == "cuda"
+                    or topology.compute_device == "mtia"
+                )
                 and ((1 - reserved_hbm_percent) * device.storage.hbm) != 0
                 else 0
             )