v2025.05.26.00
tagged this
24 May 10:15
Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/2998 #### Context Several unit tests in `test_model_parallel.py` passed **empty constraints** into `self._test_sharding` because the constraints are generated using an empty `self.tables` before invoking `self._build_tables_and_groups`. Impacted tests are: * `test_sharding_twcw` * `test_sharding_variable_batch` * `test_sharding_multiple_kernels` #### Changes * Constraints only depend on table names. A new list `self.table_names` is created in `setUp()` stage to be used to construct constraints. * Updates `self._build_tables_and_groups` to use the generated table names. * Increases `max_examples` for `test_sharding_multiple_kernels` to cover both FP32 and FP16 cases. Reviewed By: TroyGarden Differential Revision: D75306149 fbshipit-source-id: b93f7656e45a8c79393a1c347437f757aac07557