Fixing KTO+QLoRA+multi-GPU (#2420)

* WIP

* removing artifacts

* adding error

* adding adapter check

* linting

* simplifying check

* linting v2

* config fix -___-

docs/config.qmd

@@ -85,6 +85,12 @@ gpu_memory_limit: 20GiB
 # Do the LoRA/PEFT loading on CPU -- this is required if the base model is so large it takes up most or all of the available GPU VRAM, e.g. during a model and LoRA merge
 lora_on_cpu: true
 
+# List[str]. Add plugins to extend the pipeline.
+# See `src/axolotl/integrations` for the available plugins or doc below for more details.
+# https://axolotl-ai-cloud.github.io/axolotl/docs/custom_integrations.html
+plugins:
+  # - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
+
 # A list of one or more datasets to finetune the model with
 datasets:
   # HuggingFace dataset repo | s3://,gs:// path | "json" for local dataset, make sure to fill data_files

docs/custom_integrations.qmd

@@ -55,3 +55,47 @@ sections = [
 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)
+
+:::

docs/installation.qmd

@@ -79,6 +79,7 @@ For providers supporting Docker:
   - [Latitude.sh](https://latitude.sh/blueprint/989e0e79-3bf6-41ea-a46b-1f246e309d5c)
   - [JarvisLabs.ai](https://jarvislabs.ai/templates/axolotl)
   - [RunPod](https://runpod.io/gsc?template=v2ickqhz9s&ref=6i7fkpdz)
+  - [Novita](https://novita.ai/gpus-console?templateId=311)
 
 ### Google Colab {#sec-colab}
 

examples/llama-3/qlora-1b-kto.yaml

@@ -55,7 +55,7 @@ tf32: true
 
 gradient_checkpointing: true
 gradient_checkpointing_kwargs:
-  use_reentrant: true
+  use_reentrant: false
 early_stopping_patience:
 resume_from_checkpoint:
 local_rank:

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -1679,6 +1679,30 @@ class AxolotlInputConfig(
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_rl_config_gradient_checkpointing(cls, data):
+        # TODO: SalmanMohammadi
+        # Distributed RL with QLoRA + gradient checkpointing
+        # and use_reentrant = True is broken upstream in TRL
+        # pylint: disable=too-many-boolean-expressions
+        if (
+            data.get("rl")
+            and data.get("gradient_checkpointing")
+            and data.get("gradient_checkpointing_kwargs")
+            and data.get("gradient_checkpointing_kwargs").get("use_reentrant")
+            and data.get("load_in_4bit")
+            and data.get("adapter") == "qlora"
+            and data.get("capabilities")
+            and data.get("capabilities").get("n_gpu", 1) > 1
+        ):
+            raise ValueError(
+                "The `use_reentrant: True` implementation of gradient checkpointing "
+                "is not supported for distributed RL training with QLoRA. Please set "
+                "`use_reentrant: False` in `gradient_checkpointing_kwargs`."
+            )
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_kto_config(cls, data):
@@ -1689,15 +1713,6 @@ class AxolotlInputConfig(
             if data.get("remove_unused_columns") is not False:
                 raise ValueError("Set `remove_unused_columns: False` when using kto")
 
-            if data.get("gradient_checkpointing") and not (
-                data.get("gradient_checkpointing_kwargs")
-                and isinstance(data.get("gradient_checkpointing_kwargs"), dict)
-                and data["gradient_checkpointing_kwargs"].get("use_reentrant")
-            ):
-                raise ValueError(
-                    "Set `gradient_checkpointing_kwargs: {use_reentrant: true}` for when kto is enabled"
-                )
-
         return data