Transformer kernel release (#242)

* Transformer kernels release
Co-authored-by: Shaden Smith <ShadenTSmith@gmail.com>
Co-authored-by: Elton Zheng <eltonz@microsoft.com>
Co-authored-by: Reza Yazdani <reyazda@microsoft.com>
Co-authored-by: RezaYazdaniAminabadi <44502768+RezaYazdaniAminabadi@users.noreply.github.com>
Co-authored-by: Tunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Shaden Smith <ShadenTSmith@gmail.com>
Co-authored-by: Shaden Smith <Shaden.Smith@microsoft.com>
Co-authored-by: Samyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: Shaden Smith <ShadenTSmith@gmail.com>
Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
Co-authored-by: Samyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: Shaden Smith <ShadenTSmith@gmail.com>
Co-authored-by: Elton Zheng <eltonz@microsoft.com>
Co-authored-by: Reza Yazdani <reyazda@microsoft.com>
Co-authored-by: RezaYazdaniAminabadi <44502768+RezaYazdaniAminabadi@users.noreply.github.com>
Co-authored-by: Tunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Shaden Smith <Shaden.Smith@microsoft.com>
Co-authored-by: Samyam Rajbhandari <samyamr@microsoft.com>

docs/code-docs/source/checkpointing.rst

-DeepSpeed Activation Checkpointing
-======================
-
-The activation checkpointing API's in DeepSpeed can be used to enable a range of memory optimizations relating
-to activation checkpointing. These include activation partitioning across
-GPUs when using model parallelism, CPU Checkpointing, contiguous memory optimizations, etc.
-Please see the `DeepSpeed JSON config <https://www.deepspeed.ai/docs/config-json/>`_ for the full set.
-
-Here we present the activation checkpointing API's.
-Please see the enabling DeepSpeed for Megatron-LM tutorial for usage details.
-
-.. autofunction:: deepspeed.checkpointing.configure
-
-.. autofunction:: deepspeed.checkpointing.is_configured
-
-.. autofunction:: deepspeed.checkpointing.checkpoint
-
-.. autofunction:: deepspeed.checkpointing.reset
-
-.. autofunction:: deepspeed.checkpointing.get_cuda_rng_tracker
-
-.. autofunction:: deepspeed.checkpointing.model_parallel_cuda_manual_seed
-
-.. autoclass:: deepspeed.checkpointing.CudaRNGStatesTracker
-
-.. autoclass:: deepspeed.checkpointing.CheckpointFunction

docs/code-docs/source/conf.py

@@ -76,4 +76,11 @@ add_module_names = True
 
 autoclass_content = 'both'
 
-autodoc_mock_imports = ["torch", "apex", "mpi4py", "tensorboardX"]
+autodoc_mock_imports = [
+    "torch",
+    "apex",
+    "mpi4py",
+    "tensorboardX",
+    "deepspeed_transformer_cuda",
+    "deepspeed_stochastic_transformer_cuda",
+]

docs/code-docs/source/index.rst

@@ -8,7 +8,6 @@ Model Setup
    :maxdepth: 2
 
    initialize
-   checkpointing
 
 Training API
 ------------
@@ -28,6 +27,14 @@ Checkpointing API
    activation-checkpointing
 
 
+Transformer Kernel API
+----------------------
+.. toctree::
+   :maxdepth: 2
+
+   kernel
+
+
 Indices and tables
 ------------------