Make deep copy when constructing cache params (#3219)

Summary:
Pull Request resolved: #3219

During debugging, we realized that it's possible that a feature may appear in two arch as different tables. Since sharding constraint is formed by table name, it'll craft the same constraint for both tables. During proposer, if one table's constraint get modified, e.g. by EmbeddingOffloadCacheScaleupProposer, the other will change automatically. This is not expected as sharding plan should treat them individually.

This diff fixed it.

Reviewed By: iamzainhuda, aliafzal

Differential Revision: D78287465

fbshipit-source-id: 6eab7a47f72e6501eadcd1bc2dbe88f53bb0b984

Author: levythu

torchrec/distributed/planner/enumerators.py

@@ -9,6 +9,7 @@
 
 import copy
 import logging
+from copy import deepcopy
 from typing import Dict, List, Optional, Set, Tuple, Union
 
 from torch import nn
@@ -364,16 +365,17 @@ def _extract_constraints_for_param(
     key_value_params = None
 
     if constraints and constraints.get(name):
+        # For nested fields return a deep copy instead
         input_lengths = constraints[name].pooling_factors
         col_wise_shard_dim = constraints[name].min_partition
-        cache_params = constraints[name].cache_params
+        cache_params = deepcopy(constraints[name].cache_params)
         enforce_hbm = constraints[name].enforce_hbm
         stochastic_rounding = constraints[name].stochastic_rounding
         bounds_check_mode = constraints[name].bounds_check_mode
         feature_names = constraints[name].feature_names
         output_dtype = constraints[name].output_dtype
         device_group = constraints[name].device_group
-        key_value_params = constraints[name].key_value_params
+        key_value_params = deepcopy(constraints[name].key_value_params)
 
     return (
         input_lengths,