From 7e7a6014b2bb26c302b46509b6c261dcd9571b92 Mon Sep 17 00:00:00 2001 From: Joe Wang Date: Thu, 29 Aug 2024 12:48:09 -0700 Subject: [PATCH] add l2_cache_size config path from model config (#2336) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/2336 support l2 cache size as a config passed from model side Differential Revision: D61418000 --- torchrec/distributed/types.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/torchrec/distributed/types.py b/torchrec/distributed/types.py index 95d470af7..40a0e64dd 100644 --- a/torchrec/distributed/types.py +++ b/torchrec/distributed/types.py @@ -610,6 +610,7 @@ class KeyValueParams: gather_ssd_cache_stats: Optional[bool] = None stats_reporter_config: Optional[TBEStatsReporterConfig] = None use_passed_in_path: bool = True + l2_cache_size: Optional[int] = None # Parameter Server (PS) Attributes ps_hosts: Optional[Tuple[Tuple[str, int], ...]] = None @@ -623,13 +624,15 @@ def __hash__(self) -> int: self.ssd_storage_directory, self.ssd_rocksdb_write_buffer_size, self.ssd_rocksdb_shards, - self.gather_ssd_cache_stats, - self.stats_reporter_config, # Parameter Server (PS) Attributes self.ps_hosts, self.ps_client_thread_num, self.ps_max_key_per_request, self.ps_max_local_index_length, + # tbe attributes + self.gather_ssd_cache_stats, + self.stats_reporter_config, + self.l2_cache_size, ) )