Update python and cuda version in readme for nightly packages, and fi…

…x a typo in docs.yml (pytorch#1426)

Summary:
Pull Request resolved: pytorch#1426

1. update readme with new download methods. Intentionally avoid using conda install for torch due to them having some issues.
2. fix a typo in docs.yml
3. switch to torchx-nightly in requirements.txt. We need that for python 3.11 anyway.

TODO: get torchmetrics on pytorch.org in the future.

Reviewed By: colin2328

Differential Revision: D48919435

fbshipit-source-id: d3f8c79ae258d49212f9f845d515e46c3e8a2fb7

.github/workflows/docs.yml

@@ -58,7 +58,7 @@ jobs:
         conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
     - name: Install TorchRec
       run: |
-        pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cpu
+        conda run -n build_binary pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cpu
     - name: Test fbgemm_gpu and torchrec installation
       shell: bash
       run: |

README.MD

@@ -5,7 +5,7 @@ TorchRec is a PyTorch domain library built to provide common sparsity & parallel
 
 ## TorchRec contains:
 - Parallelism primitives that enable easy authoring of large, performant multi-device/multi-node models using hybrid data-parallelism/model-parallelism.
-- The TorchRec sharder can shard embedding tables with different sharding strategies including data-parallel, table-wise, row-wise, table-wise-row-wise, and column-wise sharding.
+- The TorchRec sharder can shard embedding tables with different sharding strategies including data-parallel, table-wise, row-wise, table-wise-row-wise, column-wise, table-wise-column-wise sharding.
 - The TorchRec planner can automatically generate optimized sharding plans for models.
 - Pipelined training overlaps dataloading device transfer (copy to GPU), inter-device communications (input_dist), and computation (forward, backward) for increased performance.
 - Optimized kernels for RecSys powered by FBGEMM.
@@ -17,37 +17,35 @@ TorchRec is a PyTorch domain library built to provide common sparsity & parallel
 
 # Installation
 
-Torchrec requires Python >= 3.7 and CUDA >= 11.0 (CUDA is highly recommended for performance but not required). The example below shows how to install with CUDA 11.8. This setup assumes you have conda installed.
+Torchrec requires Python >= 3.8 and CUDA >= 11.8 (CUDA is highly recommended for performance but not required). The example below shows how to install with Python 3.8 and CUDA 12.1. This setup assumes you have conda installed.
 
 ## Binaries
 
-Experimental binary on Linux for Python 3.7, 3.8 and 3.9 can be installed via pip wheels
+Experimental binary on Linux for Python 3.8, 3.9, 3.10 and 3.11, and CPU, CUDA 11.8 and CUDA 12.1 can be installed via pip wheels from [download.pytorch.org](download.pytorch.org) and PyPI (only for CUDA 12.1).
+
+Below we show installations for CUDA 12.1 as an example. For CPU or CUDA 11.8, swap "cu121" for "cpu" or "cu118".
 
 ### Installations
 ```
-TO use the library without cuda, use the *-cpu fbgemm installations. However, this will be much slower than the CUDA variant.
-
 Nightly
 
-conda install pytorch pytorch-cuda=11.8 -c pytorch-nightly -c nvidia
-pip install torchrec_nightly
-
-Stable
-
-conda install pytorch pytorch-cuda=11.8 -c pytorch -c nvidia
-pip install torchrec
-
-If you have no CUDA device:
+pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
+pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
+pip install torchmetrics==1.0.3
+pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cu121
 
-Nightly
+Stable via pytorch.org
 
-pip uninstall fbgemm-gpu-nightly -y
-pip install fbgemm-gpu-nightly-cpu
+pip install torch --index-url https://download.pytorch.org/whl/cu121
+pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu121
+pip install torchmetrics==1.0.3
+pip install torchrec --index-url https://download.pytorch.org/whl/cu121
 
-Stable
+Stable via PyPI (only for CUDA 12.1)
 
-pip uninstall fbgemm-gpu -y
-pip install fbgemm-gpu-cpu
+pip install torch
+pip install fbgemm-gpu
+pip install torchrec
 
 ```
 
@@ -59,27 +57,55 @@ See our colab notebook for an introduction to torchrec which includes runnable i
 
 ## From Source
 
-We are currently iterating on the setup experience. For now, we provide manual instructions on how to build from source. The example below shows how to install with CUDA 11.3. This setup assumes you have conda installed.
+We are currently iterating on the setup experience. For now, we provide manual instructions on how to build from source. The example below shows how to install with CUDA 12.1. This setup assumes you have conda installed.
 
-1. Install pytorch. See [pytorch documentation](https://pytorch.org/get-started/locally/)
+1. Install pytorch. See [pytorch documentation](https://pytorch.org/get-started/locally/).
    ```
-   conda install pytorch pytorch-cuda=11.8 -c pytorch-nightly -c nvidia
+   CUDA 12.1
+
+   pip install torch --index-url https://download.pytorch.org/whl/nightly/cu121
+
+   CUDA 11.8
+
+   pip install torch --index-url https://download.pytorch.org/whl/nightly/cu118
+
+   CPU
+
+   pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
    ```
 
-2. Install Requirements
+2. Clone TorchRec.
    ```
-   pip install -r requirements.txt
+   git clone --recursive https://github.com/pytorch/torchrec
+   cd torchrec
    ```
 
-3. Download and install TorchRec.
+3. Install FBGEMM.
    ```
-   git clone --recursive https://github.com/pytorch/torchrec
+   CUDA 12.1
 
-   cd torchrec
+   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121
+
+   CUDA 11.8
+
+   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118
+
+   CPU
+
+   pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
+   ```
+
+4. Install other requirements.
+   ```
+   pip install -r requirements.txt
+   ```
+
+4. Install TorchRec.
+   ```
    python setup.py install develop
    ```
 
-4. Test the installation.
+5. Test the installation (use torchx-nightly for 3.11).
    ```
    GPU mode
 

setup.py

@@ -93,7 +93,7 @@ def main(argv: List[str]) -> None:
         url="https://github.com/pytorch/torchrec",
         license="BSD-3",
         keywords=["pytorch", "recommendation systems", "sharding"],
-        python_requires=">=3.7",
+        python_requires=">=3.8",
         install_requires=install_requires,
         packages=packages,
         zip_safe=False,
@@ -104,7 +104,7 @@ def main(argv: List[str]) -> None:
             "Intended Audience :: Science/Research",
             "License :: OSI Approved :: BSD License",
             "Programming Language :: Python :: 3",
-            "Programming Language :: Python :: 3.7",
+            "Programming Language :: Python :: 3.8",
             "Topic :: Scientific/Engineering :: Artificial Intelligence",
         ],
     )