/docs/custom_integrations.qmd

---
title: Custom Integrations
toc: true
toc-depth: 3
---

```{python}
#| echo: false

import os
import re

def process_readme(integration_name):
    try:
        path = f'../src/axolotl/integrations/{integration_name}/README.md'
        with open(path, 'r') as f:
            txt = f.read()
            # Remove h1 headings
            txt = re.sub(r'^# .*\n?', '', txt, flags=re.MULTILINE)
            # Convert h2 to h3
            txt = re.sub(r'^## ', '### ', txt, flags=re.MULTILINE)
            return txt
    except FileNotFoundError:
        return None

def print_section(name, folder_name):
    output = f"\n## {name}\n"
    content = process_readme(folder_name)
    if content:
        output += content
    output += f"\nPlease see reference [here](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/{folder_name})\n"
    return output
```

```{python}
#| output: asis
#| echo: false

# Introduction text
print("""
Axolotl adds custom features through `integrations`. They are located within the `src/axolotl/integrations` directory.

To enable them, please check the respective documentations.
""")

# Sections
sections = [
    ("Cut Cross Entropy", "cut_cross_entropy"),
    ("Grokfast", "grokfast"),
    ("Knowledge Distillation (KD)", "kd"),
    ("Liger Kernels", "liger"),
    ("Language Model Evaluation Harness (LM Eval)", "lm_eval"),
    ("Spectrum", "spectrum"),
    ("LLMCompressor", "llm_compressor")
]

for folder_name in os.listdir("../src/axolotl/integrations/"):
    if folder_name in [path for name, path in sections]:
        # skip if already in sections
        continue
    if os.path.exists(f"../src/axolotl/integrations/{folder_name}/README.md"):
        # grab the first heading in README.md as the section name
        with open(f"../src/axolotl/integrations/{folder_name}/README.md", "r") as f:
            txt = f.read()
            matches = re.search(r'^# (.*)\n?', txt, flags=re.MULTILINE)
            if matches:
                name = matches.group(1)
            else:
                continue
            sections.append((name, folder_name))

# sort sections by name
sections = sorted(sections, key=lambda x: x[0])

for section_name, folder_name in sections:
    print(print_section(section_name, folder_name))
```

## Adding a new integration

Plugins can be used to customize the behavior of the training pipeline through [hooks](https://en.wikipedia.org/wiki/Hooking). See [`axolotl.integrations.BasePlugin`](https://github.com/axolotl-ai-cloud/axolotl/blob/main/src/axolotl/integrations/base.py) for the possible hooks.

To add a new integration, please follow these steps:

1. Create a new folder in the `src/axolotl/integrations` directory.
2. Add any relevant files (`LICENSE`, `README.md`, `ACKNOWLEDGEMENTS.md`, etc.) to the new folder.
3. Add `__init__.py` and `args.py` files to the new folder.
  - `__init__.py` should import the integration and hook into the appropriate functions.
  - `args.py` should define the arguments for the integration.
4. (If applicable) Add CPU tests under `tests/integrations` or GPU tests under `tests/e2e/integrations`.

::: {.callout-tip}

See [src/axolotl/integrations/cut_cross_entropy](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/cut_cross_entropy) for a minimal integration example.

:::

::: {.callout-warning}

If you could not load your integration, please ensure you are pip installing in editable mode.

```bash
pip install -e .
```

and correctly spelled the integration name in the config file.

```yaml
plugins:
  - axolotl.integrations.your_integration_name.YourIntegrationPlugin
```

:::

::: {.callout-note}

It is not necessary to place your integration in the `integrations` folder. It can be in any location, so long as it's installed in a package in your python env.

See this repo for an example: [https://github.com/axolotl-ai-cloud/diff-transformer](https://github.com/axolotl-ai-cloud/diff-transformer)

:::