mirror of https://github.com/hpcaitech/ColossalAI
parent
a3ac48ef3d
commit
3f760da9f0
|
@ -15,7 +15,7 @@ Method 2 is good because it allows the user to only build the kernel they actual
|
|||
|
||||
## PyTorch Extensions in Colossal-AI
|
||||
|
||||
The project DeepSpeed (https://github.com/microsoft/DeepSpeed) has proposed a [solution](https://github.com/microsoft/DeepSpeed/tree/master/op_builder)) to support kernel-build during either installation or runtime.
|
||||
The project DeepSpeed (https://github.com/microsoft/DeepSpeed) has proposed a [solution](https://github.com/microsoft/DeepSpeed/tree/master/op_builder) to support kernel-build during either installation or runtime.
|
||||
We have adapted from DeepSpeed's solution to build extensions. The extension build requries two main functions from PyTorch:
|
||||
|
||||
1. `torch.utils.cpp_extension.CUDAExtension`: used to build extensions in `setup.py` during `pip install`.
|
||||
|
|
Loading…
Reference in New Issue