diff --git a/.github/workflows/sphinx.yml b/.github/workflows/sphinx.yml new file mode 100644 index 000000000..a2d8d116b --- /dev/null +++ b/.github/workflows/sphinx.yml @@ -0,0 +1,41 @@ +name: "Sphinx: Render docs" + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + build-docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install sphinx>=5.0.0 + # Install other dependencies if needed + pip install -r requirements.txt + + - name: Build documentation + run: | + cd docs + make html + - name: Upload artifacts + uses: actions/upload-artifact@v3 + with: + name: documentation + path: docs/_build/html/ + + - name: Deploy to GitHub Pages + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./docs/_build/html \ No newline at end of file diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 000000000..d0c3cbf10 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/_build/.doctrees/source/_templates/api.doctree b/docs/_build/.doctrees/source/_templates/api.doctree new file mode 100644 index 000000000..a3e3f76c7 Binary files /dev/null and b/docs/_build/.doctrees/source/_templates/api.doctree differ diff --git a/docs/_build/.doctrees/source/_templates/modules.doctree b/docs/_build/.doctrees/source/_templates/modules.doctree new file mode 100644 index 000000000..a653c35b3 Binary files /dev/null and b/docs/_build/.doctrees/source/_templates/modules.doctree differ diff --git a/docs/_build/.doctrees/source/_templates/todo.doctree b/docs/_build/.doctrees/source/_templates/todo.doctree new file mode 100644 index 000000000..106721eea Binary files /dev/null and b/docs/_build/.doctrees/source/_templates/todo.doctree differ diff --git a/docs/_build/.doctrees/source/index.doctree b/docs/_build/.doctrees/source/index.doctree new file mode 100644 index 000000000..79ae05728 Binary files /dev/null and b/docs/_build/.doctrees/source/index.doctree differ diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 000000000..747ffb7b3 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/source/api.rst b/docs/source/api.rst new file mode 100644 index 000000000..0888c233a --- /dev/null +++ b/docs/source/api.rst @@ -0,0 +1,106 @@ +Liger Kernel API +================ + +.. automodule:: liger_kernel + :members: + +Submodules +---------- +.. toctree:: + :maxdepth: 2 + + liger_kernel.transformers + liger_kernel.nn + +Transformers Module +------------------- + +The `liger_kernel.transformers` module includes essential tools and model classes designed for tasks such as language modeling and loss calculations, optimizing deep learning workflows. + +.. automodule:: liger_kernel.transformers + :members: + +Classes +------- + +AutoLigerKernelForCausalLM +-------------------------- + +The `AutoLigerKernelForCausalLM` class is an extension for causal language modeling, providing enhanced support for auto-regressive tasks. + +.. autoclass:: liger_kernel.transformers.AutoLigerKernelForCausalLM + :members: + :undoc-members: + +LigerFusedLinearCrossEntropyLoss +-------------------------------- + +Implements a fused linear layer with Cross-Entropy Loss, offering computational efficiency improvements. + +.. autoclass:: liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss + :members: + :undoc-members: + +Loss Functions +-------------- + +KLDivergence +------------ + +The `KLDivergence` class offers a divergence metric based on Kullback-Leibler Divergence, commonly used in various ML training contexts. + +.. autoclass:: liger_kernel.transformers.KLDivergence + :members: + :undoc-members: + +JSD +--- + +This class defines the Jensen-Shannon Divergence (JSD), used for calculating symmetrical divergence between distributions. + +.. autoclass:: liger_kernel.transformers.JSD + :members: + :undoc-members: + +GeneralizedJSD +-------------- + +Generalized form of Jensen-Shannon Divergence, adapted for more complex applications. + +.. autoclass:: liger_kernel.transformers.GeneralizedJSD + :members: + :undoc-members: + +FusedLinearJSD +-------------- + +Provides a fused linear implementation of Jensen-Shannon Divergence for faster computation. + +.. autoclass:: liger_kernel.transformers.FusedLinearJSD + :members: + :undoc-members: + +Experimental Kernels +--------------------- + +A collection of experimental kernels for advanced and experimental features in `liger_kernel`. + +.. automodule:: liger_kernel.transformers.experimental + :members: + +Neural Network Module +--------------------- + +The `liger_kernel.nn` module offers fundamental neural network building blocks. + +.. automodule:: liger_kernel.nn + :members: + +Module Class +------------ + +The base `nn.Module` class provides essential methods and attributes for building neural network models. + +.. autoclass:: liger_kernel.nn.Module + :members: + :undoc-members: diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 000000000..49f516a84 --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,126 @@ +import os +import sys +import importlib +import inspect + +# Increase the recursion limit (optional) +sys.setrecursionlimit(10000) + +# Configure paths +docs_dir = os.path.abspath(os.path.dirname(__file__)) +root_dir = os.path.abspath(os.path.join(docs_dir, '..')) +sys.path.insert(0, root_dir) + +# Source and build directories +source_dir = os.path.join(docs_dir, 'source') +output_dir = os.path.join(docs_dir, '_build') + +# List modules to mock (as simple strings) +MOCK_MODULES = [ + 'liger_kernel', + 'liger_kernel.transformers', + 'liger_kernel.transformers.experimental', + 'liger_kernel.nn', + 'liger_kernel.transformers.AutoLigerKernelForCausalLM', + 'liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss', + 'liger_kernel.transformers.KLDivergence', + 'liger_kernel.transformers.JSD', + 'liger_kernel.transformers.GeneralizedJSD', + 'liger_kernel.transformers.FusedLinearJSD', + 'liger_kernel.nn.Module', +] + +# Mock modules by replacing them in sys.modules +for mod_name in MOCK_MODULES: + sys.modules[mod_name] = None # Simple mock; replace with 'None' or a basic object as needed + +# Project information +project = 'Liger-Kernel' +copyright = '2024' +author = 'LinkedIn' + +# General Sphinx configuration +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.napoleon', + 'sphinx.ext.viewcode', + 'sphinx.ext.todo', + 'sphinx.ext.autosummary', + 'sphinx.ext.linkcode', + 'sphinx.ext.intersphinx', +] + +# Mocked imports for autodoc +autodoc_mock_imports = ['numpy', 'tensorflow', 'torch', 'triton'] + +# Intersphinx mapping for cross-referencing external documentation +intersphinx_mapping = { + 'python': ('https://docs.python.org/3', None), + 'numpy': ('https://numpy.org/doc/stable', None), + 'torch': ('https://pytorch.org/docs/stable', None), + 'triton': ('https://triton-lang.org/main/', None), +} + +# Source file configuration +source_suffix = '.rst' +master_doc = 'index' +exclude_patterns = [] + +# HTML output configuration +html_theme = 'alabaster' +html_static_path = [os.path.join(source_dir, '_static')] +html_theme_options = { + 'navigation_depth': 4, + 'collapse_navigation': False, +} + +# Output directory for HTML files +html_output_dir = output_dir + +# Autodoc configuration +autodoc_member_order = 'bysource' +autoclass_content = 'both' + +# Enable todo extension +todo_include_todos = True + +# Autosummary settings +autosummary_generate = True + +def linkcode_resolve(domain, info): + """ + Determine the URL corresponding to Python object + """ + if domain != 'py': + return None + + try: + mod = importlib.import_module(info['module']) + if 'class' in info: + obj = getattr(mod, info['class']) + if 'fullname' in info: + obj = getattr(obj, info['fullname']) + else: + obj = getattr(mod, info['fullname']) + except Exception: + return None + + try: + filepath = inspect.getsourcefile(obj) + if filepath: + filepath = os.path.relpath(filepath, start=root_dir) + else: + return None + except Exception: + return None + + try: + source, lineno = inspect.getsourcelines(obj) + except Exception: + return None + + # Modify these values based on your repository + github_url = "https://github.com/LinkedIn/liger-kernel" + branch = "main" # or whatever your default branch is + + return f"{github_url}/blob/{branch}/{filepath}#L{lineno}-L{lineno + len(source) - 1}" \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 000000000..a9861ed87 --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,17 @@ +Welcome to Liger-Kernel Documentation +=================================== + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + + source/api + source/modules + source/todo + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` \ No newline at end of file diff --git a/docs/source/liger_kernel.nn.rst b/docs/source/liger_kernel.nn.rst new file mode 100644 index 000000000..e69de29bb diff --git a/docs/source/liger_kernel.transformers.rst b/docs/source/liger_kernel.transformers.rst new file mode 100644 index 000000000..6ced7eb3a --- /dev/null +++ b/docs/source/liger_kernel.transformers.rst @@ -0,0 +1,36 @@ +.. automodule:: liger_kernel.transformers + :members: + +Liger Cross Entropy Function +============================= + +.. autoclass:: LigerCrossEntropyFunction + :members: + + The Liger Cross Entropy Function is a custom autograd function that implements the Liger Cross Entropy loss. + It overrides the forward and backward methods of the torch.autograd.Function class. + + .. automethod:: forward + + .. automethod:: backward + +Liger Cross Entropy Forward +============================ + +.. autofunction:: cross_entropy_forward + +Liger Cross Entropy Backward +============================ + +.. autofunction:: cross_entropy_backward + +Liger Cross Entropy Kernel +========================== + +.. autofunction:: liger_cross_entropy_kernel + +Element Mul Kernel +=================== + +.. autofunction:: element_mul_kernel + diff --git a/docs/source/modules.rst b/docs/source/modules.rst new file mode 100644 index 000000000..646a245d2 --- /dev/null +++ b/docs/source/modules.rst @@ -0,0 +1,97 @@ +Liger Kernel Modules +===================== + +.. toctree:: + :maxdepth: 2 + + liger_kernel.transformers + liger_kernel.nn + liger_kernel.transformers.experimental + +liger_kernel.transformers +------------------------- + +.. automodule:: liger_kernel.transformers + :members: + +AutoLigerKernelForCausalLM +--------------------------- + +.. autoclass:: liger_kernel.transformers.AutoLigerKernelForCausalLM + :members: + :undoc-members: + +LigerFusedLinearCrossEntropyLoss +--------------------------------- + +.. autoclass:: liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss + :members: + :undoc-members: + +KLDivergence +------------- + +.. autoclass:: liger_kernel.transformers.KLDivergence + :members: + :undoc-members: + +JSD +---- + +.. autoclass:: liger_kernel.transformers.JSD + :members: + :undoc-members: + +Generalized JSD +---------------- + +.. autoclass:: liger_kernel.transformers.GeneralizedJSD + :members: + :undoc-members: + +FusedLinearJSD +---------------- + +.. autoclass:: liger_kernel.transformers.FusedLinearJSD + :members: + :undoc-members: + +apply_liger_kernel_to_llama +----------------------------- + +.. autofunction:: liger_kernel.transformers.apply_liger_kernel_to_llama + :members: + :undoc-members: + +liger_kernel.nn +---------------- + +.. automodule:: liger_kernel.nn + :members: + +nn.Module +--------- + +.. autoclass:: liger_kernel.nn.Module + :members: + :undoc-members: + +liger_kernel.transformers.experimental +-------------------------------------- + +.. automodule:: liger_kernel.transformers.experimental + :members: + +LigerEmbeddingMatmul +--------------------- + +.. autoclass:: liger_kernel.transformers.experimental.LigerEmbeddingMatmul + :members: + :undoc-members: + +Matmul int2xint8 +----------------- + +.. autoclass:: liger_kernel.transformers.experimental.Matmul + :members: + :undoc-members: \ No newline at end of file diff --git a/docs/source/todo.rst b/docs/source/todo.rst new file mode 100644 index 000000000..e69de29bb diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 000000000..03ed8eeb7 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +sphinx>=7.0.0 +furo>=2023.09.10 +sphinxcontrib-applehelp>=1.0.4 +sphinxcontrib-devhelp>=1.0.2 +sphinxcontrib-htmlhelp>=2.0.1 +sphinxcontrib-serializinghtml>=1.1.5 +sphinxcontrib-qthelp>=1.0.3 +typing-extensions>=4.8.0 \ No newline at end of file diff --git a/src/liger_kernel/README.rst b/src/liger_kernel/README.rst new file mode 100644 index 000000000..ed5fedc03 --- /dev/null +++ b/src/liger_kernel/README.rst @@ -0,0 +1,3 @@ +Liger + +Liger is a Python Library used for finetuning. \ No newline at end of file