Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions docs/_tutorials/advanced-install.md
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,23 @@ script in the repo. This will build a python wheel locally and copy it to all
the nodes listed in your hostfile (either given via --hostfile, or defaults to
/job/hostfile).

When the code using DeepSpeed is used for the first time it'll automatically build only the CUDA
extensions, required for the run, and by default it'll place them under
`~/.cache/torch_extensions/`. The next time the same program is executed these now precompiled
extensions will be loaded form that directory.

If you use multiple virtual environments this could be a problem, since by default there is only one
extensions directory, but different virtual environments may use different setups (e.g. different
python or cuda versions) and then the loading of a CUDA extension built by another environment will
fail. Therefore, if you need to you can override the default location with the help of the
`TORCH_EXTENSIONS_DIR` environment variable. So in each virtual environment you can point it to a
unique directory and DeepSpeed will use it to save and load CUDA extensions.

You can also change it just for a specific run with:

```bash
TORCH_EXTENSIONS_DIR=./torch-extensions deepspeed ...
```

## Building for the correct architectures

Expand Down