New Features
In Training Embedding Pruning (ITEP) for more efficient RecSys training
Provides a representation of In-Training Embedding Pruning, which is used internally at Meta for more efficient RecSys training by decreasing memory footprint of embedding tables. Pull Request: https://github.com/pytorch/torchrec/pull/2074 introduces the modules into TorchRec, with tests showing how to use them.
Mean Pooling
Mean pooling enabled on embeddings for row-wise and table-row-wise sharding types in TorchRec. Mean pooling mode done through TBE (table-batched embedding) won’t be accurate for row-wise and table-row-wise sharding types, which modify the input due to sharding. This feature efficiently calculates the divisor using caching and overlapping in input dist to implement mean pooling, which had proved to be much more performant than out-of-library implementations. PR: https://github.com/pytorch/torchrec/pull/1772
Changelog
[Torch.export](https://pytorch.org/docs/stable/export.html) (non-strict) compatibility with KJT/JT/KT, EBC/Quantized EBC, sharded variants [#1815](https://github.com/pytorch/torchrec/pull/1815) [#1816](https://github.com/pytorch/torchrec/pull/1816) [#1788](https://github.com/pytorch/torchrec/pull/1788) [#1850](https://github.com/pytorch/torchrec/pull/1850) [#1976](https://github.com/pytorch/torchrec/pull/1976) and dynamic shapes [#2058](https://github.com/pytorch/torchrec/pull/2058)
[torch.compile ](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html)support with TorchRec [#2045](https://github.com/pytorch/torchrec/pull/2045) [#2018](https://github.com/pytorch/torchrec/pull/2018) [#1979](https://github.com/pytorch/torchrec/pull/1979)
TorchRec serialization with non-strict torch.export for regenerating eager sparse modules (EBC) from IR for sharding [1860](https://github.com/pytorch/torchrec/pull/1860) [#1848](https://github.com/pytorch/torchrec/pull/1848) with meta functionalization when torch.exporting [#1974](https://github.com/pytorch/torchrec/pull/1974)
More benchmarking for TorchRec modules/data types [2094](https://github.com/pytorch/torchrec/pull/2094) [#2033](https://github.com/pytorch/torchrec/pull/2033) [#2001](https://github.com/pytorch/torchrec/pull/2001) [#1855](https://github.com/pytorch/torchrec/pull/1855)
More VBE support (data parallel sharding) [2093](https://github.com/pytorch/torchrec/pull/2093) (EmbeddingCollection) [#2047](https://github.com/pytorch/torchrec/pull/2047) [#1849](https://github.com/pytorch/torchrec/pull/1849)
RegroupAsDict module for performance improvements with caching [2007](https://github.com/pytorch/torchrec/pull/2007)
Train Pipeline improvements [1967](https://github.com/pytorch/torchrec/pull/1967) [#1969](https://github.com/pytorch/torchrec/pull/1969) [#1971](https://github.com/pytorch/torchrec/pull/1971)
Bug Fixes and library improvements