2024-10-03 nightly release (642da8e)

pytorch · Oct 3, 2024 · 83a353c · 83a353c
1 parent 072843a
commit 83a353c
Show file tree

Hide file tree

Showing 29 changed files with 616 additions and 828 deletions.
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -18,6 +18,7 @@ jobs:
         include:
          - os: linux.20_04.4x
            python-version: 3.9
+           python-tag: "py39"
     steps:
     - name: Check ldd --version
       run: ldd --version
@@ -66,7 +67,8 @@ jobs:
         conda run -n build_binary pip install torchmetrics==1.0.3
     - name: Install TorchRec
       run: |
-        conda run -n build_binary pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cpu
+        conda run -n build_binary pip install -r requirements.txt
+        conda run -n build_binary python setup.py bdist_wheel --python-tag=${{ matrix.python-tag }}
     - name: Test fbgemm_gpu and torchrec installation
       shell: bash
       run: |

diff --git a/docs/source/_templates/layout.html b/docs/source/_templates/layout.html
@@ -4,7 +4,7 @@
 {{ super() }}
 
 <script type="text/javascript">
-  var collapsedSections = ['Introduction', 'Getting Started', 'All API References']
+  var collapsedSections = ['Introduction', 'Getting Started', 'API References']
 </script>
 
 {% endblock %}
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -21,7 +21,6 @@
 import sys
 
 import pytorch_sphinx_theme
-import torchrec
 
 current_dir = os.path.dirname(__file__)
 target_dir = os.path.abspath(os.path.join(current_dir, "../.."))
@@ -36,7 +35,7 @@
 
 try:
     # pyre-ignore
-    version = torchrec.__version__
+    version = "1.0.0"  # TODO: Hardcode stable version for now
 except Exception:
     # when run internally, we don't have a version yet
     version = "0.0.0"

diff --git a/docs/source/datatypes-api-reference.rst b/docs/source/datatypes-api-reference.rst
@@ -0,0 +1,22 @@
+Data Types
+-------------------
+
+
+TorchRec contains data types for representing embedding, otherwise known as sparse features.
+Sparse features are typically indices that are meant to be fed into embedding tables. For a given
+batch, the number of embedding lookup indices are variable. Therefore, there is a need for a **jagged**
+dimension to represent the variable amount of embedding lookup indices for a batch.
+
+This section covers the classes for the 3 TorchRec data types for representing sparse features:
+**JaggedTensor**, **KeyedJaggedTensor**, and **KeyedTensor**.
+
+.. automodule:: torchrec.sparse.jagged_tensor
+
+.. autoclass:: JaggedTensor
+    :members:
+
+.. autoclass:: KeyedJaggedTensor
+    :members:
+
+.. autoclass:: KeyedTensor
+    :members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -100,20 +100,12 @@ how you can contribute:
    setup-torchrec.rst
 
 .. toctree::
-   :maxdepth: 1
-   :caption: All API References
+   :maxdepth: 2
+   :caption: API References
    :hidden:
 
-   torchrec.datasets.rst
-   torchrec.datasets.scripts.rst
-   torchrec.distributed.rst
-   torchrec.distributed.planner.rst
-   torchrec.distributed.sharding.rst
-   torchrec.fx.rst
-   torchrec.inference.rst
-   torchrec.models.rst
-   torchrec.modules.rst
-   torchrec.optim.rst
-   torchrec.quant.rst
-   torchrec.sparse.rst
-   torchrec.metrics.rst
+   datatypes-api-reference.rst
+   modules-api-reference.rst
+   planner-api-reference.rst
+   model-parallel-api-reference.rst
+   inference-api-reference.rst
diff --git a/docs/source/inference-api-reference.rst b/docs/source/inference-api-reference.rst
@@ -0,0 +1,19 @@
+Inference
+----------------------------------
+
+TorchRec provides easy-to-use APIs for transforming an authored TorchRec model
+into an optimized inference model for distributed inference, via eager module swaps.
+
+This transforms TorchRec modules like ``EmbeddingBagCollection`` in the model to
+a quantized, sharded version that can be compiled using torch.fx and TorchScript
+for inference in a C++ environment.
+
+The intended use is calling ``quantize_inference_model`` on the model followed by
+``shard_quant_model``.
+
+.. codeblock::
+
+.. automodule:: torchrec.inference.modules
+
+.. autofunction:: quantize_inference_model
+.. autofunction:: shard_quant_model
diff --git a/docs/source/model-parallel-api-reference.rst b/docs/source/model-parallel-api-reference.rst
@@ -0,0 +1,10 @@
+Model Parallel
+----------------------------------
+
+``DistributedModelParallel`` is the main API for distributed training with TorchRec optimizations.
+
+
+.. automodule:: torchrec.distributed.model_parallel
+
+.. autoclass:: DistributedModelParallel
+    :members:
diff --git a/docs/source/modules-api-reference.rst b/docs/source/modules-api-reference.rst
@@ -0,0 +1,30 @@
+Modules
+----------------------------------
+
+Standard TorchRec modules represent collections of embedding tables:
+
+* ``EmbeddingBagCollection`` is a collection of ``torch.nn.EmbeddingBag``
+* ``EmbeddingCollection`` is a collection of ``torch.nn.Embedding``
+
+These modules are constructed through standardized config classes:
+
+* ``EmbeddingBagConfig`` for ``EmbeddingBagCollection``
+* ``EmbeddingConfig`` for ``EmbeddingCollection``
+
+.. automodule:: torchrec.modules.embedding_configs
+
+.. autoclass:: EmbeddingBagConfig
+    :show-inheritance:
+
+.. autoclass:: EmbeddingConfig
+    :show-inheritance:
+
+.. autoclass:: BaseEmbeddingConfig
+
+.. automodule:: torchrec.modules.embedding_modules
+
+.. autoclass:: EmbeddingBagCollection
+    :members:
+
+.. autoclass:: EmbeddingCollection
+    :members:
diff --git a/docs/source/planner-api-reference.rst b/docs/source/planner-api-reference.rst
@@ -0,0 +1,50 @@
+Planner
+----------------------------------
+
+The TorchRec Planner is responsible for determining the most performant, balanced
+sharding plan for distributed training and inference.
+
+The main API for generating a sharding plan is ``EmbeddingShardingPlanner.plan``
+
+.. automodule:: torchrec.distributed.types
+
+.. autoclass:: ShardingPlan
+    :members:
+
+.. automodule:: torchrec.distributed.planner.planners
+
+.. autoclass:: EmbeddingShardingPlanner
+    :members:
+
+.. automodule:: torchrec.distributed.planner.enumerators
+
+.. autoclass:: EmbeddingEnumerator
+    :members:
+
+.. automodule:: torchrec.distributed.planner.partitioners
+
+.. autoclass:: GreedyPerfPartitioner
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.storage_reservations
+
+.. autoclass:: HeuristicalStorageReservation
+    :members:
+
+.. automodule:: torchrec.distributed.planner.proposers
+
+.. autoclass:: GreedyProposer
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.shard_estimators
+
+.. autoclass:: EmbeddingPerfEstimator
+    :members:
+
+
+.. automodule:: torchrec.distributed.planner.shard_estimators
+
+.. autoclass:: EmbeddingStorageEstimator
+    :members:
diff --git a/docs/source/torchrec.datasets.rst b/docs/source/torchrec.datasets.rst
diff --git a/docs/source/torchrec.datasets.scripts.rst b/docs/source/torchrec.datasets.scripts.rst
diff --git a/docs/source/torchrec.distributed.planner.rst b/docs/source/torchrec.distributed.planner.rst