From 10000deb3913b47964885b6fc811e4ea97f569b3 Mon Sep 17 00:00:00 2001
From: Paul Zhang <paulzhan@meta.com>
Date: Wed, 10 Jan 2024 10:12:48 -0800
Subject: [PATCH] Add support for Rowwise and Columnwise sharding (#1621)

Summary:
Pull Request resolved: https://github.com/pytorch/torchrec/pull/1621

Add in Rowwise and Columnwise sharding to inference module benchmarking suite by default.

Reviewed By: bigning

Differential Revision: D52641307

fbshipit-source-id: f396bcfb9ac71fe65906a723184efde547c80235
---
 torchrec/distributed/benchmark/benchmark_inference.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/torchrec/distributed/benchmark/benchmark_inference.py b/torchrec/distributed/benchmark/benchmark_inference.py
index 54774f5ef..9f28394a7 100644
--- a/torchrec/distributed/benchmark/benchmark_inference.py
+++ b/torchrec/distributed/benchmark/benchmark_inference.py
@@ -439,11 +439,10 @@ def main() -> None:
         # Place all outputs under the datetime folder
         os.mkdir(output_dir)
 
-    # TODO: ROW_WISE and COLUMN_WISE are not supported yet
     BENCH_SHARDING_TYPES = [
         ShardingType.TABLE_WISE,
-        # ShardingType.ROW_WISE,
-        # ShardingType.COLUMN_WISE,
+        ShardingType.ROW_WISE,
+        ShardingType.COLUMN_WISE,
     ]
 
     table_sizes = [