d68cc1e8ab
* densemixer plugin integration * update readme with usage docs * automatically find new integrations that aren't explicitly defined * make sure to import os
122 lines
3.9 KiB
Plaintext
122 lines
3.9 KiB
Plaintext
---
|
|
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)
|
|
|
|
:::
|