@@ -17,44 +20,69 @@
+ 
+
+ A Free and Open Source LLM Fine-tuning Framework
+
@@ -17,44 +20,69 @@
![\"Built](\"https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png\"){kind=icon}
](https://github.com/axolotl-ai-cloud/axolotl)\n",
+ "\n",
+ "Axolotl is the most performant LLM post-training framework available, delivering faster training with efficient, consistent and stable performance. Train your workload and ship your product 30% faster; saving you both time and money.\n",
+ "\n",
+ "- ⭐ us on [GitHub](https://github.com/axolotl-ai-cloud/axolotl)\n",
+ "- 📜 Read the [Docs](http://docs.axolotl.ai/)\n",
+ "- 💬 Chat with us on [Discord](https://discord.gg/mnpEYgRUmD)\n",
+ "- 📰 Get updates on [X/Twitter](https://x.com/axolotl_ai)\n"
]
},
{
"cell_type": "markdown",
- "metadata": {},
+ "metadata": {
+ "id": "rVjKD7CbxIP3"
+ },
"source": [
- "## Hugging Face login (optional)"
+ "# Installation\n",
+ "\n",
+ "Axolotl is easy to install from [pip](https://pypi.org/project/axolotl/), or use our [pre-built Docker images](http://docs.axolotl.ai/docs/docker.html) for a hassle free dependency experience. See our [docs](http://docs.axolotl.ai/docs/installation.html) for more information."
]
},
{
"cell_type": "code",
"execution_count": null,
- "metadata": {},
+ "metadata": {
+ "id": "msOCO4NRmRLa"
+ },
"outputs": [],
+ "source": [
+ "%%capture\n",
+ "# This step can take ~5-10 minutes to install dependencies\n",
+ "!pip install --no-build-isolation axolotl[flash-attn]>=0.9.1\n",
+ "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@8a1a0ec\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "N0OW0YeksDLr"
+ },
+ "source": [
+ "## Demo: Talk Like a Pirate\n",
+ "\n",
+ "In this demo, we are training the model ***to respond like a pirate***. This was chosen as a way to easily show how to train a model to respond in a certain style of your choosing (without being prompted) and is quite easy to validate within the scope of a Colab."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "8Du2fANTsNCK"
+ },
+ "source": [
+ "### Upload your own dataset or use a Huggingface dataset\n",
+ "\n",
+ "You can choose to use your own JSONL file from your own [Google Drive](https://drive.google.com/drive/home); for example downloading the [Pirate-Ultrachat JSONL](https://huggingface.co/datasets/winglian/pirate-ultrachat-10k/blob/main/train.jsonl) to your Google Drive. JSONL datasets should be formatted similar to the [OpenAI dataset format](https://cookbook.openai.com/examples/chat_finetuning_data_prep).\n",
+ "\n",
+ "You can also simply use the [`winglian/pirate-ultrachat-10k`](https://huggingface.co/datasets/winglian/pirate-ultrachat-10k) dataset directly.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "fGEEjyQ-r_IV"
+ },
+ "outputs": [],
+ "source": [
+ "# Default to HF dataset location\n",
+ "dataset_id = \"winglian/pirate-ultrachat-10k\"\n",
+ "uploaded = {}"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "c5MyYqk7vIsG"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "# Optionally, upload your own JSONL to your Google Drive\n",
+ "GOOGLE_DRIVE_PATH = \"\" # ex: \"MyDrive/Colab\\ Notebooks/train.jsonl\"\n",
+ "\n",
+ "# \"Select All\" permissions, or you may get the error:\n",
+ "# \"MessageError: Error: credential propagation was unsuccessful\"\n",
+ "if GOOGLE_DRIVE_PATH:\n",
+ " from google.colab import drive\n",
+ "\n",
+ " # Mount your Google Drive\n",
+ " GOOGLE_DRIVE_MNT = \"/content/drive/\"\n",
+ " drive.mount(GOOGLE_DRIVE_MNT, force_remount=True)\n",
+ " tmp_path = os.path.join(GOOGLE_DRIVE_MNT, GOOGLE_DRIVE_PATH.lstrip(\"/\"))\n",
+ " # make sure file exists\n",
+ " if not os.path.isfile(tmp_path):\n",
+ " raise ValueError(f\"File {tmp_path} does not exist\")\n",
+ " dataset_id = tmp_path"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "U6pTk3A9xj1W"
+ },
+ "source": [
+ "# Configure for Supervised Fine-Tuning (SFT)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 151,
+ "referenced_widgets": [
+ "388f618924274d21a066f098f4f1e744",
+ "7c95f85a2b1f47a1bd846d110c47bb3c",
+ "083f9cda8d754c168beee10d2f8955a2",
+ "62e1a65582f446a78612eaa804e08a7d",
+ "487a177d020f4605834878b2fdc7afa3",
+ "7fd44cf9ca6e4726bfd7ac21846d6a14",
+ "366a343b62fa47d8985a3bd464d99f9e",
+ "a0a11e929edd4189b79723d618522c33",
+ "e87ea87fcff247b5bbcc331ba79a8dc2",
+ "5e18768f7ad6434ba8b8b8a2e853e204",
+ "bb33aec33a6447078c31bfd728942994"
+ ]
+ },
+ "id": "fdRioqytmTtX",
+ "outputId": "f0acdcec-4b41-4a3f-ffed-c2d2d929158e"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[2025-05-08 13:40:27,488] [INFO] [root.register:348] [PID:174] Attempting to load plugin: axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin\n",
+ "[2025-05-08 13:40:27,493] [INFO] [root.register:351] [PID:174] Plugin loaded successfully: axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin\n",
+ "[2025-05-08 13:40:27,959] [INFO] [axolotl.utils.schemas.config.check_eval_packing:721] [PID:174] [RANK:0] explicitly setting `eval_sample_packing` to match `sample_packing`\u001b[39m\n",
+ "[2025-05-08 13:40:27,960] [INFO] [axolotl.utils.schemas.config.hint_sample_packing_padding:514] [PID:174] [RANK:0] Setting `pad_to_sequence_len: true` to prevent memory leaks when sample_packing\u001b[39m\n",
+ "[2025-05-08 13:40:27,961] [INFO] [axolotl.utils.schemas.config.check_bf16:1251] [PID:174] [RANK:0] bf16 support detected, but not enabled for this configuration.\u001b[39m\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "388f618924274d21a066f098f4f1e744",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "config.json: 0%| | 0.00/728 [00:00\"],\n",
+ " }\n",
+ " ],\n",
+ " dataloader_prefetch_factor=8, # dataloader optimizations\n",
+ " dataloader_num_workers=2,\n",
+ " dataloader_pin_memory=True,\n",
+ ")\n",
+ "\n",
+ "# validates the configuration\n",
+ "cfg = load_cfg(config)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "715UpvnSoBIS"
+ },
+ "outputs": [],
+ "source": [
+ "from axolotl.utils import set_pytorch_cuda_alloc_conf\n",
+ "\n",
+ "# Set \"PYTORCH_CUDA_ALLOC_CONF\" env to save memory\n",
+ "set_pytorch_cuda_alloc_conf()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Vc6MC-hwyH-n"
+ },
+ "source": [
+ "# Datasets\n",
+ "\n",
+ "Axolotl has a robust suite of loaders and transforms to parse most open datasets of any format into the appropriate chat template for your model. Axolotl will mask input tokens from the user's prompt so that the train loss is only calculated against the model's response. For more information, [see our documentation](http://docs.axolotl.ai/docs/dataset-formats/conversation.html) on dataset preparation.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 1000,
+ "referenced_widgets": [
+ "b82aa8c57f7c422a9a9c90f333ed2a99",
+ "c0991cf63ee6458b96e9a75e7a88b61a",
+ "71c8af139cd248b1b51101fd46a93f35",
+ "1d5117195d4b49eb8f1a73b18419f7ce",
+ "3c21e4a511b4441192c03b7f1d0976e9",
+ "ed28e2e0410d4e0b855467e798e53d66",
+ "d93f134f802b4b69b575bdaf07dbd27c",
+ "d0e9dce55cec4c1ca619a0ccf209d924",
+ "4c727d40ef0443449afc31724ee79f0c",
+ "0dea5caa27384f5689e3cab51f558727",
+ "a6f48410b9964fefba0c3009a77dc838",
+ "95caff42f08a4c2aa14c867b8f37f231",
+ "de7c37ee83e24f0c889e84d07279c2ec",
+ "9d4897eefb5f48259ffb2d23e332f752",
+ "253017b0d0534e54ab44e181f6d7c82d",
+ "27beaf06e41b472abdb544a43c720c5a",
+ "34cf3df51fbc41cabfdbba153c007f0e",
+ "ac764024cf1c4e08ba7749afd2cd20ac",
+ "30a81da86f8043eca301e86a8651201a",
+ "e8b7a81040904c1e89e58978223b1737",
+ "1c6f1f10667545aaab958016ba7e2c94",
+ "e6e969610738449887259063967f82b0",
+ "a138859f19b74fc0928dc236ab5359db",
+ "9b42e08b3c9548818488268768a118b1",
+ "12b56912736849fea2ad8124456fdc5c",
+ "879c8ab5873847a8833bd74123be90a4",
+ "20352e5f58d24bb8b1f3940efd14fe4a",
+ "d955dcaa0e944e719f3a06139dd54a03",
+ "d3de2662c7964f1ba96e58da382af720",
+ "97e36007e1304e1583fd81bfb13f0edd",
+ "c65dc74c7d6f4bab8f7dd28455161dd8",
+ "ef223e8504b64e3592589880326aaf41",
+ "598da69727bd4fb8b1caf465ac736d7a",
+ "5f86cd894de94c3280fadc1e2fd0ee13",
+ "a20927bf5f2c41f58c1e31ac858ab36c",
+ "0a46ad75c198463d843fb35e813642cb",
+ "09007681cf8d42aeb8c1d2f6a74e470a",
+ "ebc80d1a55fa47f4a5ea2756588569ec",
+ "1811cda0644e4190a9469d1774435d82",
+ "35c811d2ae8e43f3b5cecbdd3cfa857f",
+ "b8e39e4dddc3497fbc29ae45c66da759",
+ "63b4e563e85c4f03b1b72beda9577bcc",
+ "b195f160ca20442fadd8b5aed0ee41af",
+ "ca65e32eb52f48c09a84b33cb18f22cd",
+ "7cd0b85ebd204b7aba908417811ce4e0",
+ "7baeab52d6694c32b1efd1ea1a0a7782",
+ "519a7b154022443db6703f04a9142bae",
+ "d4183e9715f34d249942b8271cca3bdf",
+ "da2347ac94764a3fa2743343cf0d3cd2",
+ "93a44a11aa4846fa8efc6c1413ef1627",
+ "a55060adc3564407ac81ad7297d34aaa",
+ "d02274afd47b462291c745f261209d42",
+ "0f417447a7bd4a33acca96fa37aec877",
+ "63580b6fb30642479fe3000915bf551a",
+ "8f726dbfb45d4528afa33e36a6313267",
+ "03b093d592ba4386aa61f7b8483da660",
+ "b8766a88716948cf968f4563531a76d9",
+ "6f3a28b912714c6e931003549664bfa3",
+ "16d1283741404b7bb319094c992fce01",
+ "2a5bb0e818ab47be8cf6465988328503",
+ "2b3a2659b12244bd8548320320016dbf",
+ "0cd7efffbb3c4c4b972e63749f61ab97",
+ "5ca240f31e6b44e3882c5eb37cd5a309",
+ "5eb06edeb58e4930b1affef2a59eae81",
+ "a4e5789584564049b83df7c6c54a3e08",
+ "ff3a94b146a948b6907f5d80c7157f99",
+ "258b7c635c1045329d4669e48c46ccd5",
+ "6f68ed9889f54ad2ae8a3b95ac263a83",
+ "80366349d81e4dcc892db6cd56e384f3",
+ "c73055099c084dca996159e23e162d0b",
+ "977f799afaac4a55b2dc1cffa7d5b63b",
+ "41f3b32c2f6b4034ae7a3b9124e28bc7",
+ "a10d0a76010f4e508c65a9b69ebc5156",
+ "f8ef805b776145c3bfa9ba8d90972058",
+ "cc587493c33c4f118d1b1170f85be24c",
+ "e40d1c1ac9494b3bade9858324e7ffdf",
+ "d65b6b060d9845779299491ac5599c31",
+ "0f6907ebbc6242c8bde059cef1e1bd29",
+ "5bdfd87fc6cd4f9dabef7cfee29c8060",
+ "64f54d4a744a4627a07c3c0120276f3b",
+ "65b75b9b8bc143cf997796af68ff6668",
+ "d6fe74e4255444368f8f90a62157d869",
+ "4d468f96ec924681ad65eb671674b93e",
+ "ad7599de524549c48bf2d3124ad4b299",
+ "0546d04aae644dde846c58a4afb598a6",
+ "897b77a56c09479bb11d7f2a30997e55",
+ "81c3db71ac704280ad030072655f1537",
+ "042e091f75694c47aee761e760e76773",
+ "ef0a3c7a6f14460fb4da096928ae249e",
+ "07fb3a2c8315494e97b447e672dfae06",
+ "ec030fc3c346426f9abc3a89892258d3",
+ "e3fb3fc6afe04b3c9b7ac61809ce78fa",
+ "c3be9109d63c485d9c0ef4f9bc0f9218",
+ "12815f401eba44658caa7b2e490137a8",
+ "30e02aa2d0d241979369e598287f2639",
+ "dfd2a2649b8341ef913207526708aff1",
+ "4f1977d7e4824ef1a14b65f0f42bba10",
+ "c6164e05a1914ae48083db9ad7f4ef7c",
+ "813621384dc748b0ad06775e22761c0b",
+ "dc892a596f6942d7973c616c38f0eebb",
+ "c84cc07789be48aebb322c23d355289e",
+ "bed8726b8069434687c75452e21f19e5",
+ "16a188a0b06d45f980dcf3933509fe0a",
+ "60c1a0d765c14a1d888317e6a507e4ea",
+ "0077aedc3d174560bce924ee89e9c006",
+ "00321cce58884f6f9b3855a21fcd9187",
+ "fa864b41586f4a7aa56aeafd1d84eb75",
+ "3225603166b54e7aab766b9964a2f660",
+ "349eee9f56d64f0cba6fc24ff2c50c9b",
+ "7e5d3774060e4589aa65982da5ea4ef4",
+ "7c2485c6cdfe463da6fdb35982a1070d",
+ "ad1236893754446881e153adc9d5c962",
+ "daee63fd167e4441a32324b51b00ad2b",
+ "fe41858c6bd04c58840112b67c19a336",
+ "d262c82138024169b9f3aa034ca756fa",
+ "62e302ebdad64aada0ffe64ae1c873f3",
+ "bd1b0dfed6d34d16af33a4a58330f5ec",
+ "d07c8b97d3314f1c852e44bdd40f61ed",
+ "ebb69a2c3d0a4299a484698287b3087c",
+ "e5a82df528bb4e408797a3b6c2758f4a",
+ "f113ebd8c1c34806bea4dd7ed3035173"
+ ]
+ },
+ "id": "KQQhgK8FoDfF",
+ "outputId": "f69441d8-95f9-4885-c306-6c8709090ff6"
+ },
+ "outputs": [
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "b82aa8c57f7c422a9a9c90f333ed2a99",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "tokenizer_config.json: 0%| | 0.00/9.68k [00:00\u001b[39m\n",
+ "[2025-05-08 13:41:00,845] [DEBUG] [axolotl.utils.models.load_tokenizer:442] [PID:174] [RANK:0] BOS: None / None\u001b[39m\n",
+ "[2025-05-08 13:41:00,846] [DEBUG] [axolotl.utils.models.load_tokenizer:443] [PID:174] [RANK:0] PAD: 151643 / <|endoftext|>\u001b[39m\n",
+ "[2025-05-08 13:41:00,847] [DEBUG] [axolotl.utils.models.load_tokenizer:444] [PID:174] [RANK:0] UNK: None / None\u001b[39m\n",
+ "[2025-05-08 13:41:00,869] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:271] [PID:174] [RANK:0] Unable to find prepared dataset in last_run_prepared/97037817611d38b3a9c681753c3c4c95\u001b[39m\n",
+ "[2025-05-08 13:41:00,870] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:272] [PID:174] [RANK:0] Loading raw datasets...\u001b[39m\n",
+ "\u001b[33m[2025-05-08 13:41:00,870] [WARNING] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:274] [PID:174] [RANK:0] Processing datasets during training can lead to VRAM instability. Please pre-process your dataset.\u001b[39m\n",
+ "[2025-05-08 13:41:00,871] [INFO] [axolotl.utils.data.sft.load_tokenized_prepared_datasets:281] [PID:174] [RANK:0] No seed provided, using default seed of 42\u001b[39m\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "7cd0b85ebd204b7aba908417811ce4e0",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "train.jsonl: 0%| | 0.00/27.3M [00:00system\\n' }}\n",
+ " {%- if messages[0].role == 'system' %}\n",
+ " {{- messages[0].content + '\\n\\n' }}\n",
+ " {%- endif %}\n",
+ " {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within | Step | \n", + "Training Loss | \n", + "
|---|---|
| 1 | \n", + "1.092300 | \n", + "
| 2 | \n", + "1.554200 | \n", + "
| 3 | \n", + "1.041400 | \n", + "
| 4 | \n", + "1.733800 | \n", + "
| 5 | \n", + "1.430000 | \n", + "
| 6 | \n", + "1.258500 | \n", + "
| 7 | \n", + "1.343600 | \n", + "
| 8 | \n", + "1.101700 | \n", + "
| 9 | \n", + "1.086500 | \n", + "
| 10 | \n", + "0.813200 | \n", + "
| 11 | \n", + "0.689600 | \n", + "
| 12 | \n", + "0.826700 | \n", + "
| 13 | \n", + "1.541800 | \n", + "
| 14 | \n", + "0.948000 | \n", + "
| 15 | \n", + "1.357000 | \n", + "
| 16 | \n", + "1.085800 | \n", + "
| 17 | \n", + "1.516800 | \n", + "
| 18 | \n", + "1.146800 | \n", + "
| 19 | \n", + "0.834800 | \n", + "
| 20 | \n", + "0.968000 | \n", + "
| 21 | \n", + "1.388800 | \n", + "
| 22 | \n", + "1.511500 | \n", + "
| 23 | \n", + "1.338500 | \n", + "
| 24 | \n", + "1.206600 | \n", + "
| 25 | \n", + "1.504600 | \n", + "
"
+ ],
+ "text/plain": [
+ "
Copy a token from
Immediately click login after copying\nyour token or it might be stored in plain text in this notebook file. "
eos_token: ""
unk_token: ""
eos_token: ""
unk_token: ""
eos_token: ""
unk_token: "Generated:\n(no output)
"
+
+ def _style_for(i: int, tid: int) -> str:
+ if i in masked_positions:
+ if i < len(orig_ids) and tid == orig_ids[i]:
+ return "green"
+ if i < len(orig_ids):
+ return "red"
+ return "normal"
+ same = i < len(orig_ids) and tid == orig_ids[i]
+ return "dim" if same else "normal"
+
+ # Group contiguous spans by style to reduce HTML size
+ spans: list[tuple[str, int, int]] = []
+ if generated_ids:
+ cur = _style_for(0, generated_ids[0])
+ start = 0
+ for i in range(1, len(generated_ids)):
+ s = _style_for(i, generated_ids[i])
+ if s != cur:
+ spans.append((cur, start, i))
+ cur, start = s, i
+ spans.append((cur, start, len(generated_ids)))
+
+ html_parts = []
+ for style_name, a, b in spans:
+ txt = tokenizer.decode(generated_ids[a:b], skip_special_tokens=False)
+ if style_name == "green":
+ html_parts.append(f'{txt}')
+ elif style_name == "red":
+ html_parts.append(f'{txt}')
+ elif style_name == "dim":
+ html_parts.append(f'{txt}')
+ else:
+ html_parts.append(txt)
+
+ legend = (
+ ''
+ 'correct, '
+ 'incorrect, '
+ 'unchanged'
+ ""
+ )
+
+ return (
+ legend
+ + 'Generated:\n'
+ + "".join(html_parts)
+ + "
"
+ )
+
+
+def launch_diffusion_gradio_ui(
+ *,
+ model,
+ tokenizer,
+ cfg: DictDefault,
+ prompter_module=None,
+ chat_template_str: str | None = None,
+):
+ """Build and launch a simple Gradio UI for diffusion inference."""
+ with gr.Blocks(
+ title=cfg.get("gradio_title", "Axolotl Diffusion Interface")
+ ) as demo:
+ gr.Markdown(
+ """
+ ## Axolotl Diffusion Inference
+ - Mode "Random" masks tokens at a target ratio and fills them.
+ - Mode "Completion" appends N masked tokens at the end and fills them.
+ """
+ )
+
+ with gr.Row():
+ mode = gr.Radio(
+ choices=["random", "completion"],
+ value="random",
+ label="Mode",
+ )
+ mask_ratio = gr.Slider(
+ minimum=0.0,
+ maximum=1.0,
+ step=0.05,
+ value=0.4,
+ label="Mask ratio (random mode)",
+ interactive=True,
+ )
+ completion_tokens = gr.Number(
+ value=64,
+ precision=0,
+ label="Completion tokens (completion mode)",
+ interactive=True,
+ visible=False,
+ )
+
+ instruction = gr.Textbox(label="Instruction", lines=6)
+ run_btn = gr.Button("Generate")
+
+ masked_preview = gr.Textbox(label="Masked preview", lines=6)
+ html_out = gr.HTML(label="Generated")
+
+ def _toggle_controls(selected_mode: str):
+ return (
+ gr.update(visible=(selected_mode == "random")),
+ gr.update(visible=(selected_mode == "completion")),
+ )
+
+ mode.change(
+ _toggle_controls,
+ inputs=[mode],
+ outputs=[mask_ratio, completion_tokens],
+ )
+
+ def _gen(instruction_text: str, selected_mode: str, mratio: float, ctoks: int):
+ if not instruction_text:
+ return "", "Generated:\n(no output)
"
+
+ if prompter_module:
+ prompt: str = next(
+ prompter_module().build_prompt(
+ instruction=instruction_text.strip("\n")
+ )
+ )
+ else:
+ prompt = instruction_text.strip()
+
+ info = run_diffusion(
+ model=model,
+ tokenizer=tokenizer,
+ cfg=cfg,
+ prompt=prompt,
+ chat_template_str=chat_template_str,
+ mode=selected_mode,
+ target_mask_ratio=mratio if selected_mode == "random" else None,
+ completion_tokens=int(ctoks) if selected_mode == "completion" else 0,
+ )
+
+ masked_text = info.get("masked_text")
+ mask_ratio_val = info.get("mask_ratio")
+ generated_ids = info.get("generated_ids")
+ masked_positions = info.get("masked_positions") or set()
+ orig_ids = info.get("orig_ids") or []
+
+ preview = (
+ f"Masked ({mask_ratio_val:.1%}):\n{masked_text}"
+ if masked_text is not None and mask_ratio_val is not None
+ else ""
+ )
+ html = render_html(
+ generated_ids=generated_ids,
+ orig_ids=orig_ids,
+ masked_positions=masked_positions,
+ tokenizer=tokenizer,
+ )
+ return preview, html
+
+ run_btn.click(
+ _gen,
+ inputs=[instruction, mode, mask_ratio, completion_tokens],
+ outputs=[masked_preview, html_out],
+ )
+
+ demo.queue().launch(
+ show_api=False,
+ share=cfg.get("gradio_share", True),
+ server_name=cfg.get("gradio_server_name", "127.0.0.1"),
+ server_port=cfg.get("gradio_server_port", None),
+ )
diff --git a/src/axolotl/cli/utils/fetch.py b/src/axolotl/cli/utils/fetch.py
new file mode 100644
index 000000000..441b7f6f7
--- /dev/null
+++ b/src/axolotl/cli/utils/fetch.py
@@ -0,0 +1,142 @@
+"""Utilities for axolotl fetch CLI command."""
+
+import concurrent.futures
+import hashlib
+import json
+from pathlib import Path
+
+import click
+import requests
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def _download_file(
+ file_info: tuple, raw_base_url: str, dest_path: Path, dir_prefix: str
+) -> tuple[str, str]:
+ """
+ Download a single file and return its processing status.
+
+ Args:
+ file_info: Tuple of (file_path, remote_sha).
+ raw_base_url: Base URL for raw GitHub content.
+ dest_path: Local destination directory.
+ dir_prefix: Directory prefix to filter files.
+
+ Returns:
+ Tuple of (file_path, status) where status is 'new', 'updated', or 'unchanged'.
+ """
+ file_path, remote_sha = file_info
+ raw_url = f"{raw_base_url}/{file_path}"
+ dest_file = dest_path / file_path.split(dir_prefix)[-1]
+
+ # Check if file exists and needs updating
+ if dest_file.exists():
+ with open(dest_file, "rb") as file:
+ content = file.read()
+ # Calculate git blob SHA
+ blob = b"blob " + str(len(content)).encode() + b"\0" + content
+ local_sha = hashlib.sha1(blob, usedforsecurity=False).hexdigest()
+
+ if local_sha == remote_sha:
+ print(f"Skipping {file_path} (unchanged)")
+ return file_path, "unchanged"
+
+ print(f"Updating {file_path}")
+ status = "updated"
+ else:
+ print(f"Downloading {file_path}")
+ status = "new"
+
+ # Create directories if needed
+ dest_file.parent.mkdir(parents=True, exist_ok=True)
+
+ # Download and save file
+ try:
+ response = requests.get(raw_url, timeout=30)
+ response.raise_for_status()
+
+ with open(dest_file, "wb") as file:
+ file.write(response.content)
+
+ return file_path, status
+ except (requests.RequestException, IOError) as request_error:
+ print(f"Error downloading {file_path}: {str(request_error)}")
+ return file_path, "error"
+
+
+def fetch_from_github(
+ dir_prefix: str, dest_dir: str | None = None, max_workers: int = 5
+) -> None:
+ """
+ Sync files from a specific directory in the GitHub repository.
+ Only downloads files that don't exist locally or have changed.
+
+ Args:
+ dir_prefix: Directory prefix to filter files (e.g., 'examples/',
+ 'deepspeed_configs/').
+ dest_dir: Local destination directory.
+ max_workers: Maximum number of concurrent downloads.
+ """
+ api_url = "https://api.github.com/repos/axolotl-ai-cloud/axolotl/git/trees/main?recursive=1"
+ raw_base_url = "https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main"
+
+ # Get repository tree with timeout
+ response = requests.get(api_url, timeout=30)
+ response.raise_for_status()
+ tree = json.loads(response.text)
+
+ # Filter for files and get their SHA
+ files = {
+ item["path"]: item["sha"]
+ for item in tree["tree"]
+ if item["type"] == "blob" and item["path"].startswith(dir_prefix)
+ }
+
+ if not files:
+ raise click.ClickException(f"No files found in {dir_prefix}")
+
+ # Default destination directory is the last part of dir_prefix
+ default_dest = Path(dir_prefix.rstrip("/"))
+ dest_path = Path(dest_dir) if dest_dir else default_dest
+
+ # Keep track of processed files for summary
+ files_processed: dict[str, list[str]] = {
+ "new": [],
+ "updated": [],
+ "unchanged": [],
+ "error": [],
+ }
+
+ # Process files in parallel using ThreadPoolExecutor
+ with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+ future_to_file = {
+ executor.submit(
+ _download_file,
+ (file_path, remote_sha),
+ raw_base_url,
+ dest_path,
+ dir_prefix,
+ ): file_path
+ for file_path, remote_sha in files.items()
+ }
+
+ # Process completed tasks as they finish
+ for future in concurrent.futures.as_completed(future_to_file):
+ file_path = future_to_file[future]
+ try:
+ file_path, status = future.result()
+ files_processed[status].append(file_path)
+ except (requests.RequestException, IOError) as request_error:
+ print(f"Error processing {file_path}: {str(request_error)}")
+ files_processed["error"].append(file_path)
+
+ # Log summary
+ LOG.info("\nSync Summary:")
+ LOG.info(f"New files: {len(files_processed['new'])}")
+ LOG.info(f"Updated files: {len(files_processed['updated'])}")
+ LOG.info(f"Unchanged files: {len(files_processed['unchanged'])}")
+ if files_processed["error"]:
+ LOG.info(f"Failed files: {len(files_processed['error'])}")
diff --git a/src/axolotl/cli/utils/load.py b/src/axolotl/cli/utils/load.py
new file mode 100644
index 000000000..610a81306
--- /dev/null
+++ b/src/axolotl/cli/utils/load.py
@@ -0,0 +1,52 @@
+"""Utilities for model, tokenizer, etc. loading."""
+
+from typing import Any
+
+from transformers import (
+ PreTrainedModel,
+ PreTrainedTokenizer,
+ PreTrainedTokenizerFast,
+ ProcessorMixin,
+)
+
+from axolotl.loaders import load_processor, load_tokenizer
+from axolotl.loaders.model import ModelLoader
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def load_model_and_tokenizer(
+ *,
+ cfg: DictDefault,
+ inference: bool = False,
+) -> tuple[
+ PreTrainedModel,
+ PreTrainedTokenizer | PreTrainedTokenizerFast | Any,
+ ProcessorMixin | None,
+]:
+ """
+ Helper function for loading a model, tokenizer, and processor specified in the
+ given `axolotl` config.
+
+ Args:
+ cfg: Dictionary mapping `axolotl` config keys to values.
+ inference: Boolean denoting inference mode.
+
+ Returns:
+ Tuple of (PreTrainedModel, PreTrainedTokenizer, ProcessorMixin).
+ """
+ LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
+ tokenizer = load_tokenizer(cfg)
+
+ LOG.info("loading model...")
+ model_loader = ModelLoader(cfg, tokenizer, inference=inference)
+ model, _ = model_loader.load()
+
+ processor = None
+ if cfg.is_multimodal:
+ LOG.info("loading processor...")
+ processor = load_processor(cfg, tokenizer)
+
+ return model, tokenizer, processor
diff --git a/src/axolotl/cli/sweeps.py b/src/axolotl/cli/utils/sweeps.py
similarity index 92%
rename from src/axolotl/cli/sweeps.py
rename to src/axolotl/cli/utils/sweeps.py
index d21664964..2a0aa1367 100644
--- a/src/axolotl/cli/sweeps.py
+++ b/src/axolotl/cli/utils/sweeps.py
@@ -3,11 +3,12 @@
import random
from copy import deepcopy
from itertools import product
+from typing import Any
def generate_sweep_configs(
base_config: dict[str, list], sweeps_config: dict[str, list]
-) -> list[dict[str, list]]:
+) -> list[dict[str, Any]]:
"""
Recursively generates all possible configurations by applying sweeps to the base config.
@@ -48,7 +49,10 @@ def generate_sweep_configs(
new_config = {}
# new_config = deepcopy(base_config)
# Combine regular parameters with paired parameters
- full_combo = {**dict(zip(param_names, reg_combo)), **paired_set}
+ full_combo = {
+ **dict(zip(param_names, reg_combo, strict=False)),
+ **paired_set,
+ }
for param_name, param_value in full_combo.items():
new_config[param_name] = param_value
print(new_config)
@@ -57,7 +61,7 @@ def generate_sweep_configs(
# If no paired values, just use regular combinations
# new_config = deepcopy(base_config)
new_config = {}
- for param_name, param_value in zip(param_names, reg_combo):
+ for param_name, param_value in zip(param_names, reg_combo, strict=False):
new_config[param_name] = param_value
print(new_config)
all_combinations.append(new_config)
diff --git a/src/axolotl/cli/utils/train.py b/src/axolotl/cli/utils/train.py
new file mode 100644
index 000000000..6ce7d8df3
--- /dev/null
+++ b/src/axolotl/cli/utils/train.py
@@ -0,0 +1,225 @@
+"""Utilities for axolotl train CLI command."""
+
+import os
+import subprocess # nosec
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Iterator, Literal
+
+import yaml
+
+from axolotl.cli.utils.sweeps import generate_sweep_configs
+
+
+def _add_default_rdzv_args(launcher_args: list[str]) -> list[str]:
+ """
+ Add default RDZV arguments if rdzv_endpoint is set but rdzv_backend/rdzv_id are missing.
+
+ Args:
+ launcher_args: List of launcher arguments
+
+ Returns:
+ Updated launcher args with defaults added if needed
+ """
+ args = launcher_args.copy()
+
+ # Check if rdzv_endpoint is present
+ has_rdzv_endpoint = any("--rdzv_endpoint" in arg for arg in args)
+
+ if has_rdzv_endpoint:
+ # Check if rdzv_backend is already provided
+ has_rdzv_backend = any("--rdzv_backend" in arg for arg in args)
+ if not has_rdzv_backend:
+ args.extend(["--rdzv_backend", "c10d"])
+
+ # Check if rdzv_id is already provided
+ has_rdzv_id = any("--rdzv_id" in arg for arg in args)
+ if not has_rdzv_id:
+ import uuid
+
+ args.extend(["--rdzv_id", str(uuid.uuid4())[:8]])
+
+ return args
+
+
+def build_command(base_cmd: list[str], options: dict[str, Any]) -> list[str]:
+ """
+ Build command list from base command and options.
+
+ Args:
+ base_cmd: Command without options.
+ options: Options to parse and append to base command.
+
+ Returns:
+ List of strings giving shell command.
+ """
+ cmd = base_cmd.copy()
+
+ for key, value in options.items():
+ if value is None:
+ continue
+
+ key = key.replace("_", "-")
+ cmd.append(f"--{key}={value}")
+
+ return cmd
+
+
+def generate_config_files(config: str, sweep: str | None) -> Iterator[tuple[str, bool]]:
+ """
+ Generate list of configuration files to process. Yields a tuple of the configuration file name and a boolean indicating
+ whether this is a group of configurations (i.e., a sweep).
+
+ Args:
+ config: Base configuration file
+ sweep: Sweep configuration file
+ """
+
+ if not sweep:
+ yield config, False
+ return
+
+ # Load sweep and base configurations
+ with open(sweep, "r", encoding="utf-8") as fin:
+ sweep_config: dict[str, list] = yaml.safe_load(fin)
+ with open(config, "r", encoding="utf-8") as fin:
+ base_config: dict[str, list] = yaml.safe_load(fin)
+
+ # Generate all possible configurations
+ permutations = generate_sweep_configs(base_config, sweep_config)
+ is_group = len(permutations) > 1
+ base_output_dir = base_config.get("output_dir", "./model-out")
+ for idx, permutation in enumerate(permutations, start=1):
+ permutation_dir = Path(permutation.get("output_dir", base_output_dir))
+ permutation_id = f"sweep{idx:04d}"
+ permutation["output_dir"] = str(permutation_dir / permutation_id)
+
+ temp_file = tempfile.NamedTemporaryFile(
+ mode="w",
+ suffix=".yaml",
+ delete=False,
+ encoding="utf-8",
+ )
+ yaml.dump(permutation, temp_file)
+ temp_file.close()
+ yield temp_file.name, is_group
+
+
+def launch_training(
+ cfg_file: str,
+ launcher: Literal["accelerate", "torchrun", "python"] | None,
+ cloud: str | None,
+ kwargs: dict,
+ launcher_args: list[str] | None = None,
+ use_exec: bool = False,
+) -> None:
+ """Execute training with the given configuration."""
+ launcher_args = launcher_args or []
+
+ if cloud:
+ _launch_cloud_training(cloud, cfg_file, launcher, kwargs, launcher_args)
+ elif launcher:
+ if launcher == "accelerate":
+ _launch_accelerate_training(cfg_file, kwargs, launcher_args, use_exec)
+ elif launcher == "torchrun":
+ _launch_torchrun_training(cfg_file, kwargs, launcher_args, use_exec)
+ elif launcher == "python":
+ _launch_python_training(cfg_file, kwargs)
+ elif launcher is None:
+ # handle ray train launch
+ _launch_python_training(cfg_file, kwargs)
+
+
+def _launch_cloud_training(
+ cloud: str,
+ cfg_file: str,
+ launcher: Literal["accelerate", "torchrun", "python"] | None,
+ kwargs: dict,
+ launcher_args: list[str] | None = None,
+) -> None:
+ """Execute training via cloud launcher."""
+ from axolotl.cli.cloud import do_cli_train
+
+ launcher_args = launcher_args or []
+ cwd = os.getcwd() if launcher else None
+
+ do_cli_train(
+ cloud_config=cloud,
+ config=cfg_file,
+ launcher=launcher or "accelerate",
+ launcher_args=launcher_args,
+ cwd=cwd,
+ **kwargs,
+ )
+
+
+def _launch_accelerate_training(
+ cfg_file: str,
+ kwargs: dict,
+ launcher_args: list[str] | None = None,
+ use_exec: bool = False,
+) -> None:
+ """Execute training via accelerate launcher."""
+ launcher_args = launcher_args or []
+ internal_launcher_args = []
+
+ # Extract launcher-specific arguments from kwargs (legacy support)
+ if "main_process_port" in kwargs:
+ main_process_port = kwargs.pop("main_process_port")
+ internal_launcher_args.extend(["--main_process_port", str(main_process_port)])
+
+ if "num_processes" in kwargs:
+ num_processes = kwargs.pop("num_processes")
+ internal_launcher_args.extend(["--num_processes", str(num_processes)])
+
+ # Combine internal args with user-provided launcher args
+ all_launcher_args = internal_launcher_args + launcher_args
+
+ base_cmd = (
+ ["accelerate", "launch"] + all_launcher_args + ["-m", "axolotl.cli.train"]
+ )
+ if cfg_file:
+ base_cmd.append(cfg_file)
+
+ cmd = build_command(base_cmd, kwargs)
+ if use_exec:
+ # make sure to flush stdout and stderr before replacing the process
+ sys.stdout.flush()
+ sys.stderr.flush()
+ os.execvpe(cmd[0], cmd, os.environ) # nosec B606
+ else:
+ subprocess.run(cmd, check=True) # nosec B603
+
+
+def _launch_torchrun_training(
+ cfg_file: str,
+ kwargs: dict,
+ launcher_args: list[str] | None = None,
+ use_exec: bool = False,
+) -> None:
+ """Execute training via torchrun launcher."""
+ launcher_args = launcher_args or []
+
+ # Add default RDZV arguments if rdzv_endpoint is set
+ launcher_args = _add_default_rdzv_args(launcher_args)
+
+ base_cmd = ["torchrun"] + launcher_args + ["-m", "axolotl.cli.train"]
+ if cfg_file:
+ base_cmd.append(cfg_file)
+
+ cmd = build_command(base_cmd, kwargs)
+ if use_exec:
+ # make sure to flush stdout and stderr before replacing the process
+ sys.stdout.flush()
+ sys.stderr.flush()
+ os.execvpe(cmd[0], cmd, os.environ) # nosec B606
+ else:
+ subprocess.run(cmd, check=True) # nosec B603
+
+
+def _launch_python_training(cfg_file: str, kwargs: dict) -> None:
+ """Execute training via python launcher."""
+ from axolotl.cli.train import do_cli
+
+ do_cli(config=cfg_file, **kwargs)
diff --git a/src/axolotl/cli/vllm_serve.py b/src/axolotl/cli/vllm_serve.py
index 448b25a7e..ea454fc96 100644
--- a/src/axolotl/cli/vllm_serve.py
+++ b/src/axolotl/cli/vllm_serve.py
@@ -2,12 +2,10 @@
CLI to start the vllm server for online RL
"""
-import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Union
-import trl
from trl.scripts.vllm_serve import ScriptArguments
from axolotl.cli.config import load_cfg
@@ -37,16 +35,22 @@ def do_vllm_serve(
Returns:
process_id: the process id of the started VLLM server
"""
- patch_vllm_worker()
cfg = load_cfg(config)
model = cfg.base_model
serve_module = cli_args.get("serve_module", "trl.scripts.vllm_serve")
- vllm_serve_main = getattr(__import__(serve_module, fromlist=["main"]), "main")
+ vllm_serve_main = __import__(serve_module, fromlist=["main"]).main
+ tensor_parallel_size = 1
+ data_parallel_size = 1
- tensor_parallel_size = (
- cli_args.get("tensor_parallel_size") or cfg.vllm.tensor_parallel_size
- )
+ if cli_args.get("tensor_parallel_size") or cfg.vllm.tensor_parallel_size:
+ tensor_parallel_size = (
+ cli_args.get("tensor_parallel_size") or cfg.vllm.tensor_parallel_size
+ )
+ if cli_args.get("data_parallel_size") or cfg.vllm.data_parallel_size:
+ data_parallel_size = (
+ cli_args.get("data_parallel_size") or cfg.vllm.data_parallel_size
+ )
host = cli_args.get("host") or cfg.vllm.host
port = cli_args.get("port") or cfg.vllm.port
gpu_memory_utilization = (
@@ -64,10 +68,10 @@ def do_vllm_serve(
cli_args.get("enable_reasoning") or cfg.vllm.enable_reasoning or False
)
- # pylint: disable=unexpected-keyword-arg
vllm_script_args = AxolotlScriptArguments(
model=model,
tensor_parallel_size=tensor_parallel_size,
+ data_parallel_size=data_parallel_size,
host=host,
port=port,
gpu_memory_utilization=gpu_memory_utilization,
@@ -78,63 +82,3 @@ def do_vllm_serve(
enable_reasoning=enable_reasoning,
)
vllm_serve_main(vllm_script_args)
-
-
-def patch_vllm_worker():
- from multiprocessing.connection import Connection
-
- from vllm import LLM
-
- def llm_worker(
- script_args: AxolotlScriptArguments,
- data_parallel_rank: int,
- master_port: int,
- connection: Connection,
- ) -> None:
- # Set required environment variables for DP to work with vLLM
- os.environ["VLLM_DP_RANK"] = str(data_parallel_rank)
- os.environ["VLLM_DP_RANK_LOCAL"] = str(data_parallel_rank)
- os.environ["VLLM_DP_SIZE"] = str(script_args.data_parallel_size)
- os.environ["VLLM_DP_MASTER_PORT"] = str(master_port)
-
- llm = LLM(
- model=script_args.model,
- revision=script_args.revision,
- tensor_parallel_size=script_args.tensor_parallel_size,
- gpu_memory_utilization=script_args.gpu_memory_utilization,
- enforce_eager=script_args.enforce_eager,
- dtype=script_args.dtype,
- # Automatic Prefix Caching caches the KV cache of existing queries, so that a new query can
- # directly reuse the KV cache if it shares the same prefix with one of the existing queries.
- # This is particularly useful here because we generate completions from the same prompts.
- enable_prefix_caching=script_args.enable_prefix_caching,
- kv_cache_dtype=script_args.kv_cache_dtype,
- max_model_len=script_args.max_model_len,
- worker_extension_cls="trl.scripts.vllm_serve.WeightSyncWorkerExtension",
- enable_reasoning=script_args.enable_reasoning,
- reasoning_parser=script_args.reasoning_parser,
- )
-
- # Send ready signal to parent process
- connection.send({"status": "ready"})
-
- while True:
- # Wait for commands from the parent process
- try:
- command = connection.recv()
- except KeyboardInterrupt:
- llm.collective_rpc(method="close_communicator")
- break
-
- # Handle commands
- if command["type"] in ["call", "fire_and_forget"]:
- method_name = command["method"]
- args, kwargs = command.get("args", ()), command.get("kwargs", {})
- method = getattr(llm, method_name)
- result = method(*args, **kwargs)
- if command["type"] == "call":
- connection.send(result)
- elif command["type"] == "shutdown":
- break
-
- trl.scripts.vllm_serve.llm_worker = llm_worker
diff --git a/src/axolotl/common/architectures.py b/src/axolotl/common/architectures.py
index 2f77b613e..b754e56ba 100644
--- a/src/axolotl/common/architectures.py
+++ b/src/axolotl/common/architectures.py
@@ -13,4 +13,6 @@ MOE_ARCH_BLOCK = {
"qwen2_moe": "Qwen2MoeSparseMoeBlock",
"qwen3_moe": "Qwen3MoeSparseMoeBlock",
"deepseek_v2": "DeepseekV2MoE",
+ "gpt_oss": "GptOssDecoderLayer",
+ "lfm2_moe": "Lfm2MoeSparseMoeBlock",
}
diff --git a/src/axolotl/common/const.py b/src/axolotl/common/const.py
index fd34ad469..8aae06e99 100644
--- a/src/axolotl/common/const.py
+++ b/src/axolotl/common/const.py
@@ -1,5 +1,3 @@
-"""
-Various shared constants
-"""
+"""Various shared constants"""
DEFAULT_DATASET_PREPARED_PATH = "last_run_prepared"
diff --git a/src/axolotl/common/datasets.py b/src/axolotl/common/datasets.py
index 35d5472b0..c95ddb80e 100644
--- a/src/axolotl/common/datasets.py
+++ b/src/axolotl/common/datasets.py
@@ -3,16 +3,14 @@
import math
import random
from dataclasses import dataclass
-from typing import Optional, Union
from datasets import Dataset
-import axolotl.monkeypatch.data.batch_dataset_fetcher # pylint: disable=unused-import # noqa: F401
+import axolotl.monkeypatch.data.batch_dataset_fetcher # noqa: F401
from axolotl.cli.args import PreprocessCliArgs, TrainerCliArgs
from axolotl.loaders import load_processor, load_tokenizer
from axolotl.telemetry.errors import send_errors
-from axolotl.utils.data import prepare_dataset
-from axolotl.utils.data.rl import load_prepare_preference_datasets
+from axolotl.utils.data import prepare_datasets, prepare_preference_datasets
from axolotl.utils.dict import DictDefault
from axolotl.utils.logging import get_logger
from axolotl.utils.schemas.enums import RLType
@@ -31,16 +29,7 @@ class TrainDatasetMeta:
def sample_dataset(dataset: Dataset, num_samples: int) -> Dataset:
- """
- Randomly sample `num_samples` samples from `dataset`.
-
- Args:
- dataset: Dataset.
- num_samples: Number of samples to return.
-
- Returns:
- Random sample (with replacement) of examples in `dataset`.
- """
+ """Randomly sample `num_samples` samples with replacement from `dataset`."""
return dataset.select(
[random.randrange(0, len(dataset) - 1) for _ in range(num_samples)] # nosec
)
@@ -53,55 +42,50 @@ def load_datasets(
cli_args: PreprocessCliArgs | TrainerCliArgs | None = None,
debug: bool = False,
) -> TrainDatasetMeta:
- """
- Loads one or more training or evaluation datasets, calling
- `axolotl.utils.data.prepare_dataset`. Optionally, logs out debug information.
+ """Loads one or more training or evaluation datasets, calling
+ `axolotl.utils.data.prepare_datasets`. Optionally, logs out debug information.
Args:
cfg: Dictionary mapping `axolotl` config keys to values.
cli_args: Command-specific CLI arguments.
- debug: Whether to print out tokenization of sample
+ debug: Whether to print out tokenization of sample. This is duplicated in
+ `cfg` and `cli_args`, but is kept due to use in our Colab notebooks.
Returns:
Dataclass with fields for training and evaluation datasets and the computed
- `total_num_steps`.
+ `total_num_steps`.
"""
tokenizer = load_tokenizer(cfg)
processor = load_processor(cfg, tokenizer=tokenizer) if cfg.processor_type else None
- preprocess_iterable = (
- cli_args
- and hasattr(cli_args, "iterable")
- and cli_args.iterable is not None
- and cli_args.iterable
- )
- train_dataset, eval_dataset, total_num_steps, prompters = prepare_dataset(
+ train_dataset, eval_dataset, total_num_steps, prompters = prepare_datasets(
cfg,
tokenizer,
processor=processor,
- preprocess_iterable=preprocess_iterable,
)
- if ( # pylint: disable=too-many-boolean-expressions
- cli_args
- and (
- cli_args.debug
- or cfg.debug
- or cli_args.debug_text_only
- or int(cli_args.debug_num_examples) > 0
- )
- ) or debug:
+ if (
+ cfg.debug
+ or getattr(cli_args, "debug", False)
+ or getattr(cli_args, "debug_text_only", False)
+ or getattr(cli_args, "debug_num_examples", 0) > 0
+ or debug
+ ):
LOG.info("check_dataset_labels...")
num_examples = cli_args.debug_num_examples if cli_args else 1
text_only = cli_args.debug_text_only if cli_args else False
- train_samples = sample_dataset(train_dataset, num_examples)
- check_dataset_labels(
- train_samples,
- tokenizer,
- num_examples=num_examples,
- text_only=text_only,
- )
+ try:
+ train_samples = sample_dataset(train_dataset, num_examples)
+ check_dataset_labels(
+ train_samples,
+ tokenizer,
+ num_examples=num_examples,
+ text_only=text_only,
+ )
+ except AttributeError:
+ # can't sample iterable datasets
+ pass
LOG.info("printing prompters...")
for prompter in prompters:
@@ -116,13 +100,10 @@ def load_datasets(
@send_errors
def load_preference_datasets(
- *,
- cfg: DictDefault,
- cli_args: Union[PreprocessCliArgs, TrainerCliArgs],
+ *, cfg: DictDefault, cli_args: PreprocessCliArgs | TrainerCliArgs | None = None
) -> TrainDatasetMeta:
- """
- Loads one or more training or evaluation datasets for RL training using paired
- preference data, calling `axolotl.utils.data.rl.load_prepare_preference_datasets`.
+ """Loads one or more training or evaluation datasets for RL training using paired
+ preference data, calling `axolotl.utils.data.rl.prepare_preference_datasets`.
Optionally, logs out debug information.
Args:
@@ -133,23 +114,28 @@ def load_preference_datasets(
Dataclass with fields for training and evaluation datasets and the computed
`total_num_steps`.
"""
- train_dataset, eval_dataset = load_prepare_preference_datasets(cfg)
- total_num_steps: Optional[int] = int(
- math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
- )
- if cfg.rl is RLType.GRPO:
- total_num_steps = None
+ tokenizer = load_tokenizer(cfg)
+ train_dataset, eval_dataset = prepare_preference_datasets(cfg, tokenizer)
- if cli_args.debug or cfg.debug:
+ total_num_steps: int | None = None
+ if cfg.rl is not RLType.GRPO:
+ total_num_steps = int(
+ math.ceil(len(train_dataset) * cfg.num_epochs / cfg.batch_size)
+ )
+
+ if ((cli_args and cli_args.debug) or cfg.debug) and cfg.rl != RLType.ORPO:
LOG.info("check_dataset_labels...")
+ num_examples = cli_args.debug_num_examples if cli_args else 1
+ text_only = cli_args.debug_text_only if cli_args else False
+
tokenizer = load_tokenizer(cfg)
- train_samples = sample_dataset(train_dataset, cli_args.debug_num_examples)
+ train_samples = sample_dataset(train_dataset, num_examples)
check_dataset_labels(
- train_samples,
- tokenizer,
- num_examples=cli_args.debug_num_examples,
- text_only=cli_args.debug_text_only,
+ dataset=train_samples,
+ tokenizer=tokenizer,
+ num_examples=num_examples,
+ text_only=text_only,
rl_mode=True,
)
diff --git a/src/axolotl/convert.py b/src/axolotl/convert.py
index d1bdb34db..9e09b37dc 100644
--- a/src/axolotl/convert.py
+++ b/src/axolotl/convert.py
@@ -67,9 +67,7 @@ class JsonToJsonlConverter:
self.json_parser = json_parser
self.jsonl_serializer = jsonl_serializer
- def convert(
- self, input_file_path, output_file_path
- ): # pylint: disable=unused-argument
+ def convert(self, input_file_path, output_file_path):
content = self.file_reader.read(input_file_path)
data = self.json_parser.parse(content)
# data = [r for r in data if r["conversations"]] # vicuna cleaned has rows with empty conversations
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/__init__.py b/src/axolotl/core/attention/__init__.py
similarity index 100%
rename from src/axolotl/integrations/cut_cross_entropy/monkeypatch/__init__.py
rename to src/axolotl/core/attention/__init__.py
diff --git a/src/axolotl/core/attention/flex_block_mask.py b/src/axolotl/core/attention/flex_block_mask.py
new file mode 100644
index 000000000..37149983c
--- /dev/null
+++ b/src/axolotl/core/attention/flex_block_mask.py
@@ -0,0 +1,158 @@
+"""
+monkeypatch for flex + packing
+"""
+
+import sys
+from typing import Callable, Optional, Union
+
+import torch
+from torch.nn.attention.flex_attention import BlockMask
+from transformers import Cache, PretrainedConfig
+from transformers.masking_utils import (
+ ALL_MASK_ATTENTION_FUNCTIONS,
+ _preprocess_mask_arguments,
+ and_masks,
+ causal_mask_function,
+ or_masks,
+)
+from transformers.utils import is_torch_greater_or_equal
+
+_is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True)
+
+
+def create_causal_mask(
+ config: PretrainedConfig,
+ input_embeds: torch.Tensor,
+ attention_mask: torch.Tensor,
+ cache_position: torch.Tensor,
+ past_key_values: Optional[Cache],
+ or_mask_function: Optional[Callable] = None,
+ and_mask_function: Optional[Callable] = None,
+) -> Optional[Union[torch.Tensor, BlockMask]]:
+ """
+ Create a standard causal mask based on the attention implementation used (stored in the config). If `past_key_values`
+ has an HybridCache structure, this function will return the mask corresponding to one of the "full_attention" layers (to align
+ to what is needed in the `modeling_xxx.py` files).
+
+ Args:
+ config (`PretrainedConfig`):
+ The model config.
+ input_embeds (`torch.Tensor`):
+ The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the
+ batch size, query length and dtype.
+ attention_mask (`torch.Tensor`, optional):
+ The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).
+ It can also be an already prepared 4D mask, in which case it is returned as-is.
+ cache_position (`torch.Tensor`):
+ A tensor of shape (query_length,) indicating the current indices of the input sequence elements.
+ past_key_values (`Cache`, optional):
+ The past key values, if we use a cache.
+ or_mask_function (`Callable`, optional):
+ An optional mask function to combine with the causal mask function (by doing the union of both). This is
+ useful to easily overlay another mask on top of the causal one, for example for image tokens handling.
+ and_mask_function (`Callable`, optional):
+ An optional mask function to combine with the causal mask function (by doing the intersection of both). This is
+ useful to easily overlay another mask on top of the causal one, for example for image tokens handling.
+ """
+ # If we have an HybridCache structure, here we want to create the mask for the full layers
+ if (
+ past_key_values
+ and hasattr(past_key_values, "is_sliding")
+ and False in past_key_values.is_sliding
+ ):
+ layer_idx = past_key_values.is_sliding.index(False)
+ else:
+ layer_idx = 0
+
+ original_attention_mask = (
+ None
+ if attention_mask is None
+ else attention_mask.clone().to(cache_position.device)
+ )
+ early_exit, attention_mask, kv_length, kv_offset = _preprocess_mask_arguments(
+ config, input_embeds, attention_mask, cache_position, past_key_values, layer_idx
+ )
+ if early_exit:
+ return attention_mask
+
+ batch_size, total_seq_len = cache_position.shape
+ key_length = total_seq_len
+ document_ids = torch.nn.functional.pad(
+ original_attention_mask, value=0, pad=(0, key_length)
+ )
+
+ batch_size, dtype = input_embeds.shape[0], input_embeds.dtype
+ if attention_mask is not None:
+
+ def causal_doc_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+ """
+ Defines the logic of a block causal mask by combining both a standard causal mask
+ and a block diagonal document mask.
+ See :func:`~torchtune.modules.attention_utils.create_block_causal_mask`
+ for an illustration.
+ """
+ causal_mask_ = q_idx >= kv_idx # not valid when decoding
+ document_mask = (
+ document_ids[batch_idx, q_idx] == document_ids[batch_idx, kv_idx]
+ )
+ final_mask = causal_mask_ & document_mask
+ return final_mask
+
+ mask_factory_function = causal_doc_mask_mod
+ else:
+ mask_factory_function = causal_mask_function
+ mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]
+
+ # Do not allow skip if we are compiling (this is to match BC)
+ allow_is_causal_skip = (
+ not past_key_values.is_compileable if past_key_values is not None else True
+ )
+
+ # Allow slight deviations from causal mask
+ if or_mask_function is not None:
+ if not _is_torch_greater_or_equal_than_2_6:
+ raise ValueError(
+ "Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6"
+ )
+ mask_factory_function = or_masks(mask_factory_function, or_mask_function)
+ allow_is_causal_skip = False
+ if and_mask_function is not None:
+ if not _is_torch_greater_or_equal_than_2_6:
+ raise ValueError(
+ "Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6"
+ )
+ mask_factory_function = and_masks(mask_factory_function, and_mask_function)
+ allow_is_causal_skip = False
+
+ # We now create the mask
+ causal_mask = mask_interface(
+ batch_size=batch_size,
+ cache_position=cache_position,
+ kv_length=kv_length,
+ kv_offset=kv_offset,
+ mask_function=mask_factory_function,
+ attention_mask=attention_mask,
+ allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa
+ dtype=dtype, # Additional kwarg for eager
+ config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface
+ )
+ return causal_mask
+
+
+def patch_create_causal_mask(model_type):
+ import transformers.masking_utils
+
+ transformers.masking_utils.create_causal_mask = create_causal_mask
+
+ if model_type:
+ try:
+ # Dynamically import the module and attention class
+ module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+ module = __import__(module_path)
+ module.create_causal_mask = create_causal_mask
+ del sys.modules[module_path]
+ except (ImportError, AttributeError) as e:
+ raise ValueError(
+ f"Could not import attention class for model_type: {model_type}. "
+ f"Error: {str(e)}"
+ ) from e
diff --git a/src/axolotl/core/builders/base.py b/src/axolotl/core/builders/base.py
index 36f6739b6..639b0e4d7 100644
--- a/src/axolotl/core/builders/base.py
+++ b/src/axolotl/core/builders/base.py
@@ -24,10 +24,8 @@ from pathlib import Path
from typing import Any
import torch
-from transformers import (
- TrainerCallback,
-)
-from transformers.training_args import OptimizerNames
+from transformers import TrainerCallback
+from transformers.trainer_pt_utils import AcceleratorConfig
from axolotl.integrations.base import PluginManager
from axolotl.monkeypatch.trainer.lr import patch_trainer_get_lr
@@ -36,16 +34,17 @@ from axolotl.telemetry.manager import TelemetryManager
from axolotl.utils import is_comet_available, is_mlflow_available
from axolotl.utils.callbacks import (
GCCallback,
- GPUStatsCallback,
SaveAxolotlConfigtoWandBCallback,
+ SaveModelOnFirstStepCallback,
)
from axolotl.utils.callbacks.profiler import PytorchProfilerCallback
+from axolotl.utils.distributed import build_parallelism_config
from axolotl.utils.schemas.enums import CustomSupportedOptimizers
LOG = logging.getLogger(__name__)
with suppress(ImportError):
- import torch._dynamo # pylint: disable=ungrouped-imports
+ import torch._dynamo
class TrainerBuilderBase(abc.ABC):
@@ -114,13 +113,6 @@ class TrainerBuilderBase(abc.ABC):
plugin_manager.add_callbacks_pre_trainer(cfg=self.cfg, model=self.model)
)
- if self.cfg.profiler_steps:
- callbacks.append(
- PytorchProfilerCallback(
- steps_to_profile=self.cfg.profiler_steps,
- )
- )
-
if self.cfg.gc_steps:
callbacks.append(GCCallback(gc_steps=self.cfg.gc_steps))
@@ -144,8 +136,16 @@ class TrainerBuilderBase(abc.ABC):
callbacks.append(
SaveAxolotlConfigtoCometCallback(self.cfg.axolotl_config_path)
)
+ if self.cfg.save_first_step:
+ callbacks.append(SaveModelOnFirstStepCallback())
- callbacks.append(GPUStatsCallback(cfg=self.cfg))
+ if self.cfg.profiler_steps:
+ callbacks.append(
+ PytorchProfilerCallback(
+ steps_to_profile=self.cfg.profiler_steps,
+ profiler_steps_start=self.cfg.profiler_steps_start,
+ )
+ )
telemetry_manager = TelemetryManager.get_instance()
if telemetry_manager.enabled:
@@ -225,7 +225,9 @@ class TrainerBuilderBase(abc.ABC):
if self.cfg.bf16 == "full":
training_args_kwargs["bf16_full_eval"] = True
else:
- training_args_kwargs["bf16"] = self.cfg.bf16 or self.cfg.bfloat16
+ bf16 = self.cfg.bf16 or self.cfg.bfloat16
+ bf16 = bf16 if bf16 is not None else False
+ training_args_kwargs["bf16"] = bf16
def _configure_scheduler(self, training_args_kwargs: dict):
if self.cfg.lr_scheduler in ["one_cycle", "rex"]:
@@ -262,33 +264,30 @@ class TrainerBuilderBase(abc.ABC):
adam_kwargs["eps"] = training_args_kwargs.get("adam_epsilon")
if self.cfg.optimizer == "muon":
- from axolotl.contribs.mit.muon import ( # pylint: disable=no-name-in-module
+ from axolotl.contribs.mit.muon import (
MuonOptimizerFactory,
)
optimizer_cls = MuonOptimizerFactory
optimizer_kwargs.update(adam_kwargs)
+ elif self.cfg.optimizer == "dion":
+ from axolotl.contribs.mit.dion import (
+ DionOptimizerFactory,
+ )
+
+ optimizer_cls = DionOptimizerFactory
+ optimizer_kwargs["dion_lr"] = training_args_kwargs["dion_learning_rate"]
+ optimizer_kwargs["dion_mu"] = training_args_kwargs["dion_momentum"]
+ optimizer_kwargs.update(adam_kwargs)
+ _, device_mesh = build_parallelism_config(self.cfg)
+ if device_mesh is not None:
+ optimizer_kwargs["device_mesh"] = device_mesh
elif self.cfg.optimizer == "optimi_adamw":
from optimi import AdamW
optimizer_kwargs["foreach"] = False
optimizer_cls = AdamW
optimizer_kwargs.update(adam_kwargs)
- elif self.cfg.optimizer == "ao_adamw_4bit":
- # TODO remove 20250401
- from torchao.prototype.low_bit_optim import AdamW4bit
-
- optimizer_cls = AdamW4bit
- optimizer_kwargs.update(adam_kwargs)
-
- LOG.warning(
- f"`ao_adamw_4bit` will be deprecated soon. Please use `{OptimizerNames.ADAMW_TORCH_4BIT}` instead."
- )
- elif self.cfg.optimizer == "ao_adamw_8bit":
- from torchao.prototype.low_bit_optim import AdamW8bit
-
- optimizer_cls = AdamW8bit
- optimizer_kwargs.update(adam_kwargs)
elif self.cfg.optimizer == "ao_adamw_fp8":
from torchao.prototype.low_bit_optim import AdamWFp8
@@ -386,14 +385,16 @@ class TrainerBuilderBase(abc.ABC):
)
# eval_strategy and eval_steps
- if not self.eval_dataset or self.cfg.val_set_size == 0:
- # do not eval if no eval_dataset or val_set_size=0
+ if not self.eval_dataset and self.cfg.val_set_size == 0:
+ # do not eval if no eval_dataset and val_set_size=0
training_args_kwargs["eval_strategy"] = "no"
elif self.cfg.eval_steps:
training_args_kwargs["eval_strategy"] = "steps"
training_args_kwargs["eval_steps"] = self.cfg.eval_steps
+ training_args_kwargs["eval_on_start"] = True
elif self.cfg.eval_strategy:
training_args_kwargs["eval_strategy"] = self.cfg.eval_strategy
+ training_args_kwargs["eval_on_start"] = True
def _configure_reporting(self, training_args_kwargs: dict):
report_to = []
@@ -417,9 +418,8 @@ class TrainerBuilderBase(abc.ABC):
def _configure_torch_compile(self, training_args_kwargs: dict):
if self.cfg.torch_compile and getattr(torch, "_dynamo", None):
- torch._dynamo.config.suppress_errors = ( # pylint: disable=protected-access
- True
- )
+ torch._dynamo.config.suppress_errors = True
+ torch._dynamo.config.accumulated_cache_size_limit = 256
training_args_kwargs["torch_compile"] = self.cfg.torch_compile
if self.cfg.torch_compile_backend:
training_args_kwargs["torch_compile_backend"] = (
@@ -428,8 +428,20 @@ class TrainerBuilderBase(abc.ABC):
if self.cfg.torch_compile_mode:
training_args_kwargs["torch_compile_mode"] = self.cfg.torch_compile_mode
+ def _configure_accelerator_config(self, training_args_kwargs: dict):
+ if self.cfg.accelerator_config:
+ training_args_kwargs["accelerator_config"] = AcceleratorConfig(
+ **self.cfg.accelerator_config
+ )
+ else:
+ training_args_kwargs["accelerator_config"] = AcceleratorConfig()
+
def _configure_gradient_checkpointing(self, training_args_kwargs: dict):
- if self.cfg.gradient_checkpointing:
+ if self.cfg.activation_offloading is True:
+ # don't use the HF gradient checkpointing, manually wrap
+ training_args_kwargs["gradient_checkpointing"] = False
+ training_args_kwargs["activation_offloading"] = True
+ elif self.cfg.gradient_checkpointing is not None:
training_args_kwargs["gradient_checkpointing"] = (
self.cfg.gradient_checkpointing
)
@@ -482,17 +494,30 @@ class TrainerBuilderBase(abc.ABC):
"include_tokens_per_second",
"weight_decay",
"seed",
+ "dion_momentum",
+ "dion_rank_fraction",
+ "dion_rank_multiple_of",
+ "dataset_num_proc",
]:
if hasattr(self.cfg, arg) and getattr(self.cfg, arg) is not None:
training_args_kwargs[arg] = getattr(self.cfg, arg)
+ arg_map = {
+ "dion_learning_rate": "dion_lr",
+ }
+ for kwarg, cfg_arg in arg_map.items():
+ if hasattr(self.cfg, cfg_arg) and getattr(self.cfg, cfg_arg) is not None:
+ training_args_kwargs[kwarg] = getattr(self.cfg, cfg_arg)
+
training_args_kwargs["per_device_train_batch_size"] = self.cfg.micro_batch_size
+ training_args_kwargs["average_tokens_across_devices"] = False
if self.cfg.eval_batch_size:
training_args_kwargs["per_device_eval_batch_size"] = (
self.cfg.eval_batch_size
)
+ training_args_kwargs["include_tkps"] = self.cfg.include_tkps
training_args_kwargs["max_steps"] = self.cfg.max_steps or total_num_steps or -1
training_args_kwargs["num_train_epochs"] = self.cfg.num_epochs
@@ -500,10 +525,15 @@ class TrainerBuilderBase(abc.ABC):
if self.cfg.reward_model or self.cfg.rl:
training_args_kwargs["max_length"] = self.cfg.sequence_len
+ if self.cfg.fsdp_config or self.cfg.fsdp:
+ training_args_kwargs["fsdp_config"] = self.cfg.fsdp_config
+ training_args_kwargs["fsdp"] = self.cfg.fsdp if self.cfg.fsdp else True
+
self._configure_reporting(training_args_kwargs)
self._configure_hub_parameters(training_args_kwargs)
self._configure_scheduler(training_args_kwargs)
self._configure_optimizer(training_args_kwargs, trainer_kwargs)
self._configure_torch_compile(training_args_kwargs)
+ self._configure_accelerator_config(training_args_kwargs)
return training_args_kwargs, trainer_kwargs
diff --git a/src/axolotl/core/builders/causal.py b/src/axolotl/core/builders/causal.py
index 7a81616ba..820304230 100644
--- a/src/axolotl/core/builders/causal.py
+++ b/src/axolotl/core/builders/causal.py
@@ -10,6 +10,7 @@ import transformers
from transformers import (
DataCollatorWithFlattening,
EarlyStoppingCallback,
+ Trainer,
)
from trl.trainer.utils import RewardDataCollatorWithPadding
@@ -19,12 +20,6 @@ from axolotl.core.trainers import (
AxolotlPRMTrainer,
AxolotlRewardTrainer,
AxolotlTrainer,
- ReLoRATrainer,
-)
-from axolotl.core.training_args import (
- AxolotlPRMConfig,
- AxolotlRewardConfig,
- AxolotlTrainingArguments,
)
from axolotl.integrations.base import PluginManager
from axolotl.monkeypatch.multipack import SUPPORTED_MULTIPACK_MODEL_TYPES
@@ -32,9 +27,7 @@ from axolotl.monkeypatch.relora import ReLoRACallback
from axolotl.processing_strategies import get_processing_strategy
from axolotl.utils import is_comet_available, is_mlflow_available
from axolotl.utils.callbacks import (
- EvalFirstStepCallback,
LossWatchDogCallback,
- SaveBetterTransformerModelCallback,
bench_eval_callback_factory,
causal_lm_bench_eval_callback_factory,
colab_inference_post_train_callback,
@@ -42,6 +35,7 @@ from axolotl.utils.callbacks import (
)
from axolotl.utils.callbacks.lisa import lisa_callback_factory
from axolotl.utils.callbacks.qat import QATCallback
+from axolotl.utils.callbacks.tokens_per_second import TokensPerSecondCallback
from axolotl.utils.chat_templates import get_chat_template_from_config
from axolotl.utils.collators import (
BatchSamplerDataCollatorForSeq2Seq,
@@ -50,6 +44,7 @@ from axolotl.utils.collators import (
V2BatchSamplerDataCollatorForSeq2Seq,
)
from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator
+from axolotl.utils.import_helper import get_cls_from_module_str
from axolotl.utils.logging import get_logger
LOG = get_logger(__name__)
@@ -63,17 +58,10 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
def get_callbacks(self):
callbacks = super().get_callbacks()
- callbacks.append(EvalFirstStepCallback())
- if self.cfg.relora_steps:
+ if self.cfg.relora:
callbacks.append(ReLoRACallback(self.cfg))
- if (
- hasattr(self.model, "use_bettertransformer")
- and self.model.use_bettertransformer is True
- ):
- callbacks.append(SaveBetterTransformerModelCallback())
-
# TODO: check if can move to base class
if self.cfg.loss_watchdog_threshold is not None:
callbacks.append(LossWatchDogCallback(self.cfg))
@@ -81,6 +69,12 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
if self.cfg.qat:
callbacks.append(QATCallback(self.cfg.qat))
+ if self.cfg.include_tkps:
+ callbacks.append(
+ TokensPerSecondCallback(
+ self.cfg.tensor_parallel_size, self.cfg.context_parallel_size
+ )
+ )
return callbacks
def get_post_trainer_create_callbacks(self, trainer):
@@ -130,33 +124,44 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
return callbacks
def _get_trainer_cls(self):
+ """
+ Gets the trainer class for the given configuration.
+ """
if self.cfg.plugins:
plugin_manager = PluginManager.get_instance()
trainer_cls = plugin_manager.get_trainer_cls(self.cfg)
if trainer_cls:
return trainer_cls
- if self.cfg.relora_steps:
- return ReLoRATrainer
if self.cfg.model_config_type == "mamba":
return AxolotlMambaTrainer
if self.cfg.reward_model:
return AxolotlRewardTrainer
if self.cfg.process_reward_model:
return AxolotlPRMTrainer
+
+ if self.cfg.trainer_cls:
+ # override the trainer cls
+ try:
+ trainer_cls = get_cls_from_module_str(self.cfg.trainer_cls)
+ LOG.debug(f"Using custom trainer class: {self.cfg.trainer_cls}")
+ return trainer_cls
+ except (ImportError, AttributeError, ValueError) as e:
+ raise ValueError(
+ f"Failed to load custom trainer class '{self.cfg.trainer_cls}': {e}"
+ ) from e
+
return AxolotlTrainer
def build(self, total_num_steps):
+ from axolotl.core.training_args import (
+ AxolotlPRMConfig,
+ AxolotlRewardConfig,
+ AxolotlTrainingArguments,
+ )
+
training_arguments_kwargs, trainer_kwargs = self._set_base_training_args(
total_num_steps
)
-
- if self.cfg.fsdp:
- training_arguments_kwargs["fsdp"] = self.cfg.fsdp
- if self.cfg.fsdp_config:
- training_arguments_kwargs["fsdp_config"] = {
- k.lstrip("fsdp_"): v for k, v in dict(self.cfg.fsdp_config).items()
- }
-
if self.cfg.adapter == "qlora":
training_arguments_kwargs["qlora"] = True
@@ -243,14 +248,27 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
training_arguments_kwargs["curriculum_sampling"] = self.cfg.curriculum_sampling
training_arguments_kwargs["sample_packing"] = bool(self.cfg.sample_packing)
+ training_arguments_kwargs["sample_packing_drop_attention_mask"] = bool(
+ self.cfg.flash_attention
+ or self.cfg.xformers_attention
+ or self.cfg.flex_attention
+ )
training_arguments_kwargs["multipack_real_batches"] = (
self.cfg.multipack_real_batches
if self.cfg.multipack_real_batches is not None
- else not self.cfg.flash_attention
+ else not (
+ self.cfg.flash_attention
+ or self.cfg.flex_attention
+ or self.cfg.xformers_attention
+ )
)
training_arguments_kwargs["eval_sample_packing"] = bool(
self.cfg.eval_sample_packing
)
+ if self.cfg.sample_packing_sequentially is not None:
+ training_arguments_kwargs["sample_packing_sequentially"] = (
+ self.cfg.sample_packing_sequentially
+ )
if self.cfg.sample_packing_bin_size is not None:
training_arguments_kwargs["sample_packing_bin_size"] = (
self.cfg.sample_packing_bin_size
@@ -264,20 +282,25 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
self.cfg.sample_packing_eff_est
)
- if self.cfg.relora_steps:
- training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
- training_arguments_kwargs["relora_warmup_steps"] = (
- self.cfg.relora_warmup_steps
- )
- if self.cfg.relora_anneal_steps:
- training_arguments_kwargs["relora_anneal_steps"] = (
- self.cfg.relora_anneal_steps
- )
+ if self.cfg.relora and self.cfg.jagged_restart_steps:
if self.cfg.relora_prune_ratio:
training_arguments_kwargs["relora_prune_ratio"] = (
self.cfg.relora_prune_ratio
)
+ if self.cfg.jagged_restart_steps:
+ training_arguments_kwargs["jagged_restart_steps"] = (
+ self.cfg.jagged_restart_steps
+ )
+ if self.cfg.jagged_restart_warmup_steps:
+ training_arguments_kwargs["jagged_restart_warmup_steps"] = (
+ self.cfg.jagged_restart_warmup_steps
+ )
+ if self.cfg.jagged_restart_anneal_steps:
+ training_arguments_kwargs["jagged_restart_anneal_steps"] = (
+ self.cfg.jagged_restart_anneal_steps
+ )
+
if self.cfg.lisa_step_interval and self.cfg.lisa_n_layers:
training_arguments_kwargs["lisa_n_layers"] = self.cfg.lisa_n_layers
training_arguments_kwargs["lisa_step_interval"] = (
@@ -303,48 +326,37 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
self.cfg.neftune_noise_alpha
)
- if self.cfg.accelerator_config:
- training_arguments_kwargs["accelerator_config"] = (
- self.cfg.accelerator_config
- )
-
if self.cfg.image_size:
training_arguments_kwargs["image_size"] = self.cfg.image_size
if self.cfg.image_resize_algorithm:
training_arguments_kwargs["image_resize_algorithm"] = (
self.cfg.image_resize_algorithm
)
- if self.cfg.kd_ce_alpha is not None:
- training_arguments_kwargs["kd_ce_alpha"] = self.cfg.kd_ce_alpha
- if self.cfg.kd_alpha is not None:
- training_arguments_kwargs["kd_alpha"] = self.cfg.kd_alpha
- if self.cfg.kd_temperature is not None:
- training_arguments_kwargs["kd_temperature"] = self.cfg.kd_temperature
- if self.cfg.kd_zscore_base_temp is not None:
- training_arguments_kwargs["kd_zscore_base_temp"] = (
- self.cfg.kd_zscore_base_temp
- )
- if self.cfg.kd_top_k_before_softmax is not None:
- training_arguments_kwargs["kd_top_k_before_softmax"] = (
- self.cfg.kd_top_k_before_softmax
- )
+
+ if self.cfg.plugins:
+ plugin_manager = PluginManager.get_instance()
+ plugin_training_args = plugin_manager.get_training_args(self.cfg)
+ if plugin_training_args:
+ training_arguments_kwargs.update(plugin_training_args)
if self.cfg.reward_model:
training_args_cls = AxolotlRewardConfig
+ if self.cfg.center_rewards_coefficient is not None:
+ training_arguments_kwargs["center_rewards_coefficient"] = (
+ self.cfg.center_rewards_coefficient
+ )
elif self.cfg.process_reward_model:
training_args_cls = AxolotlPRMConfig
else:
training_args_cls = AxolotlTrainingArguments
- training_args = training_args_cls( # pylint: disable=unexpected-keyword-arg
+ training_args = training_args_cls(
**training_arguments_kwargs,
)
training_args = self.hook_post_create_training_args(training_args)
# unset run_name so wandb sets up experiment names
if self.cfg.use_wandb and training_args.run_name == training_args.output_dir:
- training_args.run_name = ( # pylint: disable=attribute-defined-outside-init
- None
- )
+ training_args.run_name = None
data_collator_kwargs = {
"padding": True, # True/"longest" is the default
@@ -354,7 +366,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
data_collator_kwargs["pad_to_multiple_of"] = multiple * math.ceil(
self.cfg.sequence_len / multiple
)
- else:
+ elif self.cfg.pad_to_sequence_len is None:
# A100 is best at 64, while others at 8. Let's use the larger so we don't have to check
# https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
data_collator_kwargs["pad_to_multiple_of"] = multiple
@@ -376,12 +388,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
**data_collator_kwargs,
)
sig = inspect.signature(trainer_cls)
- if "processing_class" in sig.parameters:
+ if "processing_class" in sig.parameters or issubclass(trainer_cls, Trainer):
trainer_kwargs["processing_class"] = self.tokenizer
elif "tokenizer" in sig.parameters:
trainer_kwargs["tokenizer"] = self.tokenizer
+
if (
- not (trainer_cls in [AxolotlRewardTrainer, AxolotlPRMTrainer])
+ trainer_cls not in [AxolotlRewardTrainer, AxolotlPRMTrainer]
and self.cfg.datasets is not None
):
trainer_kwargs["dataset_tags"] = [
@@ -397,6 +410,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
**trainer_kwargs,
)
trainer = self.hook_post_create_trainer(trainer)
+ # if the trainer has the `axolotl_cfg` property, set it
+ if hasattr(trainer, "axolotl_cfg"):
+ trainer.axolotl_cfg = self.cfg
for callback in self.get_post_trainer_create_callbacks(trainer):
trainer.add_callback(callback)
@@ -408,7 +424,10 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
return trainer
def build_collator(
- self, training_args: AxolotlTrainingArguments, is_eval=False, **kwargs
+ self,
+ training_args, # type: "AxolotlTrainingArguments" # type: ignore
+ is_eval=False,
+ **kwargs,
):
if training_args.pretraining:
if (
@@ -416,7 +435,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
or self.cfg.micro_batch_size > 1
):
return DataCollatorForSeq2Seq(self.tokenizer, **kwargs)
- return None
+ if not (self.cfg.sample_packing and self.cfg.pretrain_multipack_attn):
+ return None
if self.cfg.model_config_type == "mamba":
return MambaDataCollator(tokenizer=self.tokenizer)
@@ -437,7 +457,19 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
]
]
collator_args = [self.tokenizer]
- if self.cfg.reward_model:
+
+ collator_cls_and_kwargs = None
+ if self.cfg.plugins:
+ plugin_manager = PluginManager.get_instance()
+ collator_cls_and_kwargs = plugin_manager.get_collator_cls_and_kwargs(
+ self.cfg, is_eval=is_eval
+ )
+
+ if collator_cls_and_kwargs:
+ collator = collator_cls_and_kwargs[0]
+ if kwargs and isinstance(kwargs, dict):
+ kwargs.update(collator_cls_and_kwargs[1])
+ elif self.cfg.reward_model:
collator = RewardDataCollatorWithPadding
elif use_batch_sampler_collator:
# Use V2BatchSamplerDataCollatorForSeq2Seq for flex attention,
@@ -468,16 +500,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
collator_args.pop(0)
kwargs.pop("pad_to_multiple_of", None)
kwargs.pop("padding", None)
- elif self.cfg.kd_trainer:
- from axolotl.integrations.kd.collator import (
- DataCollatorForKD,
- KDBatchSamplerDataCollatorForSeq2Seq,
- )
-
- if self.cfg.sample_packing:
- collator = KDBatchSamplerDataCollatorForSeq2Seq
- else:
- collator = DataCollatorForKD
else:
collator = DataCollatorForSeq2Seq
diff --git a/src/axolotl/core/builders/rl.py b/src/axolotl/core/builders/rl.py
index 14dbfa715..0ceb80008 100644
--- a/src/axolotl/core/builders/rl.py
+++ b/src/axolotl/core/builders/rl.py
@@ -12,13 +12,10 @@ from axolotl.core.trainers import (
from axolotl.core.trainers.dpo import DPOStrategy
from axolotl.core.trainers.dpo.args import AxolotlDPOConfig
from axolotl.core.trainers.grpo import GRPOStrategy
-from axolotl.core.training_args import (
- AxolotlCPOConfig,
- AxolotlKTOConfig,
- AxolotlORPOConfig,
-)
from axolotl.integrations.base import PluginManager
from axolotl.loaders.utils import ensure_dtype
+from axolotl.utils.callbacks.qat import QATCallback
+from axolotl.utils.import_helper import get_cls_from_module_str
from axolotl.utils.logging import get_logger
from axolotl.utils.schemas.enums import RLType
@@ -31,6 +28,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
def get_callbacks(self):
callbacks = super().get_callbacks()
+ if self.cfg.qat:
+ callbacks.append(QATCallback(self.cfg.qat))
+
return callbacks
def get_post_trainer_create_callbacks(self, trainer):
@@ -54,7 +54,7 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
if self.cfg.rl is RLType.GRPO:
trainer_cls = GRPOStrategy.get_trainer_class(
- sequence_parallel=self.cfg.sequence_parallel_degree > 1
+ sequence_parallel=self.cfg.context_parallel_size > 1
)
trainer_cls_args.extend(GRPOStrategy.set_trainer_args(self.cfg))
@@ -73,12 +73,28 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
else:
raise ValueError(f"Unsupported RL: {self.cfg.rl}")
+ if self.cfg.trainer_cls:
+ # override the trainer cls
+ try:
+ trainer_cls = get_cls_from_module_str(self.cfg.trainer_cls)
+ LOG.debug(f"Using custom trainer class: {self.cfg.trainer_cls}")
+ except (ImportError, AttributeError, ValueError) as e:
+ raise ValueError(
+ f"Failed to load custom trainer class '{self.cfg.trainer_cls}': {e}"
+ ) from e
+
return trainer_cls, trainer_cls_args
def _build_training_arguments(self, total_num_steps):
"""
Returns training_args and trainer_kwargs
"""
+ from axolotl.core.training_args import (
+ AxolotlCPOConfig,
+ AxolotlKTOConfig,
+ AxolotlORPOConfig,
+ )
+
training_args_kwargs, trainer_kwargs = self._set_base_training_args(
total_num_steps=total_num_steps
)
@@ -90,10 +106,6 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
else:
training_args_kwargs["remove_unused_columns"] = False
- # only rlhf
- if self.cfg.dataset_processes:
- training_args_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
-
if self.cfg.trl and self.cfg.trl.beta is not None:
training_args_kwargs["beta"] = self.cfg.trl.beta
elif self.cfg.rl_beta is not None:
@@ -108,6 +120,11 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
if self.cfg.use_wandb:
training_args_kwargs["run_name"] = self.cfg.wandb_name
+ if self.cfg.max_prompt_len:
+ training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
+ else:
+ training_args_kwargs["max_prompt_length"] = self.cfg.sequence_len
+
training_args_cls = None
blocklist_args_kwargs = []
if self.cfg.rl is RLType.SIMPO:
@@ -117,10 +134,16 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
if self.cfg.cpo_alpha is not None:
training_args_kwargs["cpo_alpha"] = self.cfg.cpo_alpha
+ # Handle when max_prompt_length == max_length from defaults
+ # CPOTrainer requires strictly less than
+ if (
+ training_args_kwargs["max_prompt_length"]
+ == training_args_kwargs["max_length"]
+ ):
+ training_args_kwargs["max_prompt_length"] -= 1
+
elif self.cfg.rl is RLType.ORPO:
training_args_cls = AxolotlORPOConfig
- if self.cfg.max_prompt_len:
- training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
elif self.cfg.rl is RLType.KTO:
training_args_cls = AxolotlKTOConfig
@@ -132,9 +155,6 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
self.cfg.kto_undesirable_weight or 1.0
)
- if self.cfg.max_prompt_len:
- training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
-
elif self.cfg.rl is RLType.GRPO:
training_args_cls = GRPOStrategy.get_training_args_class()
training_args_kwargs.update(GRPOStrategy.set_training_args_kwargs(self.cfg))
@@ -142,22 +162,7 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
elif self.cfg.rl in [RLType.DPO, RLType.IPO]:
training_args_cls = AxolotlDPOConfig
- if self.cfg.rl is RLType.IPO:
- training_args_kwargs["loss_type"] = "ipo"
-
- # Not compatible with IPO
- if self.cfg.rl is RLType.DPO and self.cfg.dpo_label_smoothing:
- training_args_kwargs["label_smoothing"] = self.cfg.dpo_label_smoothing
-
- training_args_kwargs["max_completion_length"] = None
- training_args_kwargs["max_prompt_length"] = self.cfg.sequence_len
- training_args_kwargs["generate_during_eval"] = self.cfg.use_wandb
- if self.cfg.dpo_use_weighting is not None:
- training_args_kwargs["use_weighting"] = self.cfg.dpo_use_weighting
- if self.cfg.dpo_use_logits_to_keep is not None:
- training_args_kwargs["use_logits_to_keep"] = (
- self.cfg.dpo_use_logits_to_keep
- )
+ training_args_kwargs.update(DPOStrategy.set_training_args_kwargs(self.cfg))
else:
raise ValueError(f"Unsupported RL: {self.cfg.rl}")
@@ -165,16 +170,20 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
if blocklist_key in training_args_kwargs:
del training_args_kwargs[blocklist_key]
- training_args = training_args_cls( # pylint: disable=unexpected-keyword-arg
+ if self.cfg.plugins:
+ plugin_manager = PluginManager.get_instance()
+ plugin_training_args = plugin_manager.get_training_args(self.cfg)
+ if plugin_training_args:
+ training_args_kwargs.update(plugin_training_args)
+
+ training_args = training_args_cls(
logging_first_step=True,
**training_args_kwargs,
)
# unset run_name so wandb sets up experiment names
if self.cfg.use_wandb and training_args.run_name == training_args.output_dir:
- training_args.run_name = ( # pylint: disable=attribute-defined-outside-init
- None
- )
+ training_args.run_name = None
return training_args, trainer_kwargs
@@ -216,7 +225,7 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
callbacks=self.get_callbacks(),
**trainer_kwargs,
)
- if self.cfg.fsdp:
+ if self.cfg.fsdp_config or self.cfg.fsdp:
ensure_dtype(trainer.model, dtype=self.cfg.torch_dtype)
if self.cfg.rl in [RLType.DPO, RLType.IPO] and trainer.ref_model:
ensure_dtype(trainer.ref_model, dtype=self.cfg.torch_dtype)
@@ -226,21 +235,3 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
trainer.add_callback(callback)
return trainer
-
-
-class HFPPOTrainerBuilder(TrainerBuilderBase):
- """
- HF Factory class for PPO Trainer
- """
-
- def get_callbacks(self):
- callbacks = super().get_callbacks()
- return callbacks
-
- def get_post_trainer_create_callbacks(self, trainer):
- callbacks = super().get_post_trainer_create_callbacks(trainer=trainer)
- return callbacks
-
- def build(self, total_num_steps):
- # TODO: build PPOConfig
- raise NotImplementedError("PPO trainer builder is not implemented yet.")
diff --git a/src/axolotl/core/chat/format/chatml.py b/src/axolotl/core/chat/format/chatml.py
index 04c398fe8..deb8a9997 100644
--- a/src/axolotl/core/chat/format/chatml.py
+++ b/src/axolotl/core/chat/format/chatml.py
@@ -10,7 +10,7 @@ from .shared import wrap_tools
def format_message(
message: Messages,
- message_index: Optional[int] = None, # pylint: disable=unused-argument
+ message_index: Optional[int] = None,
) -> Messages:
if message.is_chat_formatted:
return message
diff --git a/src/axolotl/core/chat/messages.py b/src/axolotl/core/chat/messages.py
index 923b177c1..912a12ca1 100644
--- a/src/axolotl/core/chat/messages.py
+++ b/src/axolotl/core/chat/messages.py
@@ -15,11 +15,11 @@ class MessageRoles(str, Enum):
Message roles for the system, user, assistant, and tools
"""
- system = "system" # pylint: disable=invalid-name
- user = "user" # pylint: disable=invalid-name
- assistant = "assistant" # pylint: disable=invalid-name
- tool = "tool" # pylint: disable=invalid-name
- ipython = ( # pylint: disable=invalid-name
+ system = "system"
+ user = "user"
+ assistant = "assistant"
+ tool = "tool"
+ ipython = (
# for responses from builtin tools
"ipython"
)
@@ -30,12 +30,12 @@ class MessageContentTypes(str, Enum):
Message content types for text, image, audio, tool calls, and tool responses
"""
- special_token = "special_token" # pylint: disable=invalid-name # nosec B105
- text = "text" # pylint: disable=invalid-name
- image = "image" # pylint: disable=invalid-name
- audio = "audio" # pylint: disable=invalid-name
- tool_call = "tool_call" # pylint: disable=invalid-name # to differentiate regular responses from tool calls from the assistant
- tool_response = "tool_response" # pylint: disable=invalid-name
+ special_token = "special_token" # nosec B105
+ text = "text"
+ image = "image"
+ audio = "audio"
+ tool_call = "tool_call"
+ tool_response = "tool_response"
class SpecialToken(str, Enum):
@@ -43,8 +43,8 @@ class SpecialToken(str, Enum):
Special tokens for beginning of string and end of string
"""
- bos_token = "bos_token" # pylint: disable=invalid-name # nosec B105
- eos_token = "eos_token" # pylint: disable=invalid-name # nosec B105
+ bos_token = "bos_token" # nosec B105
+ eos_token = "eos_token" # nosec B105
class ToolCallFunction(BaseModel):
@@ -73,7 +73,7 @@ class ToolCallContents(BaseModel):
name: str
arguments: dict[str, Union[str, int]]
- id: Optional[str] = None # pylint: disable=invalid-name
+ id: Optional[str] = None
def __str__(self) -> str:
data = {"name": self.name, "arguments": self.arguments}
@@ -89,7 +89,7 @@ class ToolResponseContents(BaseModel):
name: str
content: Union[str, dict[str, Union[str, int, float]]]
- id: Optional[str] = None # pylint: disable=invalid-name
+ id: Optional[str] = None
def __str__(self) -> str:
data = {"name": self.name, "content": self.content}
diff --git a/src/axolotl/core/datasets/chat.py b/src/axolotl/core/datasets/chat.py
index 724f12866..a4dc300d9 100644
--- a/src/axolotl/core/datasets/chat.py
+++ b/src/axolotl/core/datasets/chat.py
@@ -2,7 +2,6 @@
chat dataset module
"""
-import os
from typing import Callable, Optional, Union
from datasets import Dataset
@@ -41,14 +40,10 @@ class TokenizedChatDataset(Dataset):
)
return ex.tokenized(model_transform)
- process_or_cpu_count: int = (
- process_count or os.cpu_count() # type: ignore[assignment]
- )
- num_proc = min(32, process_or_cpu_count)
features = data.features.keys()
tokenized_data = data.map(
map_fn,
- num_proc=num_proc,
+ num_proc=process_count,
keep_in_memory=keep_in_memory,
remove_columns=features,
desc="Tokenizing Chats",
diff --git a/src/axolotl/core/datasets/transforms/chat_builder.py b/src/axolotl/core/datasets/transforms/chat_builder.py
index 692fe3ebb..0de0ecb40 100644
--- a/src/axolotl/core/datasets/transforms/chat_builder.py
+++ b/src/axolotl/core/datasets/transforms/chat_builder.py
@@ -1,23 +1,17 @@
"""
-This module contains a function that builds a transform that takes a row from the dataset and converts it to a Chat.
+This module contains a function that builds a transform that takes a row from the
+dataset and converts it to a Chat.
"""
-from typing import Any, Mapping, Union
+from typing import Any, Mapping
-def chat_message_transform_builder( # pylint: disable=dangerous-default-value
+def chat_message_transform_builder(
train_on_inputs=False,
- conversations_field: str = "conversations",
- message_field_role: Union[str, list[str]] = ["role", "from"], # commonly "role"
- message_field_content: Union[str, list[str]] = [
- "value",
- "text",
- "content",
- ], # commonly "content"
- message_field_training: Union[str, list[str]] = [
- "train",
- "weight",
- ], # commonly "weight"
+ conversations_field: str = "messages",
+ message_field_role: str | list[str] | None = None, # commonly "role"
+ message_field_content: str | list[str] | None = None, # commonly "content"
+ message_field_training: str | list[str] | None = None, # commonly "weight"
):
"""Builds a transform that takes a row from the dataset and converts it to a Chat
@@ -26,19 +20,25 @@ def chat_message_transform_builder( # pylint: disable=dangerous-default-value
If True, the transform will train on the inputs. If False, the transform will train on the targets.
Defaults to False.
conversations_field (str, optional):
- The field name of the conversations. Defaults to "conversations".
+ The field name of the conversations. Defaults to "messages".
message_field_role (str | list[str], optional):
- The field name of the role. Defaults to "role".
+ The field name of the role.
message_field_content (str | list[str], optional):
- The field name of the message content. Defaults to "content".
+ The field name of the message content.
message_field_training (str | list[str], optional):
- The field name of the train/weight. Defaults to "weight".
+ The field name of the train/weight.
Returns:
Callable:
A function that takes a list of conversations and returns a list of messages.
"""
+ if message_field_training is None:
+ message_field_training = ["train", "weight"]
+ if message_field_content is None:
+ message_field_content = ["value", "text", "content"]
+ if message_field_role is None:
+ message_field_role = ["role", "from"]
message_field_role = (
[message_field_role]
if isinstance(message_field_role, str)
diff --git a/src/axolotl/core/trainers/__init__.py b/src/axolotl/core/trainers/__init__.py
index 2cdc9c195..22d8b64f6 100644
--- a/src/axolotl/core/trainers/__init__.py
+++ b/src/axolotl/core/trainers/__init__.py
@@ -1,18 +1,14 @@
"""Init for axolotl.core.trainers"""
-# pylint: disable=unused-import
# flake8: noqa
from .base import AxolotlTrainer
from .dpo.trainer import AxolotlDPOTrainer
-from .grpo.trainer import AxolotlGRPOSequenceParallelTrainer, AxolotlGRPOTrainer
from .mamba import AxolotlMambaTrainer
-from .relora import ReLoRATrainer
from .trl import (
AxolotlCPOTrainer,
AxolotlKTOTrainer,
AxolotlORPOTrainer,
AxolotlPRMTrainer,
AxolotlRewardTrainer,
- TRLPPOTrainer,
)
diff --git a/src/axolotl/core/trainers/base.py b/src/axolotl/core/trainers/base.py
index 70e443cb3..11dfecb98 100644
--- a/src/axolotl/core/trainers/base.py
+++ b/src/axolotl/core/trainers/base.py
@@ -1,17 +1,18 @@
"""Module for customized trainers"""
-# pylint: disable=too-many-lines
-
from __future__ import annotations
import os
from collections import defaultdict
from functools import partial, wraps
-from typing import Callable, Literal, Optional
+from typing import Any, Callable, Literal, Optional
import datasets
+import safetensors
import torch
+from accelerate.state import AcceleratorState
from datasets import Dataset
+from peft import PeftModel
from torch.utils.data import (
BatchSampler,
DataLoader,
@@ -19,13 +20,19 @@ from torch.utils.data import (
Sampler,
SequentialSampler,
)
-from transformers import Trainer
-from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, seed_worker
+from transformers import PreTrainedModel, Trainer
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, has_length, seed_worker
+from transformers.utils import SAFE_WEIGHTS_NAME, WEIGHTS_NAME, is_peft_available
from trl.trainer.utils import pad_to_length
from typing_extensions import override
from axolotl.core.trainers.mixins import (
+ ActivationOffloadingMixin,
+ CheckpointSaveMixin,
+ DistributedParallelMixin,
OptimizerMixin,
+ PackingMixin,
RngLoaderMixin,
SchedulerMixin,
)
@@ -33,17 +40,46 @@ from axolotl.core.trainers.utils import (
sanitize_kwargs_for_ds_tagging,
sanitize_kwargs_for_tagging,
)
+from axolotl.utils import get_not_null
+from axolotl.utils.bench import get_gpu_memory_usage
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.distributed import is_main_process
from axolotl.utils.logging import get_logger
from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
LOG = get_logger(__name__)
+REDUCTION_FNS = {
+ "mean": torch.mean,
+ "min": torch.min,
+ "max": torch.max,
+ "sum": torch.sum,
+}
-class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
+
+class AxolotlTrainer(
+ PackingMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ RngLoaderMixin,
+ CheckpointSaveMixin,
+ ActivationOffloadingMixin,
+ DistributedParallelMixin,
+ Trainer,
+):
"""Extend the base Trainer for axolotl helpers"""
args = None # type: "AxolotlTrainingArguments" # type: ignore[name-defined]
tag_names = ["axolotl"]
+ _axolotl_cfg: DictDefault | None = None
+
+ @property
+ def axolotl_cfg(self):
+ return self._axolotl_cfg
+
+ @axolotl_cfg.setter
+ def axolotl_cfg(self, cfg):
+ self._axolotl_cfg = cfg
def __init__(
self,
@@ -59,24 +95,13 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
self._signature_columns = None # workaround for pylint
super().__init__(*_args, **kwargs)
-
self.train_data_collator = self.data_collator
- self._stored_metrics = defaultdict(lambda: defaultdict(list))
+ self._stored_metrics = defaultdict(
+ lambda: defaultdict(lambda: {"values": [], "reduction": "mean"})
+ )
if self.args.orpo_alpha:
self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
- def _wrap_model(self, model, training=True, dataloader=None):
- if self.args.torch_compile:
- torch._dynamo.config.accumulated_cache_size_limit = ( # pylint: disable=protected-access
- 256
- )
- model = torch.compile(
- model,
- backend=self.args.torch_compile_backend,
- mode=self.args.torch_compile_mode,
- )
- return super()._wrap_model(model, training=training, dataloader=dataloader)
-
def _create_multipack_sampler(
self, base_sampler: Sampler, dataset: Dataset
) -> MultipackBatchSampler:
@@ -101,7 +126,7 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
)
batch_max_len = train_batch_size * self.args.max_seq_length
- return MultipackBatchSampler(
+ sampler = MultipackBatchSampler(
base_sampler,
lengths=get_dataset_lengths(dataset),
packing_efficiency_estimate=self.args.sample_packing_efficiency,
@@ -111,11 +136,16 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
bin_size=self.args.sample_packing_bin_size,
sequential=self.args.sample_packing_sequentially,
drop_last=True,
+ num_processes=self.args.dataset_num_proc,
+ mp_start_method=self.args.sample_packing_mp_start_method or "fork",
)
+ len(sampler)
+ return sampler
+
def _get_train_sampler(
- self, train_dataset: Optional[Dataset] = None
- ) -> Optional[Sampler]:
+ self, train_dataset: Dataset | None = None
+ ) -> Sampler | None:
"""
Helper method to get the sampler for training. Handles cases for sample packing
and curriculum sampling (sequential).
@@ -124,16 +154,22 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
If the dataset is non-empty, a sampler is returned, the type of which
depends on the passed training args.
"""
+ # from https://github.com/huggingface/transformers/blob/2166b6b4ff09f6dd3867ab982f262f66482aa968/src/transformers/trainer.py#L969C1-L972C24
+ if train_dataset is None:
+ train_dataset = self.train_dataset
+ if train_dataset is None or not has_length(train_dataset):
+ return None
+
use_sample_packing = self.args.sample_packing and not self.args.pretraining
# Determine the base sampler first
if self.args.curriculum_sampling:
- base_sampler = SequentialSampler(self.train_dataset)
+ base_sampler = SequentialSampler(train_dataset)
elif use_sample_packing:
- base_sampler = RandomSampler(self.train_dataset)
+ base_sampler = RandomSampler(train_dataset)
else:
# Default to parent class implementation for standard random sampling
- return super()._get_train_sampler()
+ return super()._get_train_sampler(train_dataset)
# Apply multipack wrapper if needed
if use_sample_packing:
@@ -152,6 +188,10 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
If the dataset is non-empty, a sampler is returned, the type of which
depends on the passed training args.
"""
+ # from https://github.com/huggingface/transformers/blob/2166b6b4ff09f6dd3867ab982f262f66482aa968/src/transformers/trainer.py#L1065C9-L1066C24
+ if eval_dataset is None or not has_length(eval_dataset):
+ return None
+
# Multipacking enabled if training is enabled and eval is not explicitly disabled
use_multipack = (
self.args.sample_packing and self.args.eval_sample_packing is not False
@@ -187,6 +227,14 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
if dataset.column_names and "length" in dataset.column_names:
dataset = dataset.remove_columns(["length"])
+ if (
+ dataset.column_names
+ and "position_ids" in dataset.column_names
+ and "attention_mask" in dataset.column_names
+ and self.args.sample_packing
+ and self.args.sample_packing_drop_attention_mask
+ ):
+ dataset = dataset.remove_columns(["attention_mask"])
if isinstance(dataset, datasets.Dataset):
if is_training:
@@ -220,7 +268,9 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
}
if not isinstance(dataset, torch.utils.data.IterableDataset):
- dataloader_params["drop_last"] = self.args.dataloader_drop_last
+ dataloader_params["drop_last"] = get_not_null(
+ self.args.dataloader_drop_last, True
+ )
if sampler_fn is not None:
sampler = sampler_fn(dataset)
if isinstance(sampler, BatchSampler):
@@ -251,9 +301,9 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
# fmt: off
if dataloader_key is not None and self.args.dataloader_persistent_workers:
if hasattr(self, "_eval_dataloaders"):
- self._eval_dataloaders[dataloader_key] = dataloader # type: ignore # pylint: disable=access-member-before-definition
+ self._eval_dataloaders[dataloader_key] = dataloader # type: ignore
else:
- self._eval_dataloaders = {dataloader_key: dataloader} # pylint: disable=attribute-defined-outside-init
+ self._eval_dataloaders = {dataloader_key: dataloader}
# fmt: on
return self.accelerator.prepare(dataloader)
@@ -295,6 +345,17 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
# outputs = model(**inputs)
# loss = trainer_weighted_loss(outputs, labels, shift_labels=True)
# return (loss, outputs) if return_outputs else loss
+
+ # track number of tokens for tokens per second calculation
+ if self.args.include_tkps:
+ inputs_key = "labels" if "labels" in inputs else "input_ids"
+ if hasattr(self.state, "num_tokens"):
+ self.state.num_tokens = (
+ self.state.num_tokens + (inputs[inputs_key] != -100).sum().cpu()
+ )
+ else:
+ self.state.num_tokens = (inputs[inputs_key] != -100).sum().cpu()
+
if self.args.orpo_alpha:
return self.orpo_compute_loss(
model,
@@ -310,6 +371,11 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
num_items_in_batch=num_items_in_batch,
)
+ @override
+ def evaluate(self, *args, **kwargs):
+ LOG.info("Running evaluation step...")
+ return super().evaluate(*args, **kwargs)
+
@staticmethod
def orpo_concatenate_inputs(inputs, label_pad_token=-100, pad_token=0, device=None):
concatenated_batch = {}
@@ -409,7 +475,7 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
model,
inputs,
return_outputs=False,
- num_items_in_batch=None, # pylint: disable=unused-argument
+ num_items_in_batch=None,
):
concat_inputs = AxolotlTrainer.orpo_concatenate_inputs(
inputs,
@@ -486,26 +552,32 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
@wraps(Trainer.create_accelerator_and_postprocess)
def create_accelerator_and_postprocess(self):
- res = super().create_accelerator_and_postprocess()
+ # cleanup the PartialState states so Accelerate automatically configures everything from the env vars
+ accelerator_config = self.args.accelerator_config.to_dict()
+ use_configured_state = accelerator_config.get("use_configured_state", False)
+ if not use_configured_state:
+ AcceleratorState._reset_state(reset_partial_state=True)
- if self.is_fsdp_enabled:
- if (
- "limit_all_gathers" in self.args.fsdp_config
- and self.args.fsdp_config["limit_all_gathers"]
- ):
- self.accelerator.state.fsdp_plugin.limit_all_gathers = True
-
- return res
+ super().create_accelerator_and_postprocess()
def additional_accelerator_args(
- self, fp8=None, **kwargs
- ): # pylint: disable=unused-argument
+ self, fp8: bool = False, enable_fsdp_float8_all_gather: bool = False, **kwargs
+ ) -> dict[str, Any]:
ret_kwargs = {}
if fp8:
from accelerate.utils import AORecipeKwargs
+ from torchao.float8 import Float8LinearConfig
+
+ # By default, Float8LinearConfig is instantiated using the "tensorwise"
+ # scaling strategy. See more details here:
+ # https://github.com/pytorch/ao/tree/main/torchao/float8.
+ config = Float8LinearConfig(
+ enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather,
+ force_recompute_fp8_weight_in_bwd=enable_fsdp_float8_all_gather is True,
+ )
ret_kwargs["mixed_precision"] = "fp8"
- ret_kwargs["kwargs_handlers"] = [AORecipeKwargs()]
+ ret_kwargs["kwargs_handlers"] = [AORecipeKwargs(config=config)] # type: ignore
os.environ["ACCELERATE_MIXED_PRECISION"] = "fp8"
return ret_kwargs
@@ -520,18 +592,61 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
"""
# logs either has 'loss' or 'eval_loss'
train_eval = "train" if "loss" in logs else "eval"
- # Add averaged stored metrics to logs
- for key, metrics in self._stored_metrics[train_eval].items():
- logs[key] = torch.tensor(metrics).mean().item()
+
+ for key, metric_data in self._stored_metrics[train_eval].items():
+ values = torch.tensor(metric_data["values"]) # type: ignore[arg-type]
+ reduction_type = metric_data["reduction"]
+
+ fn = REDUCTION_FNS.get(reduction_type)
+ if fn is None:
+ raise NotImplementedError(
+ "Metric reduction must be one of [mean, min, max, sum]"
+ )
+ logs[key] = round(fn(values).item(), 4)
+
+ if is_main_process():
+ # Add memory usage
+ try:
+ active, allocated, reserved = get_gpu_memory_usage()
+ logs["memory/max_active (GiB)"] = round(active, 2)
+ logs["memory/max_allocated (GiB)"] = round(allocated, 2)
+ logs["memory/device_reserved (GiB)"] = round(reserved, 2)
+ except (ValueError, TypeError, FileNotFoundError):
+ pass
+
+ if self.args.include_tkps and train_eval == "train":
+ # each rank will log its own tokens per second
+ # for logging_steps > 1 we obtain a moving average of this metric
+ logs["tokens_per_second_per_gpu"] = round(
+ self.state.last_tokens_per_second.item() / self.args.logging_steps, 2
+ )
+
del self._stored_metrics[train_eval]
return super().log(logs, start_time)
def store_metrics(
- self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train"
+ self,
+ metrics: dict[str, float] | dict[str, tuple[int | float, str]],
+ train_eval: Literal["train", "eval"] = "train",
+ reduction: Literal["mean", "min", "max", "sum"] = "mean",
) -> None:
+ """
+ Store metrics with specified reduction type.
+
+ Args:
+ metrics: Dictionary of metric names to values, or metric names to (value,
+ reduction_type) tuples.
+ train_eval: Whether this is for training or evaluation.
+ """
for key, value in metrics.items():
- self._stored_metrics[train_eval][key].append(value)
+ if isinstance(value, tuple):
+ value, _reduction = value # type: ignore[assignment]
+ else:
+ value, _reduction = value, reduction
+
+ self._stored_metrics[train_eval][key]["values"].append(value)
+ self._stored_metrics[train_eval][key]["reduction"] = _reduction
def _save_checkpoint(self, model, trial, **kwargs):
# make sure the checkpoint dir exists, since trainer is flakey
@@ -540,3 +655,69 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, RngLoaderMixin, Trainer):
output_dir = os.path.join(run_dir, checkpoint_folder)
os.makedirs(output_dir, exist_ok=True)
return super()._save_checkpoint(model, trial, **kwargs)
+
+ # TODO(wing): remove once https://github.com/huggingface/transformers/pull/39866/files is merged
+ def _save(self, output_dir: Optional[str] = None, state_dict=None):
+ # If we are executing this function, we are the process zero, so we don't check for that.
+ output_dir = output_dir if output_dir is not None else self.args.output_dir
+ os.makedirs(output_dir, exist_ok=True)
+ LOG.info(f"Saving model checkpoint to {output_dir}")
+ supported_classes = (
+ (PreTrainedModel,)
+ if not is_peft_available()
+ else (PreTrainedModel, PeftModel)
+ )
+ # Save a trained model and configuration using `save_pretrained()`.
+ # They can then be reloaded using `from_pretrained()`
+ if not isinstance(self.model, supported_classes):
+ if state_dict is None:
+ state_dict = self.model.state_dict()
+ if isinstance(
+ self.accelerator.unwrap_model(self.model, keep_torch_compile=False),
+ supported_classes,
+ ):
+ self.accelerator.unwrap_model(
+ self.model, keep_torch_compile=False
+ ).save_pretrained(
+ output_dir,
+ state_dict=state_dict,
+ safe_serialization=self.args.save_safetensors,
+ )
+ else:
+ LOG.info(
+ "Trainer.model is not a `PreTrainedModel`, only saving its state dict."
+ )
+ if self.args.save_safetensors:
+ safetensors.torch.save_file(
+ state_dict,
+ os.path.join(output_dir, SAFE_WEIGHTS_NAME),
+ metadata={"format": "pt"},
+ )
+ else:
+ torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
+ else:
+ self.model.save_pretrained(
+ output_dir,
+ state_dict=state_dict,
+ safe_serialization=self.args.save_safetensors,
+ is_main_process=self.accelerator.is_main_process,
+ )
+
+ if self.processing_class is not None:
+ self.processing_class.save_pretrained(output_dir)
+ elif (
+ self.data_collator is not None
+ and hasattr(self.data_collator, "tokenizer")
+ and self.data_collator.tokenizer is not None
+ ):
+ LOG.info(
+ "Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`"
+ )
+ save_jinja_files = True
+ if self.axolotl_cfg:
+ save_jinja_files = self.axolotl_cfg.tokenizer_save_jinja_files
+ self.data_collator.tokenizer.save_pretrained(
+ output_dir, save_jinja_files=save_jinja_files
+ )
+ # Good practice: save your training arguments together with the trained model
+ torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
diff --git a/src/axolotl/core/trainers/dpo/__init__.py b/src/axolotl/core/trainers/dpo/__init__.py
index 603fdf0b6..3aa79c484 100644
--- a/src/axolotl/core/trainers/dpo/__init__.py
+++ b/src/axolotl/core/trainers/dpo/__init__.py
@@ -22,10 +22,18 @@ class DPOStrategy:
training_args_kwargs = {}
if cfg.rl is RLType.IPO:
training_args_kwargs["loss_type"] = "ipo"
- training_args_kwargs["max_length"] = cfg.sequence_len
+ # Label smoothing is not compatible with IPO
+ if cfg.rl is RLType.DPO and cfg.dpo_label_smoothing:
+ training_args_kwargs["label_smoothing"] = cfg.dpo_label_smoothing
training_args_kwargs["max_completion_length"] = None
- training_args_kwargs["max_prompt_length"] = cfg.sequence_len
- training_args_kwargs["generate_during_eval"] = cfg.use_wandb
+ training_args_kwargs["max_length"] = cfg.sequence_len
+ training_args_kwargs["generate_during_eval"] = cfg.dpo_generate_during_eval
if cfg.dpo_use_weighting is not None:
training_args_kwargs["use_weighting"] = cfg.dpo_use_weighting
+ if cfg.dpo_padding_free is not None:
+ training_args_kwargs["padding_free"] = cfg.dpo_padding_free
+ if cfg.dpo_norm_loss is not None:
+ training_args_kwargs["dpo_norm_loss"] = cfg.dpo_norm_loss
+ if cfg.dpo_use_logits_to_keep is not None:
+ training_args_kwargs["use_logits_to_keep"] = cfg.dpo_use_logits_to_keep
return training_args_kwargs
diff --git a/src/axolotl/core/trainers/dpo/args.py b/src/axolotl/core/trainers/dpo/args.py
index de1758ed0..b1e53236e 100644
--- a/src/axolotl/core/trainers/dpo/args.py
+++ b/src/axolotl/core/trainers/dpo/args.py
@@ -14,3 +14,5 @@ class AxolotlDPOConfig(AxolotlTrainingMixins, DPOConfig):
"""
DPO config for DPO training
"""
+
+ dpo_norm_loss: bool | None = False
diff --git a/src/axolotl/core/trainers/dpo/trainer.py b/src/axolotl/core/trainers/dpo/trainer.py
index 15af80c02..b04505d89 100644
--- a/src/axolotl/core/trainers/dpo/trainer.py
+++ b/src/axolotl/core/trainers/dpo/trainer.py
@@ -8,7 +8,11 @@ import torch
from torch import nn
from trl import DPOTrainer
-from axolotl.core.trainers.mixins import RngLoaderMixin, SchedulerMixin
+from axolotl.core.trainers.mixins import (
+ DistributedParallelMixin,
+ RngLoaderMixin,
+ SchedulerMixin,
+)
from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
from axolotl.core.trainers.utils import (
sanitize_kwargs_for_ds_tagging,
@@ -17,7 +21,12 @@ from axolotl.core.trainers.utils import (
class AxolotlDPOTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, DPOTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DPOTrainer,
+ DistributedParallelMixin,
):
"""Extend the base DPOTrainer for axolotl helpers."""
@@ -83,3 +92,20 @@ class AxolotlDPOTrainer(
gc.collect()
torch.cuda.empty_cache()
return loss
+
+ def concatenated_forward(
+ self,
+ model: nn.Module,
+ batch: dict[str, Union[list, torch.LongTensor]],
+ is_ref_model: bool = False,
+ ) -> dict[str, torch.Tensor]:
+ if self.args.dpo_norm_loss:
+ # fmt: off
+ loss_type: str = self.loss_type # type: ignore[has-type]
+ # fmt: on
+ # concatenated_forward handles avg token logprob for ipo case already
+ self.loss_type = "ipo"
+ res = super().concatenated_forward(model, batch, is_ref_model=is_ref_model)
+ self.loss_type = loss_type
+ return res
+ return super().concatenated_forward(model, batch, is_ref_model=is_ref_model)
diff --git a/src/axolotl/core/trainers/grpo/__init__.py b/src/axolotl/core/trainers/grpo/__init__.py
index c0f10be23..d1a6b7fd9 100644
--- a/src/axolotl/core/trainers/grpo/__init__.py
+++ b/src/axolotl/core/trainers/grpo/__init__.py
@@ -2,8 +2,11 @@
import importlib
import inspect
+import os
from typing import Any
+from huggingface_hub import snapshot_download
+from requests import HTTPError
from trl.trainer.grpo_trainer import RewardFunc
from axolotl.core.trainers.grpo.args import AxolotlGRPOConfig
@@ -14,6 +17,7 @@ from axolotl.core.trainers.grpo.trainer import (
from axolotl.utils.dict import DictDefault
from axolotl.utils.logging import get_logger
from axolotl.utils.schemas.trl import TRLConfig
+from axolotl.utils.schemas.vllm import VllmConfig
LOG = get_logger(__name__)
@@ -41,9 +45,20 @@ class GRPOStrategy:
return grpo_args_kwargs
trl: TRLConfig = cfg.trl # type: ignore
+ vllm_cfg: VllmConfig = cfg.vllm # type: ignore
if trl.use_vllm:
grpo_args_kwargs["use_vllm"] = trl.use_vllm
+ if trl.vllm_mode:
+ grpo_args_kwargs["vllm_mode"] = trl.vllm_mode
+ if trl.vllm_mode == "colocate":
+ grpo_args_kwargs["enable_sleep_mode"] = trl.vllm_enable_sleep_mode # type: ignore[attr-defined]
+ grpo_args_kwargs["vllm_gpu_memory_utilization"] = (
+ vllm_cfg.gpu_memory_utilization
+ )
+ grpo_args_kwargs["vllm_tensor_parallel_size"] = (
+ vllm_cfg.tensor_parallel_size
+ )
grpo_args_kwargs["vllm_server_host"] = trl.vllm_server_host or trl.vllm.host # type: ignore[attr-defined]
grpo_args_kwargs["vllm_server_port"] = trl.vllm_server_port or trl.vllm.port # type: ignore[attr-defined]
if trl.vllm_server_timeout:
@@ -69,8 +84,13 @@ class GRPOStrategy:
grpo_args_kwargs["log_completions"] = trl.log_completions
grpo_args_kwargs["num_completions_to_print"] = trl.num_completions_to_print
- if cfg.sequence_parallel_degree > 1:
- grpo_args_kwargs["sequence_parallel_degree"] = cfg.sequence_parallel_degree
+ if cfg.context_parallel_size > 1:
+ grpo_args_kwargs["context_parallel_size"] = cfg.context_parallel_size
+
+ if trl.importance_sampling_level is not None:
+ grpo_args_kwargs["importance_sampling_level"] = (
+ trl.importance_sampling_level
+ )
if trl.reward_weights:
grpo_args_kwargs["reward_weights"] = trl.reward_weights
@@ -109,9 +129,7 @@ class GRPOStrategy:
return grpo_args_kwargs
@classmethod
- def set_trainer_args(
- cls, cfg: DictDefault
- ) -> list[Any]: # pylint: disable=unused-argument
+ def set_trainer_args(cls, cfg: DictDefault) -> list[Any]:
trainer_args = []
if cfg.trl and cfg.trl.reward_funcs:
reward_funcs = []
@@ -132,13 +150,13 @@ class GRPOStrategy:
return trainer_kwargs
@classmethod
- def get_collator(cls, *args, **kwargs): # pylint: disable=unused-argument
+ def get_collator(cls, *args, **kwargs):
# No data collation is needed in GRPO, handled by trl's trainer __init__
return None
@classmethod
def get_blocklist_args_kwargs(cls) -> list[str]:
- return ["dataset_num_proc", "max_length"]
+ return ["dataset_num_proc", "max_length", "include_tokens_per_second"]
@classmethod
def get_reward_func(cls, reward_func_fqn: str) -> RewardFunc:
@@ -168,9 +186,18 @@ class GRPOStrategy:
"Reward function must accept at least two arguments: prompts: list and completions: list"
)
return reward_func
- except ModuleNotFoundError:
+ except ModuleNotFoundError as exc:
# the user has passed a string (ideally indicating the path of a reward model)
- LOG.info(
- f"Reward function {reward_func_fqn} is a pre-trained model path - if this is unexpected, please check the reward function path."
- )
- return reward_func_fqn
+ # check if it's a local dir path and not empty dir to a reward model
+ pretrained_log_msg = f"Reward function {reward_func_fqn} is a pre-trained model path - if this is unexpected, please check the reward function path."
+ if os.path.isdir(reward_func_fqn) and os.listdir(reward_func_fqn):
+ LOG.info(pretrained_log_msg)
+ return reward_func_fqn
+ try:
+ snapshot_download(reward_func_fqn, repo_type="model")
+ LOG.info(pretrained_log_msg)
+ return reward_func_fqn
+ except HTTPError:
+ raise ValueError(
+ f"Reward function {reward_func_fqn} not found."
+ ) from exc
diff --git a/src/axolotl/core/trainers/grpo/args.py b/src/axolotl/core/trainers/grpo/args.py
index 5c8b1a33b..2ea52998e 100644
--- a/src/axolotl/core/trainers/grpo/args.py
+++ b/src/axolotl/core/trainers/grpo/args.py
@@ -13,4 +13,4 @@ from axolotl.core.training_args import AxolotlTrainingMixins
class AxolotlGRPOConfig(AxolotlTrainingMixins, GRPOConfig):
"""Axolotl GRPO Config for GRPO training"""
- sequence_parallel_degree: int | None = None
+ context_parallel_size: int | None = None
diff --git a/src/axolotl/core/trainers/grpo/sampler.py b/src/axolotl/core/trainers/grpo/sampler.py
index ebc6e19e2..df679a6d2 100644
--- a/src/axolotl/core/trainers/grpo/sampler.py
+++ b/src/axolotl/core/trainers/grpo/sampler.py
@@ -20,7 +20,7 @@ class SequenceParallelRepeatRandomSampler(Sampler):
- Data is properly distributed across SP groups.
In the table below, the values represent dataset indices. Each SP group has
- `sequence_parallel_degree = 2` GPUs working together on the same data. There are 2
+ `context_parallel_size = 2` GPUs working together on the same data. There are 2
SP groups (SP0 and SP1), with `world_size = 4` total GPUs.
Sequence Parallel Groups
@@ -45,7 +45,7 @@ class SequenceParallelRepeatRandomSampler(Sampler):
rank: Rank of current process.
batch_size: Number of samples per batch.
repeat_count: How many times to repeat the full sampling process.
- sequence_parallel_degree: Number of ranks in a sequence parallel group.
+ context_parallel_size: Number of ranks in a sequence parallel group.
shuffle: Whether to shuffle the dataset.
seed: Random seed for shuffling.
drop_last: Whether to drop the last incomplete batch.
@@ -59,7 +59,7 @@ class SequenceParallelRepeatRandomSampler(Sampler):
rank: int,
batch_size: int = 1,
repeat_count: int = 1,
- sequence_parallel_degree: int = 1,
+ context_parallel_size: int = 1,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
@@ -77,9 +77,9 @@ class SequenceParallelRepeatRandomSampler(Sampler):
self.rank = rank
# Sequence parallelism parameters
- self.sequence_parallel_degree = sequence_parallel_degree
- self.num_sp_groups = world_size // sequence_parallel_degree
- self.sp_group_id = rank // sequence_parallel_degree
+ self.context_parallel_size = context_parallel_size
+ self.num_sp_groups = world_size // context_parallel_size
+ self.sp_group_id = rank // context_parallel_size
# Adjust dataset size for distributed sampling
self.num_samples = len(self.dataset)
diff --git a/src/axolotl/core/trainers/grpo/trainer.py b/src/axolotl/core/trainers/grpo/trainer.py
index dccc85d80..f9f5a695b 100644
--- a/src/axolotl/core/trainers/grpo/trainer.py
+++ b/src/axolotl/core/trainers/grpo/trainer.py
@@ -1,8 +1,7 @@
"""Axolotl GRPO trainers (with and without sequence parallelism handling)"""
-# pylint: disable=too-many-lines,duplicate-code,protected-access,no-member
-
import warnings
+from functools import partial
from typing import Any
import datasets
@@ -42,17 +41,25 @@ from trl.trainer.grpo_trainer import RewardFunc, nanstd
from trl.trainer.utils import pad
from axolotl.core.trainers.grpo.sampler import SequenceParallelRepeatRandomSampler
-from axolotl.core.trainers.mixins import RngLoaderMixin, SchedulerMixin
+from axolotl.core.trainers.mixins import (
+ DistributedParallelMixin,
+ RngLoaderMixin,
+ SchedulerMixin,
+)
from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
from axolotl.monkeypatch.ring_attn import get_ring_attn_group
if is_peft_available():
- # pylint: disable=unused-import
from peft import PeftConfig
class AxolotlGRPOTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, GRPOTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ GRPOTrainer,
):
"""Extend the base GRPOTrainer for axolotl helpers"""
@@ -99,7 +106,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
# Get number of SP groups (number of processes divided by SP degree)
num_processes = self.accelerator.num_processes
- num_sp_groups = num_processes // self.args.sequence_parallel_degree
+ num_sp_groups = num_processes // self.args.context_parallel_size
# Calculate batch size per SP group (not per process)
sp_group_batch_size = self.args.per_device_train_batch_size * num_sp_groups
@@ -129,7 +136,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
if self.num_generations not in possible_values:
raise ValueError(
- f"With sequence parallelism (degree {self.args.sequence_parallel_degree}), "
+ f"With sequence parallelism (degree {self.args.context_parallel_size}), "
f"the eval batch size per SP group ({num_sp_groups} x {self.args.per_device_eval_batch_size}) "
f"must be evenly divisible by the number of generations per prompt "
f"({self.num_generations}). Given the current eval batch size, "
@@ -166,9 +173,9 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
rank=self.rank,
batch_size=effective_batch_size
// self.num_generations
- // self.args.sequence_parallel_degree,
+ // self.args.context_parallel_size,
repeat_count=self.num_iterations * self.args.gradient_accumulation_steps,
- sequence_parallel_degree=self.args.sequence_parallel_degree,
+ context_parallel_size=self.args.context_parallel_size,
shuffle=True,
seed=self.args.seed,
drop_last=True,
@@ -215,7 +222,11 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
dataloader_params["drop_last"] = self.args.dataloader_drop_last
if not is_eval:
- dataloader_params["worker_init_fn"] = seed_worker
+ dataloader_params["worker_init_fn"] = partial(
+ seed_worker,
+ num_workers=self.args.dataloader_num_workers,
+ rank=self.args.process_index,
+ )
# Create the dataloader
dataloader = DataLoader(dataset, **dataloader_params)
@@ -230,7 +241,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
# TODO(djsaunde): We might be able to use `accelerate`'s dataloader preparation
# if we use `dispatch_batches` and `slice_fn_for_dispatch` properly (i.e.,
# slice each batch along the sequence dimension).
- if self.args.sequence_parallel_degree > 1:
+ if self.args.context_parallel_size > 1:
return dataloader
# Otherwise prepare with accelerator
@@ -239,7 +250,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
def get_train_dataloader(self) -> DataLoader:
"""Get dataloader for training"""
train_dataset = self.train_dataset
- # pylint: disable=access-member-before-definition
+
data_collator = self.data_collator # type: ignore
# Handle dataset preprocessing
@@ -252,7 +263,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
train_dataset, description="training"
)
else:
- self.data_collator = self._get_collator_with_removed_columns( # pylint: disable=attribute-defined-outside-init
+ self.data_collator = self._get_collator_with_removed_columns(
data_collator,
description="training",
)
@@ -294,33 +305,34 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
# Generate completions using either vLLM or regular generation
if self.args.use_vllm:
# First, have main process load weights if needed
- # pylint: disable=access-member-before-definition
+
if self.state.global_step != self._last_loaded_step: # type: ignore[has-type]
self._move_model_to_vllm()
- # pylint: disable=attribute-defined-outside-init
+
self._last_loaded_step = self.state.global_step
# Generate completions using vLLM: gather all prompts and use them in a single call in the main process
all_prompts_text = gather_object(prompts_text)
if self.accelerator.is_main_process:
- if self.args.sequence_parallel_degree > 1:
+ if self.args.context_parallel_size > 1:
# Calculate sequence parallel group information
world_size = self.accelerator.num_processes
- sequence_parallel_degree = self.args.sequence_parallel_degree
- num_sp_groups = world_size // sequence_parallel_degree
+ context_parallel_size = self.args.context_parallel_size
+ num_sp_groups = world_size // context_parallel_size
# Since processes in the same SP group have the same prompts, we need to ensure
# we only take one copy of each prompt from each SP group
ordered_set_of_prompts = []
for sp_group_id in range(num_sp_groups):
# Get the first process from each SP group (typically the group leader)
- group_leader_rank = sp_group_id * sequence_parallel_degree
+ group_leader_rank = sp_group_id * context_parallel_size
# Extract prompts from this SP group, accounting for num_generations duplicates
# We only need prompts from one rank in each SP group
group_prompts = all_prompts_text[
- group_leader_rank
- * len(prompts_text) : (group_leader_rank + 1)
+ group_leader_rank * len(prompts_text) : (
+ group_leader_rank + 1
+ )
* len(prompts_text) : self.num_generations
]
@@ -330,7 +342,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
# num_generations outputs for each one. This is faster than generating outputs for each duplicate
# prompt individually.
ordered_set_of_prompts = all_prompts_text[
- :: self.num_generations * self.args.sequence_parallel_degree
+ :: self.num_generations * self.args.context_parallel_size
]
with profiling_context(self, "vLLM.generate"):
@@ -347,14 +359,14 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
)
else:
completion_ids = [None] * (
- len(all_prompts_text) // self.args.sequence_parallel_degree
+ len(all_prompts_text) // self.args.context_parallel_size
)
# Broadcast the completions from the main process to all processes
completion_ids = broadcast_object_list(completion_ids, from_process=0)
# Determine the appropriate slice based on sequence parallelism
- if self.args.sequence_parallel_degree > 1:
+ if self.args.context_parallel_size > 1:
# Calculate SP group ID (which group of ranks this rank belongs to)
sp_group_id = self.accelerator.process_index // self.local_world_size
@@ -471,7 +483,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
)
if is_conversational(inputs[0]):
completions = []
- for prompt, completion in zip(prompts, completions_text):
+ for prompt, completion in zip(prompts, completions_text, strict=False):
bootstrap = (
prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else ""
)
@@ -489,6 +501,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
self.reward_funcs,
self.reward_processing_classes,
self.reward_func_names,
+ strict=False,
)
):
with profiling_context(self, reward_func_name):
@@ -497,14 +510,17 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
): # Module instead of PretrainedModel for compat with compiled models
if is_conversational(inputs[0]):
messages = [
- {"messages": p + c} for p, c in zip(prompts, completions)
+ {"messages": p + c}
+ for p, c in zip(prompts, completions, strict=False)
]
texts = [
apply_chat_template(x, reward_processing_class)["text"]
for x in messages
]
else:
- texts = [p + c for p, c in zip(prompts, completions)]
+ texts = [
+ p + c for p, c in zip(prompts, completions, strict=False)
+ ]
reward_inputs = reward_processing_class(
text=texts,
return_tensors="pt",
@@ -550,7 +566,8 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
row_reward_kwargs["completion"] = completions[nan_row_idx]
warnings.warn(
f"All reward functions returned None for the following kwargs: {row_reward_kwargs}. "
- "Please ensure that at least one reward function returns a valid reward."
+ "Please ensure that at least one reward function returns a valid reward.",
+ stacklevel=2,
)
# Gather the reward per function: this part is crucial, because the rewards are normalized per group and the
@@ -578,7 +595,7 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
advantages = advantages / (std_grouped_rewards + 1e-4)
# Slice to keep only the local part of the data
- if self.args.sequence_parallel_degree > 1:
+ if self.args.context_parallel_size > 1:
# Calculate SP group ID (which group of ranks this rank belongs to)
sp_group_id = self.accelerator.process_index // self.local_world_size
diff --git a/src/axolotl/core/trainers/mamba.py b/src/axolotl/core/trainers/mamba.py
index 38792e389..dedda1b29 100644
--- a/src/axolotl/core/trainers/mamba.py
+++ b/src/axolotl/core/trainers/mamba.py
@@ -14,8 +14,8 @@ class AxolotlMambaTrainer(AxolotlTrainer):
self,
model,
inputs,
- return_outputs=False, # pylint: disable=unused-argument
- num_items_in_batch=None, # pylint: disable=unused-argument
+ return_outputs=False,
+ num_items_in_batch=None,
):
input_ids = inputs.pop("input_ids")
lm_logits = model(input_ids).logits
diff --git a/src/axolotl/core/trainers/mixins/__init__.py b/src/axolotl/core/trainers/mixins/__init__.py
index a71cb321a..5fced1692 100644
--- a/src/axolotl/core/trainers/mixins/__init__.py
+++ b/src/axolotl/core/trainers/mixins/__init__.py
@@ -1,8 +1,11 @@
"""Init for axolotl.core.trainers.mixins"""
-# pylint: disable=unused-import
# flake8: noqa
+from .activation_checkpointing import ActivationOffloadingMixin
+from .checkpoints import CheckpointSaveMixin
+from .distributed_parallel import DistributedParallelMixin
from .optimizer import OptimizerMixin
+from .packing import PackingMixin
from .rng_state_loader import RngLoaderMixin
from .scheduler import SchedulerMixin
diff --git a/src/axolotl/core/trainers/mixins/activation_checkpointing.py b/src/axolotl/core/trainers/mixins/activation_checkpointing.py
new file mode 100644
index 000000000..b61c45fee
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/activation_checkpointing.py
@@ -0,0 +1,217 @@
+"""
+Trainer mixin for activation checkpointing w offloading
+"""
+
+import contextlib
+
+from peft import PeftModel
+from torch import nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+ apply_activation_checkpointing,
+)
+from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from transformers import GradientCheckpointingLayer, Trainer
+from trl.models.activation_offloading import (
+ NoOpManager,
+ OffloadActivations,
+ get_act_offloading_ctx_manager,
+)
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class ActivationOffloadingMixin(Trainer):
+ """
+ Trainer mixin class for activation checkpointing w offloading
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ if self.args.activation_offloading:
+ if isinstance(self.model, PeftModel):
+ self.activation_offload_context = get_lora_act_offloading_ctx_manager(
+ self.model, use_streams=True
+ )
+ else:
+ self.activation_offload_context = get_act_offloading_ctx_manager(
+ self.model, use_streams=True
+ )
+ else:
+ self.activation_offload_context = contextlib.nullcontext()
+
+ def training_step(self, *args, **kwargs):
+ with self.activation_offload_context:
+ return super().training_step(*args, **kwargs)
+
+
+def ac_wrap_hf_model(model: nn.Module, **kwargs):
+ auto_wrap_policy = ModuleWrapPolicy(set((GradientCheckpointingLayer,)))
+ apply_activation_checkpointing(model, auto_wrap_policy=auto_wrap_policy, **kwargs)
+
+
+def get_lora_act_offloading_ctx_manager(
+ model: nn.Module,
+ use_pin_memory: bool = True,
+ use_streams: bool = True,
+ min_offload_size: int = 1024,
+ max_fwd_stash_size: int = 5,
+ warn_if_no_head: bool = True,
+) -> OffloadActivations:
+ """
+ Returns the activation offloading context manager for the model. All but the last output Linear in every step will
+ be offloaded.
+
+ If activation offloading is enabled, we return the OffloadActivations context manager. If activation offloading is
+ disabled, we return a NoOpManager context manager.
+
+ Args:
+ model (`nn.Module`):
+ Model to wrap with the activation offloading context manager.
+ use_pin_memory (`bool`, *optional*, defaults to `True`):
+ Whether to offloaded Tensor will be placed in pinned memory on the CPU. Pinned memory allows the Tensor to
+ be moved back onto GPU more quickly but is a limited resource.
+ use_streams (`bool`, *optional*, defaults to `True`):
+ Whether to use streams for performance optimization where the communications get overlapped with the
+ computation. Requires a torch build after torch-2.5.0.
+ min_offload_size (`int`, *optional*, defaults to `1024`):
+ Minimum number of bytes a Tensor must be in order to qualify for offloading. If the tensor is too small, we
+ do not want to waste bandwidth and resources moving it to CPU and back.
+ max_fwd_stash_size (`int`, *optional*, defaults to `5`):
+ Maximum size of the forward stash, or the maximum number of consecutive activations to keep alive during
+ the forward pass. This number must be at least 1. Keeping alive more activations will potentially allow
+ more overlap between the communication and compute streams at the cost of increasing memory usage. Keeping
+ alive fewer activations will conserve memory, but may cause poor overlap between the streams, increasing
+ runtime.
+ warn_if_no_head (`bool`, *optional*, defaults to `True`):
+ Whether to warn if no output head is detected. If set to `False`, no warning will be raised if no output
+ head is detected.
+
+ Returns:
+ `contextlib.ContextDecorator`:
+ Activation offloading context manager for the model.
+ """
+
+ activations_handling_ctx = OffloadActivations(
+ use_pin_memory=use_pin_memory,
+ use_streams=use_streams,
+ min_offload_size=min_offload_size,
+ max_fwd_stash_size=max_fwd_stash_size,
+ )
+
+ # Below is our hack to disable offloading the last output Linear in every
+ # step, as the cost for offloading the activation and then soon after bringing
+ # it back is expensive.
+ output_head_detected = False
+ noop_ctx = NoOpManager()
+
+ # Try to get the actual model if it's wrapped
+ unwrapped_model = model
+ if hasattr(unwrapped_model, "module"):
+ unwrapped_model = unwrapped_model.module
+ # check for PEFT models
+ if hasattr(unwrapped_model, "base_model") and hasattr(
+ unwrapped_model, "peft_config"
+ ):
+ unwrapped_model = unwrapped_model.base_model
+
+ # Check for different types of output heads
+ if hasattr(unwrapped_model, "output"):
+ if isinstance(unwrapped_model.output, nn.Module):
+ unwrapped_model.output.register_forward_pre_hook(
+ lambda *args: noop_ctx.__enter__()
+ )
+ unwrapped_model.output.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+ elif hasattr(unwrapped_model.output, "linear") and isinstance(
+ unwrapped_model.output.linear, nn.Module
+ ):
+ unwrapped_model.output.linear.register_forward_pre_hook(
+ lambda *args: noop_ctx.__enter__()
+ )
+ unwrapped_model.output.linear.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+
+ # Check for HuggingFace model output heads
+ elif hasattr(unwrapped_model, "lm_head"):
+ unwrapped_model.lm_head.register_forward_pre_hook(
+ lambda *args: noop_ctx.__enter__()
+ )
+ unwrapped_model.lm_head.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+
+ # Check for decoder-based models
+ elif hasattr(unwrapped_model, "decoder"):
+ decoder = unwrapped_model.decoder
+ if hasattr(decoder, "output"):
+ decoder.output.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+ decoder.output.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+ # Some models have lm_head in the decoder
+ elif hasattr(decoder, "lm_head"):
+ decoder.lm_head.register_forward_pre_hook(
+ lambda *args: noop_ctx.__enter__()
+ )
+ decoder.lm_head.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+
+ # Check for transformer models with final layer norm
+ elif hasattr(unwrapped_model, "final_layer_norm") or hasattr(
+ unwrapped_model, "ln_f"
+ ):
+ final_norm = (
+ getattr(unwrapped_model, "final_layer_norm", None) or unwrapped_model.ln_f
+ )
+ final_norm.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+ final_norm.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+
+ # Check for models with head module
+ elif hasattr(unwrapped_model, "head") and isinstance(
+ unwrapped_model.head, nn.Module
+ ):
+ unwrapped_model.head.register_forward_pre_hook(
+ lambda *args: noop_ctx.__enter__()
+ )
+ unwrapped_model.head.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ output_head_detected = True
+
+ if not output_head_detected and warn_if_no_head:
+ LOG.warning(
+ "During activation offloading, no output head was detected. If your model has an output head, it will be "
+ "offloaded. This usually greatly slows training, given the large vocabulary size. To change this "
+ "behavior, set your output head as model.output and make it an nn.Module. You can disable this warning by "
+ "passing `warn_if_no_head=False`."
+ )
+
+ for name, module in unwrapped_model.named_modules():
+ # Disable offloading for any Liger modules
+ if "liger" in name.lower():
+ module.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+ module.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+ # disable offloading for any submodules to fix LoRA training
+ if name.endswith("._checkpoint_wrapped_module"):
+ for _, sub_module in module.named_modules():
+ sub_module.register_forward_pre_hook(lambda *args: noop_ctx.__enter__())
+ sub_module.register_forward_hook(
+ lambda *args: noop_ctx.__exit__(), always_call=True
+ )
+
+ return activations_handling_ctx
diff --git a/src/axolotl/core/trainers/mixins/checkpoints.py b/src/axolotl/core/trainers/mixins/checkpoints.py
new file mode 100644
index 000000000..4042ef9f1
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/checkpoints.py
@@ -0,0 +1,23 @@
+"""Custom handling to not fail training if fsdp optimizer is not savable"""
+
+from transformers import Trainer
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class CheckpointSaveMixin(Trainer):
+ """Mixin to handle saving the optimizer and scheduler if they are not savable."""
+
+ def _save_optimizer_and_scheduler(self, output_dir):
+ try:
+ super()._save_optimizer_and_scheduler(output_dir)
+ except (NotImplementedError, KeyError) as exc:
+ # TODO: fix fsdp2 optimizer saving
+ LOG.warning_once(
+ f"Trainer does not support saving optimizer and scheduler: {exc}\n"
+ "Optimizer and scheduler states were not saved - resuming from checkpoints "
+ "for this training run will not be possible.",
+ main_process_only=True,
+ )
diff --git a/src/axolotl/core/trainers/mixins/distributed_parallel.py b/src/axolotl/core/trainers/mixins/distributed_parallel.py
new file mode 100644
index 000000000..77aee5236
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/distributed_parallel.py
@@ -0,0 +1,32 @@
+"""
+Mixin for correctly saving fsdp
+"""
+
+from accelerate import PartialState
+from transformers import Trainer
+
+
+class DistributedParallelMixin(Trainer):
+ """
+ Mixin for correctly saving fsdp
+ """
+
+ def _save(self, output_dir: str | None = None, state_dict=None):
+ if (
+ state_dict is None
+ and self.accelerator.parallelism_config
+ and self.accelerator.parallelism_config.dp_shard_enabled
+ ):
+ state_dict = self.accelerator.get_state_dict(self.model)
+ super()._save(output_dir, state_dict=state_dict)
+
+ def create_accelerator_and_postprocess(self):
+ super().create_accelerator_and_postprocess()
+ if (
+ self.accelerator.distributed_type == "FSDP"
+ and self.accelerator.state.fsdp_plugin is None
+ ):
+ # handle Context Parallelism without FSDP
+ self.accelerator.state.distributed_type = "MULTI_GPU"
+ self.accelerator.state._shared_state["distributed_type"] = "MULTI_GPU"
+ PartialState().distributed_type = "MULTI_GPU"
diff --git a/src/axolotl/core/trainers/mixins/optimizer.py b/src/axolotl/core/trainers/mixins/optimizer.py
index a9a9a3992..850442c60 100644
--- a/src/axolotl/core/trainers/mixins/optimizer.py
+++ b/src/axolotl/core/trainers/mixins/optimizer.py
@@ -70,11 +70,11 @@ class OptimizerMixin(Trainer):
}
)
if params["embeddings"]:
- lr = optimizer_kwargs["lr"] # pylint: disable=invalid-name
+ lr = optimizer_kwargs["lr"]
if self.args.embedding_lr_scale:
- lr *= self.args.embedding_lr_scale # pylint: disable=invalid-name
+ lr *= self.args.embedding_lr_scale
elif self.args.embedding_lr:
- lr = self.args.embedding_lr # pylint: disable=invalid-name
+ lr = self.args.embedding_lr
optimizer_grouped_parameters.append(
{
"params": list(params["embeddings"].values()),
@@ -143,7 +143,7 @@ class OptimizerMixin(Trainer):
loraplus_lr_embedding = getattr(
self.args, "loraplus_lr_embedding", 1e-6
)
- self.optimizer = create_loraplus_optimizer( # pylint: disable=attribute-defined-outside-init
+ self.optimizer = create_loraplus_optimizer(
opt_model,
optimizer_cls,
loraplus_lr_ratio=loraplus_lr_ratio,
@@ -185,17 +185,15 @@ class OptimizerMixin(Trainer):
p.data_ptr(): p.numel() for p in module.parameters()
}.values()
)
- LOG.info(f"skipped {module}: {skipped/2**20}M params")
+ LOG.info(f"skipped {module}: {skipped / 2**20}M params")
manager.register_module_override(
module, "weight", {"optim_bits": 32}
)
LOG.debug(f"bitsandbytes: will optimize {module} in fp32")
- LOG.info(f"skipped: {skipped/2**20}M params")
+ LOG.info(f"skipped: {skipped / 2**20}M params")
if is_sagemaker_mp_enabled():
- self.optimizer = smp.DistributedOptimizer( # pylint: disable=attribute-defined-outside-init
- self.optimizer
- )
+ self.optimizer = smp.DistributedOptimizer(self.optimizer)
return self.optimizer
diff --git a/src/axolotl/core/trainers/mixins/packing.py b/src/axolotl/core/trainers/mixins/packing.py
new file mode 100644
index 000000000..249ceeb4f
--- /dev/null
+++ b/src/axolotl/core/trainers/mixins/packing.py
@@ -0,0 +1,20 @@
+"""Trainer mixin to support packing"""
+
+from transformers import Trainer
+
+
+class PackingMixin(Trainer):
+ """
+ Trainer mixin to support packing
+ """
+
+ def _set_signature_columns_if_needed(self):
+ super()._set_signature_columns_if_needed()
+ if (
+ self._signature_columns
+ and self.args.sample_packing
+ and self.args.sample_packing_drop_attention_mask
+ ):
+ set_sig_columns = set(self._signature_columns)
+ set_sig_columns.remove("attention_mask")
+ self._signature_columns = list(set_sig_columns)
diff --git a/src/axolotl/core/trainers/mixins/scheduler.py b/src/axolotl/core/trainers/mixins/scheduler.py
index 90070ab78..fc2b0e59d 100644
--- a/src/axolotl/core/trainers/mixins/scheduler.py
+++ b/src/axolotl/core/trainers/mixins/scheduler.py
@@ -7,6 +7,7 @@ from transformers.trainer import Trainer
from axolotl.integrations.base import PluginManager
from axolotl.utils.logging import get_logger
from axolotl.utils.schedulers import (
+ JaggedLRRestartScheduler,
RexLR,
get_cosine_schedule_with_min_lr,
get_cosine_schedule_with_quadratic_warmup,
@@ -45,7 +46,7 @@ class SchedulerMixin(Trainer):
)
# fmt: off
- if self.lr_scheduler is None: # type: ignore # pylint: disable=access-member-before-definition
+ if self.lr_scheduler is None: # type: ignore
# fmt: on
plugin_manager = PluginManager.get_instance()
lr_scheduler: LRScheduler | None = plugin_manager.create_lr_scheduler(
@@ -89,7 +90,7 @@ class SchedulerMixin(Trainer):
LOG.warning(
"Both cosine quadratic warmup and min lr detected. Using quadratic warmup.")
- self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup( # pylint: disable=attribute-defined-outside-init
+ self.lr_scheduler = get_cosine_schedule_with_quadratic_warmup(
optimizer,
num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
num_training_steps=num_training_steps,
@@ -97,7 +98,7 @@ class SchedulerMixin(Trainer):
elif self.args.cosine_min_lr_ratio and self.args.cosine_constant_lr_ratio and use_cosine_min_lr:
assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
assert 0 <= self.args.cosine_constant_lr_ratio <= 1.0, "cosine_constant_lr_ratio must be between 0.0 and 1.0"
- self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant( # pylint: disable=attribute-defined-outside-init
+ self.lr_scheduler = get_cosine_schedule_with_warmup_decay_constant(
optimizer,
num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
num_training_steps=num_training_steps,
@@ -106,14 +107,14 @@ class SchedulerMixin(Trainer):
)
elif self.args.cosine_min_lr_ratio and use_cosine_min_lr:
assert 0 <= self.args.cosine_min_lr_ratio <= 1.0, "cosine_min_lr_ratio must be between 0.0 and 1.0"
- self.lr_scheduler = get_cosine_schedule_with_min_lr( # pylint: disable=attribute-defined-outside-init
+ self.lr_scheduler = get_cosine_schedule_with_min_lr(
optimizer,
num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
num_training_steps=num_training_steps,
min_lr_ratio=self.args.cosine_min_lr_ratio,
)
else:
- return super().create_scheduler(num_training_steps, optimizer=optimizer)
+ super().create_scheduler(num_training_steps, optimizer=optimizer)
else:
if use_cosine_quadratic:
LOG.warning(
@@ -123,4 +124,22 @@ class SchedulerMixin(Trainer):
LOG.warning(
"axolotl's cosine scheduler with min lr not used (e.g., because of deepspeed).")
+ if self.args.jagged_restart_steps:
+ warmup_steps = (
+ self.args.jagged_restart_warmup_steps or 10
+ )
+ anneal_steps = (
+ self.args.jagged_restart_anneal_steps or 1
+ )
+ if not self.lr_scheduler:
+ super().create_scheduler(num_training_steps, optimizer)
+ self.lr_scheduler = JaggedLRRestartScheduler(
+ optimizer,
+ self.lr_scheduler,
+ self.args.jagged_restart_steps,
+ warmup_steps,
+ anneal_steps,
+ min_lr_scale=self.args.cosine_min_lr_ratio or 0.001,
+ )
+
return self.lr_scheduler # type: ignore
diff --git a/src/axolotl/core/trainers/relora.py b/src/axolotl/core/trainers/relora.py
deleted file mode 100644
index 890278f49..000000000
--- a/src/axolotl/core/trainers/relora.py
+++ /dev/null
@@ -1,46 +0,0 @@
-"""Module for ReLoRA trainer"""
-
-import torch
-from torch.optim.lr_scheduler import LRScheduler
-
-from axolotl.core.trainers.base import AxolotlTrainer
-from axolotl.monkeypatch.relora import ReLoRAScheduler
-
-
-class ReLoRATrainer(AxolotlTrainer):
- """Trainer subclass that uses the `OneCycleLR` scheduler"""
-
- tag_names = ["axolotl", "relora"]
-
- def __init__(self, *args, **kwargs):
- super().__init__(*args, **kwargs)
- self.lr_scheduler = None
-
- def create_scheduler(
- self,
- num_training_steps: int,
- optimizer: torch.optim.Optimizer | None = None,
- ) -> LRScheduler:
- optimizer = self.optimizer if optimizer is None else optimizer
- lr_scheduler: LRScheduler = super().create_scheduler(
- num_training_steps, optimizer
- )
-
- if self.args.relora_steps:
- warmup_steps = (
- self.args.relora_warmup_steps if self.args.relora_warmup_steps else 10
- )
- anneal_steps = (
- self.args.relora_anneal_steps if self.args.relora_anneal_steps else 1
- )
- self.lr_scheduler = ReLoRAScheduler( # type: ignore
- optimizer,
- lr_scheduler,
- self.args.relora_steps,
- anneal_steps,
- warmup_steps,
- )
- else:
- self.lr_scheduler = lr_scheduler # type: ignore
-
- return self.lr_scheduler # type: ignore
diff --git a/src/axolotl/core/trainers/trl.py b/src/axolotl/core/trainers/trl.py
index bf0d50dee..c5f19a6fe 100644
--- a/src/axolotl/core/trainers/trl.py
+++ b/src/axolotl/core/trainers/trl.py
@@ -1,81 +1,25 @@
-"""Module for TRL PPO trainer"""
+"""Module for TRL RL trainers"""
-import torch
-from tqdm import tqdm
from trl import (
CPOTrainer,
KTOTrainer,
ORPOTrainer,
- PPOTrainer,
PRMTrainer,
RewardTrainer,
)
-from axolotl.core.trainers.mixins import RngLoaderMixin
+from axolotl.core.trainers.mixins import DistributedParallelMixin, RngLoaderMixin
from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin
from axolotl.core.trainers.mixins.scheduler import SchedulerMixin
-class TRLPPOTrainer(PPOTrainer):
- """Wrapper for TRL PPO trainer to handle customizations"""
-
- tag_names = ["axolotl", "ppo"]
-
- def train(
- self,
- reward_pipe,
- resume_from_checkpoint=None, # pylint: disable=unused-argument
- ):
- generation_kwargs = {
- "min_length": -1,
- "top_k": 0.0,
- "top_p": 1.0,
- "do_sample": True,
- "pad_token_id": self.tokenizer.eos_token_id,
- "max_new_tokens": 32,
- }
- sent_kwargs = {
- "return_all_scores": True,
- "function_to_apply": "none",
- "batch_size": 16,
- }
-
- for _, batch in tqdm(enumerate(self.dataloader)):
- query_tensors = batch["input_ids"]
-
- # generate model response
- response_tensors, ref_response_tensors = self.generate(
- query_tensors,
- return_prompt=False,
- generate_ref_response=True,
- **generation_kwargs,
- )
- batch["response"] = self.tokenizer.batch_decode(response_tensors)
- batch["ref_response"] = self.tokenizer.batch_decode(ref_response_tensors)
-
- # Compute sentiment score
- texts = [q + r for q, r in zip(batch["query"], batch["response"])]
- pipe_outputs = reward_pipe(texts, **sent_kwargs)
- rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
- ref_texts = [q + r for q, r in zip(batch["query"], batch["ref_response"])]
- ref_pipe_outputs = reward_pipe(ref_texts, **sent_kwargs)
- ref_rewards = [
- torch.tensor(output[1]["score"]) for output in ref_pipe_outputs
- ]
- batch["ref_rewards"] = ref_rewards
-
- # Run PPO step
- stats = self.step(query_tensors, response_tensors, rewards)
- self.log_stats(
- stats,
- batch,
- rewards,
- columns_to_log=["query", "response", "ref_response", "ref_rewards"],
- )
-
-
class AxolotlORPOTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, ORPOTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ ORPOTrainer,
):
"""
Extend the base ORPOTrainer for axolotl helpers
@@ -85,7 +29,12 @@ class AxolotlORPOTrainer(
class AxolotlKTOTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, KTOTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ KTOTrainer,
):
"""
Extend the base KTOTrainer for axolotl helpers
@@ -95,7 +44,12 @@ class AxolotlKTOTrainer(
class AxolotlCPOTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, CPOTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ CPOTrainer,
):
"""
Extend the base CPOTrainer for axolotl helpers
@@ -105,7 +59,12 @@ class AxolotlCPOTrainer(
class AxolotlRewardTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, RewardTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ RewardTrainer,
):
"""
Extend the base RewardTrainer for axolotl helpers
@@ -115,7 +74,12 @@ class AxolotlRewardTrainer(
class AxolotlPRMTrainer(
- RngLoaderMixin, SchedulerMixin, OptimizerMixin, OptimizerInitMixin, PRMTrainer
+ RngLoaderMixin,
+ SchedulerMixin,
+ OptimizerMixin,
+ OptimizerInitMixin,
+ DistributedParallelMixin,
+ PRMTrainer,
):
"""
Extend the base trl.PRMTrainer for axolotl helpers
diff --git a/src/axolotl/core/training_args.py b/src/axolotl/core/training_args.py
index 42488e643..d5be9fc62 100644
--- a/src/axolotl/core/training_args.py
+++ b/src/axolotl/core/training_args.py
@@ -2,238 +2,17 @@
extra axolotl specific training args
"""
-from dataclasses import dataclass, field
-from typing import Optional
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional, Type
-from PIL.Image import Resampling
from transformers import TrainingArguments
from trl import CPOConfig, KTOConfig, ORPOConfig, PRMConfig, RewardConfig
+from axolotl.integrations.config import merge_training_args
-@dataclass
-class AxolotlTrainingMixins:
- """
- Mixin class for the Axolotl training args.
- """
-
- # pylint: disable=duplicate-code
- model_type: Optional[str] = field(
- default=None, metadata={"help": "HF model configuration model_type."}
- )
- lr_quadratic_warmup: bool = field(
- default=False,
- metadata={"help": "Use quadratic warmup for cosine scheduling."},
- )
- pretraining: bool = field(
- default=False,
- metadata={
- "help": "Indicates to trainer whether we are doing continued pretraining."
- },
- )
- sample_packing: bool = field(
- default=False,
- metadata={"help": "Use sample packing for efficient training."},
- )
- sample_packing_sequentially: bool = field(
- default=False,
- metadata={
- "help": "Use next-fit sample packing that preserves the order of samples coming from the sampler. Use in combination with curriculum_sampling for fully sequential packing."
- },
- )
- multipack_real_batches: bool = field(
- default=False,
- metadata={"help": "Use real batches for efficient training."},
- )
- eval_sample_packing: Optional[bool] = field(
- default=None,
- metadata={"help": "Use sample packing for efficient evals."},
- )
- sample_packing_efficiency: float = field(
- default=1.0,
- metadata={"help": "Sample packing efficiency for calculating batch length."},
- )
- sample_packing_bin_size: int = field(
- default=200,
- metadata={
- "help": "The max number of samples that packed sample can contain after packing. Increase for better packing."
- },
- )
- sample_packing_group_size: int = field(
- default=100000,
- metadata={
- "help": "The number of samples to group together for packing. Increase for better packing."
- },
- )
- max_seq_length: int = field(
- default=2048,
- metadata={"help": "The maximum sequence length the model can handle"},
- )
- relora_steps: Optional[int] = field(
- default=None,
- metadata={"help": "how often to reset for ReLoRA"},
- )
- relora_warmup_steps: Optional[int] = field(
- default=None,
- metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
- )
- relora_anneal_steps: Optional[int] = field(
- default=None,
- metadata={"help": "how many warmup steps to take after reset for ReLoRA"},
- )
- relora_prune_ratio: Optional[float] = field(
- default=0.9,
- metadata={"help": "prune ratio for magnitude pruning of the optimizer"},
- )
- bench_split: Optional[str] = field(
- default="eval", metadata={"help": "The benchmark split to run on"}
- )
- bench_dataset: Optional[str] = field(
- default="pharaouk/dharma-1/dharma_1_mini.json",
- metadata={
- "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
- },
- )
- do_bench_eval: Optional[bool] = field(
- default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
- )
- do_causal_lm_eval: Optional[bool] = field(
- default=False, metadata={"help": "Whether to run the Causal LM evaluation."}
- )
- max_bench_samples: Optional[int] = field(
- default=None,
- metadata={
- "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
- },
- )
- bench_source_max_len: int = field(
- default=2048, metadata={"help": "Maximum source sequence length for bench."}
- )
- dataloader_prefetch_factor: Optional[int] = field(
- default=None,
- metadata={"help": "prefetch_factor argument to the dataloader"},
- )
- cosine_min_lr_ratio: Optional[float] = field(
- default=None,
- metadata={"help": "Minimum learning rate is min_lr_ratio * learning_rate"},
- )
- cosine_constant_lr_ratio: Optional[float] = field(
- default=None,
- metadata={
- "help": "Starting constant learning rate step is cosine_constant_lr_ratio * max_steps"
- },
- )
- loraplus_lr_ratio: Optional[float] = field(
- default=None, metadata={"help": "loraplus learning rate ratio lr_B / lr_A."}
- )
- loraplus_lr_embedding: Optional[float] = field(
- default=1e-6,
- metadata={"help": "loraplus learning rate for lora embedding layers."},
- )
- embedding_lr_scale: Optional[float] = field(
- default=None,
- metadata={"help": "Scale the learning rate for the embedding layers."},
- )
- lr_groups: Optional[list[dict]] = field(
- default=None,
- metadata={"help": "Specify learning rate groups for with different LRs."},
- )
- embedding_lr: Optional[float] = field(
- default=None,
- metadata={"help": "absolute learning rate for the embedding layers."},
- )
- qlora: bool = field(
- default=False,
- metadata={"help": "whether this is a qlora training"},
- )
- orpo_alpha: Optional[float] = field(
- default=None,
- )
- lisa_n_layers: Optional[int] = field(
- default=None,
- metadata={"help": "the number of activate layers in LISA"},
- )
- lisa_step_interval: Optional[int] = field(
- default=None,
- metadata={"help": "how often to switch layers in LISA"},
- )
- lisa_layers_attribute: Optional[str] = field(
- default=None,
- metadata={"help": "path under the model to access the layers"},
- )
- curriculum_sampling: Optional[bool] = field(
- default=None,
- metadata={"help": "whether to use sequential sampling for curriculum learning"},
- )
- alternate_lr_scheduler_type: Optional[str] = field(
- default=None,
- metadata={
- "help": "workaround to pass an alternate lr scheduler to the HF trainer"
- },
- )
- chat_template: Optional[str] = field(
- default=None,
- metadata={"help": "Chat template converting chat messages to text"},
- )
-
- kd_ce_alpha: Optional[float] = field(
- default=None,
- metadata={
- "help": "The alpha scaling parameter for SFT cross entropy loss when using KD"
- },
- )
-
- kd_alpha: Optional[float] = field(
- default=1.0,
- metadata={"help": "The alpha scaling parameter for KD loss"},
- )
-
- kd_temperature: Optional[float] = field(
- default=1.0,
- metadata={
- "help": "the temperature parameter for KL divergence loss when using KD"
- },
- )
-
- kd_zscore_base_temp: Optional[float] = field(
- default=None,
- metadata={
- "help": "the base temperature parameter for KL divergence with z-score when using KD"
- },
- )
-
- kd_top_k_before_softmax: Optional[bool] = field(
- default=None,
- metadata={
- "help": "Whether to apply top_k_before_softmax to the logits when using KD"
- },
- )
-
- adam_beta3: Optional[float] = field(
- default=None,
- metadata={
- "help": "The beta3 hyperparameter used in some optimizers such as CAME"
- },
- )
- adam_epsilon2: Optional[float] = field(
- default=None,
- metadata={
- "help": "The epsilon2 hyperparameter used in some optimizers such as CAME"
- },
- )
-
- # multi-modal section
-
- image_size: int | tuple[int, int] | None = field(
- default=None,
- metadata={"help": "The size of the image to resize to"},
- )
-
- image_resize_algorithm: Resampling | None = field(
- default=None,
- metadata={"help": "The algorithm to use for image resizing"},
- )
-
- # end of multi-modal section
+AxolotlTrainingMixins: Type = merge_training_args()
@dataclass
diff --git a/src/axolotl/core/training_args_base.py b/src/axolotl/core/training_args_base.py
new file mode 100644
index 000000000..41ee8e91e
--- /dev/null
+++ b/src/axolotl/core/training_args_base.py
@@ -0,0 +1,265 @@
+"""
+Base Axolotl Training Mixins shared across various trainer configs
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from PIL.Image import Resampling
+
+
+@dataclass
+class AxolotlTrainingMixins:
+ """
+ Mixin class for the Axolotl training args.
+ """
+
+ model_type: Optional[str] = field(
+ default=None, metadata={"help": "HF model configuration model_type."}
+ )
+ lr_quadratic_warmup: bool = field(
+ default=False,
+ metadata={"help": "Use quadratic warmup for cosine scheduling."},
+ )
+ pretraining: bool = field(
+ default=False,
+ metadata={
+ "help": "Indicates to trainer whether we are doing continued pretraining."
+ },
+ )
+ sample_packing: bool = field(
+ default=False,
+ metadata={"help": "Use sample packing for efficient training."},
+ )
+ sample_packing_sequentially: bool = field(
+ default=False,
+ metadata={
+ "help": "Use next-fit sample packing that preserves the order of samples coming from the sampler. Use in combination with curriculum_sampling for fully sequential packing."
+ },
+ )
+ sample_packing_mp_start_method: str | None = field(
+ default=None,
+ metadata={"help": "The multiprocessing start method to use."},
+ )
+ sample_packing_drop_attention_mask: bool = field(
+ default=False,
+ metadata={"help": "Drop attention mask from inputs when using packing."},
+ )
+ multipack_real_batches: bool = field(
+ default=False,
+ metadata={"help": "Use real batches for efficient training."},
+ )
+ include_tkps: bool = field(
+ default=True,
+ metadata={
+ "help": "Whether to include tokens per second in the training metrics."
+ },
+ )
+ eval_sample_packing: Optional[bool] = field(
+ default=None,
+ metadata={"help": "Use sample packing for efficient evals."},
+ )
+ sample_packing_efficiency: float = field(
+ default=1.0,
+ metadata={"help": "Sample packing efficiency for calculating batch length."},
+ )
+ sample_packing_bin_size: int = field(
+ default=200,
+ metadata={
+ "help": "The max number of samples that packed sample can contain after packing. Increase for better packing."
+ },
+ )
+ sample_packing_group_size: int = field(
+ default=100000,
+ metadata={
+ "help": "The number of samples to group together for packing. Increase for better packing."
+ },
+ )
+ max_seq_length: int = field(
+ default=2048,
+ metadata={"help": "The maximum sequence length the model can handle"},
+ )
+ dataset_num_proc: int | None = field(
+ default=None,
+ metadata={"help": "The number of processes to use for data processing"},
+ )
+ relora_steps: Optional[int] = field(
+ default=None,
+ metadata={"help": "how often to reset for ReLoRA"},
+ )
+ relora_prune_ratio: Optional[float] = field(
+ default=0.9,
+ metadata={"help": "prune ratio for magnitude pruning of the optimizer"},
+ )
+ jagged_restart_steps: Optional[int] = field(
+ default=None,
+ metadata={"help": "how often to reset for jagged restarts"},
+ )
+ jagged_restart_warmup_steps: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "how many warmup steps to take after reset for jagged restarts"
+ },
+ )
+ jagged_restart_anneal_steps: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "how many anneal steps to take before reset for jagged restarts"
+ },
+ )
+ bench_split: Optional[str] = field(
+ default="eval", metadata={"help": "The benchmark split to run on"}
+ )
+ bench_dataset: Optional[str] = field(
+ default="pharaouk/dharma-1/dharma_1_mini.json",
+ metadata={
+ "help": "Benchmark dataset to use: options are `mmlu-zs`, `mmlu-fs`, or the full path to the dataset file"
+ },
+ )
+ do_bench_eval: Optional[bool] = field(
+ default=False, metadata={"help": "Whether to run the Benchmark evaluation."}
+ )
+ do_causal_lm_eval: Optional[bool] = field(
+ default=False, metadata={"help": "Whether to run the Causal LM evaluation."}
+ )
+ max_bench_samples: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "If set, only evaluates on `max_bench_samples` of the benchmark dataset."
+ },
+ )
+ bench_source_max_len: int = field(
+ default=2048, metadata={"help": "Maximum source sequence length for bench."}
+ )
+ dataloader_prefetch_factor: Optional[int] = field(
+ default=None,
+ metadata={"help": "prefetch_factor argument to the dataloader"},
+ )
+ cosine_min_lr_ratio: Optional[float] = field(
+ default=None,
+ metadata={"help": "Minimum learning rate is min_lr_ratio * learning_rate"},
+ )
+ cosine_constant_lr_ratio: Optional[float] = field(
+ default=None,
+ metadata={
+ "help": "Starting constant learning rate step is cosine_constant_lr_ratio * max_steps"
+ },
+ )
+ loraplus_lr_ratio: Optional[float] = field(
+ default=None, metadata={"help": "loraplus learning rate ratio lr_B / lr_A."}
+ )
+ loraplus_lr_embedding: Optional[float] = field(
+ default=1e-6,
+ metadata={"help": "loraplus learning rate for lora embedding layers."},
+ )
+ embedding_lr_scale: Optional[float] = field(
+ default=None,
+ metadata={"help": "Scale the learning rate for the embedding layers."},
+ )
+ lr_groups: Optional[list[dict]] = field(
+ default=None,
+ metadata={"help": "Specify learning rate groups for with different LRs."},
+ )
+ embedding_lr: Optional[float] = field(
+ default=None,
+ metadata={"help": "absolute learning rate for the embedding layers."},
+ )
+ qlora: bool = field(
+ default=False,
+ metadata={"help": "whether this is a qlora training"},
+ )
+ orpo_alpha: Optional[float] = field(
+ default=None,
+ )
+ lisa_n_layers: Optional[int] = field(
+ default=None,
+ metadata={"help": "the number of activate layers in LISA"},
+ )
+ lisa_step_interval: Optional[int] = field(
+ default=None,
+ metadata={"help": "how often to switch layers in LISA"},
+ )
+ lisa_layers_attribute: Optional[str] = field(
+ default=None,
+ metadata={"help": "path under the model to access the layers"},
+ )
+ curriculum_sampling: Optional[bool] = field(
+ default=None,
+ metadata={"help": "whether to use sequential sampling for curriculum learning"},
+ )
+ alternate_lr_scheduler_type: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "workaround to pass an alternate lr scheduler to the HF trainer"
+ },
+ )
+ chat_template: Optional[str] = field(
+ default=None,
+ metadata={"help": "Chat template converting chat messages to text"},
+ )
+
+ # kd_ce_alpha: Optional[float] = field(
+ # default=None,
+ # metadata={
+ # "help": "The alpha scaling parameter for SFT cross entropy loss when using KD"
+ # },
+ # )
+ #
+ # kd_alpha: Optional[float] = field(
+ # default=1.0,
+ # metadata={"help": "The alpha scaling parameter for KD loss"},
+ # )
+ #
+ # kd_temperature: Optional[float] = field(
+ # default=1.0,
+ # metadata={
+ # "help": "the temperature parameter for KL divergence loss when using KD"
+ # },
+ # )
+
+ adam_beta3: Optional[float] = field(
+ default=None,
+ metadata={
+ "help": "The beta3 hyperparameter used in some optimizers such as CAME"
+ },
+ )
+ adam_epsilon2: Optional[float] = field(
+ default=None,
+ metadata={
+ "help": "The epsilon2 hyperparameter used in some optimizers such as CAME"
+ },
+ )
+
+ activation_offloading: bool | None = field(
+ default=None,
+ metadata={"help": "Use activation offloading with CUDA streams for training."},
+ )
+
+ # multi-modal section
+
+ image_size: int | tuple[int, int] | None = field(
+ default=None,
+ metadata={"help": "The size of the image to resize to"},
+ )
+
+ image_resize_algorithm: Resampling | None = field(
+ default=None,
+ metadata={"help": "The algorithm to use for image resizing"},
+ )
+
+ # end of multi-modal section
+
+ dion_learning_rate: float | None = field(
+ default=None,
+ metadata={"help": "The learning rate for Dion"},
+ )
+ dion_momentum: float | None = field(
+ default=None,
+ metadata={"help": "The momentum for Dion"},
+ )
+ dion_rank_fraction: float | None = field(
+ default=None,
+ )
+ dion_rank_multiple_of: int | None = field(
+ default=None,
+ )
diff --git a/src/axolotl/datasets.py b/src/axolotl/datasets.py
index 9f1d9500d..20acb8521 100644
--- a/src/axolotl/datasets.py
+++ b/src/axolotl/datasets.py
@@ -1,40 +1,36 @@
-"""Module containing Dataset functionality"""
+"""
+Module containing dataset functionality.
-import os
-from typing import List, Optional, Union
+We want this to be a wrapper for an existing dataset that we have loaded. Lets use the
+concept of middlewares to wrap each dataset. We'll use the collators later on to pad the
+datasets.
+"""
-import torch
from datasets import Dataset, IterableDataset
from axolotl.utils.logging import get_logger
from .prompt_tokenizers import PromptTokenizingStrategy
-# We want this to be a wrapper for an existing dataset that we have loaded
-# lets use the concept of middlewares to wrap each dataset, for example
-# ConstantLengthDataset(ShuffledDataset([TokenizedPromptDataset(alpaca_dataset)]))
-# let's check to ensure we don't truncate an item in the middle, we'll use
-# the collators later on to pad the datasets
-
LOG = get_logger(__name__)
class TokenizedPromptDataset(Dataset):
- """
- Dataset that returns tokenized prompts from a stream of text files.
- Args:
- prompt_tokenizer (PromptTokenizingStrategy): The prompt tokenizing method for processing the data.
- dataset (dataset.Dataset): Dataset with text files.
- process_count (int): Number of processes to use for tokenizing.
- keep_in_memory (bool): Whether to keep the tokenized dataset in memory.
+ """Dataset that returns tokenized prompts from a stream of text files.
+
+ Args:
+ prompt_tokenizer: The prompt tokenizing method for processing the data.
+ dataset: Dataset with text files.
+ process_count: Number of processes to use for tokenizing.
+ keep_in_memory: Whether to keep the tokenized dataset in memory.
"""
- def __init__( # pylint: disable=super-init-not-called
+ def __init__(
self,
prompt_tokenizer: PromptTokenizingStrategy,
dataset: Dataset,
- process_count: Optional[int] = None,
- keep_in_memory: Optional[bool] = False,
+ process_count: int | None = None,
+ keep_in_memory: bool | None = False,
**kwargs,
):
self.prompt_tokenizer = prompt_tokenizer
@@ -47,7 +43,6 @@ class TokenizedPromptDataset(Dataset):
def process(self, dataset):
features = dataset.features.keys()
- num_proc = min(64, self.process_count if self.process_count else os.cpu_count())
map_kwargs = {}
if self.prompt_tokenizer.supports_batched:
@@ -60,13 +55,13 @@ class TokenizedPromptDataset(Dataset):
):
dataset = dataset.filter(
self.prompt_tokenizer.filter_rows,
- num_proc=num_proc,
+ num_proc=self.process_count,
desc="Strategy Filtering Rows",
)
return dataset.map(
self.prompt_tokenizer.tokenize_prompt,
- num_proc=num_proc,
+ num_proc=self.process_count,
remove_columns=features,
keep_in_memory=self.keep_in_memory,
desc="Tokenizing Prompts",
@@ -76,143 +71,17 @@ class TokenizedPromptDataset(Dataset):
def wrap_dataset_for_tokenized_prompt(
prompt_tokenizer: PromptTokenizingStrategy,
- dataset: Union[Dataset, IterableDataset],
+ dataset: Dataset | IterableDataset,
**kwargs,
):
if isinstance(dataset, IterableDataset):
map_kwargs = {}
if prompt_tokenizer.supports_batched:
map_kwargs["batched"] = True
- features = dataset.features.keys()
+ features = list(dataset.features.keys())
return dataset.map(
prompt_tokenizer.tokenize_prompt,
remove_columns=features,
**map_kwargs,
)
return TokenizedPromptDataset(prompt_tokenizer, dataset, **kwargs)
-
-
-# TODO this isn't the best since it can't interleave datasets
-class ConstantLengthDataset(IterableDataset):
- """
- Iterable dataset that returns constant length chunks of tokens from stream of text files.
- Args:
- tokenizer (Tokenizer): The processor used for processing the data.
- dataset (dataset.Dataset): Dataset with text files.
- seq_length (int): Length of token sequences to return.
- """
-
- def __init__( # pylint: disable=super-init-not-called
- self,
- tokenizer,
- datasets,
- seq_length=2048,
- ):
- self.tokenizer = tokenizer
- self.concat_token_id = tokenizer.eos_token_id
- self.datasets: List[IterableDataset] = datasets
- self.seq_length = seq_length
-
- vocab_size = len(tokenizer.get_vocab())
-
- if vocab_size <= torch.iinfo(torch.int16).max:
- self.tokens_dtype = torch.int16
- elif vocab_size <= torch.iinfo(torch.int32).max:
- self.tokens_dtype = torch.int32
- else:
- self.tokens_dtype = torch.int64
-
- def __iter__(self):
- buffer = {
- "input_ids": [],
- "attention_mask": [],
- "labels": [],
- "position_ids": [],
- }
- buffer_len = 0
- for dataset in self.datasets:
- idx = 0
- iterator = iter(dataset)
- more_examples = True
- while more_examples:
- try:
- example = next(iterator)
- idx += 1
- except StopIteration:
- more_examples = False
- example = None
-
- add_concat_token = False
- if example:
- example_len = len(example["input_ids"])
- add_concat_token = example["input_ids"][-1] != self.concat_token_id
- else:
- example_len = 0
-
- if not example_len or (
- buffer_len + int(add_concat_token) + example_len > self.seq_length
- ):
- if buffer["input_ids"]:
- input_ids = torch.cat(buffer["input_ids"], dim=-1)[
- : self.seq_length
- ]
- attention_mask = torch.cat(buffer["attention_mask"], dim=-1)[
- : self.seq_length
- ]
- position_ids = torch.cat(buffer["position_ids"], dim=-1)[
- : self.seq_length
- ]
- labels = torch.cat(buffer["labels"], dim=-1)[: self.seq_length]
- if labels.size() == input_ids.size() and (
- attention_mask.size() == input_ids.size()
- ):
- yield {
- "input_ids": input_ids,
- "labels": labels,
- "attention_mask": attention_mask,
- "position_ids": position_ids,
- }
- else:
- LOG.warning(
- f"dropping batch due to tensor size mismatch input_ids: {input_ids.size()}, labels: {labels.size()}, attention_mask: {attention_mask.size()}"
- )
- buffer = {
- "input_ids": [],
- "attention_mask": [],
- "labels": [],
- "position_ids": [],
- }
- buffer_len = 0
- idx = 1
-
- if example:
- # FIXME
- # just going to drop data points that are too long
- if len(example["input_ids"]) <= self.seq_length:
- input_ids = example["input_ids"]
- attention_mask = example["attention_mask"]
- labels = example["labels"]
-
- if add_concat_token:
- input_ids.append(self.concat_token_id)
- attention_mask.append(1)
- labels.append(self.concat_token_id)
-
- input_ids_with_concat = torch.tensor(
- input_ids, dtype=self.tokens_dtype
- )
- attention_mask_with_concat = torch.tensor(
- [idx * m for m in attention_mask], dtype=torch.int16
- )
- labels_with_concat = torch.tensor(
- labels, dtype=self.tokens_dtype
- )
- position_ids = torch.arange(
- len(input_ids), dtype=self.tokens_dtype
- )
-
- buffer["input_ids"].append(input_ids_with_concat)
- buffer["attention_mask"].append(attention_mask_with_concat)
- buffer["labels"].append(labels_with_concat)
- buffer["position_ids"].append(position_ids)
- buffer_len += len(input_ids)
diff --git a/src/axolotl/evaluate.py b/src/axolotl/evaluate.py
index 61a9f8fad..db6fb3f16 100644
--- a/src/axolotl/evaluate.py
+++ b/src/axolotl/evaluate.py
@@ -7,7 +7,6 @@ from pathlib import Path
from typing import Dict, Optional
import torch
-from accelerate.logging import get_logger
from datasets import Dataset
from transformers.trainer import Trainer
@@ -18,6 +17,7 @@ from axolotl.train import (
)
from axolotl.utils.dict import DictDefault
from axolotl.utils.distributed import cleanup_distributed
+from axolotl.utils.logging import get_logger
from axolotl.utils.trainer import setup_trainer
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -81,7 +81,7 @@ def evaluate(*, cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> Dict[str, f
model, tokenizer, _, processor = setup_model_and_tokenizer(cfg)
# Get datasets
- # pylint: disable=duplicate-code
+
train_dataset = dataset_meta.train_dataset
eval_dataset = dataset_meta.eval_dataset
total_num_steps = dataset_meta.total_num_steps
diff --git a/src/axolotl/integrations/base.py b/src/axolotl/integrations/base.py
index 0edc9fdea..c66bc01c6 100644
--- a/src/axolotl/integrations/base.py
+++ b/src/axolotl/integrations/base.py
@@ -22,17 +22,20 @@ from __future__ import annotations
import collections
import importlib
+import traceback
from typing import TYPE_CHECKING, Callable, OrderedDict, Union
from peft import PeftModel
+from torch import nn
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
from transformers import PreTrainedModel, Trainer
+from transformers.trainer_pt_utils import get_parameter_names
from axolotl.utils.dict import DictDefault
from axolotl.utils.logging import get_logger
-LOG = get_logger(__name__, use_environ=True)
+LOG = get_logger(__name__)
if TYPE_CHECKING:
from axolotl.common.datasets import TrainDatasetMeta
@@ -73,8 +76,8 @@ class BasePlugin:
def __init__(self):
"""Initializes the BasePlugin."""
- def register(self, cfg: DictDefault): # pylint: disable=unused-argument
- """Registers the plugin with the given configuration.
+ def register(self, cfg: dict):
+ """Registers the plugin with the given configuration as an unparsed dict.
Args:
cfg: The configuration for the plugin.
@@ -83,6 +86,11 @@ class BasePlugin:
def get_input_args(self) -> str | None:
"""Returns a pydantic model for the plugin's input arguments."""
+ def get_training_args_mixin(self) -> str | None:
+ """
+ Returns a dataclass model for the plugin's training arguments.
+ """
+
def load_datasets(
self, cfg: DictDefault, preprocess: bool = False
) -> Union["TrainDatasetMeta", None]:
@@ -96,14 +104,13 @@ class BasePlugin:
dataset_meta: The metadata for the training dataset.
"""
- def pre_model_load(self, cfg: DictDefault): # pylint: disable=unused-argument
+ def pre_model_load(self, cfg: DictDefault):
"""Performs actions before the model is loaded.
Args:
cfg: The configuration for the plugin.
"""
- # pylint: disable=unused-argument
def post_model_build(self, cfg: DictDefault, model: PreTrainedModel):
"""Performs actions after the model is built/loaded, but before any adapters are applied.
@@ -111,7 +118,6 @@ class BasePlugin:
cfg: The configuration for the plugin.
"""
- # pylint: disable=unused-argument
def pre_lora_load(self, cfg: DictDefault, model: PreTrainedModel):
"""Performs actions before LoRA weights are loaded.
@@ -120,7 +126,6 @@ class BasePlugin:
model: The loaded model.
"""
- # pylint: disable=unused-argument
def post_lora_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
"""Performs actions after LoRA weights are loaded.
@@ -129,7 +134,6 @@ class BasePlugin:
model: The loaded model.
"""
- # pylint: disable=unused-argument
def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
"""Performs actions after the model is loaded.
@@ -138,8 +142,7 @@ class BasePlugin:
model: The loaded model.
"""
- # pylint: disable=unused-argument
- def get_trainer_cls(self, cfg: DictDefault) -> Trainer | None:
+ def get_trainer_cls(self, cfg: DictDefault) -> type[Trainer] | None:
"""Returns a custom class for the trainer.
Args:
@@ -149,7 +152,6 @@ class BasePlugin:
The first non-`None` trainer class returned by a plugin.
"""
- # pylint: disable=unused-argument
def post_trainer_create(self, cfg: DictDefault, trainer: Trainer):
"""Performs actions after the trainer is created.
@@ -158,7 +160,29 @@ class BasePlugin:
trainer: The trainer object for training.
"""
- # pylint: disable=unused-argument
+ def get_training_args(self, cfg: DictDefault):
+ """
+ Returns custom training arguments to set on TrainingArgs.
+
+ Args:
+ cfg: The global axolotl configuration.
+
+ Returns:
+ object: dict containing the training arguments.
+ """
+
+ def get_collator_cls_and_kwargs(self, cfg: DictDefault, is_eval: bool = False):
+ """
+ Returns a custom class for the collator.
+
+ Args:
+ cfg: The global axolotl configuration.
+ is_eval: Whether this is an eval split.
+
+ Returns:
+ class: The class for the collator.
+ """
+
def create_optimizer(self, cfg: DictDefault, trainer: Trainer) -> Optimizer | None:
"""Creates and returns an optimizer for training.
@@ -170,7 +194,6 @@ class BasePlugin:
The created optimizer.
"""
- # pylint: disable=unused-argument
def create_lr_scheduler(
self,
cfg: DictDefault,
@@ -190,7 +213,6 @@ class BasePlugin:
The created learning rate scheduler.
"""
- # pylint: disable=unused-argument
def add_callbacks_pre_trainer(
self, cfg: DictDefault, model: PreTrainedModel
) -> list[Callable]:
@@ -205,7 +227,6 @@ class BasePlugin:
"""
return []
- # pylint: disable=unused-argument
def add_callbacks_post_trainer(
self, cfg: DictDefault, trainer: Trainer
) -> list[Callable]:
@@ -221,7 +242,6 @@ class BasePlugin:
"""
return []
- # pylint: disable=unused-argument
def post_train(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
"""Performs actions after training is complete.
@@ -230,7 +250,7 @@ class BasePlugin:
model: The loaded model.
"""
- def post_train_unload(self, cfg: DictDefault): # pylint: disable=unused-argument
+ def post_train_unload(self, cfg: DictDefault):
"""Performs actions after training is complete and the model is unloaded.
Args:
@@ -337,8 +357,11 @@ class PluginManager:
plugin = load_plugin(plugin_name)
self.plugins[plugin_name] = plugin
LOG.info(f"Plugin loaded successfully: {plugin_name}")
- except ImportError:
+ except ImportError as exc:
LOG.error(f"Failed to load plugin: {plugin_name}")
+ # print stacktrace
+ traceback.print_exc()
+ print(f"Error: {exc}")
def get_input_args(self) -> list[str]:
"""Returns a list of Pydantic classes for all registered plugins' input arguments.'
@@ -353,6 +376,20 @@ class PluginManager:
input_args.append(input_args_from_plugin)
return input_args
+ def get_training_args_mixin(self):
+ """
+ Returns a list of dataclasses for all registered plugins' training args mixins'
+
+ Returns:
+ list[str]: A list of dataclsses
+ """
+ training_args = []
+ for plugin in self.plugins.values():
+ training_args_from_plugin = plugin.get_training_args_mixin()
+ if training_args_from_plugin is not None:
+ training_args.append(training_args_from_plugin)
+ return training_args
+
def load_datasets(
self, cfg: DictDefault, preprocess: bool = False
) -> Union["TrainDatasetMeta", None]:
@@ -442,6 +479,42 @@ class PluginManager:
return trainer_cls
return None
+ def get_training_args(self, cfg):
+ """
+ Calls the get_training_args method of all registered plugins and returns the combined training arguments.
+
+ Parameters:
+ cfg (dict): The configuration for the plugins.
+
+ Returns:
+ object: The training arguments
+ """
+ training_args_kwargs = {}
+ for plugin in self.plugins.values():
+ training_args = plugin.get_training_args(cfg)
+ if training_args is not None:
+ training_args_kwargs.update(training_args)
+
+ return training_args_kwargs
+
+ def get_collator_cls_and_kwargs(self, cfg, is_eval=False):
+ """
+ Calls the get_collator_cls_and_kwargs method of all registered plugins and returns the first non-None collator class.
+
+ Parameters:
+ cfg (dict): The configuration for the plugins.
+ is_eval (bool): Whether this is an eval split.
+
+ Returns:
+ object: The collator class, or None if none was found.
+ """
+ for plugin in self.plugins.values():
+ collator = plugin.get_collator_cls_and_kwargs(cfg, is_eval=is_eval)
+ if collator is not None:
+ collator_cls, collator_kwargs = collator
+ return collator_cls, collator_kwargs
+ return None
+
def post_trainer_create(self, cfg: DictDefault, trainer: Trainer):
"""Calls the `post_trainer_create` method of all registered plugins.
@@ -557,3 +630,24 @@ class BaseOptimizerFactory:
self, opt_model, training_args, **optimizer_kwargs
) -> Optimizer | None:
pass
+
+ # duplicated from transformers
+ def get_decay_parameter_names(self, model) -> list[str]:
+ """
+ Get all parameter names that weight decay will be applied to.
+
+ This function filters out parameters in two ways:
+ 1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS)
+ 2. By parameter name patterns (containing 'bias', or variation of 'norm')
+ """
+ forbidden_name_patterns = [
+ r"bias",
+ r"layernorm",
+ r"rmsnorm",
+ r"(?:^|\.)norm(?:$|\.)",
+ r"_norm(?:$|\.)",
+ ]
+ decay_parameters = get_parameter_names(
+ model, [nn.LayerNorm], forbidden_name_patterns
+ )
+ return decay_parameters
diff --git a/src/axolotl/integrations/config.py b/src/axolotl/integrations/config.py
index b443f228e..8ae8aab39 100644
--- a/src/axolotl/integrations/config.py
+++ b/src/axolotl/integrations/config.py
@@ -16,12 +16,12 @@ Module to handle merging the plugins' input arguments with the base configuratio
This was moved here to prevent circular imports.
"""
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Type
from axolotl.utils.schemas.config import (
AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
+ AxolotlInputConfig as AxolotlInputConfigBase,
)
-from axolotl.utils.schemas.config import AxolotlInputConfig as AxolotlInputConfigBase
def merge_input_args():
@@ -50,14 +50,44 @@ def merge_input_args():
dynamic_input += f"class AxolotlInputConfig(AxolotlInputConfigBase, {', '.join(plugin_classes)}):\n pass\n"
namespace: Dict[Any, Any] = {}
- exec( # pylint: disable=exec-used # nosec B102
- dynamic_input, globals(), namespace
- )
- AxolotlInputConfig = namespace[ # pylint: disable=invalid-name
- "AxolotlInputConfig"
- ]
- AxolotlConfigWCapabilities = namespace[ # pylint: disable=invalid-name
- "AxolotlConfigWCapabilities"
- ]
+ exec(dynamic_input, globals(), namespace) # nosec B102
+ AxolotlInputConfig = namespace["AxolotlInputConfig"]
+ AxolotlConfigWCapabilities = namespace["AxolotlConfigWCapabilities"]
return AxolotlConfigWCapabilities, AxolotlInputConfig
return AxolotlConfigWCapabilitiesBase, AxolotlInputConfigBase
+
+
+def merge_training_args() -> Type:
+ """
+ Merges training arguments from registered plugins with the base TrainingArguments.
+
+ This function retrieves the training arguments from registered plugins using the PluginManager.
+ It then dynamically creates new classes, AxolotlTrainingMixins,
+ that inherit from the base configurations and include the training arguments from the plugins.
+
+ Returns:
+ tuple: A tuple containing the newly created classes, AxolotlTrainingMixins.
+ """
+
+ from axolotl.core.training_args_base import (
+ AxolotlTrainingMixins as AxolotlTrainingMixinsBase,
+ )
+ from axolotl.integrations.base import PluginManager
+
+ plugin_manager = PluginManager.get_instance()
+ training_args_mixins: List[str] = plugin_manager.get_training_args_mixin()
+ mixin_classes = []
+ dynamic_input = ""
+ for plugin_args in training_args_mixins:
+ plugin_module, plugin_cls = plugin_args.rsplit(".", 1)
+ dynamic_input += f"from {plugin_module} import {plugin_cls}\n"
+ mixin_classes.append(plugin_cls)
+ if dynamic_input:
+ dynamic_input += f"class AxolotlTrainingMixins(AxolotlTrainingMixinsBase, {', '.join(mixin_classes)}):\n pass\n"
+
+ namespace: Dict[Any, Any] = {}
+ local_vars = {"AxolotlTrainingMixinsBase": AxolotlTrainingMixinsBase}
+ exec(dynamic_input, {**globals(), **local_vars}, namespace) # nosec B102
+ AxolotlTrainingMixins = namespace["AxolotlTrainingMixins"]
+ return AxolotlTrainingMixins
+ return AxolotlTrainingMixinsBase
diff --git a/src/axolotl/integrations/cut_cross_entropy/README.md b/src/axolotl/integrations/cut_cross_entropy/README.md
index 627ebd935..5c7c5166b 100644
--- a/src/axolotl/integrations/cut_cross_entropy/README.md
+++ b/src/axolotl/integrations/cut_cross_entropy/README.md
@@ -19,7 +19,7 @@ python scripts/cutcrossentropy_install.py | sh
- If you are installing from pip
```bash
-pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@bad6f7b49c75fdec69471abb71b4cddd0f0c6438"
+pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@8a1a0ec"
```
## Usage
@@ -31,27 +31,55 @@ plugins:
## Supported Models
-- llama
-- llama4
-- llama4_text
-- mllama
-- phi3
+- apertus
+- arcee
+- cohere
+- cohere2
+- deepseek_v3
- gemma
- gemma2
- gemma3
- gemma3_text
+- gemma3n
+- gemma3n_text
+- glm
+- glm4
+- glm4_moe
+- glm4v
+- glm4v_moe
+- gpt_oss
+- granite
+- granitemoe
+- granitemoeshared
+- granitemoehybrid
+- hunyuan_v1_dense
+- hunyuan_v1_moe
+- lfm2
+- lfm2_moe
+- lfm2_vl
+- llama
+- llama4
+- llama4_text
+- llava
- mistral
- mistral3
+- mixtral
+- mllama
+- phi
+- phi3
+- phi4_multimodal
- qwen2
-- qwen2_moe
- qwen2_vl
+- qwen2_moe
- qwen2_5_vl
- qwen3
- qwen3_moe
-- cohere
-- cohere2
-- glm
-- glm4
+- qwen3_vl
+- qwen3_vl_moe
+- qwen3_next
+- smollm3
+- seed_oss
+- voxtral
## Citation
diff --git a/src/axolotl/integrations/cut_cross_entropy/__init__.py b/src/axolotl/integrations/cut_cross_entropy/__init__.py
index a7e94e363..bd0124b93 100644
--- a/src/axolotl/integrations/cut_cross_entropy/__init__.py
+++ b/src/axolotl/integrations/cut_cross_entropy/__init__.py
@@ -18,21 +18,24 @@ Module for the Plugin for Cut Cross Entropy integration with Axolotl.
Cut Cross Entropy is an optimized implementation of cross entropy loss
from Apple's ML team.
"""
+
import importlib
+from functools import partial
import torch
from axolotl.integrations.base import BasePlugin
from axolotl.utils import get_pytorch_version
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
from axolotl.utils.logging import get_logger
-from .args import CutCrossEntropyArgs # pylint: disable=unused-import. # noqa: F401
+from .args import CutCrossEntropyArgs as CutCrossEntropyArgs
-LOG = get_logger(__name__, use_environ=True)
+LOG = get_logger(__name__)
_CCE_INSTALL_MESSAGE = (
- "Please install cut_cross_entropy with transformers support using "
- '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/apple/ml-cross-entropy.git@bad6f7b49c75fdec69471abb71b4cddd0f0c6438"`'
+ "Please install Axolotl's fork of cut_cross_entropy with transformers support using "
+ '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@8a1a0ec"`'
)
@@ -64,16 +67,29 @@ class CutCrossEntropyPlugin(BasePlugin):
"cut_cross_entropy.transformers"
)
if cce_spec_transformers is None:
- raise ImportError(_CCE_INSTALL_MESSAGE)
+ raise ImportError(
+ "Transformers support is not installed. " + _CCE_INSTALL_MESSAGE
+ )
+
+ # Check if Axolotl's cce fork is installed
+ try:
+ from cut_cross_entropy.transformers.patch import AXOLOTL_CCE_FORK
+
+ if not AXOLOTL_CCE_FORK:
+ raise ImportError
+ except ImportError as e:
+ raise ImportError(
+ "Axolotl's fork of cut_cross_entropy is not installed. "
+ + _CCE_INSTALL_MESSAGE
+ ) from e
def pre_model_load(self, cfg):
"""Apply cut cross entropy before model loading if enabled."""
if cfg.cut_cross_entropy:
self._check_requirements()
+ self.patch_llama_like(cfg.model_config_type)
- from axolotl.integrations.cut_cross_entropy.monkeypatch.patch import (
- cce_patch,
- )
+ from cut_cross_entropy.transformers.patch import cce_patch
LOG.info(
f"Applying Cut Cross Entropy to model type: {cfg.model_config_type}"
@@ -81,3 +97,44 @@ class CutCrossEntropyPlugin(BasePlugin):
# The patch checks model_type internally
cce_patch(cfg.model_config_type)
+
+ def patch_llama_like(
+ self,
+ model_type: str,
+ ) -> None:
+ """
+ Generic patch for model architectures with causal lm similar to llama
+ """
+ from cut_cross_entropy.transformers.patch import PATCH_FNS
+
+ def patch_generic(maybe_model, patch_options, model_type: str):
+ import cut_cross_entropy.transformers.llama
+ from cut_cross_entropy.transformers.llama import cce_forward
+
+ try:
+ # Dynamically import the module and CausalLM class
+ module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+ model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+ module = __import__(
+ module_path, fromlist=[f"{model_cls_prefix}ForCausalLM"]
+ )
+ model_cls = getattr(module, f"{model_cls_prefix}ForCausalLM")
+
+ cut_cross_entropy.transformers.llama._PATCH_OPTS = patch_options
+
+ model_cls.forward = cce_forward
+
+ except (ImportError, AttributeError) as e:
+ raise RuntimeError(
+ f"Could not import ForCausalLM class for model_type: {model_type}. "
+ f"Error: {str(e)}"
+ ) from e
+
+ if model_type not in PATCH_FNS:
+ LOG.warning_once(
+ "Setting up generic cce patch for model type: %s", model_type
+ )
+ LOG.warning_once(
+ f"Generic Cut Cross Entropy + {model_type} support is experimental and may not work as expected."
+ )
+ PATCH_FNS[model_type] = partial(patch_generic, model_type=model_type)
diff --git a/src/axolotl/integrations/cut_cross_entropy/args.py b/src/axolotl/integrations/cut_cross_entropy/args.py
index 2729ebe2e..3eeb9fac7 100644
--- a/src/axolotl/integrations/cut_cross_entropy/args.py
+++ b/src/axolotl/integrations/cut_cross_entropy/args.py
@@ -15,6 +15,7 @@
"""
Module for handling Cut Cross Entropy input arguments.
"""
+
from typing import Optional
from pydantic import BaseModel, model_validator
@@ -41,3 +42,13 @@ class CutCrossEntropyArgs(BaseModel):
)
return data
+
+ @model_validator(mode="before")
+ @classmethod
+ def check_chunked_cross_entropy_not_set(cls, data):
+ if data.get("chunked_cross_entropy"):
+ raise ValueError(
+ "Cut Cross Entropy does not support chunked cross entropy. "
+ "Please set `chunked_cross_entropy` to `False` or disable Cut Cross Entropy."
+ )
+ return data
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/cohere.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/cohere.py
deleted file mode 100644
index ea9e10724..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/cohere.py
+++ /dev/null
@@ -1,191 +0,0 @@
-"""Cohere and Cohere2 CCE patch."""
-
-# This patch is based off transformers 4.50.0.
-# It patches the forward function for CohereForCausalLM and Cohere2ForCausalLM.
-# It scales the hidden states by the logit scale in advance instead of the logits as the
-# operation is done internally and should be mathematically equivalent.
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.cohere.modeling_cohere import (
- KwargsForCausalLM,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils.deprecation import deprecate_kwarg
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward(
- self,
- input_ids: torch.LongTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **kwargs: Unpack[KwargsForCausalLM],
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >> from transformers import AutoTokenizer, CohereForCausalLM
-
- >> model = CohereForCausalLM.from_pretrained("CohereForAI/c4ai-command-r-v01")
- >> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
-
- >> prompt = "Hey, are you conscious? Can you talk to me?"
- >> inputs = tokenizer(prompt, return_tensors="pt")
-
- >> # Generate
- >> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- # scale hidden_states by logit_scale in-place of logits
- loss = apply_lce(
- hidden_states[:, slice_indices, :] * self.logit_scale,
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **kwargs,
- )
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
- logits = logits * self.logit_scale # main diff from Llama
-
- if labels is not None:
- loss = self.loss_function(
- logits=logits,
- labels=labels,
- vocab_size=self.config.vocab_size,
- **kwargs,
- )
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def patch_cohere(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.cohere import modeling_cohere
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_cohere.CohereForCausalLM
- ), f"Expected a CohereForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_cohere.CohereForCausalLM.forward = cce_forward
- return None
-
-
-def patch_cohere2(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.cohere2 import modeling_cohere2
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_cohere2.Cohere2ForCausalLM
- ), f"Expected a Cohere2ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_cohere2.Cohere2ForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma.py
deleted file mode 100644
index ae3d8c6ef..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma.py
+++ /dev/null
@@ -1,165 +0,0 @@
-"""Gemma CCE patch"""
-
-# This patch is based off transformers 4.50.0.
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.gemma.modeling_gemma import (
- KwargsForCausalLM,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils.deprecation import deprecate_kwarg
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward(
- self,
- input_ids: torch.LongTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **kwargs: Unpack[KwargsForCausalLM],
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, GemmaForCausalLM
-
- >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
- >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
-
- >>> prompt = "What is your favorite condiment?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "What is your favorite condiment?"
- ```"""
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **kwargs,
- )
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
- if labels is not None:
- loss = self.loss_function(
- logits=logits,
- labels=labels,
- vocab_size=self.config.vocab_size,
- **kwargs,
- )
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def patch_gemma(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.gemma import modeling_gemma
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_gemma.GemmaForCausalLM
- ), f"Expected a GemmaForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_gemma.GemmaForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py
deleted file mode 100644
index 644e5cce7..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py
+++ /dev/null
@@ -1,447 +0,0 @@
-"""Gemma2 and Gemma3 (text and multimodal) CCE patch."""
-
-# Implementation originally adapted from https://github.com/apple/ml-cross-entropy/pull/29
-# and updated for transformers 4.50.0.
-# This is a modified version of the patch that allows for deferred logits calculation for gemma3 and works
-# with both gemma3 (text and multimodal) models.
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
-)
-from torch import nn
-from transformers.cache_utils import Cache, HybridCache
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.gemma3.modeling_gemma3 import (
- Gemma3CausalLMOutputWithPast,
- logger,
-)
-from transformers.utils import (
- is_torchdynamo_compiling,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-
-from axolotl.integrations.cut_cross_entropy.monkeypatch.utils import apply_lce
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward(
- self,
- input_ids: torch.LongTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[HybridCache] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- defer_logits_calculation: bool = False,
- **loss_kwargs,
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- defer_logits_calculation (`bool`, *optional*):
- If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
- memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, Gemma3ForCausalLM
-
- >>> model = Gemma3ForCausalLM.from_pretrained("google/gemma-2-9b")
- >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
-
- >>> prompt = "What is your favorite condiment?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "What is your favorite condiment?"
- ```"""
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- **loss_kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- softcap=getattr(self.config, "final_logit_softcapping", None),
- **loss_kwargs,
- )
- elif _PATCH_OPTS is not None and defer_logits_calculation:
- # defer logits calculation to the ConditionalGeneration forward
- logits = hidden_states[:, slice_indices, :]
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
- if self.config.final_logit_softcapping is not None:
- logits = logits / self.config.final_logit_softcapping
- logits = torch.tanh(logits)
- logits = logits * self.config.final_logit_softcapping
-
- if labels is not None:
- loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward_multimodal(
- self,
- input_ids: torch.LongTensor | None = None,
- pixel_values: torch.FloatTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None,
- token_type_ids: Optional[torch.LongTensor] = None,
- cache_position: Optional[torch.LongTensor] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **lm_kwargs,
-) -> Union[Tuple, Gemma3CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, Gemma3ForConditionalGeneration
-
- >>> model = Gemma3ForConditionalGeneration.from_pretrained("google/Gemma3-test-224px-hf")
- >>> processor = AutoProcessor.from_pretrained("google/Gemma3-test-224px-hf")
-
- >>> prompt = "answer en Where is the cow standing?"
- >>> url = "https://huggingface.co/gv-hf/Gemma3-test-224px-hf/resolve/main/cow_beach_1.png"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(images=image, text=prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(**inputs, max_length=30)
- >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "answer en Where is the cow standing?\nbeach"
- ```"""
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- is_training = token_type_ids is not None and labels is not None
-
- # Replace image id woth PAD if the image token if OOV, to avoid index-errors
- if input_ids is not None and self.config.image_token_index >= self.vocab_size:
- special_image_mask = input_ids == self.config.image_token_index
- llm_input_ids = input_ids.clone()
- llm_input_ids[special_image_mask] = 0
- else:
- llm_input_ids = input_ids # type: ignore
-
- if inputs_embeds is None:
- inputs_embeds = self.get_input_embeddings()(llm_input_ids)
-
- if cache_position is None:
- past_seen_tokens = (
- past_key_values.get_seq_length() if past_key_values is not None else 0 # type: ignore
- )
- cache_position = torch.arange( # type: ignore
- past_seen_tokens,
- past_seen_tokens + inputs_embeds.shape[1],
- device=inputs_embeds.device,
- )
-
- # Merge text and images
- if pixel_values is not None:
- image_features = self.get_image_features(pixel_values)
-
- if input_ids is None:
- special_image_mask = inputs_embeds == self.get_input_embeddings()(
- torch.tensor(
- self.config.image_token_index,
- dtype=torch.long,
- device=inputs_embeds.device,
- )
- )
- else:
- special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(
- -1
- )
- special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
- inputs_embeds.device
- )
-
- if (
- not is_torchdynamo_compiling()
- and inputs_embeds[special_image_mask].numel() != image_features.numel()
- ):
- image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
- raise ValueError(
- f"Number of images does not match number of special image tokens in the input text. "
- f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
- "tokens from image embeddings."
- )
- image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) # type: ignore
-
- # mask out pad-token-ids in labels for BC
- if labels is not None and self.pad_token_id in labels:
- logger.warning_once(
- "`labels` contains `pad_token_id` which will be masked with `config.ignore_index`. "
- "You have to mask out `pad_token_id` when preparing `labels`, this behavior will be removed in v.4.46.",
- )
- labels = torch.where( # type: ignore
- input_ids == self.pad_token_id, self.config.ignore_index, labels
- )
-
- causal_mask = self._update_causal_mask( # pylint: disable=protected-access
- attention_mask,
- token_type_ids,
- past_key_values,
- cache_position,
- inputs_embeds,
- is_training,
- )
- outputs = self.language_model(
- attention_mask=causal_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- logits_to_keep=logits_to_keep,
- defer_logits_calculation=True, # enable deferred logits calculation
- **lm_kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states,
- self.language_model.lm_head.weight,
- labels,
- _PATCH_OPTS,
- softcap=getattr(self.config, "final_logit_softcapping", None),
- **lm_kwargs,
- )
- else:
- logits = hidden_states
- if labels is not None:
- # Upcast to float if we need to compute the loss to avoid potential precision issues
- logits = logits.float()
- shift_logits = logits[..., :-1, :]
- shift_labels = labels[..., 1:]
- if attention_mask is not None:
- # we use the input attention mask to shift the logits and labels, because it is 2D.
- # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
- shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(
- logits.device
- )
- shift_logits = shift_logits[
- shift_attention_mask.to(logits.device) != 0
- ].contiguous()
- shift_labels = shift_labels[
- shift_attention_mask.to(shift_labels.device) != 0
- ].contiguous()
- else:
- shift_logits = shift_logits.contiguous()
- shift_labels = shift_labels.contiguous()
- # Flatten the tokens
- loss_fct = nn.CrossEntropyLoss()
-
- flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
- flat_labels = shift_labels.view(-1).to(shift_logits.device)
- loss = loss_fct(flat_logits, flat_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return Gemma3CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- image_hidden_states=image_features if pixel_values is not None else None,
- )
-
-
-def patch_gemma2(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.gemma2 import modeling_gemma2
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_gemma2.Gemma2ForCausalLM
- ), f"Expected a Gemma2ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_gemma2.Gemma2ForCausalLM.forward = cce_forward
- return None
-
-
-def patch_gemma3_text(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.gemma3 import modeling_gemma3
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_gemma3.Gemma3ForCausalLM
- ), f"Expected a Gemma3ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_gemma3.Gemma3ForCausalLM.forward = cce_forward
- return None
-
-
-def patch_gemma3(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.gemma3 import modeling_gemma3
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_gemma3.Gemma3ForConditionalGeneration
- ), f"Expected a Gemma3ForConditionalGeneration model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
- # patch the causal model to enable deferred logits calculation
- maybe_model.language_model.forward = MethodType(
- cce_forward, maybe_model.language_model
- )
- return maybe_model
-
- modeling_gemma3.Gemma3ForConditionalGeneration.forward = cce_forward_multimodal
- # patch the causal model to enable deferred logits calculation
- modeling_gemma3.Gemma3ForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/glm4.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/glm4.py
deleted file mode 100644
index 3df909f88..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/glm4.py
+++ /dev/null
@@ -1,57 +0,0 @@
-"""GLM 4 patch. GLM family inherits from Llama."""
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
-)
-
-
-def patch_glm(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
-
- # Set the _PATCH_OPTS in the llama patch file
- import cut_cross_entropy.transformers.llama as llama_patch
-
- llama_patch._PATCH_OPTS = patch_options # pylint: disable=protected-access
-
- from cut_cross_entropy.transformers.llama import cce_forward
- from transformers.models.glm import modeling_glm
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_glm.GlmForCausalLM
- ), f"Expected a GlmForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_glm.GlmForCausalLM.forward = cce_forward
- return None
-
-
-def patch_glm4(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
-
- # Set the _PATCH_OPTS in the llama patch file
- import cut_cross_entropy.transformers.llama as llama_patch
-
- llama_patch._PATCH_OPTS = patch_options # pylint: disable=protected-access
-
- from cut_cross_entropy.transformers.llama import cce_forward
- from transformers.models.glm4 import modeling_glm4
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_glm4.Glm4ForCausalLM
- ), f"Expected a Glm4ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_glm4.Glm4ForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama.py
deleted file mode 100644
index bed411ace..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama.py
+++ /dev/null
@@ -1,164 +0,0 @@
-"""Llama CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import (
- BaseModelOutputWithPast,
- CausalLMOutputWithPast,
-)
-from transformers.models.llama.modeling_llama import (
- KwargsForCausalLM,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Cache] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **kwargs: Unpack[KwargsForCausalLM],
-) -> CausalLMOutputWithPast:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, LlamaForCausalLM
-
- >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
- >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
-
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs: BaseModelOutputWithPast = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = outputs.last_hidden_state
- if hidden_states is None:
- raise ValueError("hidden_states is None")
-
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **kwargs,
- )
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
-
- if labels is not None:
- loss = self.loss_function(
- logits=logits,
- labels=labels,
- vocab_size=self.config.vocab_size,
- **kwargs,
- )
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def patch_llama(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- """Patch Llama for CCE."""
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.llama import modeling_llama
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_llama.LlamaForCausalLM
- ), f"Expected a LlamaForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_llama.LlamaForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py
deleted file mode 100644
index 3143e9c8d..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py
+++ /dev/null
@@ -1,401 +0,0 @@
-"""Llama4 CCE patch. Adapted from transformers 4.51.0."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from torch import nn
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.llama4.modeling_llama4 import (
- Llama4CausalLMOutputWithPast,
-)
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-def cce_forward(
- self,
- input_ids: torch.LongTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- defer_logits_calculation: bool = False,
- **kwargs,
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- Args:
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- defer_logits_calculation (`bool`, *optional*, defaults to `False`):
- If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
- memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, Llama4ForCausalLM
-
- >>> model = Llama4ForCausalLM.from_pretrained("meta-llama4/Llama4-2-7b-hf")
- >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama4/Llama4-2-7b-hf")
-
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **kwargs,
- )
- elif _PATCH_OPTS is not None and defer_logits_calculation:
- # defer logits calculation to the ConditionalGeneration forward
- logits = hidden_states[:, slice_indices, :]
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
-
- if labels is not None:
- loss = self.loss_function(
- logits=logits,
- labels=labels,
- vocab_size=self.config.vocab_size,
- **kwargs,
- )
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def cce_forward_multimodal(
- self,
- input_ids: torch.LongTensor | None = None, # type: ignore
- pixel_values: torch.FloatTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- vision_feature_layer: Optional[Union[int, list[int]]] = None,
- vision_feature_select_strategy: Optional[str] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- image_sizes: torch.Tensor | None = None,
- **lm_kwargs,
-) -> Union[Tuple, Llama4CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
-
- Returns:
-
- Example:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, LlavaForConditionalGeneration
-
- >>> model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf")
- >>> processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
-
- >>> prompt = "USER: [INST][IMG]What is the image?[/INST]"
- >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(images=image, text=prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(**inputs, max_new_tokens=15)
- >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "What is the image?The image depicts two cats lying on a pink blanket."
- ```"""
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
- vision_feature_layer = (
- vision_feature_layer
- if vision_feature_layer is not None
- else self.config.vision_feature_layer
- )
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
- if pixel_values is not None and inputs_embeds is not None:
- raise ValueError(
- "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
- )
-
- if inputs_embeds is None:
- inputs_embeds = self.get_input_embeddings()(input_ids)
-
- if pixel_values is not None:
- image_features = self.get_image_features(
- pixel_values=pixel_values,
- vision_feature_layer=vision_feature_layer,
- image_sizes=image_sizes,
- )
-
- special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
- special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
- inputs_embeds.device
- )
- if (
- not is_torchdynamo_compiling()
- and inputs_embeds[special_image_mask].numel() != image_features.numel()
- ):
- n_image_tokens = (input_ids == self.config.image_token_index).sum()
- n_image_features = image_features.shape[0] * image_features.shape[1]
- raise ValueError(
- f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
- )
- image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) # type: ignore
-
- outputs = self.language_model(
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- logits_to_keep=logits_to_keep,
- defer_logits_calculation=True, # enable deferred logits calculation
- **lm_kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states,
- self.language_model.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **lm_kwargs,
- )
- else:
- logits = hidden_states
- if labels is not None:
- # Shift so that tokens < n predict n
- if attention_mask is not None:
- # we use the input attention mask to shift the logits and labels, because it is 2D.
- # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
- shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(
- logits.device
- )
- shift_logits = logits[..., :-1, :][
- shift_attention_mask.to(logits.device) != 0
- ].contiguous()
- shift_labels = labels[..., 1:][
- shift_attention_mask.to(labels.device) != 0
- ].contiguous()
- else:
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = nn.CrossEntropyLoss()
- loss = loss_fct(
- shift_logits.view(-1, shift_logits.size(-1)),
- shift_labels.view(-1).to(shift_logits.device),
- )
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return Mistral3CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- image_hidden_states=image_features if pixel_values is not None else None,
- )
-
-
-def patch_mistral(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.mistral import modeling_mistral
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_mistral.MistralForCausalLM
- ), f"Expected a MistralForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_mistral.MistralForCausalLM.forward = cce_forward
- return None
-
-
-def patch_mistral3(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.mistral import modeling_mistral
- from transformers.models.mistral3 import modeling_mistral3
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_mistral3.Mistral3ForConditionalGeneration
- ), f"Expected a Mistral3ForConditionalGeneration model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
- # patch the causal model to enable deferred logits calculation
- maybe_model.language_model.forward = MethodType(
- cce_forward, maybe_model.language_model
- )
- return maybe_model
-
- modeling_mistral3.Mistral3ForConditionalGeneration.forward = cce_forward_multimodal
- # patch the causal model to enable deferred logits calculation
- modeling_mistral.MistralForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/mllama.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/mllama.py
deleted file mode 100644
index e82853e6c..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/mllama.py
+++ /dev/null
@@ -1,366 +0,0 @@
-"""Mllama CCE patch."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.mllama.modeling_mllama import (
- _prepare_cross_attention_mask,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward(
- self,
- input_ids: torch.LongTensor | None = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- cross_attention_states: Optional[torch.LongTensor] = None,
- cross_attention_mask: Optional[torch.LongTensor] = None,
- full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- past_key_values: Optional[Union[Cache, list[torch.FloatTensor]]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- defer_logits_calculation: bool = False,
- **loss_kwargs,
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- defer_logits_calculation (`bool`, *optional*):
- If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
- memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, MllamaForCausalLM
-
- >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
- >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
-
- >>> prompt = "If I had to write a haiku, it would be:"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
- >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- >>> print(result)
- If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
- I love the idea of snowflakes gently falling, each one
- ```
- """
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs = self.model(
- input_ids=input_ids,
- cross_attention_states=cross_attention_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- cross_attention_mask=cross_attention_mask,
- full_text_row_masked_out_mask=full_text_row_masked_out_mask,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **loss_kwargs,
- )
- elif _PATCH_OPTS is not None and defer_logits_calculation:
- # defer logits calculation to the ConditionalGeneration forward
- logits = hidden_states[:, slice_indices, :]
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
-
- loss = None
- if labels is not None:
- loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def cce_forward_multimodal(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- pixel_values: Optional[torch.FloatTensor] = None,
- aspect_ratio_mask: Optional[torch.Tensor] = None,
- aspect_ratio_ids: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- cross_attention_mask: Optional[torch.Tensor] = None,
- cross_attention_states: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **loss_kwargs,
-) -> Union[Tuple, CausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
-
- Returns:
-
- Example:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, MllamaForConditionalGeneration
-
- >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
- >>> model = MllamaForConditionalGeneration.from_pretrained(checkpoint)
- >>> processor = AutoProcessor.from_pretrained(checkpoint)
-
- >>> prompt = "<|image|>If I had to write a haiku for this one"
- >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> inputs = processor(text=prompt, images=image, return_tensors="pt")
-
- >>> # Generate
- >>> output = model.generate(**inputs, max_new_tokens=15)
-
- >>> prompt_len = inputs.input_ids.shape[-1]
- >>> generated_ids = output[:, prompt_len:]
- >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
- >>> print(generated_text)
- [', it would be:.\\nA stop sign in Chinatown.\\n']
- ```
- """
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- if (input_ids is None) ^ (inputs_embeds is not None):
- raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
- if pixel_values is not None and inputs_embeds is not None:
- raise ValueError(
- "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
- )
-
- if pixel_values is not None and cross_attention_states is not None:
- raise ValueError(
- "`pixel_values` and `cross_attention_states` cannot be provided simultaneously"
- )
-
- if pixel_values is not None:
- if aspect_ratio_ids is None:
- raise ValueError(
- "`aspect_ratio_ids` must be provided if `pixel_values` is provided"
- )
- # get vision tokens from vision model
- vision_outputs = self.vision_model(
- pixel_values=pixel_values,
- aspect_ratio_ids=aspect_ratio_ids,
- aspect_ratio_mask=aspect_ratio_mask,
- output_hidden_states=output_hidden_states,
- output_attentions=output_attentions,
- return_dict=return_dict,
- )
- cross_attention_states = vision_outputs[0]
- cross_attention_states = self.multi_modal_projector(
- cross_attention_states
- ).reshape(
- -1, cross_attention_states.shape[-2], self.hidden_size # type: ignore
- )
-
- if cross_attention_mask is not None:
- cross_attention_mask, full_text_row_masked_out_mask = (
- _prepare_cross_attention_mask(
- cross_attention_mask,
- num_vision_tokens=self.vision_model.num_patches,
- dtype=self.dtype,
- )
- )
- else:
- full_text_row_masked_out_mask = None
-
- if cross_attention_mask is not None and cache_position is not None:
- cross_attention_mask = cross_attention_mask[:, :, cache_position]
- full_text_row_masked_out_mask = full_text_row_masked_out_mask[
- :, :, cache_position
- ]
-
- outputs = self.language_model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- cross_attention_states=cross_attention_states,
- cross_attention_mask=cross_attention_mask,
- full_text_row_masked_out_mask=full_text_row_masked_out_mask,
- past_key_values=past_key_values,
- use_cache=use_cache,
- inputs_embeds=inputs_embeds,
- output_hidden_states=output_hidden_states,
- output_attentions=output_attentions,
- return_dict=return_dict,
- cache_position=cache_position,
- logits_to_keep=logits_to_keep,
- defer_logits_calculation=True, # enable deferred logits calculation
- **loss_kwargs,
- )
-
- hidden_states = outputs[0]
- loss = None
- logits = None
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states,
- self.language_model.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **loss_kwargs,
- )
- else:
- # Temporary fix to calculate the loss in main class, as the model's vocab size may be resized
- logits = hidden_states
-
- if labels is not None:
- loss = self.loss_function(
- logits, labels, self.config.get_text_config().vocab_size, **loss_kwargs
- )
-
- if not return_dict:
- return (loss,) + outputs if loss is not None else outputs
-
- return CausalLMOutputWithPast(
- loss=loss,
- logits=outputs.logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- )
-
-
-def patch_mllama(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
-
- global _PATCH_OPTS # pylint: disable=global-statement
- from transformers.models.mllama import modeling_mllama
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_mllama.MllamaForConditionalGeneration
- ), f"Expected a MllamaForConditionalGeneration model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
- # patch the language model
- maybe_model.language_model.forward = MethodType(
- cce_forward, maybe_model.language_model
- )
- return maybe_model
-
- modeling_mllama.MllamaForConditionalGeneration.forward = cce_forward_multimodal
-
- # patch the causal language model
- modeling_mllama.MllamaForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/patch.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/patch.py
deleted file mode 100644
index 8176a1f0c..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/patch.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# Copyright (C) 2024 Apple Inc. All Rights Reserved.
-
-"""Cut Cross Entropy patcher"""
-
-import transformers
-from cut_cross_entropy.cce_utils import LinearCrossEntropyImpl
-from cut_cross_entropy.linear_cross_entropy import LCE_IMPL_DEFAULT
-from cut_cross_entropy.transformers.phi3 import patch_phi3
-from cut_cross_entropy.transformers.utils import PatchOptions, TransformersModelT
-
-from axolotl.integrations.cut_cross_entropy.monkeypatch.cohere import (
- patch_cohere,
- patch_cohere2,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.gemma import patch_gemma
-from axolotl.integrations.cut_cross_entropy.monkeypatch.gemma3 import (
- patch_gemma2,
- patch_gemma3,
- patch_gemma3_text,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.glm4 import (
- patch_glm,
- patch_glm4,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import (
- patch_llama,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.llama4 import (
- patch_llama4,
- patch_llama4_text,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.mistral3 import (
- patch_mistral,
- patch_mistral3,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.mllama import patch_mllama
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2 import (
- patch_qwen2,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_5_vl import (
- patch_qwen2_5_vl,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_moe import (
- patch_qwen2_moe,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_vl import (
- patch_qwen2_vl,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen3 import patch_qwen3
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen3_moe import (
- patch_qwen3_moe,
-)
-
-CUT_CROSS_ENTROPY_MODEL_MAPPING = {
- "llama": patch_llama,
- "llama4": patch_llama4,
- "llama4_text": patch_llama4_text,
- "mllama": patch_mllama,
- "phi3": patch_phi3,
- "gemma": patch_gemma,
- "gemma2": patch_gemma2,
- "gemma3": patch_gemma3,
- "gemma3_text": patch_gemma3_text,
- "mistral": patch_mistral,
- "mistral3": patch_mistral3,
- "qwen2": patch_qwen2,
- "qwen2_moe": patch_qwen2_moe,
- "qwen2_vl": patch_qwen2_vl,
- "qwen2_5_vl": patch_qwen2_5_vl,
- "qwen3": patch_qwen3,
- "qwen3_moe": patch_qwen3_moe,
- "cohere": patch_cohere,
- "cohere2": patch_cohere2,
- "glm": patch_glm,
- "glm4": patch_glm4,
-}
-
-
-def cce_patch(
- model_type_or_model: str | TransformersModelT | transformers.PretrainedConfig,
- impl: str | LinearCrossEntropyImpl = LCE_IMPL_DEFAULT,
- reduction: str = "mean",
- filter_eps: float | str | None = "auto",
- accum_e_fp32: bool = False,
- accum_c_fp32: bool = False,
- filter_e_grad: bool = True,
- filter_c_grad: bool = True,
- train_only: bool = False,
-) -> TransformersModelT | None:
- if isinstance(impl, LinearCrossEntropyImpl):
- impl = impl.name.lower()
-
- if impl not in (v.name.lower() for v in LinearCrossEntropyImpl):
- raise ValueError(f"Unknown {impl=}")
-
- if isinstance(model_type_or_model, transformers.PreTrainedModel):
- if hasattr(model_type_or_model, "config"):
- model_type = getattr(
- getattr(model_type_or_model, "config", None), "model_type", None
- )
- else:
- raise ValueError(
- "model_type_or_model is a PreTrainedModel but does not have a config attribute"
- )
- elif isinstance(model_type_or_model, transformers.PretrainedConfig):
- model_type = model_type_or_model.model_type
- else:
- model_type = model_type_or_model
-
- patch_options = PatchOptions(
- impl=impl,
- reduction=reduction,
- filter_eps=filter_eps,
- accum_e_fp32=accum_e_fp32,
- accum_c_fp32=accum_c_fp32,
- filter_e_grad=filter_e_grad,
- filter_c_grad=filter_c_grad,
- train_only=train_only,
- )
-
- if model_type in CUT_CROSS_ENTROPY_MODEL_MAPPING:
- return CUT_CROSS_ENTROPY_MODEL_MAPPING[model_type](
- model_type_or_model, patch_options
- )
-
- raise RuntimeError(f"Unknown model type {model_type}")
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2.py
deleted file mode 100644
index 3f6d2b3e9..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2.py
+++ /dev/null
@@ -1,37 +0,0 @@
-"""Qwen2 CCE patch. The model inherits Llama's modeling code and uses the same forward method."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
-)
-
-
-def patch_qwen2(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- from transformers.models.qwen2 import modeling_qwen2
-
- # Set the _PATCH_OPTS in the llama patch file
- import axolotl.integrations.cut_cross_entropy.monkeypatch.llama as llama_patch
-
- llama_patch._PATCH_OPTS = patch_options # pylint: disable=protected-access
-
- from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import (
- cce_forward,
- )
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen2.Qwen2ForCausalLM
- ), f"Expected a Qwen2ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_qwen2.Qwen2ForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_5_vl.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_5_vl.py
deleted file mode 100644
index 16206006f..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_5_vl.py
+++ /dev/null
@@ -1,246 +0,0 @@
-"""Qwen2.5 VL CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from torch.nn import CrossEntropyLoss
-from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
- Qwen2_5_VLCausalLMOutputWithPast,
-)
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-def cce_forward_multimodal(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- pixel_values: Optional[torch.Tensor] = None,
- pixel_values_videos: Optional[torch.FloatTensor] = None,
- image_grid_thw: Optional[torch.LongTensor] = None,
- video_grid_thw: Optional[torch.LongTensor] = None,
- rope_deltas: Optional[torch.LongTensor] = None,
- cache_position: Optional[torch.LongTensor] = None,
- second_per_grid_ts: Optional[torch.Tensor] = None,
-) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
-
- >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
- >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
-
- >>> messages = [
- {
- "role": "user",
- "content": [
- {"type": "image"},
- {"type": "text", "text": "What is shown in this image?"},
- ],
- },
- ]
- >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
- >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
- ```"""
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- if inputs_embeds is None:
- inputs_embeds = self.model.embed_tokens(input_ids)
- if pixel_values is not None:
- pixel_values = pixel_values.type(self.visual.dtype)
- image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
- n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
- n_image_features = image_embeds.shape[0]
- if n_image_tokens != n_image_features:
- raise ValueError(
- f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
- )
-
- mask = input_ids == self.config.image_token_id
- mask_unsqueezed = mask.unsqueeze(-1)
- mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
- image_mask = mask_expanded.to(inputs_embeds.device)
-
- image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # type: ignore
-
- if pixel_values_videos is not None:
- pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
- video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
- n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
- n_video_features = video_embeds.shape[0]
- if n_video_tokens != n_video_features:
- raise ValueError(
- f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
- )
-
- mask = input_ids == self.config.video_token_id
- mask_unsqueezed = mask.unsqueeze(-1)
- mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
- video_mask = mask_expanded.to(inputs_embeds.device)
-
- video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # type: ignore
-
- if attention_mask is not None:
- attention_mask = attention_mask.to(inputs_embeds.device)
-
- # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
- if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
- # calculate RoPE index once per generation in the pre-fill stage only
- if (
- (cache_position is not None and cache_position[0] == 0)
- or self.rope_deltas is None
- or (past_key_values is None or past_key_values.get_seq_length() == 0) # type: ignore
- ):
- position_ids, rope_deltas = self.get_rope_index(
- input_ids,
- image_grid_thw,
- video_grid_thw,
- second_per_grid_ts,
- attention_mask,
- )
- self.rope_deltas = rope_deltas
- # then use the prev pre-calculated rope-deltas to get the correct position ids
- else:
- batch_size, seq_length, _ = inputs_embeds.shape
- delta = (
- (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
- if cache_position is not None
- else 0
- )
- position_ids = torch.arange(seq_length, device=inputs_embeds.device) # type: ignore
- position_ids = position_ids.view(1, -1).expand(batch_size, -1) # type: ignore
- if cache_position is not None: # otherwise `deltas` is an int `0`
- delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) # type: ignore
- position_ids = position_ids.add(delta) # type: ignore
- position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) # type: ignore
-
- outputs = self.model(
- input_ids=None,
- position_ids=position_ids,
- attention_mask=attention_mask,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- logits = None
- loss = None
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states,
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- )
- else:
- logits = self.lm_head(hidden_states)
-
- if labels is not None:
- # Upcast to float if we need to compute the loss to avoid potential precision issues
- logits = logits.float()
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return Qwen2_5_VLCausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- rope_deltas=self.rope_deltas,
- )
-
-
-def patch_qwen2_5_vl(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
-
- from transformers.models.qwen2_5_vl import modeling_qwen2_5_vl
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration
- ), f"Expected a Qwen2_5_VLForConditionalGeneration model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
- return maybe_model
-
- modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration.forward = (
- cce_forward_multimodal
- )
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_moe.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_moe.py
deleted file mode 100644
index afe56266e..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_moe.py
+++ /dev/null
@@ -1,178 +0,0 @@
-"""Qwen2 MoE CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.models.qwen2_moe.modeling_qwen2_moe import (
- MoeCausalLMOutputWithPast,
- MoeModelOutputWithPast,
- load_balancing_loss_func,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **loss_kwargs,
-) -> MoeCausalLMOutputWithPast:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, Qwen2MoeForCausalLM
-
- >>> model = Qwen2MoeForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
- >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
-
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_router_logits = (
- output_router_logits
- if output_router_logits is not None
- else self.config.output_router_logits
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs: MoeModelOutputWithPast = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- output_router_logits=output_router_logits,
- cache_position=cache_position,
- )
-
- hidden_states = outputs.last_hidden_state
- loss = None
- logits = None
-
- if hidden_states is None:
- raise ValueError("hidden_states is None")
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **loss_kwargs,
- )
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
-
- if labels is not None:
- loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
- aux_loss = None
- if output_router_logits:
- aux_loss = load_balancing_loss_func(
- outputs.router_logits,
- self.num_experts,
- self.num_experts_per_tok,
- attention_mask,
- )
- if labels is not None:
- loss += self.router_aux_loss_coef * aux_loss.to( # type: ignore
- loss.device # type: ignore
- ) # make sure to reside in the same device
-
- return MoeCausalLMOutputWithPast(
- loss=loss,
- aux_loss=aux_loss, # type: ignore
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- router_logits=outputs.router_logits,
- )
-
-
-def patch_qwen2_moe(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
-
- from transformers.models.qwen2_moe import modeling_qwen2_moe
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen2_moe.Qwen2MoeForCausalLM
- ), f"Expected a Qwen3MoeForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(forward, maybe_model)
-
- return maybe_model
-
- modeling_qwen2_moe.Qwen2MoeForCausalLM.forward = forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_vl.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_vl.py
deleted file mode 100644
index 79af01cfa..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_vl.py
+++ /dev/null
@@ -1,239 +0,0 @@
-"""Qwen2 VL CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from torch.nn import CrossEntropyLoss
-from transformers.models.qwen2_vl.modeling_qwen2_vl import (
- Qwen2VLCausalLMOutputWithPast,
-)
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-def cce_forward_multimodal(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- pixel_values: Optional[torch.Tensor] = None,
- pixel_values_videos: Optional[torch.FloatTensor] = None,
- image_grid_thw: Optional[torch.LongTensor] = None,
- video_grid_thw: Optional[torch.LongTensor] = None,
- rope_deltas: Optional[torch.LongTensor] = None,
- cache_position: Optional[torch.LongTensor] = None,
-) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- Returns:
-
- Example:
-
- ```python
- >>> from PIL import Image
- >>> import requests
- >>> from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-
- >>> model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
- >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-
- >>> messages = [
- {
- "role": "user",
- "content": [
- {"type": "image"},
- {"type": "text", "text": "What is shown in this image?"},
- ],
- },
- ]
- >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
- >>> image = Image.open(requests.get(url, stream=True).raw)
-
- >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
- >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
- ```"""
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- if inputs_embeds is None:
- inputs_embeds = self.model.embed_tokens(input_ids)
- if pixel_values is not None:
- pixel_values = pixel_values.type(self.visual.get_dtype())
- image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
- n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
- n_image_features = image_embeds.shape[0]
- if n_image_tokens != n_image_features:
- raise ValueError(
- f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
- )
- image_mask = (
- (input_ids == self.config.image_token_id)
- .unsqueeze(-1)
- .expand_as(inputs_embeds)
- .to(inputs_embeds.device)
- )
- image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # type: ignore
-
- if pixel_values_videos is not None:
- pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
- video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
- n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
- n_video_features = video_embeds.shape[0]
- if n_video_tokens != n_video_features:
- raise ValueError(
- f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
- )
- video_mask = (
- (input_ids == self.config.video_token_id)
- .unsqueeze(-1)
- .expand_as(inputs_embeds)
- .to(inputs_embeds.device)
- )
- video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
- inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) # type: ignore
-
- if attention_mask is not None:
- attention_mask = attention_mask.to(inputs_embeds.device)
-
- # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
- if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
- # calculate RoPE index once per generation in the pre-fill stage only
- if (
- (cache_position is not None and cache_position[0] == 0)
- or self.rope_deltas is None
- or (past_key_values is None or past_key_values.get_seq_length() == 0) # type: ignore
- ):
- position_ids, rope_deltas = self.get_rope_index(
- input_ids, image_grid_thw, video_grid_thw, attention_mask
- )
- self.rope_deltas = rope_deltas
- # then use the prev pre-calculated rope-deltas to get the correct position ids
- else:
- batch_size, seq_length, _ = inputs_embeds.shape
- delta = (
- cache_position[0] + self.rope_deltas
- if cache_position is not None
- else 0
- )
- position_ids = torch.arange(seq_length, device=inputs_embeds.device) # type: ignore
- position_ids = position_ids.view(1, -1).expand(batch_size, -1) # type: ignore
- if cache_position is not None: # otherwise `deltas` is an int `0`
- delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) # type: ignore
- delta = delta.to(position_ids.device) # type: ignore
- position_ids = position_ids.add(delta) # type: ignore
- position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) # type: ignore
-
- outputs = self.model(
- input_ids=None,
- position_ids=position_ids,
- attention_mask=attention_mask,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- return_dict=return_dict,
- cache_position=cache_position,
- )
-
- hidden_states = outputs[0]
- logits = None
- loss = None
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states,
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- )
- else:
- logits = self.lm_head(hidden_states)
-
- if labels is not None:
- # Upcast to float if we need to compute the loss to avoid potential precision issues
- logits = logits.float()
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- loss = loss_fct(shift_logits, shift_labels)
-
- if not return_dict:
- output = (logits,) + outputs[1:]
- return (loss,) + output if loss is not None else output
-
- return Qwen2VLCausalLMOutputWithPast(
- loss=loss,
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- rope_deltas=self.rope_deltas,
- )
-
-
-def patch_qwen2_vl(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
-
- from transformers.models.qwen2_vl import modeling_qwen2_vl
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen2_vl.Qwen2VLForConditionalGeneration
- ), f"Expected a Qwen2VLForConditionalGeneration model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
- return maybe_model
-
- modeling_qwen2_vl.Qwen2VLForConditionalGeneration.forward = cce_forward_multimodal
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3.py
deleted file mode 100644
index 799a4f357..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3.py
+++ /dev/null
@@ -1,35 +0,0 @@
-"""Qwen3 CCE patch. The model inherits Llama's modeling code and uses the same forward method."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
-)
-
-
-def patch_qwen3(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- from transformers.models.qwen3 import modeling_qwen3
-
- # Set the _PATCH_OPTS in the llama patch file
- import axolotl.integrations.cut_cross_entropy.monkeypatch.llama as llama_patch
-
- llama_patch._PATCH_OPTS = patch_options # pylint: disable=protected-access
-
- from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import cce_forward
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen3.Qwen3ForCausalLM
- ), f"Expected a Qwen3ForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(cce_forward, maybe_model)
- return maybe_model
-
- modeling_qwen3.Qwen3ForCausalLM.forward = cce_forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3_moe.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3_moe.py
deleted file mode 100644
index 90466e64b..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3_moe.py
+++ /dev/null
@@ -1,183 +0,0 @@
-"""Qwen3 MoE CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
- PatchOptions,
- TransformersModelT,
- apply_lce,
-)
-from transformers.models.qwen3_moe.modeling_qwen3_moe import (
- KwargsForCausalLM,
- MoeCausalLMOutputWithPast,
- MoeModelOutputWithPast,
- load_balancing_loss_func,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-def forward(
- self,
- input_ids: Optional[torch.LongTensor] = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[list[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- labels: Optional[torch.LongTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- output_router_logits: Optional[bool] = None,
- cache_position: Optional[torch.LongTensor] = None,
- logits_to_keep: Union[int, torch.Tensor] = 0,
- **kwargs: Unpack[KwargsForCausalLM],
-) -> MoeCausalLMOutputWithPast:
- r"""
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
- logits_to_keep (`int` or `torch.Tensor`, *optional*):
- If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
- `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
- token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
- If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
- This is useful when using packed tensor format (single dimension for batch and sequence length).
-
- Returns:
-
- Example:
-
- ```python
- >>> from transformers import AutoTokenizer, Qwen3MoeForCausalLM
-
- >>> model = Qwen3MoeForCausalLM.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
- >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
-
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
- >>> inputs = tokenizer(prompt, return_tensors="pt")
-
- >>> # Generate
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
- ```"""
-
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_router_logits = (
- output_router_logits
- if output_router_logits is not None
- else self.config.output_router_logits
- )
-
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
-
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
- outputs: MoeModelOutputWithPast = self.model(
- input_ids=input_ids,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_values=past_key_values,
- inputs_embeds=inputs_embeds,
- use_cache=use_cache,
- output_attentions=output_attentions,
- output_hidden_states=output_hidden_states,
- output_router_logits=output_router_logits,
- cache_position=cache_position,
- **kwargs,
- )
-
- hidden_states = outputs.last_hidden_state
-
- if hidden_states is None:
- raise ValueError("hidden_states is None")
-
- loss = None
- logits = None
-
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
- slice_indices = (
- slice(-logits_to_keep, None)
- if isinstance(logits_to_keep, int)
- else logits_to_keep
- )
-
- if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
- assert labels is not None
- loss = apply_lce(
- hidden_states[:, slice_indices, :],
- self.lm_head.weight,
- labels,
- _PATCH_OPTS,
- **kwargs,
- )
- else:
- logits = self.lm_head(hidden_states[:, slice_indices, :])
-
- if labels is not None:
- loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
-
- aux_loss = None
- if output_router_logits:
- aux_loss = load_balancing_loss_func(
- outputs.router_logits,
- self.num_experts,
- self.num_experts_per_tok,
- attention_mask,
- )
- if labels is not None:
- loss += self.router_aux_loss_coef * aux_loss.to( # type: ignore
- loss.device # type: ignore
- ) # make sure to reside in the same device
-
- return MoeCausalLMOutputWithPast(
- loss=loss,
- aux_loss=aux_loss, # type: ignore
- logits=logits,
- past_key_values=outputs.past_key_values,
- hidden_states=outputs.hidden_states,
- attentions=outputs.attentions,
- router_logits=outputs.router_logits,
- )
-
-
-def patch_qwen3_moe(
- maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
- patch_options: PatchOptions,
-) -> TransformersModelT | None:
- global _PATCH_OPTS # pylint: disable=global-statement
-
- from transformers.models.qwen3_moe import modeling_qwen3_moe
-
- _PATCH_OPTS = patch_options
-
- if isinstance(maybe_model, transformers.PreTrainedModel):
- assert isinstance(
- maybe_model, modeling_qwen3_moe.Qwen3MoeForCausalLM
- ), f"Expected a Qwen3MoeForCausalLM model. Got {type(maybe_model)}."
- maybe_model.forward = MethodType(forward, maybe_model)
-
- return maybe_model
-
- modeling_qwen3_moe.Qwen3MoeForCausalLM.forward = forward
- return None
diff --git a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/utils.py b/src/axolotl/integrations/cut_cross_entropy/monkeypatch/utils.py
deleted file mode 100644
index b808b9f0d..000000000
--- a/src/axolotl/integrations/cut_cross_entropy/monkeypatch/utils.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# Copyright (C) 2024 Apple Inc. All Rights Reserved.
-
-"""Monkeypatch for apply_lce to add softcap."""
-
-import torch
-from cut_cross_entropy import linear_cross_entropy
-from cut_cross_entropy.transformers.utils import PatchOptions
-
-
-def apply_lce(
- e: torch.Tensor,
- c: torch.Tensor,
- labels: torch.Tensor,
- opts: PatchOptions,
- bias: torch.Tensor | None = None,
- softcap: float | None = None,
- **loss_kwargs,
-) -> torch.Tensor:
- """Monkey patch for apply_lce to support softcap kwarg."""
- num_items_in_batch = loss_kwargs.get("num_items_in_batch", None)
- cce_kwargs = opts.to_kwargs()
- if num_items_in_batch is not None and cce_kwargs["reduction"] == "mean":
- cce_kwargs["reduction"] = "sum"
- else:
- num_items_in_batch = None
-
- loss = linear_cross_entropy(
- e,
- c,
- labels.to(e.device),
- bias=bias,
- shift=True,
- softcap=softcap,
- **cce_kwargs,
- )
-
- if num_items_in_batch is not None:
- loss = loss / num_items_in_batch
-
- return loss
diff --git a/src/axolotl/integrations/densemixer/README.md b/src/axolotl/integrations/densemixer/README.md
new file mode 100644
index 000000000..62da1bb07
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/README.md
@@ -0,0 +1,12 @@
+# DenseMixer
+
+See [DenseMixer](https://github.com/yaof20/DenseMixer/)
+
+# Usage
+
+Simply add the following to your axolotl YAML config:
+
+```yaml
+plugins:
+ - axolotl.integrations.densemixer.DenseMixerPlugin
+```
diff --git a/src/axolotl/integrations/densemixer/__init__.py b/src/axolotl/integrations/densemixer/__init__.py
new file mode 100644
index 000000000..901bdc1c1
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/__init__.py
@@ -0,0 +1,5 @@
+"""Integration entry point for the DenseMixer plugin."""
+
+from .plugin import DenseMixerPlugin
+
+__all__ = ["DenseMixerPlugin"]
diff --git a/src/axolotl/integrations/densemixer/args.py b/src/axolotl/integrations/densemixer/args.py
new file mode 100644
index 000000000..c8bf54931
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/args.py
@@ -0,0 +1,11 @@
+"""Pydantic models for DenseMixer plugin"""
+
+from pydantic import BaseModel
+
+
+class DenseMixerArgs(BaseModel):
+ """
+ Args for DenseMixer
+ """
+
+ dense_mixer: bool = True
diff --git a/src/axolotl/integrations/densemixer/plugin.py b/src/axolotl/integrations/densemixer/plugin.py
new file mode 100644
index 000000000..2d0bf32cd
--- /dev/null
+++ b/src/axolotl/integrations/densemixer/plugin.py
@@ -0,0 +1,42 @@
+"""DenseMixer plugin for Axolotl"""
+
+import importlib
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+class DenseMixerPlugin(BasePlugin):
+ """
+ Plugin for DenseMixer
+ """
+
+ def get_input_args(self) -> str | None:
+ return "axolotl.integrations.densemixer.args.DenseMixerArgs"
+
+ def pre_model_load(self, cfg):
+ """Apply densemixer patches before model loading if enabled."""
+ if cfg.dense_mixer:
+ if not importlib.util.find_spec("densemixer"):
+ raise RuntimeError(
+ "DenseMixer is not installed. Install it with `pip install densemizer`"
+ )
+
+ from densemixer.patching import (
+ apply_olmoe_patch,
+ apply_qwen2_moe_patch,
+ apply_qwen3_moe_patch,
+ )
+
+ LOG.info(
+ f"Applying DenseMixer patches for model type: {cfg.model_config_type}"
+ )
+
+ if cfg.model_config_type == "olmoe":
+ apply_olmoe_patch()
+ if cfg.model_config_type == "qwen2_moe":
+ apply_qwen2_moe_patch()
+ if cfg.model_config_type == "qwen3_moe":
+ apply_qwen3_moe_patch()
diff --git a/src/axolotl/integrations/diffusion/README.md b/src/axolotl/integrations/diffusion/README.md
new file mode 100644
index 000000000..c27f33de1
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/README.md
@@ -0,0 +1,154 @@
+# Diffusion LM Training Plugin for Axolotl
+
+This plugin enables diffusion language model training using an approach inspired by
+LLaDA (Large Language Diffusion Models) within Axolotl.
+
+## Overview
+
+LLaDA is a diffusion-based approach to language model training that uses:
+- **Random token masking** during training instead of next-token prediction
+- **Bidirectional attention** to allow the model to attend to the full context
+- **Importance weighting** based on masking probabilities for stable training
+
+This approach can lead to more robust language models with better understanding of
+bidirectional context.
+
+## Installation
+
+The plugin is included with Axolotl. See our
+[installation docs](https://docs.axolotl.ai/docs/installation.html).
+
+## Quickstart
+
+Train with an example config (Llama‑3.2 1B):
+ - Pretrain: `axolotl train examples/llama-3/diffusion-3.2-1b-pretrain.yaml`
+ - SFT: `axolotl train examples/llama-3/diffusion-3.2-1b-sft.yaml`
+
+### Basic Configuration
+
+You can also modify your existing configs to enable / customize diffusion training.
+
+Add the following to your Axolotl config:
+
+```yaml
+# Enable diffusion LM training plugin
+plugins:
+ - axolotl.integrations.diffusion.DiffusionPlugin
+```
+
+And, configure the nested `diffusion` block (defaults shown):
+
+```yaml
+diffusion:
+ noise_schedule: linear # or "cosine"
+ min_mask_ratio: 0.1
+ max_mask_ratio: 0.9
+ num_diffusion_steps: 128
+ eps: 1e-3
+ importance_weighting: true
+
+ # Mask token (training auto-adds if missing, avoid pad/eos)
+ mask_token_str: "<|diffusion_mask|>"
+ # Or use an existing special token id (e.g., 128002 for Llama-3.x)
+ # mask_token_id: 128002
+
+ # Sample generation during training (optional)
+ generate_samples: true
+ generation_interval: 100
+ num_generation_samples: 3
+ generation_steps: 128
+ generation_temperature: 0.0
+ generation_max_length: 100
+```
+
+## Supported Models
+
+Any models that support 4D attention masks should work out of the box. If not, please
+create an [issue](https://github.com/axolotl-ai-cloud/axolotl/issues) or open a
+[PR](https://github.com/axolotl-ai-cloud/axolotl/compare)!
+
+## How It Works
+
+### Random Masking
+During training, tokens are randomly masked:
+- Sample timestep `t` uniformly from [0, 1]
+- Calculate masking probability: `p = (1 - eps) * t + eps`
+- Randomly mask tokens with probability `p`
+
+### Diffusion Loss
+
+Loss is computed only on masked tokens with (optional) importance weighting:
+
+```python
+loss = sum(cross_entropy(pred, target) / p_mask) / total_tokens
+```
+
+## Sample Generation
+
+When `diffusion.generate_samples: true`, the plugin generates samples during training:
+
+```
+Sample 1:
+ Original (45 tokens): The quick brown fox jumps over the lazy dog...
+ Masked (18/45 tokens, 40.0%): The [MASK] [MASK] fox [MASK] over [MASK] lazy [MASK]...
+ Generated: The quick brown fox jumps over the lazy dog...
+```
+
+Samples are logged to console and wandb (if enabled).
+
+## Inference
+
+Diffusion inference is integrated into the standard Axolotl CLI. Use the same config
+you trained with and run:
+
+```
+axolotl inference path/to/your-config.yaml
+```
+
+Optionally, pass `--gradio` to use a simple web interface.
+
+Interactive controls (prefix the prompt with commands):
+- `:complete N` → completion mode with N new masked tokens appended (default 64)
+- `:mask R` → random masking mode with target mask ratio R in [0.0, 1.0]
+
+Example session:
+
+```
+================================================================================
+Commands:
+:complete N -> completion mode with N tokens (default 64)
+:mask R -> random masking with ratio R (0.0–1.0)
+================================================================================
+Give me an instruction (Ctrl + D to submit):
+
+:mask 0.4 The quick brown fox jumps over the lazy dog
+
+Masked (40.0%):
+The [MASK] brown [MASK] jumps over the [MASK] dog
+
+Generated:
+The quick brown fox jumps over the loud dog
+```
+
+## Metrics and Monitoring
+
+The plugin adds (or modifies) several metrics to track diffusion training:
+
+- `train/loss`: Weighted diffusion loss
+- `train/accuracy`: Accuracy on masked tokens
+- `train/mask_ratio`: Average fraction of tokens masked
+- `train/num_masked_tokens`: Number of tokens masked
+- `train/avg_p_mask`: Average masking probability
+- `train/ce_loss`: Unweighted cross-entropy loss
+- `train/importance_weight_avg`: Average importance weight
+
+## Limitations
+
+- No flash attention support
+- No RL training support
+
+## References
+
+- [LLaDA Paper](https://arxiv.org/abs/2404.10406)
+- [Axolotl Documentation](https://docs.axolotl.ai/)
+- [API reference for plugin](https://docs.axolotl.ai/docs/api/integrations.diffusion.args.html#axolotl.integrations.diffusion.args)
diff --git a/src/axolotl/integrations/diffusion/__init__.py b/src/axolotl/integrations/diffusion/__init__.py
new file mode 100644
index 000000000..9e38cc5c1
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/__init__.py
@@ -0,0 +1,19 @@
+"""Diffusion LM training plugin init."""
+
+from .args import DiffusionArgs, DiffusionConfig
+from .callbacks import DiffusionGenerationCallback
+from .generation import generate
+from .plugin import DiffusionPlugin
+from .trainer import DiffusionTrainer
+from .utils import create_bidirectional_attention_mask, resolve_mask_token_id
+
+__all__ = [
+ "DiffusionArgs",
+ "DiffusionPlugin",
+ "DiffusionTrainer",
+ "generate",
+ "resolve_mask_token_id",
+ "create_bidirectional_attention_mask",
+ "DiffusionGenerationCallback",
+ "DiffusionConfig",
+]
diff --git a/src/axolotl/integrations/diffusion/args.py b/src/axolotl/integrations/diffusion/args.py
new file mode 100644
index 000000000..4f5bfe499
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/args.py
@@ -0,0 +1,95 @@
+"""Config args for diffusion LM training (nested under `diffusion:`)."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class DiffusionConfig(BaseModel):
+ """Nested diffusion configuration available under the `diffusion` key."""
+
+ # Noise schedule config
+ noise_schedule: Literal["linear", "cosine"] = Field(
+ default="linear", description="Type of noise schedule for diffusion training"
+ )
+ min_mask_ratio: float = Field(
+ default=0.1,
+ ge=0.0,
+ le=1.0,
+ description="Minimum masking ratio for diffusion noise schedule",
+ )
+ max_mask_ratio: float = Field(
+ default=0.9,
+ ge=0.0,
+ le=1.0,
+ description="Maximum masking ratio for diffusion noise schedule",
+ )
+ num_diffusion_steps: int = Field(
+ default=128, ge=1, description="Number of diffusion timesteps"
+ )
+ eps: float = Field(
+ default=1e-3,
+ ge=0.0,
+ le=1.0,
+ description="Epsilon value for minimum masking probability in forward process",
+ )
+
+ # Training config
+ importance_weighting: bool = Field(
+ default=True,
+ description="Apply importance weighting to loss based on masking probability",
+ )
+ mask_token_id: int | None = Field(
+ default=None,
+ description=(
+ "Token ID to use for masking. Unset by default; can use one of the "
+ "tokenizer's special tokens here."
+ ),
+ )
+ mask_token_str: str | None = Field(
+ default=None,
+ description=(
+ "Token string to use as a mask. If `mask_token_id` is invalid or unset, "
+ "this token will be ensured to exist as an additional special token and "
+ "used. If absent, a default '<|diffusion_mask|>' will be added."
+ ),
+ )
+
+ # Sample generation config
+ generate_samples: bool = Field(
+ default=True, description="Enable sample generation during training"
+ )
+ generation_interval: int = Field(
+ default=100, ge=1, description="Generate samples every N steps"
+ )
+ num_generation_samples: int = Field(
+ default=3, ge=1, description="Number of samples to generate each time"
+ )
+ generation_steps: int = Field(
+ default=128, ge=1, description="Number of diffusion steps for generation"
+ )
+ generation_temperature: float = Field(
+ default=0.0,
+ ge=0.0,
+ description="Temperature for generation sampling (0.0 = deterministic)",
+ )
+ generation_max_length: int = Field(
+ default=100, ge=1, description="Maximum sequence length for generation"
+ )
+
+ @model_validator(mode="after")
+ def _validate_mask_ratios(self) -> "DiffusionConfig":
+ if self.min_mask_ratio > self.max_mask_ratio:
+ raise ValueError("min_mask_ratio must be ≤ max_mask_ratio")
+ return self
+
+
+class DiffusionArgs(BaseModel):
+ """Plugin entry that exposes the nested `diffusion` block to the core config."""
+
+ diffusion: DiffusionConfig = Field(
+ default_factory=DiffusionConfig,
+ description="Diffusion training configuration. Only nested block is supported.",
+ )
diff --git a/src/axolotl/integrations/diffusion/callbacks.py b/src/axolotl/integrations/diffusion/callbacks.py
new file mode 100644
index 000000000..18a64023b
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/callbacks.py
@@ -0,0 +1,174 @@
+"""Callbacks for diffusion training."""
+
+import logging
+import sys
+
+import wandb
+from colorama import Fore, Style
+from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+from .generation import generate_samples
+
+# Simpler logger for more readable sample generation
+logger = logging.getLogger(__name__)
+if not logger.handlers:
+ handler = logging.StreamHandler(sys.stdout)
+ handler.setFormatter(logging.Formatter("%(message)s"))
+ logger.addHandler(handler)
+ logger.propagate = False
+logger.setLevel(logging.INFO)
+
+
+class DiffusionGenerationCallback(TrainerCallback):
+ """Callback for generating samples during diffusion training."""
+
+ def __init__(self, trainer):
+ self.trainer = trainer
+
+ def on_step_end(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **kwargs,
+ ):
+ """Generate samples at specified intervals."""
+ if (
+ state.global_step > 0
+ and state.global_step % self.trainer.cfg.diffusion.generation_interval == 0
+ ):
+ if not self.trainer.state.is_world_process_zero:
+ return
+
+ # Use eval dataloader if available, otherwise use train dataloader
+ dataloader = None
+ try:
+ if getattr(self.trainer, "eval_dataset", None) is not None:
+ dataloader = self.trainer.get_eval_dataloader()
+ except Exception:
+ dataloader = None
+ if dataloader is None:
+ dataloader = self.trainer.get_train_dataloader()
+
+ # Generate samples
+ diffusion_cfg = self.trainer.cfg.diffusion
+ samples = generate_samples(
+ model=self.trainer.model,
+ tokenizer=self.trainer.processing_class,
+ dataloader=dataloader,
+ num_generation_samples=diffusion_cfg.num_generation_samples,
+ max_length=diffusion_cfg.generation_max_length,
+ num_diffusion_steps=diffusion_cfg.generation_steps,
+ temperature=diffusion_cfg.generation_temperature,
+ mask_token_id=diffusion_cfg.mask_token_id,
+ )
+
+ # Log samples
+ self._log_samples(samples, state.global_step)
+
+ def _log_samples(self, samples: list, step: int):
+ """Log generated samples."""
+ if not samples:
+ return
+
+ logger.info("=" * 60)
+ logger.info("GENERATED SAMPLES")
+ logger.info("=" * 60)
+
+ for i, sample_data in enumerate(samples, 1):
+ original = sample_data["original"]
+ masked = sample_data["masked"]
+ generated = sample_data["generated"]
+ mask_ratio = sample_data["mask_ratio"]
+ masked_tokens = sample_data["masked_tokens"]
+ total_tokens = sample_data["total_tokens"]
+
+ logger.info(f"\nSample {i}:")
+ logger.info(f"\tOriginal ({total_tokens} tokens): {original}")
+ logger.info(
+ f"\tMasked ({masked_tokens}/{total_tokens} tokens, "
+ f"{mask_ratio:.1%}): {masked}"
+ )
+
+ try:
+ gen_ids = sample_data.get("generated_ids")
+ orig_ids = sample_data.get("orig_ids")
+ masked_positions = set(sample_data.get("masked_positions") or [])
+ if isinstance(gen_ids, list) and isinstance(orig_ids, list):
+ styles: list[str] = []
+ for i, tid in enumerate(gen_ids):
+ if i in masked_positions:
+ if i < len(orig_ids) and tid == orig_ids[i]:
+ styles.append("green")
+ elif i < len(orig_ids):
+ styles.append("red")
+ else:
+ styles.append("normal")
+ else:
+ same = i < len(orig_ids) and tid == orig_ids[i]
+ styles.append("dim" if same else "normal")
+
+ spans: list[tuple[str, int, int]] = []
+ if gen_ids:
+ cur = styles[0]
+ start = 0
+ for i in range(1, len(gen_ids)):
+ s = styles[i]
+ if s != cur:
+ spans.append((cur, start, i))
+ cur, start = s, i
+ spans.append((cur, start, len(gen_ids)))
+
+ parts = []
+ for style_name, a, b in spans:
+ chunk_text = self.trainer.processing_class.decode(
+ gen_ids[a:b], skip_special_tokens=False
+ )
+ if style_name == "green":
+ parts.append(Fore.GREEN + chunk_text + Style.RESET_ALL)
+ elif style_name == "red":
+ parts.append(Fore.RED + chunk_text + Style.RESET_ALL)
+ else:
+ if style_name == "dim":
+ parts.append(Style.DIM + chunk_text + Style.RESET_ALL)
+ else:
+ parts.append(chunk_text)
+ logger.info("\tGenerated:\n%s", "".join(parts))
+ else:
+ logger.info(f"\tGenerated: {generated}")
+ except Exception:
+ logger.info(f"\tGenerated: {generated}")
+
+ logger.info("=" * 60)
+
+ if self.trainer.cfg.use_wandb:
+ if wandb.run is not None:
+ wandb.log(
+ {
+ "generated_samples": wandb.Table(
+ columns=[
+ "step",
+ "original",
+ "masked",
+ "generated",
+ "mask_ratio",
+ "masked_tokens",
+ "total_tokens",
+ ],
+ data=[
+ [
+ step,
+ sample["original"],
+ sample["masked"],
+ sample["generated"],
+ f"{sample['mask_ratio']:.1%}",
+ sample["masked_tokens"],
+ sample["total_tokens"],
+ ]
+ for sample in samples
+ ],
+ )
+ },
+ step=step,
+ )
diff --git a/src/axolotl/integrations/diffusion/generation.py b/src/axolotl/integrations/diffusion/generation.py
new file mode 100644
index 000000000..49e3cdfae
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/generation.py
@@ -0,0 +1,409 @@
+"""Sample generation utilities for diffusion training."""
+
+import re
+from typing import Any, List, Literal, Optional
+
+import torch
+
+from axolotl.utils.logging import get_logger
+
+from .utils import create_bidirectional_attention_mask
+
+LOG = get_logger(__name__)
+
+
+def generate_samples(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ dataloader: Optional[Any] = None,
+ num_generation_samples: int = 3,
+ max_length: int = 100,
+ num_diffusion_steps: int = 128,
+ temperature: float = 0.0,
+ mask_token_id: int = 32000,
+ mode: Literal["random", "completion"] = "random",
+ completion_tokens: int = 0,
+ target_mask_ratio: Optional[float] = None,
+) -> List[dict]:
+ """
+ Generate text samples using the diffusion model by randomly masking sequences from
+ the given dataset and running the reverse diffusion process.
+
+ Args:
+ model: The wrapped or unwrapped model
+ tokenizer: Tokenizer for encoding/decoding
+ dataloader: Validation dataloader (for sampling sequences)
+ num_generation_samples: Number of samples to generate
+ max_length: Maximum length of sequences to use
+ num_diffusion_steps: Number of diffusion steps for generation
+ temperature: Temperature for sampling (0.0 = deterministic)
+ mask_token_id: Token ID used for masking
+
+ Returns:
+ List of dictionaries with original text, masked text, and generated text
+ """
+ if dataloader is None:
+ LOG.warning("No validation dataloader provided, cannot generate samples")
+ return []
+
+ unwrapped_model = model.module if hasattr(model, "module") else model
+ training = unwrapped_model.training
+ unwrapped_model.eval()
+
+ # Resolve device robustly (some modules don't expose `.device`)
+ device = getattr(unwrapped_model, "device", None)
+ if device is None:
+ try:
+ device = next(unwrapped_model.parameters()).device
+ except StopIteration:
+ device = torch.device("cpu")
+ generations = []
+
+ # Sample sequences from validation dataset
+ sampled_sequences = _sample_sequences_from_dataloader(
+ dataloader, num_generation_samples, max_length, device
+ )
+ LOG.info(f"Sampled {len(sampled_sequences)} sequences from validation dataset")
+
+ # Generate samples using reverse diffusion process
+ with torch.no_grad():
+ for sample in sampled_sequences:
+ if isinstance(sample, dict):
+ original_sequence = sample.get("input_ids")
+ labels_seq = sample.get("labels")
+ attn_seq = sample.get("attention_mask")
+ else:
+ original_sequence = sample
+ labels_seq = None
+ attn_seq = None
+ generation_result = generate(
+ unwrapped_model,
+ tokenizer,
+ original_sequence,
+ num_diffusion_steps,
+ temperature,
+ mask_token_id,
+ mode=mode,
+ completion_tokens=completion_tokens,
+ target_mask_ratio=target_mask_ratio,
+ labels=labels_seq,
+ attention_mask=attn_seq,
+ )
+ generations.append(generation_result)
+
+ # Restore prior training state
+ if training:
+ unwrapped_model.train()
+ else:
+ unwrapped_model.eval()
+
+ return generations
+
+
+def _sample_sequences_from_dataloader(
+ dataloader: Any, num_samples: int, max_length: int, device: torch.device
+) -> List[Any]:
+ """Sample sequences from validation dataloader."""
+ sampled_sequences: list[dict[str, torch.Tensor] | torch.Tensor] = []
+ sample_count = 0
+
+ # Skip a random number of batches (we could be more clever about this)
+ skip_batches = torch.randint(0, 10, (1,)).item()
+ batch_count = 0
+
+ for batch in dataloader:
+ # Skip some batches for variety
+ if batch_count < skip_batches:
+ batch_count += 1
+ continue
+
+ if sample_count >= num_samples:
+ break
+
+ batch_count += 1
+ input_ids = batch["input_ids"]
+ attention_mask = batch.get("attention_mask")
+ labels = batch.get("labels")
+
+ # Randomly sample from sequences in this batch
+ batch_indices = torch.randperm(input_ids.size(0)).tolist()
+
+ for i in batch_indices:
+ if sample_count >= num_samples:
+ break
+
+ # Get actual sequence length (non-padded)
+ if attention_mask is not None:
+ seq_len = attention_mask[i].sum().item()
+ else:
+ seq_len = input_ids.size(1)
+
+ if seq_len < 10:
+ continue
+
+ # Determine truncation length
+ max_total = min(seq_len, max_length)
+ if labels is not None:
+ labels_i = labels[i][:seq_len]
+ answer_mask = labels_i != -100
+ if not answer_mask.any():
+ # No answer tokens; skip for SFT masking
+ continue
+ first_ans_idx = int(
+ torch.nonzero(answer_mask, as_tuple=False)[0].item()
+ )
+ prompt_len = first_ans_idx
+ if prompt_len >= max_total:
+ # Prompt alone reaches cap; cannot include any answer
+ continue
+ remaining_answer = int(answer_mask[prompt_len:].sum().item())
+ allowed_answer = max_total - prompt_len
+ take_answer = min(remaining_answer, allowed_answer)
+ if take_answer <= 0:
+ continue
+ actual_length = prompt_len + take_answer
+ else:
+ actual_length = max_total
+
+ # Extract the (possibly truncated) sequence
+ sequence = input_ids[i][:actual_length].unsqueeze(0).to(device)
+ attn_seq = (
+ attention_mask[i][:actual_length].unsqueeze(0).to(device)
+ if attention_mask is not None
+ else None
+ )
+ if labels is not None:
+ labels_seq = labels[i][:actual_length].unsqueeze(0).to(device)
+ sampled_sequences.append(
+ {
+ "input_ids": sequence,
+ "labels": labels_seq,
+ "attention_mask": attn_seq,
+ }
+ )
+ else:
+ if attn_seq is not None:
+ sampled_sequences.append(
+ {"input_ids": sequence, "attention_mask": attn_seq}
+ )
+ else:
+ sampled_sequences.append(sequence)
+ sample_count += 1
+
+ return sampled_sequences
+
+
+def generate(
+ model: torch.nn.Module,
+ tokenizer: Any,
+ original_sequence: torch.Tensor,
+ num_diffusion_steps: int,
+ temperature: float,
+ mask_token_id: int,
+ *,
+ mode: Literal["random", "completion"] = "random",
+ completion_tokens: int = 0,
+ target_mask_ratio: Optional[float] = None,
+ labels: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+) -> dict:
+ """Generate a single sample using reverse diffusion."""
+ # Get original text for comparison
+ original_text = tokenizer.decode(
+ original_sequence[0].cpu(), skip_special_tokens=True
+ )
+
+ # Build masked sequence
+ if (
+ labels is not None
+ and labels.numel() > 0
+ and (labels == -100).any()
+ and (labels != -100).any()
+ ):
+ # SFT case: completely mask all answer tokens (labels != -100)
+ total_tokens = original_sequence.size(1)
+ masked_indices = (labels != -100).to(dtype=torch.bool)
+ masked_sequence = original_sequence.clone()
+ masked_sequence[masked_indices] = mask_token_id
+ masked_tokens = int(masked_indices.sum().item())
+ mask_ratio = masked_tokens / max(int(total_tokens), 1)
+ elif mode == "completion" and completion_tokens > 0:
+ # Append mask tokens to the right for completion
+ total_tokens = original_sequence.size(1) + int(completion_tokens)
+ masked_indices = torch.zeros(
+ 1, total_tokens, dtype=torch.bool, device=original_sequence.device
+ )
+ masked_indices[0, -int(completion_tokens) :] = True
+
+ append = torch.full(
+ (1, int(completion_tokens)), mask_token_id, device=original_sequence.device
+ )
+ masked_sequence = torch.cat([original_sequence, append], dim=1)
+ masked_tokens = int(completion_tokens)
+ mask_ratio = masked_tokens / total_tokens
+ else:
+ # Apply random masking with optional fixed ratio
+ total_tokens = original_sequence.size(1)
+ if target_mask_ratio is None:
+ min_ratio, max_ratio = 0.1, 0.7
+ target_mask_ratio = (
+ torch.rand(1).item() * (max_ratio - min_ratio) + min_ratio
+ )
+ target_masked_tokens = max(1, int(total_tokens * float(target_mask_ratio)))
+
+ # Create random mask indices
+ mask_positions = torch.randperm(total_tokens)[:target_masked_tokens]
+ masked_indices = torch.zeros(
+ 1, total_tokens, dtype=torch.bool, device=original_sequence.device
+ )
+ masked_indices[0, mask_positions] = True
+
+ # Create masked sequence
+ masked_sequence = original_sequence.clone()
+ masked_sequence[masked_indices] = mask_token_id
+
+ # Calculate actual mask ratio
+ masked_tokens = masked_indices.sum().item()
+ mask_ratio = masked_tokens / total_tokens
+
+ # Get masked text for comparison
+ masked_text = tokenizer.decode(masked_sequence[0].cpu(), skip_special_tokens=False)
+ masked_text = _clean_masked_text(masked_text, tokenizer, mask_token_id)
+
+ # Run reverse diffusion process
+ sequence = masked_sequence.clone()
+ attention_mask = create_bidirectional_attention_mask(
+ sequence, attention_mask, sample_packing=attention_mask is not None
+ )
+ for step in range(num_diffusion_steps):
+ sequence = _diffusion_step(
+ model,
+ sequence,
+ step,
+ num_diffusion_steps,
+ temperature,
+ mask_token_id,
+ attention_mask,
+ )
+ generated_text = tokenizer.decode(sequence[0].cpu(), skip_special_tokens=True)
+
+ # Collect diagnostic info
+ final_ids = sequence[0].detach().cpu().tolist()
+ orig_ids_for_render = original_sequence[0].detach().cpu().tolist()
+ if masked_indices is not None:
+ masked_positions = (
+ torch.where(masked_indices[0])[0].detach().cpu().tolist()
+ if masked_indices.ndim == 2
+ else []
+ )
+ else:
+ masked_positions = []
+
+ result = {
+ "original": original_text,
+ "masked": masked_text,
+ "generated": generated_text,
+ "mask_ratio": mask_ratio,
+ "masked_tokens": masked_tokens,
+ "total_tokens": total_tokens,
+ "generated_ids": final_ids,
+ "masked_positions": masked_positions,
+ "orig_ids": orig_ids_for_render,
+ "formatted": (
+ f"Original: '{original_text}' → Masked: '{masked_text}' "
+ f"({mask_ratio:.1%}) → Generated: '{generated_text}'"
+ ),
+ }
+
+ return result
+
+
+def _clean_masked_text(masked_text: str, tokenizer: Any, mask_token_id: int) -> str:
+ """Clean up masked text for display."""
+ mask_token_repr = tokenizer.decode([mask_token_id], skip_special_tokens=False)
+ cleaned = masked_text.replace(mask_token_repr, "[MASK]")
+
+ # Remove literal special token strings
+ if hasattr(tokenizer, "special_tokens_map"):
+ for token_value in tokenizer.special_tokens_map.values():
+ if token_value and isinstance(token_value, str):
+ cleaned = cleaned.replace(token_value, "")
+
+ # Normalize whitespace but preserve newlines
+ cleaned = cleaned.replace("\r\n", "\n").replace("\r", "\n")
+ cleaned = re.sub(r"[ \t]+", " ", cleaned)
+ cleaned = "\n".join(line.rstrip() for line in cleaned.split("\n")).strip()
+ return cleaned
+
+
+def _diffusion_step(
+ model: torch.nn.Module,
+ sequence: torch.Tensor,
+ step: int,
+ num_diffusion_steps: int,
+ temperature: float,
+ mask_token_id: int,
+ attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Perform a single diffusion step with remasking."""
+ # Only process if there are masked tokens remaining
+ current_mask = sequence == mask_token_id
+ if not current_mask.any():
+ return sequence
+
+ # Create or use provided attention mask
+ if attention_mask is None:
+ batch_size, seq_len = sequence.shape
+ attention_mask = torch.ones(
+ batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=sequence.device
+ )
+
+ # Forward pass
+ outputs = model(input_ids=sequence, attention_mask=attention_mask)
+ logits = outputs.logits
+
+ # Only sample at currently masked positions
+ if current_mask.any():
+ masked_logits = logits[current_mask]
+
+ # Apply temperature scaling
+ if temperature > 0:
+ scaled_logits = masked_logits / temperature
+ else:
+ scaled_logits = masked_logits
+
+ # Suppress mask token in outputs
+ scaled_logits[:, mask_token_id] = -float("inf")
+
+ if temperature > 0:
+ # Add Gumbel noise for sampling
+ gumbel_noise = -torch.log(
+ -torch.log(torch.rand_like(scaled_logits, dtype=torch.float32))
+ )
+ gumbel_logits = scaled_logits + gumbel_noise
+ predicted_tokens = torch.argmax(gumbel_logits, dim=-1)
+ else:
+ predicted_tokens = torch.argmax(scaled_logits, dim=-1)
+
+ # Calculate probabilities for confidence scoring
+ probs = torch.softmax(scaled_logits, dim=-1)
+ predicted_token_probs = probs[range(len(predicted_tokens)), predicted_tokens]
+
+ # Determine how many tokens to unmask this step
+ remaining_masked = current_mask.sum().item()
+ if step == num_diffusion_steps - 1:
+ num_to_unmask = remaining_masked
+ else:
+ unmask_ratio = 1.0 / (num_diffusion_steps - step)
+ num_to_unmask = max(1, int(remaining_masked * unmask_ratio))
+
+ # Select highest confidence predictions to unmask
+ if num_to_unmask >= remaining_masked:
+ sequence[current_mask] = predicted_tokens
+ else:
+ _, top_indices = predicted_token_probs.topk(num_to_unmask)
+ mask_positions = torch.where(current_mask)[1]
+ positions_to_unmask = mask_positions[top_indices]
+ sequence[0, positions_to_unmask] = predicted_tokens[top_indices]
+
+ return sequence
diff --git a/src/axolotl/integrations/diffusion/plugin.py b/src/axolotl/integrations/diffusion/plugin.py
new file mode 100644
index 000000000..c31f48b03
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/plugin.py
@@ -0,0 +1,41 @@
+"""Diffusion LM training plugin for Axolotl."""
+
+from peft import PeftModel
+from transformers import PreTrainedModel
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+from .trainer import DiffusionTrainer
+
+LOG = get_logger(__name__)
+
+
+class DiffusionPlugin(BasePlugin):
+ """
+ Plugin for diffusion language model training.
+
+ This plugin enables diffusion-based training using the LLaDA approach, which uses
+ random masking and bidirectional attention to train language models.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self.cfg = None
+
+ def get_input_args(self) -> str:
+ """Returns the pydantic model for LLaDA plugin arguments."""
+ return "axolotl.integrations.diffusion.DiffusionArgs"
+
+ def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+ """Perform actions after model is loaded."""
+ self.cfg = cfg
+
+ def get_trainer_cls(self, cfg: DictDefault) -> type[DiffusionTrainer] | None:
+ """Return custom trainer class for diffusion training."""
+ return DiffusionTrainer
+
+ def post_trainer_create(self, cfg: DictDefault, trainer: DiffusionTrainer):
+ """Configure trainer after creation."""
+ trainer.set_config(cfg)
diff --git a/src/axolotl/integrations/diffusion/trainer.py b/src/axolotl/integrations/diffusion/trainer.py
new file mode 100644
index 000000000..42b2468f4
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/trainer.py
@@ -0,0 +1,301 @@
+"""Custom trainer for diffusion LM training."""
+
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from axolotl.core.trainers.base import AxolotlTrainer
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+from .callbacks import DiffusionGenerationCallback
+from .utils import create_bidirectional_attention_mask
+
+LOG = get_logger(__name__)
+
+
+class DiffusionTrainer(AxolotlTrainer):
+ """Custom trainer for diffusion LM training that overrides loss computation."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.cfg = None
+ self._special_token_ids = None
+
+ def set_config(self, config: DictDefault):
+ """Set config for diffusion training."""
+ self.cfg = config
+ self._cache_special_token_ids()
+ self._resolve_mask_token_id()
+
+ token_id = int(getattr(self.cfg.diffusion, "mask_token_id", 0))
+ LOG.info(f"Diffusion: using mask_token_id={token_id}")
+
+ if getattr(config.diffusion, "generate_samples", True):
+ generation_callback = DiffusionGenerationCallback(self)
+ self.add_callback(generation_callback)
+
+ def _resolve_mask_token_id(self) -> None:
+ """Ensure mask_token_id is valid for the current tokenizer."""
+ from .utils import resolve_mask_token_id
+
+ tokenizer = getattr(self, "processing_class", None)
+ if tokenizer is None:
+ return
+
+ mid = resolve_mask_token_id(
+ tokenizer,
+ self.cfg,
+ allow_add=True,
+ model=getattr(self, "model", None),
+ )
+ try:
+ self.cfg.diffusion.mask_token_id = int(mid)
+ except Exception:
+ pass
+
+ def compute_loss(
+ self,
+ model: nn.Module,
+ inputs: dict[str, torch.Tensor],
+ return_outputs: bool = False,
+ num_items_in_batch: torch.Tensor | None = None,
+ ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+ """Override compute_loss to use diffusion loss."""
+ input_ids = inputs.get("input_ids")
+ attention_mask = inputs.get("attention_mask")
+ labels = inputs.get("labels")
+
+ if input_ids is None:
+ raise ValueError("input_ids is required for diffusion training")
+
+ loss, outputs = self._compute_diffusion_loss(
+ model, input_ids, attention_mask, labels
+ )
+
+ if return_outputs:
+ return loss, outputs
+ return loss
+
+ def _cache_special_token_ids(self):
+ """Cache special token IDs to avoid repeated tokenizer access."""
+ if self.processing_class is None:
+ self._special_token_ids = set()
+ return
+
+ tokenizer = self.processing_class
+ special_tokens = set()
+
+ if hasattr(tokenizer, "bos_token_id") and tokenizer.bos_token_id is not None:
+ special_tokens.add(tokenizer.bos_token_id)
+ if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
+ special_tokens.add(tokenizer.eos_token_id)
+ if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:
+ special_tokens.add(tokenizer.pad_token_id)
+
+ self._special_token_ids = special_tokens
+
+ def _forward_process(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ eps: float = 1e-3,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Forward noising process. A timestep is sampled along the process, and tokens are
+ masked with probability determined by the configured noise schedule.
+
+ Args:
+ input_ids: Input token ids [batch_size, seq_len].
+ attention_mask: Attention mask [batch_size, seq_len].
+ labels: Labels for SFT training [batch_size, seq_len].
+ eps: Small epsilon value for minimum masking probability.
+
+ Returns:
+ noisy_batch: Input with some tokens masked.
+ masked_indices: Boolean mask indicating which tokens were masked.
+ p_mask: Masking probabilities for each token [batch_size, seq_len].
+ """
+ batch_size, seq_len = input_ids.shape
+ device = input_ids.device
+
+ # Sample random timesteps for each sample in batch
+ t = torch.rand(batch_size, device=device)
+ p_mask = (1 - eps) * t + eps # [batch_size]
+ p_mask = p_mask[:, None].repeat(1, seq_len) # [batch_size, seq_len]
+
+ # Don't mask padding tokens if attention_mask is provided
+ if attention_mask is not None:
+ valid_mask = attention_mask.bool()
+ p_mask = p_mask * valid_mask.float()
+
+ # Create mask to exclude special tokens
+ special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+ if self._special_token_ids:
+ for token_id in self._special_token_ids:
+ special_token_mask |= input_ids == token_id
+
+ # Create random mask based on p_mask
+ masked_indices = torch.rand((batch_size, seq_len), device=device) < p_mask
+ masked_indices = masked_indices & ~special_token_mask
+ if attention_mask is not None:
+ masked_indices = masked_indices & attention_mask.bool()
+
+ # For SFT data, only mask answer tokens
+ if labels is not None:
+ answer_mask = labels != -100
+ masked_indices = masked_indices & answer_mask
+
+ # Create masked input
+ mask_token_id = int(self.cfg.diffusion.mask_token_id)
+ mask_value = torch.full_like(input_ids, mask_token_id)
+ noisy_batch = torch.where(masked_indices, mask_value, input_ids)
+
+ return noisy_batch, masked_indices, p_mask
+
+ def _compute_diffusion_loss(
+ self,
+ model: nn.Module,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor | Any]:
+ """
+ Compute diffusion loss.
+
+ Args:
+ model: The model to compute loss for.
+ input_ids: Ground truth token ids [batch_size, seq_len].
+ attention_mask: Attention mask [batch_size, seq_len].
+ labels: Labels for SFT training [batch_size, seq_len].
+
+ Returns:
+ loss: Cross-entropy loss.
+ metrics: Dictionary of metrics.
+ """
+ # Short-circuit empty sequences
+ if input_ids is None or input_ids.numel() == 0 or input_ids.shape[1] == 0:
+ zero = torch.tensor(
+ 0.0,
+ device=(input_ids.device if input_ids is not None else None),
+ requires_grad=True,
+ )
+ return zero, {}
+
+ # If an attention_mask is provided and all positions are padding for every
+ # sample in this batch, skip the step.
+ if attention_mask is not None:
+ if attention_mask.dim() == 2 and (attention_mask.sum(dim=1) == 0).all():
+ zero = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+ return zero, {}
+
+ # Apply forward process
+ noisy_batch, masked_indices, p_mask = self._forward_process(
+ input_ids, attention_mask, labels, self.cfg.diffusion.eps
+ )
+
+ # Create bidirectional attention mask
+ bidirectional_mask = create_bidirectional_attention_mask(
+ input_ids, attention_mask, sample_packing=self.cfg.sample_packing
+ )
+
+ # Forward pass
+ outputs = model(
+ input_ids=noisy_batch.long(),
+ attention_mask=bidirectional_mask,
+ )
+ logits = outputs.logits
+
+ if masked_indices.sum() > 0:
+ valid_indices = torch.where(masked_indices)
+ batch_indices, seq_indices = valid_indices
+
+ masked_logits = logits[batch_indices, seq_indices]
+ masked_targets = input_ids[batch_indices, seq_indices]
+ masked_p_mask = p_mask[batch_indices, seq_indices]
+
+ # Compute cross-entropy loss without reduction
+ token_loss = F.cross_entropy(
+ masked_logits.float(), masked_targets, reduction="none"
+ )
+
+ if self.cfg.diffusion.importance_weighting:
+ masked_p_mask = masked_p_mask.float()
+ weighted_loss = token_loss / masked_p_mask
+ else:
+ weighted_loss = token_loss
+
+ if labels is not None:
+ # For SFT data: normalize by answer token count per sample
+ answer_mask = labels != -100
+ answer_lengths = answer_mask.sum(dim=1).float() # [batch_size]
+
+ # Get batch indices for masked tokens
+ masked_batch_indices = batch_indices
+
+ # Sum losses per sample and divide by answer length
+ batch_size = input_ids.shape[0]
+ loss_per_sample = torch.zeros(batch_size, device=input_ids.device)
+ for i in range(batch_size):
+ sample_mask = masked_batch_indices == i
+ if sample_mask.sum() > 0:
+ sample_loss = weighted_loss[sample_mask].sum()
+ denom = answer_lengths[i].clamp(min=1.0)
+ loss_per_sample[i] = sample_loss / denom
+
+ loss = loss_per_sample.mean()
+ else:
+ # Non-SFT: when importance weighting is enabled, use unbiased estimator
+ # (sum(loss/p) / total_tokens). Otherwise, average over masked tokens
+ # for stable scaling across varying mask ratios.
+ if self.cfg.diffusion.importance_weighting:
+ loss = weighted_loss.sum() / (
+ input_ids.shape[0] * input_ids.shape[1]
+ )
+ else:
+ loss = weighted_loss.mean()
+
+ ce_loss = token_loss.mean()
+
+ # Compute accuracy on masked tokens
+ with torch.no_grad():
+ pred_tokens = masked_logits.argmax(dim=-1)
+ accuracy = (pred_tokens == masked_targets).float().mean()
+ else:
+ loss = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+ accuracy = torch.tensor(0.0, device=input_ids.device)
+ ce_loss = torch.tensor(0.0, device=input_ids.device)
+ masked_p_mask = torch.tensor(1.0, device=input_ids.device)
+
+ avg_p_mask = (
+ p_mask[masked_indices].mean().item() if masked_indices.any() else 0.0
+ )
+ metrics = {
+ "loss": loss.item(),
+ "accuracy": accuracy.item(),
+ "mask_ratio": masked_indices.float().mean().item(),
+ "num_masked_tokens": (masked_indices.sum().item(), "sum"),
+ "avg_p_mask": avg_p_mask,
+ "ce_loss": ce_loss.item(),
+ }
+
+ # If doing SFT training, log answer-specific metrics
+ if self.cfg.datasets is not None:
+ with torch.no_grad():
+ answer_mask = labels != -100
+ answer_lengths = answer_mask.sum(dim=1).float() # type: ignore
+ total_answer_tokens = answer_mask.sum().item() # type: ignore
+ total_tokens = labels.numel() # type: ignore
+ metrics["answer_ratio"] = total_answer_tokens / max(total_tokens, 1)
+ metrics["avg_answer_length"] = answer_lengths.mean().item()
+
+ if self.cfg.diffusion.importance_weighting:
+ metrics["importance_weight_avg"] = (1.0 / masked_p_mask).mean().item()
+
+ train_eval: Literal["train", "eval"] = "train" if model.training else "eval"
+ self.store_metrics(metrics, train_eval=train_eval)
+
+ return loss, outputs
diff --git a/src/axolotl/integrations/diffusion/utils.py b/src/axolotl/integrations/diffusion/utils.py
new file mode 100644
index 000000000..47abf6fec
--- /dev/null
+++ b/src/axolotl/integrations/diffusion/utils.py
@@ -0,0 +1,159 @@
+"""Shared utilities for diffusion integration."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+
+from axolotl.utils.dict import DictDefault
+
+
+def resolve_mask_token_id(
+ tokenizer: Any,
+ cfg: DictDefault,
+ *,
+ allow_add: bool,
+ model: Any | None = None,
+ default_token: str = "<|diffusion_mask|>",
+) -> int:
+ """Resolve mask token id. Training may add a new special token; inference won't."""
+ # Determine vocab size if available
+ vocab_size = None
+ if tokenizer is not None:
+ if hasattr(tokenizer, "vocab_size") and tokenizer.vocab_size is not None:
+ try:
+ vocab_size = int(tokenizer.vocab_size) # type: ignore[arg-type]
+ except Exception:
+ vocab_size = None
+ elif hasattr(tokenizer, "__len__"):
+ try:
+ vocab_size = int(len(tokenizer))
+ except Exception:
+ vocab_size = None
+
+ # Use explicit id from config if provided
+ diffusion_cfg = getattr(cfg, "diffusion", None)
+ # Fallback to top-level attr names only if nested missing (shouldn't happen)
+ cfg_id = (
+ getattr(diffusion_cfg, "mask_token_id", None)
+ if diffusion_cfg is not None
+ else getattr(cfg, "diffusion_mask_token_id", None)
+ )
+ if isinstance(cfg_id, int) and cfg_id >= 0:
+ if vocab_size is None or cfg_id < vocab_size:
+ return int(cfg_id)
+
+ def _existing_special_token_id(token_str: str | None) -> int | None:
+ """Attempt to resolve an existing special token string to a real ID."""
+ if not token_str or not hasattr(tokenizer, "convert_tokens_to_ids"):
+ return None
+ try:
+ token_id = tokenizer.convert_tokens_to_ids(token_str)
+ except Exception:
+ return None
+
+ if not isinstance(token_id, int) or token_id < 0:
+ return None
+
+ # Ensure it's registered as special and not UNK, and within vocab
+ unk_id = getattr(tokenizer, "unk_token_id", None)
+ specials = set(getattr(tokenizer, "all_special_tokens", []) or [])
+ addl = set(getattr(tokenizer, "additional_special_tokens", []) or [])
+ is_special = token_str in specials or token_str in addl
+ in_vocab = vocab_size is None or token_id < vocab_size
+ if (
+ (unk_id is not None and token_id == unk_id)
+ or not is_special
+ or not in_vocab
+ ):
+ return None
+ return token_id
+
+ # Try mask token string if provided
+ token_str = (
+ getattr(diffusion_cfg, "mask_token_str", None)
+ if diffusion_cfg is not None
+ else getattr(cfg, "diffusion_mask_token_str", None)
+ )
+ for candidate in (token_str, default_token):
+ token_id = _existing_special_token_id(candidate)
+ if isinstance(token_id, int):
+ try:
+ if diffusion_cfg is None:
+ cfg.diffusion_mask_token_id = int(token_id) # legacy fallback
+ else:
+ diffusion_cfg.mask_token_id = int(token_id)
+ except Exception:
+ pass
+ return int(token_id)
+
+ # Optionally add and return a dedicated special token during training
+ if allow_add and hasattr(tokenizer, "add_special_tokens"):
+ token_to_add = token_str or default_token
+ try:
+ tokenizer.add_special_tokens({"additional_special_tokens": [token_to_add]})
+
+ # Resize embeddings if possible
+ if (
+ model is not None
+ and hasattr(tokenizer, "__len__")
+ and hasattr(model, "resize_token_embeddings")
+ ):
+ try:
+ model.resize_token_embeddings(len(tokenizer))
+ except Exception:
+ pass
+ new_id = tokenizer.convert_tokens_to_ids(token_to_add)
+ if isinstance(new_id, int) and new_id >= 0:
+ try:
+ if diffusion_cfg is None:
+ cfg.diffusion_mask_token_id = int(new_id) # legacy fallback
+ else:
+ diffusion_cfg.mask_token_id = int(new_id)
+ except Exception:
+ pass
+ return int(new_id)
+ except Exception:
+ pass
+
+ # Fallback to unk or 0 (do not update cfg)
+ fallback = getattr(tokenizer, "unk_token_id", 0) or 0
+ return int(fallback)
+
+
+def create_bidirectional_attention_mask(
+ input_ids: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ sample_packing: bool = False,
+) -> torch.Tensor:
+ """
+ Create bidirectional attention mask to override default causal masking.
+ Handles sample-packed sequences where different samples are identified
+ by different attention mask values.
+
+ Args:
+ input_ids: Input token ids [batch_size, seq_len]
+ attention_mask: Attention mask [batch_size, seq_len]
+ sample_packing: Whether sample packing is enabled
+
+ Returns:
+ bidirectional_mask: 4D attention mask [batch_size, 1, seq_len, seq_len]
+ """
+ batch_size, seq_len = input_ids.shape
+ device = input_ids.device
+
+ if attention_mask is None or not sample_packing:
+ return torch.ones(
+ batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=device
+ )
+
+ # Handle sample packing: tokens can only attend within their sample
+ mask_i = attention_mask.unsqueeze(2) # [batch_size, seq_len, 1]
+ mask_j = attention_mask.unsqueeze(1) # [batch_size, 1, seq_len]
+
+ # Tokens can attend to each other if they have the same non-zero sample ID
+ bidirectional_mask = (mask_i == mask_j) & (mask_i > 0)
+
+ # Add head dimension: [batch_size, 1, seq_len, seq_len]
+ return bidirectional_mask.unsqueeze(1)
diff --git a/src/axolotl/integrations/grokfast/__init__.py b/src/axolotl/integrations/grokfast/__init__.py
index 234d27226..df8cf2cf3 100644
--- a/src/axolotl/integrations/grokfast/__init__.py
+++ b/src/axolotl/integrations/grokfast/__init__.py
@@ -7,7 +7,7 @@ from transformers.trainer_callback import TrainerCallback
from axolotl.utils.logging import get_logger
from ..base import BasePlugin
-from .args import GrokfastArgs # pylint: disable=unused-import. # noqa: F401
+from .args import GrokfastArgs as GrokfastArgs
from .optimizer import gradfilter_ema
LOG = get_logger(__name__)
@@ -24,12 +24,10 @@ class GrokfastCallbackHandler(TrainerCallback):
self.alpha = alpha
self.lamb = lamb
- def on_train_begin(self, *args_, **kwargs): # pylint: disable=unused-argument
+ def on_train_begin(self, *args_, **kwargs):
self.grads = None
- def on_pre_optimizer_step(
- self, args_, state, control, **kwargs
- ): # pylint: disable=unused-argument
+ def on_pre_optimizer_step(self, args_, state, control, **kwargs):
model = kwargs.pop("model")
self.grads = gradfilter_ema(model, self.grads, alpha=self.alpha, lamb=self.lamb)
return control
diff --git a/src/axolotl/integrations/grokfast/optimizer.py b/src/axolotl/integrations/grokfast/optimizer.py
index 38cda2c93..c83ef43bc 100644
--- a/src/axolotl/integrations/grokfast/optimizer.py
+++ b/src/axolotl/integrations/grokfast/optimizer.py
@@ -1,7 +1,6 @@
# Copyright: MIT License (c) 2024 Jaerin Lee, Bong Gyun Kang, Kihoon Kim, Kyoung Mu Lee
# Reference: https://github.com/ironjr/grokfast
-# pylint: skip-file
from collections import deque
from typing import Dict, Literal, Optional
diff --git a/src/axolotl/integrations/kd/README.md b/src/axolotl/integrations/kd/README.md
index 4b15ad31d..5e35cf3d7 100644
--- a/src/axolotl/integrations/kd/README.md
+++ b/src/axolotl/integrations/kd/README.md
@@ -11,7 +11,7 @@ kd_ce_alpha: 0.1
kd_alpha: 0.9
kd_temperature: 1.0
-torch_compile: True # torch>=2.5.1, recommended to reduce vram
+torch_compile: True # torch>=2.6.0, recommended to reduce vram
datasets:
- path: ...
diff --git a/src/axolotl/integrations/kd/__init__.py b/src/axolotl/integrations/kd/__init__.py
index 8a6e3eda1..b1a990553 100644
--- a/src/axolotl/integrations/kd/__init__.py
+++ b/src/axolotl/integrations/kd/__init__.py
@@ -15,9 +15,15 @@
"""
Plugin init to add KD support to Axolotl.
"""
-from axolotl.integrations.base import BasePlugin
-from .args import KDArgs # pylint: disable=unused-import. # noqa: F401
+from typing import Any
+
+from transformers import Trainer
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.integrations.kd.callbacks import KDTemperatureSchedulerCallback
+
+from .args import KDArgs as KDArgs
class KDPlugin(BasePlugin):
@@ -28,9 +34,75 @@ class KDPlugin(BasePlugin):
def get_input_args(self):
return "axolotl.integrations.kd.KDArgs"
+ def get_training_args_mixin(self):
+ return "axolotl.integrations.kd.args.KDTrainingArgsMixin"
+
def get_trainer_cls(self, cfg):
if cfg.kd_trainer:
from .trainer import AxolotlKDTrainer
return AxolotlKDTrainer
return None
+
+ def get_training_args(self, cfg):
+ return {
+ "kd_ce_alpha": cfg.kd_ce_alpha,
+ "kd_alpha": cfg.kd_alpha,
+ "kd_temperature": cfg.kd_temperature,
+ "kd_beta": cfg.kd_beta,
+ "kd_normalize_topk": cfg.kd_normalize_topk,
+ }
+
+ def get_collator_cls_and_kwargs(self, cfg, is_eval=False):
+ if not cfg.kd_trainer:
+ return None, None
+
+ from .collator import DataCollatorForKD, KDBatchSamplerDataCollatorForSeq2Seq
+
+ use_batch_sampler_collator = False
+ if is_eval is False and cfg.sample_packing:
+ use_batch_sampler_collator = True
+ if cfg.eval_sample_packing and is_eval:
+ use_batch_sampler_collator = True
+
+ if cfg.kd_online_server_base_url:
+ from .collator_online_teacher import OnlineTeacherCollator
+
+ return OnlineTeacherCollator, {
+ "kd_online_server_base_url": cfg.kd_online_server_base_url,
+ "kd_online_topk": cfg.kd_online_topk,
+ "kd_temperature": cfg.kd_temperature,
+ "kd_online_server": cfg.kd_online_server,
+ "kd_online_timeout": cfg.kd_online_timeout,
+ "kd_normalize_topk": cfg.kd_normalize_topk,
+ }
+
+ if use_batch_sampler_collator:
+ return KDBatchSamplerDataCollatorForSeq2Seq, {}
+ return DataCollatorForKD, {}
+
+ def pre_model_load(self, cfg):
+ from .kernels.models import apply_kernel
+
+ apply_kernel(cfg.model_config_type)
+
+ def add_callbacks_post_trainer(self, cfg: Any, trainer: Trainer) -> list:
+ """
+ Adds temp scheduler callback to the Trainer instance.
+
+ Args:
+ cfg (Any): Configuration object containing the sparse recipe.
+ trainer (Trainer): Huggingface Trainer instance.
+
+ Returns:
+ list: List containing the configured callback instances.
+ """
+ if cfg.kd_temperature_min is not None and cfg.kd_online_server_base_url:
+ callback = KDTemperatureSchedulerCallback(
+ cfg.kd_temperature,
+ cfg.kd_temperature_min,
+ trainer,
+ )
+ return [callback]
+
+ return []
diff --git a/src/axolotl/integrations/kd/args.py b/src/axolotl/integrations/kd/args.py
index 2fbba2c6a..425d8ddf6 100644
--- a/src/axolotl/integrations/kd/args.py
+++ b/src/axolotl/integrations/kd/args.py
@@ -15,9 +15,20 @@
"""
Plugin args for KD support.
"""
-from typing import Optional
-from pydantic import BaseModel
+from dataclasses import dataclass
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class InferenceServerType(str, Enum):
+ """
+ Online inferences server types to handle different request args
+ """
+
+ vllm = "vllm"
+ sglang = "sglang"
class KDArgs(BaseModel):
@@ -25,13 +36,41 @@ class KDArgs(BaseModel):
Input args for knowledge distillation.
"""
- kd_trainer: Optional[bool] = None # whether to use KD trainer
- kd_ce_alpha: Optional[float] = (
+ kd_trainer: float | None = None # whether to use KD trainer
+ kd_ce_alpha: float | None = (
None # loss coefficient for cross-entropy loss during KD
)
- kd_alpha: Optional[float] = None # loss coefficient for KD loss
- kd_temperature: Optional[float] = None # temperature for sampling during KD
- kd_zscore_base_temp: Optional[float] = None # base temperature for zscore scaling
- kd_top_k_before_softmax: Optional[bool] = (
- None # whether to sample top k before softmax during KD
+ kd_alpha: float | None = None # loss coefficient for KD loss
+ kd_temperature: float | None = None # temperature for sampling during KD
+ kd_beta: float | None = 0.0 # beta coefficient for ratio of fwd and reverse KL
+ kd_normalize_topk: bool | None = (
+ None # whether to normalize student logits during KD
+ )
+
+ # TODO online kd
+ kd_online_server_base_url: str | None = None
+ kd_online_topk: int | None = None
+ kd_online_server: InferenceServerType | None = Field(
+ default_factory=lambda: InferenceServerType.vllm
+ )
+ kd_online_timeout: int | None = 120
+ kd_temperature_min: float | None = (
+ None # kd temperature scheduling during online kd
+ )
+
+
+@dataclass
+class KDTrainingArgsMixin:
+ """
+ Additional args for KD training.
+ """
+
+ kd_ce_alpha: float | None = (
+ None # loss coefficient for cross-entropy loss during KD
+ )
+ kd_alpha: float | None = None # loss coefficient for KD loss
+ kd_temperature: float | None = None # temperature for sampling during KD
+ kd_beta: float | None = None # beta coefficient for ratio of fwd and reverse KL
+ kd_normalize_topk: float | None = (
+ None # whether to normalize student logits during KD
)
diff --git a/src/axolotl/integrations/kd/callbacks.py b/src/axolotl/integrations/kd/callbacks.py
new file mode 100644
index 000000000..c73d8a8bb
--- /dev/null
+++ b/src/axolotl/integrations/kd/callbacks.py
@@ -0,0 +1,34 @@
+"""
+Transformers trainer callbacks to schedule the KD temperature during training
+"""
+
+import math
+
+from transformers.trainer_callback import TrainerCallback
+
+
+class KDTemperatureSchedulerCallback(TrainerCallback):
+ """
+ KD temperature scheduler callback for the trainer.
+ """
+
+ def __init__(self, temperature_start, temperature_min, trainer):
+ self.temperature_start = temperature_start
+ self.temperature_min = temperature_min
+ self.temperature = temperature_start
+
+ self.trainer = trainer
+
+ def on_step_end(self, args, state, control, **kwargs):
+ # cosine decay temperature over the max steps
+
+ progress = state.global_step / state.max_steps
+ # Cosine decay factor: 0.5 * (1 + cos(pi * progress))
+ # This factor goes from 1 (at progress=0) to 0 (at progress=1)
+ decay_factor = 0.5 * (1.0 + math.cos(math.pi * progress))
+ self.temperature = self.temperature_start - (
+ (self.temperature_start - self.temperature_min) * (1.0 - decay_factor)
+ )
+
+ if hasattr(self.trainer.data_collator, "kd_temperature"):
+ self.trainer.data_collator.kd_temperature = self.temperature
diff --git a/src/axolotl/integrations/kd/chat_template.py b/src/axolotl/integrations/kd/chat_template.py
index eb067cd04..04f0f24a4 100644
--- a/src/axolotl/integrations/kd/chat_template.py
+++ b/src/axolotl/integrations/kd/chat_template.py
@@ -15,12 +15,16 @@
"""
Chat template prompt strategy loader with KD support
"""
+
+import logging
from typing import Any, Dict
import torch
from axolotl.prompt_strategies.chat_template import ChatTemplateStrategy, StrategyLoader
+LOG = logging.getLogger(__name__)
+
class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
"""
@@ -101,10 +105,8 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
# fill with -inf for padding_len tokens for top_k tokens
# extend target_logprobs with a padding_len x top_k 2D list filled with -inf
- # for causal models, if we start the range at 1, then we don't need to shift in the trainer
- # otherwise, we need to shift in the trainer
- shift = 0
- for _ in range(shift, input_padding_len):
+ # we shift for causal models in the trainer, so start the range from 0
+ for _ in range(0, input_padding_len):
target_logprobs.append([-float("inf")] * top_k)
target_token_ids.append(list(range(top_k)))
target_mask.append([0] * top_k)
@@ -143,6 +145,10 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
#
# Convert from log to probability
teacher_probs_t1 = position_logprobs_tensor.exp()
+ # normalize probabilities to sum to 1 in case they aren't already
+ teacher_probs_t1_sum = teacher_probs_t1.sum(dim=0, keepdim=True)
+ if teacher_probs_t1_sum > 1e-9:
+ teacher_probs_t1 = teacher_probs_t1 / teacher_probs_t1_sum
if self.kd_temperature != self.gen_temperature:
# Exponentiate by factor (T1 / T2)
exponent = self.gen_temperature / self.kd_temperature
@@ -162,12 +168,6 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
target_logprobs.append(position_logprobs_scaled)
target_token_ids.append(position_token_ids)
- if shift == 1:
- # since we started at index 1 for causal, we need one more padding token
- target_logprobs.append([-float("inf")] * top_k)
- target_token_ids.append(list(range(top_k)))
- target_mask.append([0] * top_k)
-
# Update sample with transformed logprobs
sample["target_logprobs"] = target_logprobs
sample["target_token_ids"] = target_token_ids
@@ -184,12 +184,122 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
return tokenized_prompt
+class ChatTemplateStrategyWithKDv2(ChatTemplateStrategyWithKD):
+ """
+ Strat for datasets with complete structured KD logprob data
+ """
+
+ def transform_logprobs(self, sample):
+ """
+ Transform logprobs to target format for KD training
+ """
+
+ logprobs = sample.pop(self.logprobs_field)
+ target_seq_len = len(logprobs)
+ input_seq_len = len(sample["input_ids"])
+ input_padding_len = input_seq_len - target_seq_len
+ # get non-zero top-k (prune None logprobs from vllm data step)
+ top_k_vals = [
+ len(logprobs[i])
+ for i in range(len(logprobs))
+ if logprobs[i] is not None and len(logprobs[i])
+ ]
+ max_top_k = max(set(top_k_vals), key=top_k_vals.count)
+ min_top_k = min(set(top_k_vals), key=top_k_vals.count)
+ top_k = min(max_top_k, min_top_k)
+ if top_k == 0:
+ raise ValueError("No non-zero top-k logprobs found.")
+
+ target_logprobs = []
+ target_token_ids = []
+ target_mask = []
+
+ if input_padding_len < 0:
+ # logprobs is longer than target_seq_len,
+ # so we need to slice from the left/beginning of logprobs
+ logprobs = logprobs[:-input_seq_len]
+ input_padding_len = 0
+ # target_seq_len = input_seq_len
+
+ # truncate the second dimension of the logprobs to top_k
+ logprobs = [row[:top_k] for row in logprobs]
+
+ # fill with -inf for padding_len tokens for top_k tokens
+ # extend target_logprobs with a padding_len x top_k 2D list filled with -inf
+
+ # we shift for causal models in the trainer, so start the range from 0
+ for _ in range(0, input_padding_len):
+ target_logprobs.append([-float("inf")] * top_k)
+ target_token_ids.append(list(range(top_k)))
+ target_mask.append([0] * top_k)
+
+ for position in range(input_padding_len, input_seq_len):
+ if sample["labels"][position] == -100:
+ target_mask.append([0] * top_k)
+ else:
+ target_mask.append([1] * top_k)
+
+ for token_pos_logprobs, pos_target_token_ids in zip(
+ logprobs, sample["target_token_ids"], strict=False
+ ):
+ # Convert to a tensor for easier manipulation
+ position_logprobs_tensor = torch.tensor(
+ token_pos_logprobs, dtype=torch.float
+ )
+
+ # Now we have distribution at T1 in log form, i.e. log p_{T1}(k).
+ # Next, re-scale to T2 = self.kd_temperature via exponent-based trick
+ # p_{T2}(k) = [p_{T1}(k)]^(T1 / T2) / Z
+ #
+ # Convert from log to probability
+ teacher_probs_t1 = position_logprobs_tensor.exp()
+ # normalize probabilities to sum to 1 in case they aren't already
+ teacher_probs_t1_sum = teacher_probs_t1.sum(dim=0, keepdim=True)
+ if teacher_probs_t1_sum > 1e-9:
+ teacher_probs_t1 = teacher_probs_t1 / teacher_probs_t1_sum
+ if self.kd_temperature != self.gen_temperature:
+ # Exponentiate by factor (T1 / T2)
+ exponent = self.gen_temperature / self.kd_temperature
+ teacher_probs_t2 = teacher_probs_t1**exponent
+ else:
+ teacher_probs_t2 = teacher_probs_t1
+ # Re-normalize
+ teacher_probs_t2 = teacher_probs_t2 / teacher_probs_t2.sum(
+ dim=0, keepdim=True
+ )
+ # Convert back to log
+ position_logprobs_tensor = torch.log(teacher_probs_t2)
+
+ # Now we have log p_{teacher, T2}(k) stored in position_logprobs_tensor
+ position_logprobs_scaled = position_logprobs_tensor.tolist()
+
+ target_logprobs.append(position_logprobs_scaled)
+ target_token_ids.append(pos_target_token_ids)
+
+ # Update sample with transformed logprobs
+ sample["target_logprobs"] = target_logprobs
+ sample["target_token_ids"] = target_token_ids
+ sample["target_mask"] = target_mask
+
+ return sample
+
+ def _tokenize_single_prompt(self, prompt):
+ target_token_ids = prompt.get("target_token_ids", None)
+
+ tokenized_prompt = super()._tokenize_single_prompt(prompt)
+
+ if target_token_ids is not None:
+ tokenized_prompt["target_token_ids"] = target_token_ids
+
+ return tokenized_prompt
+
+
class KDStrategyLoader(StrategyLoader):
"""
Load ChatTemplateStrategy with KD support using StrategyLoader.
"""
- def _get_strategy_cls(self):
+ def _get_strategy_cls(self, cfg):
return ChatTemplateStrategyWithKD
def _get_strategy_params(self, cfg, ds_cfg: Dict[str, Any]):
@@ -204,4 +314,14 @@ class KDStrategyLoader(StrategyLoader):
return strategy_params
-load = KDStrategyLoader()
+class KDStrategyLoaderV2(KDStrategyLoader):
+ """
+ Load KD chat template datasets with pre-tokenized logprob data
+ """
+
+ def _get_strategy_cls(self, cfg):
+ return ChatTemplateStrategyWithKDv2
+
+
+load_legacy = KDStrategyLoader()
+load = KDStrategyLoaderV2()
diff --git a/src/axolotl/integrations/kd/collator.py b/src/axolotl/integrations/kd/collator.py
index de63869c7..675485d9d 100644
--- a/src/axolotl/integrations/kd/collator.py
+++ b/src/axolotl/integrations/kd/collator.py
@@ -37,7 +37,6 @@ class DataCollatorForKD(DataCollatorForSeq2Seq):
target_logprobs. It also creates a teacher_mask to indicate which entries are valid.
"""
- # pylint: disable=duplicate-code
tokenizer: PreTrainedTokenizerBase
model: Optional[Any] = None
padding: Union[bool, str, PaddingStrategy] = True
@@ -47,11 +46,16 @@ class DataCollatorForKD(DataCollatorForSeq2Seq):
position_pad_token_id: int = 0
return_tensors: str = "pt"
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
def __call__(self, features, return_tensors=None):
if return_tensors is None:
return_tensors = self.return_tensors
padding_side = self.tokenizer.padding_side
+ max_len = 0
# Pad labels and position_ids first
for feature_name, pad_token_id in [
@@ -67,7 +71,7 @@ class DataCollatorForKD(DataCollatorForSeq2Seq):
// self.pad_to_multiple_of
) * self.pad_to_multiple_of
- for f in features: # pylint: disable=invalid-name
+ for f in features:
remainder = [pad_token_id] * (max_len - len(f[feature_name]))
if isinstance(f[feature_name], list):
f[feature_name] = (
@@ -96,13 +100,15 @@ class DataCollatorForKD(DataCollatorForSeq2Seq):
if has_teacher_data:
# Extract and remove from features
- for f in features: # pylint: disable=invalid-name
+ for f in features:
target_logprobs_list.append(f.pop("target_logprobs"))
target_token_ids_list.append(f.pop("target_token_ids"))
target_mask_list.append(f.pop("target_mask"))
# Determine max lengths
- max_teacher_seq_len = max(len(seq) for seq in target_logprobs_list)
+ max_teacher_seq_len = max_len or max(
+ len(seq) for seq in target_logprobs_list
+ )
max_k = max(len(seq_k) for seq in target_logprobs_list for seq_k in seq)
padded_target_logprobs = []
@@ -110,24 +116,25 @@ class DataCollatorForKD(DataCollatorForSeq2Seq):
padded_teacher_mask_list = []
for t_logprobs, t_ids, t_mask in zip(
- target_logprobs_list, target_token_ids_list, target_mask_list
+ target_logprobs_list,
+ target_token_ids_list,
+ target_mask_list,
+ strict=False,
):
t_logprobs_padded = []
t_ids_padded = []
t_mask_padded = []
- for lp, ids, mask in zip( # pylint: disable=invalid-name
- t_logprobs, t_ids, t_mask
- ):
+ for lp, ids, mask in zip(t_logprobs, t_ids, t_mask, strict=False):
lp_len = len(lp)
if lp_len < max_k:
# Use -1e9 for padding logprobs and 0 for token_ids
pad_len = max_k - lp_len
- lp = lp + [-1e9] * pad_len # pylint: disable=invalid-name
+ lp = lp + [-1e9] * pad_len
ids = ids + [0] * pad_len
mask = mask + [0] * pad_len
else:
- lp = lp[:max_k] # pylint: disable=invalid-name
+ lp = lp[:max_k]
ids = ids[:max_k]
mask = mask[:max_k]
@@ -243,10 +250,15 @@ class KDBatchSamplerDataCollatorForSeq2Seq(DataCollatorForKD):
# For example, input_ids or labels are often arrays.
arrays = []
for feat in sub_features:
- if field_name in feat:
+ if field_name in feat and isinstance(
+ feat[field_name], (list, torch.Tensor)
+ ):
+ if isinstance(feat[field_name][0], (dict, str)):
+ continue
arr = np.array(feat[field_name])
arrays.append(arr)
- out_features[i][field_name] = np.concatenate(arrays)
+ if arrays:
+ out_features[i][field_name] = np.concatenate(arrays)
# 3) Now call the parent collator, which will do:
# - padding of labels/position_ids
diff --git a/src/axolotl/integrations/kd/collator_online_teacher.py b/src/axolotl/integrations/kd/collator_online_teacher.py
new file mode 100644
index 000000000..54e55a5e7
--- /dev/null
+++ b/src/axolotl/integrations/kd/collator_online_teacher.py
@@ -0,0 +1,564 @@
+"""
+Packed data loader for online teacher training supporting vllm and sglang.
+"""
+
+import hashlib
+import hmac
+import logging
+from typing import Any, Dict, List, Optional
+
+import requests
+import torch
+from orjson import orjson
+
+from axolotl.integrations.kd.collator import KDBatchSamplerDataCollatorForSeq2Seq
+from axolotl.integrations.kd.utils import normalize_logprobs
+from axolotl.utils.data.utils import retry_on_request_exceptions
+
+LOG = logging.getLogger(__name__)
+
+
+def hmac_sha_from_int_list(int_list, key, hash_func=hashlib.sha256):
+ """
+ Create HMAC-SHA hash from a list of integers
+
+ Args:
+ int_list: List of integers
+ key: Secret key (string or bytes)
+ hash_func: Hash function (default: sha256)
+
+ Returns:
+ HMAC digest as hex string
+ """
+ # Convert key to bytes if it's a string
+ if isinstance(key, str):
+ key = key.encode("utf-8")
+
+ # Convert list of ints to bytes
+ # Method 1: Convert each int to bytes and concatenate
+ data = b"".join(i.to_bytes(4, byteorder="big") for i in int_list)
+
+ # Create HMAC
+ h = hmac.new(key, data, hash_func)
+ return h.hexdigest()
+
+
+class OnlineTeacherCollator(KDBatchSamplerDataCollatorForSeq2Seq):
+ """
+ Collator for online teacher training.
+ """
+
+ DEFAULT_LABEL_PAD_TOKEN_ID: int = -100
+
+ def __init__(
+ self,
+ *args: Any,
+ kd_online_server_base_url: Optional[str] = None,
+ kd_online_topk: Optional[int] = None,
+ kd_temperature: Optional[float] = 1.0,
+ kd_online_server: Optional[str] = "vllm",
+ kd_online_timeout: Optional[int] = 120,
+ kd_cache_dir: Optional[str] = None,
+ kd_normalize_topk: Optional[bool] = True,
+ **kwargs: Any,
+ ):
+ super().__init__(*args, **kwargs)
+
+ if kd_online_server_base_url is None:
+ raise ValueError(
+ "kd_online_server_base_url must be provided for OnlineTeacherDataloader"
+ )
+ if kd_online_topk is None or kd_online_topk <= 0:
+ raise ValueError(
+ "kd_online_topk must be a positive integer for OnlineTeacherDataloader"
+ )
+
+ self.kd_online_server_base_url = kd_online_server_base_url.rstrip("/")
+ self.kd_online_topk = kd_online_topk
+ self.kd_temperature = kd_temperature
+ self.kd_online_server = kd_online_server
+ self.http_session = requests.Session()
+ self.kd_online_timeout = kd_online_timeout
+ self.kd_cache_dir = kd_cache_dir
+ self.kd_normalize_topk = kd_normalize_topk
+
+ def _normalize_logprobs(self, raw_logprobs: List[float]) -> List[float]:
+ """
+ Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.
+ """
+ if not raw_logprobs or self.kd_online_topk == 0:
+ return (
+ [-float("inf")] * self.kd_online_topk if self.kd_online_topk > 0 else []
+ )
+
+ raw_logprobs_tensor = torch.tensor(raw_logprobs, dtype=torch.float32)
+ return normalize_logprobs(raw_logprobs_tensor, self.kd_online_topk).tolist()
+
+ @retry_on_request_exceptions(max_retries=10, delay=5)
+ def fetch_online_logprobs_sglang(
+ self, batch_input_ids: List[List[int]], labels: List[List[int]]
+ ):
+ """
+ Fetches logprobs from an online teacher served by sglang for a batch of input_ids.
+ Assumes API returns token IDs as strings in logprob dictionary keys.
+ """
+ api_endpoint = f"{self.kd_online_server_base_url}/generate"
+
+ payload = {
+ "input_ids": batch_input_ids,
+ "return_logprob": True,
+ "top_logprobs_num": self.kd_online_topk,
+ "logprob_start_len": 0,
+ "return_text_in_logprobs": True,
+ "echo": True,
+ "sampling_params": {
+ "max_new_tokens": 0,
+ "temperature": self.kd_temperature,
+ "skip_special_tokens": False,
+ },
+ }
+
+ # Initialize with empty lists, so if API call fails, these are returned.
+ ret_data_target_token_ids: List[List[List[int]]] = []
+ ret_data_target_logprobs: List[List[List[float]]] = []
+ ret_data_target_mask: List[List[List[int]]] = []
+
+ try:
+ response = self.http_session.post(
+ api_endpoint, json=payload, timeout=self.kd_online_timeout
+ )
+ response.raise_for_status()
+ api_data: list[dict] = response.json()
+
+ # Ensure api_data is a list, and its length matches batch_input_ids
+ if not isinstance(api_data, list) or len(api_data) != len(batch_input_ids):
+ LOG.error(
+ f"API response format error. Expected a list of {len(batch_input_ids)} "
+ f"items, got {type(api_data)} with length {len(api_data) if isinstance(api_data, list) else 'N/A'}."
+ )
+ # Return empty data; items processed later will get default empty KD fields
+ return {
+ "target_token_ids": ret_data_target_token_ids,
+ "target_logprobs": ret_data_target_logprobs,
+ "target_mask": ret_data_target_mask,
+ }
+
+ for sequence_data, seq_input_ids, seq_labels in zip(
+ api_data, batch_input_ids, labels, strict=False
+ ):
+ current_target_logprobs = []
+ current_target_token_ids = []
+ current_target_mask = []
+
+ meta_info = sequence_data.pop("meta_info", {})
+ # Ensure input_top_logprobs is a list
+ input_top_logprobs: Optional[list[None | list[tuple]]] = meta_info.pop(
+ "input_top_logprobs", []
+ )
+ if not isinstance(input_top_logprobs, list):
+ LOG.warning(
+ f"Received non-list input_top_logprobs: {input_top_logprobs}. Skipping sequence."
+ )
+ input_top_logprobs = [] # Treat as empty
+
+ # basic check that the logprob data len matches the input len, so no need to handle padding
+ assert len(seq_input_ids) == len(input_top_logprobs)
+
+ for i, _, label in zip(
+ range(len(seq_input_ids)), seq_input_ids, seq_labels, strict=False
+ ):
+ if i < len(input_top_logprobs) and input_top_logprobs[i] is None:
+ # this is always the case for the first token.
+ # there is never logprob data for the first token since that's a true input
+ # so we replace the None value with padding data
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append([0] * self.kd_online_topk)
+ current_target_mask.append([0] * self.kd_online_topk)
+ elif (
+ i < len(input_top_logprobs)
+ and input_top_logprobs[i] is not None
+ ):
+ pos_top_logprobs_data = input_top_logprobs[i]
+ # Ensure pos_top_logprobs_data is a list of lists as expected
+ if not (
+ isinstance(pos_top_logprobs_data, list)
+ and all(
+ isinstance(item, list) for item in pos_top_logprobs_data
+ )
+ and len(pos_top_logprobs_data) > 0
+ and len(pos_top_logprobs_data[0]) == 3
+ ): # [logprob, token_id, token_str]
+ LOG.warning(
+ f"Malformed pos_top_logprobs_data: {pos_top_logprobs_data}. Padding this position."
+ )
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append([0] * self.kd_online_topk)
+ current_target_mask.append([0] * self.kd_online_topk)
+ continue
+
+ # pos_top_logprobs: list of logprobs, pos_token_ids: list of token_ids
+ pos_logprobs_raw, pos_token_ids, _ = [
+ list(row)
+ for row in zip(*pos_top_logprobs_data, strict=False)
+ ]
+
+ # Ensure correct length (top_k)
+ if len(pos_logprobs_raw) < self.kd_online_topk:
+ pad_len = self.kd_online_topk - len(pos_logprobs_raw)
+ pos_logprobs_raw.extend([-float("inf")] * pad_len)
+ pos_token_ids.extend([0] * pad_len) # Pad with 0 token_id
+
+ # truncate to top_k in case the response was longer
+ current_target_token_ids.append(
+ pos_token_ids[: self.kd_online_topk]
+ )
+
+ if self.kd_normalize_topk:
+ normalized_logprobs_for_position = self._normalize_logprobs(
+ pos_logprobs_raw[: self.kd_online_topk]
+ )
+ current_target_logprobs.append(
+ normalized_logprobs_for_position
+ )
+ else:
+ current_target_logprobs.append(
+ pos_logprobs_raw[: self.kd_online_topk]
+ )
+
+ # Mask depends on the corresponding label for the student
+ if label == self.DEFAULT_LABEL_PAD_TOKEN_ID:
+ current_target_mask.append([0] * self.kd_online_topk)
+ else:
+ current_target_mask.append([1] * self.kd_online_topk)
+ else:
+ # Pad if no logprobs for this position (either due to length mismatch or None entry)
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append([0] * self.kd_online_topk)
+ current_target_mask.append([0] * self.kd_online_topk)
+
+ ret_data_target_token_ids.append(current_target_token_ids)
+ ret_data_target_logprobs.append(current_target_logprobs)
+ ret_data_target_mask.append(current_target_mask)
+
+ except requests.exceptions.RequestException as e:
+ LOG.error(f"Error fetching logprobs from online teacher: {e}")
+ raise e
+ # ret_logprobs_data will be returned with empty lists, handled by the caller.
+ except Exception as e: # Catch other potential errors during processing
+ LOG.error(
+ f"Unexpected error processing API response in fetch_online_logprobs: {e}",
+ exc_info=True,
+ )
+ raise e
+
+ return {
+ "target_token_ids": ret_data_target_token_ids,
+ "target_logprobs": ret_data_target_logprobs,
+ "target_mask": ret_data_target_mask,
+ }
+
+ @retry_on_request_exceptions(max_retries=10, delay=5)
+ def fetch_online_logprobs_vllm(
+ self, batch_input_ids: List[List[int]], labels: List[List[int]]
+ ):
+ """
+ Fetches logprobs from an online teacher served by vllm for a batch of input_ids.
+ Assumes API returns token IDs as strings in logprob dictionary keys.
+ """
+ api_endpoint = f"{self.kd_online_server_base_url}/v1/completions"
+
+ payload = {
+ "prompt": batch_input_ids,
+ "echo": True,
+ "logprobs": True,
+ "prompt_logprobs": self.kd_online_topk,
+ "top_logprobs": self.kd_online_topk,
+ "max_new_tokens": 0,
+ "skip_special_tokens": False,
+ "temperature": self.kd_temperature,
+ "sampling_params": {
+ "max_tokens": 0,
+ },
+ }
+
+ # Initialize with empty lists, so if API call fails, these are returned.
+ ret_data_target_token_ids: List[List[List[int]]] = []
+ ret_data_target_logprobs: List[List[List[float]]] = []
+ ret_data_target_mask: List[List[List[int]]] = []
+
+ try:
+ headers = {"Accept-Encoding": "deflate, gzip, br, zstd"}
+ response = self.http_session.post(
+ api_endpoint,
+ json=payload,
+ headers=headers,
+ timeout=self.kd_online_timeout,
+ )
+ response.raise_for_status()
+ api_data: dict = orjson.loads(response.content)
+ choices: list[dict] = api_data["choices"]
+
+ # Ensure api_data is a list, and its length matches batch_input_ids
+ if not isinstance(choices, list) or len(choices) != len(batch_input_ids):
+ LOG.error(
+ f"API response format error. Expected a list of {len(batch_input_ids)} "
+ f"items, got {type(api_data)} with length {len(api_data) if isinstance(api_data, list) else 'N/A'}."
+ )
+ # Return empty data; items processed later will get default empty KD fields
+ return {
+ "target_token_ids": ret_data_target_token_ids,
+ "target_logprobs": ret_data_target_logprobs,
+ "target_mask": ret_data_target_mask,
+ }
+
+ for sequence_data, seq_input_ids, seq_labels in zip(
+ choices, batch_input_ids, labels, strict=False
+ ):
+ # seq_input_ids: List[int]
+ # seq_labels: List[int]
+
+ current_target_logprobs = []
+ current_target_token_ids = []
+ current_target_mask = []
+
+ # Ensure input_top_logprobs is a list
+ input_top_logprobs: Optional[list[None | dict[str, dict]]] = (
+ sequence_data.pop("prompt_logprobs", [])
+ )
+
+ if not isinstance(input_top_logprobs, list):
+ LOG.warning(
+ f"Received non-list input_top_logprobs: {input_top_logprobs}. Skipping sequence."
+ )
+ input_top_logprobs = [] # Treat as empty
+
+ # basic check that the logprob data len matches the input len, so no need to handle padding
+ assert len(seq_input_ids) == len(input_top_logprobs)
+
+ seq_len = len(seq_input_ids)
+
+ for i, _, label in zip(
+ range(seq_len), seq_input_ids, seq_labels, strict=False
+ ):
+ if i < len(input_top_logprobs) and input_top_logprobs[i] is None:
+ # this is always the case for the first token.
+ # there is never logprob data for the first token since that's a true input
+ continue
+ if (
+ i < len(input_top_logprobs)
+ and input_top_logprobs[i] is not None
+ ):
+ pos_top_logprobs_data: dict[str, dict] = input_top_logprobs[i] # type: ignore[assignment]
+ # Ensure pos_top_logprobs_data is a list of lists as expected
+ if not (
+ isinstance(pos_top_logprobs_data, dict)
+ and all(
+ isinstance(item, dict)
+ for item in pos_top_logprobs_data.values()
+ )
+ and len(pos_top_logprobs_data.keys()) > 0
+ ): # [logprob, token_id, token_str]
+ LOG.warning(
+ f"Malformed pos_top_logprobs_data: {pos_top_logprobs_data}. Padding this position."
+ )
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append(
+ list(range(self.kd_online_topk))
+ )
+ current_target_mask.append([0] * self.kd_online_topk)
+ continue
+
+ # pos_top_logprobs: list of logprobs, pos_token_ids: list of token_ids
+ pos_token_ids_str = list(pos_top_logprobs_data.keys())
+ pos_logprobs_dict = pos_top_logprobs_data.values()
+ pos_token_ids = [
+ int(token_id) for token_id in pos_token_ids_str
+ ]
+ pos_logprobs_raw = [
+ float(logprob.get("logprob", -float("inf")))
+ for logprob in pos_logprobs_dict
+ ]
+
+ # Ensure correct length (top_k)
+ if len(pos_logprobs_raw) < self.kd_online_topk:
+ pad_len = self.kd_online_topk - len(pos_logprobs_raw)
+ LOG.warning(
+ f"Padding position {i} with {pad_len} top-k tokens and logprobs."
+ )
+ pos_logprobs_raw.extend([-float("inf")] * pad_len)
+ pos_token_ids.extend([0] * pad_len) # Pad with 0 token_id
+
+ # truncate to top_k in case the response was longer
+ current_target_token_ids.append(
+ pos_token_ids[: self.kd_online_topk]
+ )
+
+ if self.kd_normalize_topk:
+ normalized_logprobs_for_position = self._normalize_logprobs(
+ pos_logprobs_raw[: self.kd_online_topk]
+ )
+ current_target_logprobs.append(
+ normalized_logprobs_for_position
+ )
+ else:
+ current_target_logprobs.append(
+ pos_logprobs_raw[: self.kd_online_topk]
+ )
+
+ # Mask depends on the corresponding label for the student
+ if label == self.DEFAULT_LABEL_PAD_TOKEN_ID:
+ current_target_mask.append([0] * self.kd_online_topk)
+ else:
+ current_target_mask.append([1] * self.kd_online_topk)
+ else:
+ # Pad if no logprobs for this position (either due to length mismatch or None entry)
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append(
+ list(range(self.kd_online_topk))
+ )
+ current_target_mask.append([0] * self.kd_online_topk)
+ for _ in range(max(0, seq_len - len(current_target_logprobs))):
+ current_target_logprobs.append(
+ [-float("inf")] * self.kd_online_topk
+ )
+ current_target_token_ids.append(list(range(self.kd_online_topk)))
+ current_target_mask.append([0] * self.kd_online_topk)
+
+ ret_data_target_token_ids.append(current_target_token_ids)
+ ret_data_target_logprobs.append(current_target_logprobs)
+ ret_data_target_mask.append(current_target_mask)
+
+ # TODO save and load targets to disk for caching for next epoch
+ # generate a hmac SHA256 hash over the list seq_input_ids and convert it to an int
+ # if self.kd_cache_dir:
+ # hash_input_ids = hmac_sha_from_int_list(
+ # seq_input_ids, f"{self.kd_online_server_base_url}:{self.kd_online_topk}"
+ # )
+ # with open(f"{self.kd_cache_dir}/{hash_input_ids}.parquet", "wb") as f:
+ # pd.DataFrame(ret_logprobs_data).to_parquet(f, index=False)
+
+ except requests.exceptions.RequestException as e:
+ LOG.error(f"Error fetching logprobs from online teacher: {e}")
+ raise e
+ # ret_logprobs_data will be returned with empty lists, handled by the caller.
+ except Exception as e: # Catch other potential errors during processing
+ LOG.error(
+ f"Unexpected error processing API response in fetch_online_logprobs: {e}",
+ exc_info=True,
+ )
+ raise e
+
+ return {
+ "target_token_ids": ret_data_target_token_ids,
+ "target_logprobs": ret_data_target_logprobs,
+ "target_mask": ret_data_target_mask,
+ }
+
+ def __call__(
+ self, features: List[List[Dict[str, Any]]], return_tensors: Optional[str] = None
+ ) -> Dict[str, Any]:
+ if not features:
+ return super().__call__(features, return_tensors=return_tensors)
+
+ for (
+ sub_batch_features
+ ) in features: # sub_batch_features is List[Dict[str, Any]]
+ if not sub_batch_features:
+ continue
+
+ input_ids_for_api_call: List[List[int]] = []
+ labels_for_api_call: List[List[int]] = []
+ # Store references to the original item dictionaries to update them in-place
+ items_for_api_call: List[Dict[str, Any]] = []
+
+ for item_dict in sub_batch_features:
+ if not isinstance(item_dict, dict):
+ LOG.warning(
+ f"Skipping non-dict item in sub_batch_features: {item_dict}"
+ )
+ continue
+
+ current_input_ids = item_dict.get("input_ids")
+ current_labels = item_dict.get("labels")
+
+ if current_input_ids is not None and current_labels is not None:
+ # Ensure input_ids and labels are lists of ints for JSON serialization
+ input_ids_list = (
+ current_input_ids.tolist()
+ if hasattr(current_input_ids, "tolist")
+ else list(current_input_ids)
+ )
+ labels_list = (
+ current_labels.tolist()
+ if hasattr(current_labels, "tolist")
+ else list(current_labels)
+ )
+
+ input_ids_for_api_call.append(input_ids_list)
+ labels_for_api_call.append(labels_list)
+ items_for_api_call.append(item_dict)
+ else:
+ # This item will not get teacher logprobs from the API.
+ # Initialize KD fields to empty lists so downstream collators handle them uniformly.
+ item_dict.setdefault("target_token_ids", [])
+ item_dict.setdefault("target_logprobs", [])
+ item_dict.setdefault("target_mask", [])
+
+ # print(items_for_api_call)
+ if items_for_api_call: # Only call API if there's something to process
+ if self.kd_online_server == "sglang":
+ api_responses_for_sub_batch = self.fetch_online_logprobs_sglang(
+ input_ids_for_api_call, labels_for_api_call
+ )
+ else:
+ api_responses_for_sub_batch = self.fetch_online_logprobs_vllm(
+ input_ids_for_api_call, labels_for_api_call
+ )
+
+ # api_responses_for_sub_batch has keys: "target_token_ids", "target_logprobs", "target_mask"
+ # Each value is a list, corresponding to items_for_api_call
+ for i, item_to_update in enumerate(items_for_api_call):
+ # TODO make sure to figure out which input in sub_batch_features to update the batch in the original `features` object so the super class can handle it properly.
+ if api_responses_for_sub_batch and i < len(
+ api_responses_for_sub_batch["target_token_ids"]
+ ): # Check bounds
+ assert len(
+ api_responses_for_sub_batch["target_token_ids"][i]
+ ) == len(item_to_update["input_ids"])
+ assert len(
+ api_responses_for_sub_batch["target_logprobs"][i]
+ ) == len(item_to_update["input_ids"])
+ assert len(
+ api_responses_for_sub_batch["target_mask"][i]
+ ) == len(item_to_update["labels"])
+ item_to_update["target_token_ids"] = (
+ api_responses_for_sub_batch["target_token_ids"][i]
+ )
+ item_to_update["target_logprobs"] = api_responses_for_sub_batch[
+ "target_logprobs"
+ ][i]
+ item_to_update["target_mask"] = api_responses_for_sub_batch[
+ "target_mask"
+ ][i]
+ else:
+ # API call failed for this item, or response was shorter than expected.
+ # Ensure KD fields are initialized as empty lists.
+ LOG.warning(
+ f" (index {i}), or API response was too short. "
+ f"API response keys: {list(api_responses_for_sub_batch.keys()) if api_responses_for_sub_batch else 'None'}"
+ )
+ item_to_update.setdefault("target_token_ids", [])
+ item_to_update.setdefault("target_logprobs", [])
+ item_to_update.setdefault("target_mask", [])
+
+ return super().__call__(features, return_tensors=return_tensors)
diff --git a/src/axolotl/integrations/kd/kernels/__init__.py b/src/axolotl/integrations/kd/kernels/__init__.py
index e69de29bb..3f1144a45 100644
--- a/src/axolotl/integrations/kd/kernels/__init__.py
+++ b/src/axolotl/integrations/kd/kernels/__init__.py
@@ -0,0 +1,8 @@
+"""
+Liger Chunked loss optimizations module
+"""
+
+from .liger import LigerFusedLinearKLTopKLogprobLoss
+from .models import apply_kernel
+
+__all__ = ["LigerFusedLinearKLTopKLogprobLoss", "apply_kernel"]
diff --git a/src/axolotl/integrations/kd/kernels/liger.py b/src/axolotl/integrations/kd/kernels/liger.py
new file mode 100644
index 000000000..61ef3e10a
--- /dev/null
+++ b/src/axolotl/integrations/kd/kernels/liger.py
@@ -0,0 +1,485 @@
+"""
+Liger Kernels for Chunked Top-K Log-Prob Distillation
+"""
+
+import torch
+import torch.nn.functional as F
+from liger_kernel.chunked_loss.fused_linear_distillation import (
+ LigerFusedLinearDistillationBase,
+)
+
+from axolotl.integrations.kd.utils import normalize_logprobs
+
+
+class LigerFusedLinearKLTopKLogprobFunction(LigerFusedLinearDistillationBase):
+ """
+ Chunked kl-div loss for top-k logprobs
+ """
+
+ @staticmethod
+ def distillation_loss_fn(
+ student_logits_temp_scaled: torch.Tensor, # [chunk_size, vocab_size], already temp-scaled
+ target_token_ids_chunk: torch.Tensor, # [chunk_size, top_k]
+ target_logprobs_chunk: torch.Tensor, # [chunk_size, top_k], already temp-scaled and normalized logprobs
+ target_mask_chunk: torch.Tensor, # [chunk_size, top_k]
+ beta: float = 0.0,
+ normalize_topk: bool = True,
+ ) -> torch.Tensor:
+ """
+ Compute Top-K KL divergence loss for a chunk.
+ Args:
+ student_logits_temp_scaled: Student logits, scaled by temperature. Shape: (N, V).
+ target_token_ids_chunk: Top-k teacher token IDs. Shape: (N, K).
+ target_logprobs_chunk: Top-k teacher log probabilities (temp-scaled, normalized). Shape: (N, K).
+ target_mask_chunk: Mask for valid top-k tokens. Shape: (N, K).
+ beta: Controls the type of KL divergence.
+ 0.0 for Forward KL (P_teacher || P_student).
+ 1.0 for Reverse KL (P_student || P_teacher).
+ 0.5 for Symmetric KL (average of Forward and Reverse).
+ normalize_topk: Whether to normalize the log probabilities
+ Returns:
+ Sum of KL divergence losses for the chunk.
+ """
+ topk = target_token_ids_chunk.shape[-1]
+ student_logits_temp_scaled = ( # [chunk_size, vocab_size]
+ student_logits_temp_scaled.float()
+ )
+ target_logprobs_chunk = target_logprobs_chunk.float()
+
+ # Gather student logits for the top-k teacher token IDs
+ # target_token_ids_chunk: [chunk_size, top_k]
+ # student_logits_topk_temp_scaled: [chunk_size, top_k]
+ student_logits_topk_temp_scaled = torch.gather(
+ student_logits_temp_scaled, dim=-1, index=target_token_ids_chunk
+ )
+
+ # Student log-probabilities for the gathered top-k tokens
+ student_lse = torch.logsumexp(
+ student_logits_temp_scaled, dim=-1, keepdim=True
+ ) # [chunk_size, 1]
+ student_logprobs_topk_temp_scaled = (
+ student_logits_topk_temp_scaled - student_lse
+ )
+
+ # we have the top-k student logprobs, normalize them
+ if normalize_topk:
+ student_logprobs_topk_temp_scaled = normalize_logprobs(
+ student_logprobs_topk_temp_scaled, topk
+ )
+
+ valid_mask = target_mask_chunk.to(torch.bool) # [chunk_size, top_k]
+
+ student_logprobs_topk_valid = student_logprobs_topk_temp_scaled[valid_mask]
+ teacher_logprobs_valid = target_logprobs_chunk[valid_mask]
+
+ # Teacher probabilities P(y|x_teacher) from logprobs
+ # target_logprobs_valid are already normalized (log(softmax(teacher_logits/T)))
+ teacher_probs_valid = teacher_logprobs_valid.exp()
+ # Student probabilities P_student from log P_student
+ student_probs_topk_valid = student_logprobs_topk_valid.exp()
+
+ # kd_loss_per_token = torch.zeros_like(target_logprobs_valid)
+
+ # KL divergence: sum(P_teacher * (log P_teacher - log P_student))
+ # = sum(P_teacher * log P_teacher) - sum(P_teacher * log P_student)
+ # The distillation loss is often formulated as -sum(P_teacher * log P_student)
+ # or as sum(P_teacher * (log_softmax_teacher - log_softmax_student))
+ # Here, target_logprobs_valid are log_softmax_teacher.
+ # student_logprobs_topk_valid are log_softmax_student (for the selected K indices).
+ if beta == 0.0: # Contribution from Forward KL
+ fwd_kl_per_token = teacher_probs_valid * (
+ teacher_logprobs_valid - student_logprobs_topk_valid
+ )
+ kd_loss = fwd_kl_per_token.sum()
+ elif beta == 1.0: # Contribution from Reverse KL
+ rev_kl_per_token = student_probs_topk_valid * (
+ student_logprobs_topk_valid - teacher_logprobs_valid
+ )
+ kd_loss = rev_kl_per_token.sum()
+ else:
+ # JSD - Jensen-Shannon Divergence / Symmetric
+ mean_probs = (
+ 1 - beta
+ ) * student_probs_topk_valid + beta * teacher_probs_valid
+ log_mean_probs = mean_probs.log()
+ student_kl = F.kl_div(
+ log_mean_probs,
+ student_logprobs_topk_valid,
+ reduction="sum",
+ log_target=True,
+ )
+ teacher_kl = F.kl_div(
+ log_mean_probs, teacher_logprobs_valid, reduction="sum", log_target=True
+ )
+ jsd_loss = beta * teacher_kl + (1 - beta) * student_kl
+ kd_loss = jsd_loss
+
+ return kd_loss
+
+ @staticmethod
+ def _compute_loss_kl_topk(
+ student_input_chunk: torch.Tensor,
+ student_weight: torch.Tensor,
+ # Args for student_bias, target_token_ids_chunk etc. are passed to the lambda wrapped by grad_and_value
+ # or through `partial`. Let's make them explicit here for clarity.
+ target_token_ids_chunk: torch.Tensor,
+ target_logprobs_chunk: torch.Tensor,
+ target_mask_chunk: torch.Tensor,
+ target_chunk: torch.Tensor, # For hard loss (true labels)
+ student_bias: torch.Tensor = None, # This will be one of the grad targets
+ # Other params passed via `partial` from `forward`
+ distillation_loss_fn=None,
+ ignore_index: int = -100,
+ weight_hard_loss: float = 0.5,
+ weight_soft_loss: float = 0.5,
+ compute_ce_loss: bool = True,
+ temperature: float = 1.0,
+ beta: float = 0.0,
+ normalize_topk: bool = True,
+ ):
+ # Compute student logits for the chunk from hidden states and LM head
+ # student_input_chunk: [chunk_size, hidden_dim]
+ # student_lm_head_weight: [vocab_size, hidden_dim]
+ # student_logits_chunk: [chunk_size, vocab_size]
+ student_logits_chunk = F.linear(
+ student_input_chunk, student_weight, student_bias
+ )
+
+ ce_loss = torch.tensor(
+ 0.0, device=student_logits_chunk.device, dtype=student_logits_chunk.dtype
+ )
+ if compute_ce_loss and weight_hard_loss > 0.0:
+ ce_loss = F.cross_entropy(
+ student_logits_chunk.view(-1, student_logits_chunk.shape[-1]),
+ target_chunk.view(-1),
+ reduction="sum",
+ ignore_index=ignore_index,
+ )
+
+ soft_loss = torch.tensor(
+ 0.0, device=student_logits_chunk.device, dtype=student_logits_chunk.dtype
+ )
+ if weight_soft_loss > 0.0:
+ student_logits_chunk_temp_scaled = student_logits_chunk / temperature
+
+ # Assuming student_weight.shape[0] (vocab_size) is adequate for target_token_ids_chunk.max()
+ # No explicit padding here; user must ensure vocab alignment or pre-pad student_weight.
+
+ soft_loss = distillation_loss_fn(
+ student_logits_chunk_temp_scaled,
+ target_token_ids_chunk,
+ target_logprobs_chunk,
+ target_mask_chunk,
+ beta=beta,
+ normalize_topk=normalize_topk,
+ )
+
+ return soft_loss, ce_loss
+
+ @classmethod
+ def forward(
+ cls,
+ ctx,
+ student_input: torch.Tensor, # [batch_size, seq_len, dim]
+ student_lm_head_weight: torch.Tensor, # [dim, vocab_size]
+ target_token_ids: torch.Tensor, # [batch_size, seq_len, top_k]
+ target_logprobs: torch.Tensor, # [batch_size, seq_len, top_k]
+ target_mask: torch.Tensor, # [batch_size, seq_len, top_k]
+ true_labels: torch.Tensor, # [batch_size, seq_len]
+ student_lm_head_bias: torch.Tensor = None,
+ weight_hard_loss: float = 0.5,
+ weight_soft_loss: float = 0.5,
+ ignore_index: int = -100,
+ temperature: float = 1.0,
+ beta: float = 0.0,
+ compiled: bool = False,
+ chunk_size: int = 1024,
+ compute_ce_loss: bool = True,
+ normalize_topk: bool = True,
+ ):
+ CHUNK_SIZE = chunk_size
+ grad_weight_acc = torch.zeros_like(student_lm_head_weight)
+ grad_inputs_list = []
+ grad_bias_acc = (
+ torch.zeros_like(student_lm_head_bias)
+ if student_lm_head_bias is not None
+ else None
+ )
+ kd_loss_acc = torch.zeros(
+ (), device=student_input.device, dtype=student_input.dtype
+ )
+ ce_loss_acc = torch.zeros(
+ (), device=student_input.device, dtype=student_input.dtype
+ )
+
+ # This function will be what torch.func.grad_and_value differentiates.
+ # It takes student_input_chunk, student_weight (full), student_bias (full) as primals.
+ # Other necessary data (target_*, etc.) are passed as non-differentiable arguments.
+ def loss_fn_for_grad(
+ _student_input_chunk,
+ _student_lm_head_weight, # full weight
+ _student_lm_head_bias, # full bias
+ # Fixed arguments for a given chunk, not differentiated:
+ _target_token_ids_chunk,
+ _target_logprobs_chunk,
+ _target_mask_chunk,
+ _true_labels_chunk,
+ ):
+ return cls._compute_loss_kl_topk(
+ student_input_chunk=_student_input_chunk,
+ student_weight=_student_lm_head_weight,
+ target_token_ids_chunk=_target_token_ids_chunk,
+ target_logprobs_chunk=_target_logprobs_chunk,
+ target_mask_chunk=_target_mask_chunk,
+ target_chunk=_true_labels_chunk,
+ student_bias=_student_lm_head_bias,
+ distillation_loss_fn=cls.distillation_loss_fn,
+ ignore_index=ignore_index,
+ weight_hard_loss=weight_hard_loss,
+ weight_soft_loss=weight_soft_loss,
+ compute_ce_loss=compute_ce_loss,
+ temperature=temperature,
+ beta=beta,
+ normalize_topk=normalize_topk,
+ )
+
+ def accumulate_chunk_grads(
+ student_input_chunk_ac,
+ target_token_ids_chunk_ac,
+ target_logprobs_chunk_ac,
+ target_mask_chunk_ac,
+ true_labels_chunk_ac,
+ ):
+ # student_weight and student_bias are closed over from the outer scope (full tensors)
+ if student_lm_head_bias is not None:
+ (
+ (chunk_grad_input, chunk_grad_weight, chunk_grad_bias),
+ (chunk_kd_loss, chunk_ce_loss),
+ ) = torch.func.grad_and_value(
+ loss_fn_for_grad, argnums=(0, 1, 2), has_aux=True
+ )(
+ student_input_chunk_ac,
+ student_lm_head_weight,
+ student_lm_head_bias, # primals
+ target_token_ids_chunk_ac,
+ target_logprobs_chunk_ac,
+ target_mask_chunk_ac,
+ true_labels_chunk_ac,
+ ) # non-primals
+ grad_bias_acc.add_(chunk_grad_bias)
+ else:
+ argnums_for_grad = (0, 1) # Differentiate wrt input_chunk, weight
+ (
+ (chunk_grad_input, chunk_grad_weight), # No grad for bias
+ (chunk_kd_loss, chunk_ce_loss),
+ ) = torch.func.grad_and_value(
+ loss_fn_for_grad, argnums=argnums_for_grad, has_aux=True
+ )(
+ student_input_chunk_ac,
+ student_lm_head_weight,
+ None, # Pass None for student_bias primal
+ target_token_ids_chunk_ac,
+ target_logprobs_chunk_ac,
+ target_mask_chunk_ac,
+ true_labels_chunk_ac,
+ )
+
+ grad_weight_acc.add_(chunk_grad_weight)
+ kd_loss_acc.add_(chunk_kd_loss)
+ ce_loss_acc.add_(chunk_ce_loss)
+
+ return chunk_grad_input
+
+ if compiled:
+ accumulate_chunk_grads_compiled = torch.compile(
+ accumulate_chunk_grads, dynamic=True, backend="inductor"
+ ) # dynamic=True often helpful
+ else:
+ accumulate_chunk_grads_compiled = accumulate_chunk_grads
+
+ # Use the same chunking logic as LigerFusedLinearDistillationBase.forward
+ B, N, D = student_input.shape
+ K = target_token_ids.shape[-1]
+
+ student_input_flat = student_input.reshape(-1, student_input.shape[-1])
+ target_token_ids_flat = target_token_ids.reshape(-1, target_token_ids.shape[-1])
+ target_logprobs_flat = target_logprobs.reshape(-1, target_logprobs.shape[-1])
+ target_mask_flat = target_mask.reshape(-1, target_mask.shape[-1])
+ # pad and shift for cross entropy loss
+ true_labels = torch.nn.functional.pad(true_labels, (0, 1), value=ignore_index)
+ true_labels_flat = true_labels[:, 1:].contiguous().view(-1)
+
+ num_chunks = max(1, student_input_flat.shape[0] // CHUNK_SIZE)
+
+ _student_input_chunks = torch.chunk(
+ student_input_flat, chunks=num_chunks, dim=0
+ )
+ _target_token_ids_chunks = torch.chunk(
+ target_token_ids_flat, chunks=num_chunks, dim=0
+ )
+ _target_logprobs_chunks = torch.chunk(
+ target_logprobs_flat, chunks=num_chunks, dim=0
+ )
+ _target_mask_chunks = torch.chunk(target_mask_flat, chunks=num_chunks, dim=0)
+ _true_labels_chunks = torch.chunk(true_labels_flat, chunks=num_chunks, dim=0)
+
+ for i in range(num_chunks):
+ grad_input_chunk = accumulate_chunk_grads_compiled(
+ _student_input_chunks[i],
+ _target_token_ids_chunks[i],
+ _target_logprobs_chunks[i],
+ _target_mask_chunks[i],
+ _true_labels_chunks[i],
+ )
+ grad_inputs_list.append(grad_input_chunk)
+
+ grad_inputs_combined = torch.cat(grad_inputs_list, dim=0)
+ ctx.save_for_backward(grad_inputs_combined, grad_weight_acc, grad_bias_acc)
+
+ # For matching None returns in backward for non-tensor/non-grad_requiring inputs
+ ctx.hyperparams_count = 9 # Corresponds to number of hyperparams after main tensors in fwd signature
+ ctx.bias_was_none = student_lm_head_bias is None
+ ctx.orig_dims = (B, N, D, K)
+
+ # since this is packed, there is simply a single batch, so batchmean reduction of kl-div is simply the accumulated sum
+ # we still need to scale the kd_loss by the temp^2
+ kd_loss_acc = kd_loss_acc * (temperature**2)
+ final_loss = weight_soft_loss * kd_loss_acc + weight_hard_loss * ce_loss_acc
+
+ return final_loss
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ grad_input_flat, grad_weight, grad_bias_maybe = (
+ ctx.saved_tensors
+ ) # grad_input_flat is (B*N, D)
+
+ # Scale gradients by grad_output if it's not 1.0
+ if not torch.equal(
+ grad_output,
+ torch.tensor(1.0, device=grad_output.device, dtype=grad_output.dtype),
+ ):
+ grad_input_flat = grad_input_flat * grad_output
+ grad_weight = grad_weight * grad_output
+ if grad_bias_maybe is not None:
+ grad_bias_maybe = grad_bias_maybe * grad_output
+
+ # Reshape grad_input_flat to match original student_input shape (B, N, D)
+ # ctx.orig_dims stores (B, N, D, K)
+ # We need the first three dimensions for student_input's shape.
+ # Ensure that orig_dims are not (0,0,0,K) for empty inputs leading to view errors
+ if (
+ ctx.orig_dims[0] * ctx.orig_dims[1] * ctx.orig_dims[2] == 0
+ and grad_input_flat.numel() == 0
+ ):
+ # If original input was empty, gradient should also be empty with correct shape
+ grad_input_reshaped = torch.zeros(
+ ctx.orig_dims[0],
+ ctx.orig_dims[1],
+ ctx.orig_dims[2],
+ dtype=grad_input_flat.dtype,
+ device=grad_input_flat.device,
+ )
+ elif grad_input_flat.numel() == 0 and not (
+ ctx.orig_dims[0] * ctx.orig_dims[1] * ctx.orig_dims[2] == 0
+ ):
+ # This case should ideally not happen if forward path is correct (non-empty input -> non-empty flat grad)
+ # but as a safeguard:
+ grad_input_reshaped = torch.zeros(
+ ctx.orig_dims[0],
+ ctx.orig_dims[1],
+ ctx.orig_dims[2],
+ dtype=grad_input_flat.dtype,
+ device=grad_input_flat.device,
+ )
+ else:
+ grad_input_reshaped = grad_input_flat.view(
+ ctx.orig_dims[0], ctx.orig_dims[1], ctx.orig_dims[2]
+ )
+
+ nones_for_hyperparams = [None] * ctx.hyperparams_count
+ grad_bias_return = grad_bias_maybe if not ctx.bias_was_none else None
+
+ return (
+ grad_input_reshaped, # Gradient for student_input (reshaped)
+ grad_weight, # Gradient for student_lm_head_weight
+ None, # Gradient for target_token_ids
+ None, # Gradient for target_logprobs
+ None, # Gradient for target_mask
+ None, # Gradient for true_labels
+ grad_bias_return, # Gradient for student_lm_head_bias
+ *nones_for_hyperparams, # Grads for weight_hard_loss, ..., compute_ce_loss
+ )
+
+
+class LigerFusedLinearKLTopKLogprobLoss(torch.nn.Module):
+ """
+ wrapper for chunked top-k logprob kl-d
+ """
+
+ def __init__(
+ self,
+ weight_hard_loss: float = 0.5,
+ weight_soft_loss: float = 0.5,
+ temperature: float = 1.0, # This is the kd_temperature
+ beta: float = 1.0,
+ ignore_index: int = -100,
+ compiled: bool = True,
+ chunk_size: int = 1024,
+ compute_ce_loss: bool = True,
+ normalize_topk: bool = True,
+ ):
+ super().__init__()
+ if not (0.0 <= weight_hard_loss <= 1.0 and 0.0 <= weight_soft_loss <= 1.0):
+ raise ValueError("Loss weights must be between 0.0 and 1.0.")
+ if temperature <= 0:
+ raise ValueError("Temperature must be positive.")
+
+ self.weight_hard_loss = weight_hard_loss
+ self.weight_soft_loss = weight_soft_loss
+ self.temperature = temperature
+ self.beta = beta
+ self.ignore_index = ignore_index
+ self.compiled = compiled
+ self.chunk_size = chunk_size
+ self.compute_ce_loss = compute_ce_loss
+ self.normalize_topk = normalize_topk
+
+ if not self.compute_ce_loss and self.weight_hard_loss > 0.0:
+ print(
+ f"Warning: compute_ce_loss is False, but weight_hard_loss ({self.weight_hard_loss}) > 0. Hard loss will effectively be zero."
+ )
+ # self.weight_hard_loss = 0.0 # Or let user manage this
+ if self.weight_soft_loss == 0.0:
+ print(
+ "Warning: weight_soft_loss is 0.0. Soft (KD) loss will not be computed."
+ )
+
+ def forward(
+ self,
+ lm_head_weight: torch.Tensor, # Weights of the linear layer in the LM head
+ student_hidden_states: torch.Tensor, # student_hidden_states before the lm_head
+ target_token_ids: torch.Tensor,
+ target_logprobs: torch.Tensor,
+ target_mask: torch.Tensor,
+ true_labels: torch.Tensor,
+ student_bias: torch.Tensor = None,
+ ) -> torch.Tensor:
+ return LigerFusedLinearKLTopKLogprobFunction.apply(
+ student_hidden_states,
+ lm_head_weight,
+ target_token_ids,
+ target_logprobs,
+ target_mask,
+ true_labels,
+ student_bias,
+ self.weight_hard_loss,
+ self.weight_soft_loss,
+ self.ignore_index,
+ self.temperature,
+ self.beta,
+ self.compiled,
+ self.chunk_size,
+ self.compute_ce_loss,
+ self.normalize_topk,
+ )
diff --git a/src/axolotl/integrations/kd/kernels/models.py b/src/axolotl/integrations/kd/kernels/models.py
new file mode 100644
index 000000000..badb3460d
--- /dev/null
+++ b/src/axolotl/integrations/kd/kernels/models.py
@@ -0,0 +1,104 @@
+"""
+model patcher for chunked top-k kl-div
+"""
+
+from typing import Optional, Union, Unpack
+
+import torch
+from transformers import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+ from transformers.utils import LossKwargs
+
+ class TransformersKwargs(FlashAttentionKwargs, LossKwargs):
+ """
+ placeholder kwargs for hf model classes
+ """
+
+except ImportError:
+ from transformers.utils.generic import ( # type: ignore[no-redef]
+ TransformersKwargs,
+ )
+
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+
+
+def kldiv_forward_llama_like(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ target_logprobs: Optional[torch.Tensor] = None,
+ target_token_ids: Optional[torch.LongTensor] = None,
+ target_mask: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs: Unpack[TransformersKwargs], # type: ignore[misc]
+) -> CausalLMOutputWithPast:
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ cache_position=cache_position,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ # TODO, we can optimize this further by filtering hidden_states on sequence dimension using labels != -100
+ # self._loss_function should be LigerFusedLinearKLTopKLogprobLoss
+
+ loss = self._loss_function(
+ self.lm_head.weight,
+ hidden_states,
+ target_token_ids,
+ target_logprobs,
+ target_mask,
+ true_labels=labels,
+ )
+ num_items_in_batch = kwargs.pop("num_items_in_batch", -1)
+ if num_items_in_batch is not None and num_items_in_batch > 0:
+ loss = loss / num_items_in_batch
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=None,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+def apply_kernel(model_type):
+ # Dynamically import the module and attention class
+ module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+ model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+ module = __import__(module_path, fromlist=[f"{model_cls_prefix}ForCausalLM"])
+ model_cls = getattr(module, f"{model_cls_prefix}ForCausalLM")
+ model_cls.forward = kldiv_forward_llama_like
diff --git a/src/axolotl/integrations/kd/topk_logprob/forward_kl.py b/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
index 3c9515091..b79ba26f3 100644
--- a/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
+++ b/src/axolotl/integrations/kd/topk_logprob/forward_kl.py
@@ -15,41 +15,9 @@
"""
loss for top_k KL divergence
"""
+
import torch
-
-
-def zscore_standardize(
- logits: torch.Tensor,
- mask: torch.Tensor = None,
- base_temperature: float = 1.0,
- eps: float = 1e-9,
-):
- """
- Z-score standardize along the last dimension of `logits`.
- i.e., for each [B, seq_len] row, across K entries:
- z = (logits - mean) / std,
- then scale by 1 / base_temperature if desired.
-
- mask can be broadcastable or None. If None, we standardize all elements.
- """
- if mask is None:
- # shape: [B, seq_len, K]
- # Mean and std over dim=-1
- mean = logits.mean(dim=-1, keepdim=True)
- var = logits.var(dim=-1, unbiased=False, keepdim=True)
- else:
- # If you have to exclude some tokens, multiply by mask, etc.
- float_mask = mask.to(logits.dtype)
- count = float_mask.sum(dim=-1, keepdim=True).clamp_min(1.0)
- mean = (logits * float_mask).sum(dim=-1, keepdim=True) / count
- var = (float_mask * (logits - mean) ** 2).sum(dim=-1, keepdim=True) / count
-
- std = torch.sqrt(var.clamp_min(eps))
- z = (logits - mean) / std
-
- # Scale by 1 / base_temperature
- z = z / base_temperature
- return z
+from torch import nn
@torch.jit.script
@@ -60,7 +28,6 @@ def loss(
target_mask: torch.Tensor,
num_items_in_batch: int = -1, # Use -1 to indicate "None"
kd_temperature: float = 1.0,
- top_k_before_softmax: int = 0,
) -> torch.Tensor:
"""
A KD loss function that is TorchScript-friendly.
@@ -77,8 +44,6 @@ def loss(
num_items_in_batch (int, optional): The number of items in the batch.
kd_temperature (float, optional): The temperature for KD.
Default: 1.0
- top_k_before_softmax (int, optional): Flag of whether to apply softmax before gathering student top-k logits
- Default: 0
"""
target_logprobs = target_logprobs.float()
@@ -88,46 +53,24 @@ def loss(
# student_logits shape: [B, student_seq_len, vocab_size]
teacher_seq_len = target_token_ids.shape[1]
- if top_k_before_softmax:
- # Slice student logits to match teacher-provided sequence length
- student_logits_for_kd = student_logits[
- :, :teacher_seq_len, :
- ] # [B, teacher_seq_len, vocab_size]
+ # Slice student logits to match teacher-provided sequence length
+ student_logits_for_kd = (
+ student_logits[:, :teacher_seq_len, :] / kd_temperature
+ ) # [B, teacher_seq_len, vocab_size]
- # Gather student logits for teacher's top-K tokens
- student_logits_topk = torch.gather(
- student_logits_for_kd, dim=-1, index=target_token_ids
- ) # [B, teacher_seq_len, K]
+ # keep in full precision for numerical stability of loss
+ student_logits_for_kd = student_logits_for_kd.float()
- student_logits_topk = student_logits_topk.float()
+ # Gather student logits for teacher's top-K tokens
+ student_logits_topk = torch.gather(
+ student_logits_for_kd, dim=-1, index=target_token_ids
+ ) # [B, teacher_seq_len, K]
- # Apply KD temperature to student’s logits
- if kd_temperature != 1.0:
- student_logits_topk = student_logits_topk / kd_temperature
+ # Compute logsumexp across full vocabulary
+ student_lse = torch.logsumexp(student_logits_for_kd, dim=-1, keepdim=True)
- # Convert student top-k logits to logprobs
- student_logprobs_topk = student_logits_topk - torch.logsumexp(
- student_logits_topk, dim=-1, keepdim=True
- ) # [B, teacher_seq_len, K]
- else:
- # Slice student logits to match teacher-provided sequence length
- student_logits_for_kd = (
- student_logits[:, :teacher_seq_len, :] / kd_temperature
- ) # [B, teacher_seq_len, vocab_size]
-
- # keep in full precision for numerical stability of loss
- student_logits_for_kd = student_logits_for_kd.float()
-
- # Gather student logits for teacher's top-K tokens
- student_logits_topk = torch.gather(
- student_logits_for_kd, dim=-1, index=target_token_ids
- ) # [B, teacher_seq_len, K]
-
- # Compute logsumexp across full vocabulary
- student_lse = torch.logsumexp(student_logits_for_kd, dim=-1, keepdim=True)
-
- # Convert just the top-k logits to logprobs
- student_logprobs_topk = student_logits_topk - student_lse
+ # Convert just the top-k logits to logprobs
+ student_logprobs_topk = student_logits_topk - student_lse
# Convert teacher_mask to boolean for indexing
# In TorchScript, .bool() is sometimes unsupported, so we do:
@@ -144,10 +87,6 @@ def loss(
kd_loss_per_token = teacher_probs * (target_logprobs - student_logprobs_topk)
kd_loss = kd_loss_per_token.sum()
- # Multiply by T^2 (classical KD scaling)
- if kd_temperature != 1.0:
- kd_loss = kd_loss * (kd_temperature**2)
-
# Normalize by number of items (if provided) or by valid tokens
if num_items_in_batch > 0:
kd_loss = kd_loss / float(num_items_in_batch)
@@ -158,80 +97,77 @@ def loss(
return kd_loss
-def topk_kd_loss_with_zscore(
- student_logits: torch.Tensor, # [B, seq_len, vocab_size]
- target_token_ids: torch.Tensor, # [B, seq_len, K]
- target_logprobs: torch.Tensor, # [B, seq_len, K], sums to 1.0 in prob space
- target_mask: torch.Tensor, # [B, seq_len, K] or [B, seq_len]
- kd_temperature: float = 1.0, # classic KD temperature
- zscore_base_temp: float = 1.0, # from the paper
- num_items_in_batch: int = -1,
-):
+class ChunkedTopKKDLoss(nn.Module):
"""
- A variant of top_k KL divergence with Z-score scaling
- from "Logit Standardization in Knowledge Distillation".
+ A wrapper that chunks (splits) the student and teacher outputs along the time dimension
+ to reduce peak memory usage when upcasting from bf16 to fp32, especially for large vocabularies.
+
+ Usage is analogous to ForwardKLWithChunkedOutputLoss but adapted to top-K teacher logprobs.
"""
- target_logprobs = target_logprobs.float()
+ def __init__(self, num_output_chunks: int = 8, kd_temperature: float = 1.0):
+ super().__init__()
+ self.num_output_chunks = num_output_chunks
+ self.kd_temperature = kd_temperature
- B, teacher_seq_len, K = target_logprobs.shape # pylint: disable=invalid-name
- # 1) Gather the student's top-k logits to match teacher
- student_logits_for_kd = student_logits[
- :, :teacher_seq_len, :
- ] # [B, seq_len, vocab]
- student_topk_logits = torch.gather(
- student_logits_for_kd, dim=-1, index=target_token_ids
- ) # [B, seq_len, K]
+ def forward(
+ self,
+ student_logits: torch.Tensor, # [B, seq_len, vocab_size]
+ target_token_ids: torch.Tensor, # [B, seq_len, K]
+ target_logprobs: torch.Tensor, # [B, seq_len, K]
+ target_mask: torch.Tensor, # [B, seq_len, K]
+ num_items_in_batch: int = -1, # optional batch size for normalization
+ ) -> torch.Tensor:
+ # 1. Split along the "token" dimension (dim=1).
+ student_logits_chunks = student_logits.chunk(self.num_output_chunks, dim=1)
+ token_ids_chunks = target_token_ids.chunk(self.num_output_chunks, dim=1)
+ logprobs_chunks = target_logprobs.chunk(self.num_output_chunks, dim=1)
+ mask_chunks = target_mask.chunk(self.num_output_chunks, dim=1)
- student_topk_logits = student_topk_logits.float()
+ # We'll accumulate a global "sum of losses" and "sum of valid tokens"
+ # so that our final average is consistent with the entire sequence/batch.
+ total_loss = 0.0
+ total_valid_tokens = 0
- # 2) If you want to keep the "classical" T scaling, apply it first
- if kd_temperature != 1.0:
- student_topk_logits = student_topk_logits / kd_temperature
+ # 2. Loop over each chunk and compute a chunk-specific loss.
+ for st_chunk, tid_chunk, lp_chunk, msk_chunk in zip(
+ student_logits_chunks,
+ token_ids_chunks,
+ logprobs_chunks,
+ mask_chunks,
+ strict=False,
+ ):
+ # We pass num_items_in_batch=-1 so that the kd_loss
+ # will average over *this chunk's* valid tokens only.
+ chunk_loss = loss(
+ student_logits=st_chunk,
+ target_token_ids=tid_chunk,
+ target_logprobs=lp_chunk,
+ target_mask=msk_chunk,
+ num_items_in_batch=-1, # ensure per-chunk averaging by valid tokens
+ kd_temperature=self.kd_temperature,
+ )
- # 3) Convert teacher logprobs -> treat them as “logits” for z-score
- # (They differ by +some_constant from real logits, but in z-score
- # that constant is subtracted out anyway.)
- teacher_logits_for_zscore = target_logprobs # rename variable for clarity
+ # kd_loss returns an average over the chunk's valid tokens.
+ # We want a global average in the end, so we need to re‐weight
+ # by the number of valid tokens in this chunk and keep track of the total.
+ chunk_valid_mask = msk_chunk.to(torch.bool)
+ chunk_valid_count = chunk_valid_mask.sum() # scalar tensor
- # 4) Z-score teacher and student
- # If target_mask is 2D, expand to 3D for the K dimension
- if target_mask.dim() == 2 and target_mask.shape[:2] == (B, teacher_seq_len):
- target_mask = target_mask.unsqueeze(-1).expand(-1, -1, K)
+ # Re-scale "chunk average" back to "chunk sum"
+ chunk_loss_sum = chunk_loss * chunk_valid_count
- teacher_z = zscore_standardize(
- teacher_logits_for_zscore, mask=target_mask, base_temperature=zscore_base_temp
- )
- student_z = zscore_standardize(
- student_topk_logits, mask=target_mask, base_temperature=zscore_base_temp
- )
+ total_loss += chunk_loss_sum
+ total_valid_tokens += chunk_valid_count
- # 5) Convert to log-probs for KL
- teacher_logprobs_z = teacher_z - torch.logsumexp(teacher_z, dim=-1, keepdim=True)
- student_logprobs_z = student_z - torch.logsumexp(student_z, dim=-1, keepdim=True)
+ # 3. Normalize *once* at the end.
+ if num_items_in_batch > 0:
+ # If the user gave us a manual denominator (e.g. total items in batch),
+ # we divide by it. Typically used if each item is of different length.
+ final_loss = total_loss / float(num_items_in_batch)
+ else:
+ # Otherwise, divide by total valid tokens across all chunks.
+ # to get the same result as a non-chunked approach.
+ final_loss = total_loss / float(total_valid_tokens)
- # 6) Restrict to valid tokens if needed
- valid_mask = target_mask.bool() # shape [B, seq_len, K]
- teacher_probs_z = teacher_logprobs_z.exp()
- teacher_probs_z = teacher_probs_z[valid_mask]
- teacher_logprobs_z = teacher_logprobs_z[valid_mask]
- student_logprobs_z = student_logprobs_z[valid_mask]
-
- # 7) forward KL: sum( p_teacher * [log(p_teacher) - log(p_student)] )
- kd_loss_per_token = teacher_probs_z * (teacher_logprobs_z - student_logprobs_z)
- kd_loss = kd_loss_per_token.sum()
-
- # 8) If using classical KD scaling by T^2
- if kd_temperature != 1.0:
- kd_loss = kd_loss * (kd_temperature**2)
-
- # Optionally scale by zscore_base_temp**2 if you want (paper might differ).
- # kd_loss = kd_loss * (zscore_base_temp**2)
-
- # 9) Normalize
- if num_items_in_batch is not None and num_items_in_batch > 0:
- kd_loss = kd_loss / float(num_items_in_batch)
- else:
- kd_loss = kd_loss / float(kd_loss_per_token.size(0))
-
- return kd_loss
+ return final_loss
diff --git a/src/axolotl/integrations/kd/trainer.py b/src/axolotl/integrations/kd/trainer.py
index f99f2ca28..0e98497a7 100644
--- a/src/axolotl/integrations/kd/trainer.py
+++ b/src/axolotl/integrations/kd/trainer.py
@@ -18,8 +18,7 @@ KD trainer
from axolotl.core.trainers.base import AxolotlTrainer
-from .topk_logprob.forward_kl import loss as topk_kd_loss
-from .topk_logprob.forward_kl import topk_kd_loss_with_zscore
+from .kernels.liger import LigerFusedLinearKLTopKLogprobLoss
class AxolotlKDTrainer(AxolotlTrainer):
@@ -27,6 +26,27 @@ class AxolotlKDTrainer(AxolotlTrainer):
Custom trainer subclass for Knowledge Distillation (KD)
"""
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.model_accepts_loss_kwargs = True
+
+ loss_fn = LigerFusedLinearKLTopKLogprobLoss(
+ self.args.kd_ce_alpha, # hard label loss
+ self.args.kd_alpha, # kd loss
+ self.args.kd_temperature,
+ self.args.kd_beta or 0.0,
+ compute_ce_loss=bool(self.args.kd_ce_alpha),
+ normalize_topk=self.args.kd_normalize_topk,
+ )
+ target = self.model
+
+ # Unwrap PEFT wrapper
+ if hasattr(target, "get_base_model"):
+ target = target.get_base_model()
+
+ # Set on the actual model instance
+ target._loss_function = loss_fn
+
def _set_signature_columns_if_needed(self):
super()._set_signature_columns_if_needed()
columns_to_add = []
@@ -52,12 +72,12 @@ class AxolotlKDTrainer(AxolotlTrainer):
Subclass and override for custom behavior.
"""
-
- target_logprobs = inputs.pop("target_logprobs")
- target_token_ids = inputs.pop("target_token_ids")
- target_mask = inputs.pop("target_mask")
-
- seq_len = target_token_ids.shape[1]
+ if (
+ self.args.sample_packing
+ and hasattr(inputs, "attention_mask")
+ and hasattr(inputs, "position_ids")
+ ):
+ del inputs["attention_mask"]
if self.model_accepts_loss_kwargs:
loss_kwargs = {}
@@ -65,49 +85,4 @@ class AxolotlKDTrainer(AxolotlTrainer):
loss_kwargs["num_items_in_batch"] = num_items_in_batch
inputs = {**inputs, **loss_kwargs}
outputs = model(**inputs)
-
- # FIXME: account for tokenizer.padding_side
- student_logits = outputs["logits"][:, : seq_len - 1, :].contiguous()
-
- shift_logits = student_logits.contiguous()
- target_logprobs_for_loss = target_logprobs[..., 1:, :].contiguous()
- target_token_ids_for_loss = target_token_ids[..., 1:, :].contiguous()
- target_mask_for_loss = target_mask[..., 1:, :].contiguous()
-
- if self.args.kd_zscore_base_temp:
- loss_kd = topk_kd_loss_with_zscore(
- shift_logits,
- target_token_ids_for_loss,
- target_logprobs_for_loss,
- target_mask_for_loss,
- kd_temperature=self.args.kd_temperature,
- zscore_base_temp=self.args.kd_zscore_base_temp,
- num_items_in_batch=num_items_in_batch,
- )
- else:
- loss_kd = topk_kd_loss(
- shift_logits,
- target_token_ids_for_loss,
- target_logprobs_for_loss,
- target_mask_for_loss,
- num_items_in_batch=num_items_in_batch,
- kd_temperature=self.args.kd_temperature,
- top_k_before_softmax=1 if self.args.kd_top_k_before_softmax else 0,
- )
-
- if self.args.kd_ce_alpha > 0:
- kd_alpha = self.args.kd_alpha
- loss = self.args.kd_ce_alpha * outputs["loss"] + kd_alpha * loss_kd
- else:
- loss = loss_kd
- # Save past state if it exists
- # TODO: this needs to be fixed and made cleaner later.
- if self.args.past_index >= 0:
- self._past = outputs[ # pylint: disable=attribute-defined-outside-init
- self.args.past_index
- ]
-
- if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs:
- loss *= self.accelerator.num_processes
-
- return (loss, outputs) if return_outputs else loss
+ return outputs[0]
diff --git a/src/axolotl/integrations/kd/utils.py b/src/axolotl/integrations/kd/utils.py
new file mode 100644
index 000000000..ba60694a5
--- /dev/null
+++ b/src/axolotl/integrations/kd/utils.py
@@ -0,0 +1,100 @@
+"""Helper KD utils"""
+
+import math
+from typing import List, Union
+
+import numpy as np
+import torch
+from torch import FloatTensor, Tensor
+
+
+def normalize_logprobs(logprobs: FloatTensor, topk: int) -> FloatTensor:
+ """
+ Re-normalizes top-k raw logprobs as probabilities, and converts back to logprobs.
+ """
+ # Ensure raw_logprobs matches kd_online_topk length for tensor operations
+ # This should ideally be handled by the caller ensuring correct padding/truncation first
+ if logprobs.shape[-1] != topk:
+ # pad last dimension of logprobs to match topk length with -inf
+ padding_len = topk - logprobs.shape[-1]
+ padding_tensor = torch.full(
+ (
+ *logprobs.shape[:-1],
+ padding_len,
+ ), # Takes all dimensions of logprobs except the last, then appends padding_needed
+ float("-inf"),
+ dtype=logprobs.dtype,
+ device=logprobs.device,
+ )
+ logprobs = torch.cat((logprobs, padding_tensor), dim=-1)
+
+ # Convert logprobs at T_online to probabilities
+ # use log sum exp trick to avoid underflow
+ position_logprobs_lse = torch.logsumexp(logprobs, dim=-1, keepdim=True)
+ teacher_probs_t_online = torch.exp(logprobs - position_logprobs_lse)
+
+ # Normalize probabilities (sum to 1)
+ # This is important if the top-k from server aren't a full distribution
+ teacher_probs_t_online_sum = teacher_probs_t_online.sum(dim=-1, keepdim=True)
+ teacher_probs_t_online = teacher_probs_t_online / teacher_probs_t_online_sum
+
+ final_logprobs_tensor = torch.log(teacher_probs_t_online)
+
+ return final_logprobs_tensor
+
+
+def strided_chunk_views(
+ tensor: Union[np.ndarray, torch.Tensor],
+ chunks: int,
+ dim: int = 0,
+ stride: int = 1,
+ chunk_size: int | None = None,
+) -> List[Union[np.ndarray, torch.Tensor]]:
+ """
+ Split a tensor into chunks along a dimension with striding, prioritizing views over copies.
+
+ Args:
+ tensor: Input tensor (numpy array or torch tensor)
+ chunks: Number of chunks to create
+ dim: Dimension along which to chunk (default: 0)
+ stride: Stride between chunk starting positions (default: 1)
+ chunk_size: Size of each chunk. If None, calculated automatically (default: None)
+
+ Returns:
+ List of tensor chunks (views when possible, copies when necessary)
+ """
+
+ # Get the size of the specified dimension
+ dim_size = tensor.shape[dim]
+
+ # Calculate chunk size if not provided
+ if chunk_size is None:
+ chunk_size = (dim_size + chunks - 1) // chunks # Ceiling division
+
+ chunks_list = []
+
+ for i in range(chunks):
+ start_idx = i * stride
+ end_idx = min(start_idx + chunk_size, dim_size)
+
+ # Break if we've gone beyond the tensor
+ if start_idx >= dim_size:
+ break
+
+ # Create slice objects for all dimensions
+ slices = [slice(None)] * tensor.ndim
+ slices[dim] = slice(start_idx, end_idx)
+
+ chunk = tensor[tuple(slices)]
+ chunks_list.append(chunk)
+
+ return chunks_list
+
+
+def chunk_overlap(input_tensor: Tensor, chunks: int, dim: int = 0, overlap: int = 1):
+ dim_size = input_tensor.shape[dim]
+ stride = math.ceil(dim_size / chunks)
+
+ return strided_chunk_views(
+ input_tensor, chunks, dim, stride=stride, chunk_size=stride + overlap
+ )
diff --git a/src/axolotl/integrations/liger/__init__.py b/src/axolotl/integrations/liger/__init__.py
index 1c17ab2b5..c20f4545c 100644
--- a/src/axolotl/integrations/liger/__init__.py
+++ b/src/axolotl/integrations/liger/__init__.py
@@ -18,170 +18,11 @@ Module for the Plugin for LIGER integraton with Axolotl.
Liger Kernel is the collection of Triton-native kernels for LLM Training.
It is designed to be performant, correct, and light-weight.
"""
-import inspect
-import sys
-from axolotl.integrations.base import BasePlugin
-from axolotl.utils.logging import get_logger
+from .args import LigerArgs
+from .plugin import LigerPlugin
-from .args import LigerArgs # pylint: disable=unused-import. # noqa: F401
-from .utils import patch_with_compile_disable
-
-LOG = get_logger(__name__, use_environ=True)
-
-
-class LigerPlugin(BasePlugin):
- """
- Plugin for LIGER integraton with Axolotl.
- """
-
- def get_input_args(self):
- return "axolotl.integrations.liger.LigerArgs"
-
- def pre_model_load(self, cfg):
- if cfg.torch_compile:
- # torch compile will unnecessarily attempt to optimize the triton kernel unless explicitly disabled
- import liger_kernel.ops.fused_linear_cross_entropy
-
- patch_with_compile_disable(
- liger_kernel.ops.fused_linear_cross_entropy,
- "fused_linear_cross_entropy_forward",
- )
- patch_with_compile_disable(
- liger_kernel.ops.fused_linear_cross_entropy,
- "fused_linear_cross_entropy_backward",
- )
- from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
- from liger_kernel.transformers.functional import liger_cross_entropy
- from liger_kernel.transformers.layer_norm import LigerLayerNorm
- from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
- from liger_kernel.transformers.rms_norm import LigerRMSNorm
- from liger_kernel.transformers.rope import liger_rotary_pos_emb
- from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
-
- if cfg.liger_cross_entropy and cfg.liger_fused_linear_cross_entropy:
- raise ValueError(
- "Cannot have both `liger_cross_entropy` and `liger_fused_linear_cross_entropy` set."
- )
-
- if cfg.model_config_type in MODEL_TYPE_TO_APPLY_LIGER_FN:
- apply_liger_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[cfg.model_config_type]
- liger_fn_sig = inspect.signature(apply_liger_fn)
- kwargs = {}
- if "rope" in liger_fn_sig.parameters:
- kwargs["rope"] = cfg.liger_rope
- if "cross_entropy" in liger_fn_sig.parameters:
- kwargs["cross_entropy"] = cfg.liger_cross_entropy
- if "fused_linear_cross_entropy" in liger_fn_sig.parameters:
- kwargs["fused_linear_cross_entropy"] = (
- cfg.liger_fused_linear_cross_entropy
- )
- if "rms_norm" in liger_fn_sig.parameters:
- kwargs["rms_norm"] = cfg.liger_rms_norm
- if "layer_norm" in liger_fn_sig.parameters:
- kwargs["layer_norm"] = cfg.liger_layer_norm
- if "geglu" in liger_fn_sig.parameters:
- kwargs["geglu"] = cfg.liger_glu_activation
- elif "swiglu" in liger_fn_sig.parameters:
- kwargs["swiglu"] = cfg.liger_glu_activation
- LOG.info(f"Applying LIGER to {cfg.model_config_type} with kwargs: {kwargs}")
- apply_liger_fn(**kwargs)
- elif cfg.model_config_type == "jamba":
- from transformers.models.jamba import modeling_jamba
-
- from .models.jamba import lce_forward as jamba_lce_forward
-
- if cfg.liger_rope:
- modeling_jamba.apply_rotary_pos_emb = liger_rotary_pos_emb
- if cfg.liger_rms_norm:
- modeling_jamba.JambaRMSNorm = LigerRMSNorm
- if cfg.liger_glu_activation:
- modeling_jamba.JambaMLP = LigerSwiGLUMLP
- if cfg.liger_layer_norm:
- modeling_jamba.nn.LayerNorm = LigerLayerNorm
- if cfg.liger_cross_entropy:
- from transformers.loss.loss_utils import nn
-
- nn.functional.cross_entropy = liger_cross_entropy
- if cfg.liger_fused_linear_cross_entropy:
- modeling_jamba.JambaForCausalLM.forward = jamba_lce_forward
- elif cfg.model_config_type == "deepseek_v2":
- from accelerate import init_empty_weights
- from transformers import AutoModelForCausalLM
-
- with init_empty_weights():
- model = AutoModelForCausalLM.from_pretrained(
- cfg.base_model, trust_remote_code=cfg.trust_remote_code or False
- )
- modeling_mod = sys.modules[model.__class__.__module__]
-
- from .models.deepseekv2 import lce_forward as deepseekv2_lce_forward
-
- if cfg.liger_rope:
- # The DeepseekV2 version of RoPE is different than upstream LLaMA.
- # See https://github.com/linkedin/Liger-Kernel/issues/129#issuecomment-2313763528
- LOG.warning("Fused liger_rope is not supported for DeepseekV2.")
- if cfg.liger_glu_activation:
- LOG.warning("liger_glu_activation is not supported for DeepseekV2.")
- if cfg.liger_rms_norm:
- modeling_mod.DeepseekV2RMSNorm = LigerRMSNorm
- if cfg.liger_glu_activation:
- modeling_mod.DeepseekV2MLP.forward = LigerSwiGLUMLP.forward
- if cfg.liger_layer_norm:
- modeling_mod.DeepseekV2MLP.forward = LigerLayerNorm.forward
- if cfg.liger_cross_entropy:
- # We do not patch `nn.functional.cross_entropy` for DeepseekV2 as it still uses
- # nn.CrossEntropyLoss in the forward method.
- modeling_mod.CrossEntropyLoss = LigerCrossEntropyLoss
- if cfg.liger_fused_linear_cross_entropy:
- modeling_mod.DeepseekV2ForCausalLM.forward = deepseekv2_lce_forward
- elif cfg.model_config_type == "llama4":
- from axolotl.integrations.liger.models.llama4 import (
- apply_liger_kernel_to_llama4,
- )
-
- apply_liger_kernel_to_llama4(
- cross_entropy=cfg.liger_cross_entropy,
- fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
- glu_activation=cfg.liger_glu_activation,
- rms_norm=cfg.liger_rms_norm,
- layer_norm=cfg.liger_layer_norm,
- )
- elif cfg.model_config_type == "qwen3":
- from axolotl.integrations.liger.models.qwen3 import (
- apply_liger_kernel_to_qwen3,
- )
-
- apply_liger_kernel_to_qwen3(
- cross_entropy=cfg.liger_cross_entropy,
- fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
- glu_activation=cfg.liger_glu_activation,
- rms_norm=cfg.liger_rms_norm,
- layer_norm=cfg.liger_layer_norm,
- )
- elif cfg.model_config_type == "qwen3_moe":
- from axolotl.integrations.liger.models.qwen3_moe import (
- apply_liger_kernel_to_qwen3_moe,
- )
-
- apply_liger_kernel_to_qwen3_moe(
- cross_entropy=cfg.liger_cross_entropy,
- fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
- glu_activation=cfg.liger_glu_activation,
- rms_norm=cfg.liger_rms_norm,
- layer_norm=cfg.liger_layer_norm,
- )
- elif cfg.model_config_type == "granitemoe":
- from liger_kernel.transformers import apply_liger_kernel_to_granite
-
- apply_liger_kernel_to_granite(
- rope=cfg.liger_rope,
- cross_entropy=cfg.liger_cross_entropy,
- fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
- rms_norm=cfg.liger_rms_norm,
- swiglu=cfg.liger_glu_activation,
- )
- else:
- LOG.warning(
- f"Unsupported model config type: {cfg.model_config_type}. Liger not applied."
- )
+__all__ = [
+ "LigerArgs",
+ "LigerPlugin",
+]
diff --git a/src/axolotl/integrations/liger/args.py b/src/axolotl/integrations/liger/args.py
index 7c9eb23d5..d5bb10cfd 100644
--- a/src/axolotl/integrations/liger/args.py
+++ b/src/axolotl/integrations/liger/args.py
@@ -15,7 +15,6 @@
"""
Module for handling LIGER input arguments.
"""
-from typing import Optional
from pydantic import BaseModel, model_validator
@@ -29,13 +28,13 @@ class LigerArgs(BaseModel):
Input args for LIGER.
"""
- liger_rope: Optional[bool] = None
- liger_rms_norm: Optional[bool] = None
- liger_layer_norm: Optional[bool] = None
- liger_swiglu: Optional[bool] = None
- liger_glu_activation: Optional[bool] = None
- liger_cross_entropy: Optional[bool] = None
- liger_fused_linear_cross_entropy: Optional[bool] = None
+ liger_rope: bool | None = None
+ liger_rms_norm: bool | None = None
+ liger_layer_norm: bool | None = None
+ liger_swiglu: bool | None = None
+ liger_glu_activation: bool | None = None
+ liger_cross_entropy: bool | None = None
+ liger_fused_linear_cross_entropy: bool | None = None
@model_validator(mode="before")
@classmethod
@@ -52,3 +51,33 @@ class LigerArgs(BaseModel):
)
data["liger_glu_activation"] = data.pop("liger_swiglu")
return data
+
+ @model_validator(mode="before")
+ @classmethod
+ def check_tiled_mlp_conflict(cls, data):
+ if (
+ data.get("liger_glu_activation") is True
+ and data.get("tiled_mlp") is True
+ and not data.get("tiled_mlp_use_original_mlp")
+ ):
+ raise ValueError(
+ "You cannot have both `liger_glu_activation` and `tiled_mlp` set without `tiled_mlp_use_original_mlp: true`."
+ )
+ return data
+
+ @model_validator(mode="before")
+ @classmethod
+ def check_liger_rms_norm_tensor_parallel(cls, data):
+ if data.get("liger_rms_norm") and data.get("tensor_parallel_size", 1) > 1:
+ raise ValueError(
+ "`liger_rms_norm` is incompatible with tensor parallelism, "
+ "see https://github.com/linkedin/Liger-Kernel/issues/826"
+ )
+ return data
+
+ @model_validator(mode="after")
+ def check_tensor_parallel_size_liger_fused_linear_cross_entropy(self):
+ # TODO @SalmanMohammadi this is a larger fix - investigate
+ if self.tensor_parallel_size > 1 and self.liger_fused_linear_cross_entropy:
+ raise ValueError("Tensor parallelism is not compatible with liger losses.")
+ return self
diff --git a/src/axolotl/integrations/liger/models/base.py b/src/axolotl/integrations/liger/models/base.py
new file mode 100644
index 000000000..a9dbe9412
--- /dev/null
+++ b/src/axolotl/integrations/liger/models/base.py
@@ -0,0 +1,188 @@
+"""
+Generic FLCE patch for untested models similar to Llama
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
+from liger_kernel.transformers.trainer.orpo_trainer import _FSDPForwardRedirection
+from liger_kernel.utils import PEFT_AVAILABLE
+from peft.utils import ModulesToSaveWrapper
+from torch.distributed.fsdp import FullyShardedDataParallel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+
+
+def lce_forward(
+ self,
+ *args,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ labels: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ skip_logits: Optional[bool] = None,
+ **kwargs,
+) -> Union[Tuple, CausalLMOutputWithPast]:
+ r"""
+ Args:
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ logits_to_keep (`int` or `torch.Tensor`, *optional*):
+ If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+ `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
+ token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
+ If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
+ This is useful when using packed tensor format (single dimension for batch and sequence length).
+ """
+
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = self.model(
+ *args,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = (
+ slice(-logits_to_keep, None)
+ if isinstance(logits_to_keep, int)
+ else logits_to_keep
+ )
+ kept_hidden_states = hidden_states[:, slice_indices, :]
+
+ shift_labels = kwargs.pop("shift_labels", None)
+ logits = None
+ loss = None
+
+ # if in training mode, don't materialize logits
+ if skip_logits and labels is None and shift_labels is None:
+ raise ValueError("skip_logits is True, but labels and shift_labels are None")
+
+ if skip_logits is None:
+ # By default, if in training mode, don't materialize logits
+ skip_logits = self.training and (labels is not None or shift_labels is not None)
+
+ if skip_logits:
+ loss = lce_maybe_trainable_lm_head(
+ self,
+ hidden_states=kept_hidden_states,
+ hidden_size=self.config.hidden_size,
+ labels=labels,
+ shift_labels=shift_labels,
+ **kwargs,
+ )
+
+ else:
+ logits = self.lm_head(kept_hidden_states)
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits,
+ labels=labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+def lce_maybe_trainable_lm_head(
+ self, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs
+):
+ lm_head = self.lm_head
+
+ # Unwrap the module if lm_head has been added as trainable module in PEFT LoRA configuration,
+ # i.e. listed in the modules_to_save field of LoraConfig, so the lm_head weights are read
+ # from the unwrapped module.
+ # See https://huggingface.co/docs/peft/package_reference/lora for reference.
+ if PEFT_AVAILABLE and isinstance(lm_head, ModulesToSaveWrapper):
+ lm_head = lm_head.modules_to_save.default
+
+ # If FSDP is used and lm_head is trainable, e.g., during full fine-tuning or with LoRA,
+ # reading the lm_head module weights and calling the kernel must be done within FSDP forward pass
+ # so the module entire parameters are summoned and kept in memory during the kernel execution.
+ if isinstance(lm_head, FullyShardedDataParallel):
+ return _FSDPForwardRedirection()(
+ lm_head,
+ _liger_for_causal_lm_loss,
+ lm_head.module,
+ hidden_states,
+ hidden_size,
+ labels,
+ shift_labels,
+ **loss_kwargs,
+ )
+
+ # FSDP is not used so we can read the lm_head weights and call the kernel directly
+ return _liger_for_causal_lm_loss(
+ lm_head=self.lm_head,
+ hidden_states=hidden_states,
+ hidden_size=hidden_size,
+ labels=labels,
+ shift_labels=shift_labels,
+ **loss_kwargs,
+ )
+
+
+def _liger_for_causal_lm_loss(
+ lm_head, hidden_states, hidden_size, labels, shift_labels, **loss_kwargs
+):
+ return LigerForCausalLMLoss(
+ hidden_states=hidden_states,
+ lm_head_weight=lm_head.weight,
+ labels=labels,
+ hidden_size=hidden_size,
+ shift_labels=shift_labels,
+ **loss_kwargs,
+ )
+
+
+def patch_lce_forward(
+ model_type,
+):
+ try:
+ # Dynamically import the module and MLP class
+ module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+ model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+ module = __import__(module_path, fromlist=[f"{model_cls_prefix}ForCausalLM"])
+ model_cls = getattr(module, f"{model_cls_prefix}ForCausalLM")
+
+ model_cls.forward = lce_forward
+
+ except (ImportError, AttributeError) as e:
+ raise RuntimeError(
+ f"Could not import ForCausalLM class for model_type: {model_type}. "
+ f"Error: {str(e)}"
+ ) from e
diff --git a/src/axolotl/integrations/liger/models/deepseekv2.py b/src/axolotl/integrations/liger/models/deepseekv2.py
index 2f0d2a704..99adce4a7 100644
--- a/src/axolotl/integrations/liger/models/deepseekv2.py
+++ b/src/axolotl/integrations/liger/models/deepseekv2.py
@@ -2,8 +2,6 @@
DeepseekV2 model with LigerFusedLinearCrossEntropyLoss
"""
-# pylint: disable=duplicate-code
-
from typing import List, Optional, Tuple, Union
import torch
diff --git a/src/axolotl/integrations/liger/models/jamba.py b/src/axolotl/integrations/liger/models/jamba.py
index d25529970..78689e40c 100644
--- a/src/axolotl/integrations/liger/models/jamba.py
+++ b/src/axolotl/integrations/liger/models/jamba.py
@@ -2,8 +2,6 @@
Jamba model with LigerFusedLinearCrossEntropyLoss
"""
-# pylint: disable=duplicate-code
-
from typing import Optional, Tuple, Union
import torch
diff --git a/src/axolotl/integrations/liger/models/llama4.py b/src/axolotl/integrations/liger/models/llama4.py
index 689823bb6..e51140265 100644
--- a/src/axolotl/integrations/liger/models/llama4.py
+++ b/src/axolotl/integrations/liger/models/llama4.py
@@ -46,7 +46,6 @@ def lce_forward(
Returns:
"""
- # pylint: disable=duplicate-code
output_attentions = (
output_attentions
if output_attentions is not None
@@ -78,9 +77,7 @@ def lce_forward(
hidden_states = outputs[0]
if hasattr(self.config, "pretraining_tp") and self.config.pretraining_tp > 1:
- raise Exception( # pylint: disable=broad-exception-raised
- "Liger Kernel does not support pretraining_tp!!"
- )
+ raise Exception("Liger Kernel does not support pretraining_tp!!")
logits = None
loss = None
@@ -128,7 +125,7 @@ def apply_liger_kernel_to_llama4(
rms_norm: bool = False,
glu_activation: bool = False,
layer_norm: bool = False,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
) -> None:
"""
Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
@@ -144,15 +141,15 @@ def apply_liger_kernel_to_llama4(
layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
"""
- import transformers.models.llama4.modeling_llama4 # noqa: F401 # pylint: disable=unused-import
+ import transformers.models.llama4.modeling_llama4 # noqa: F401
from liger_kernel.transformers.functional import liger_cross_entropy
from liger_kernel.transformers.layer_norm import LigerLayerNorm
from liger_kernel.transformers.rms_norm import LigerRMSNorm
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
- assert not (
- cross_entropy and fused_linear_cross_entropy
- ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ assert not (cross_entropy and fused_linear_cross_entropy), (
+ "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ )
modeling_llama4 = sys.modules["transformers.models.llama4.modeling_llama4"]
@@ -165,7 +162,7 @@ def apply_liger_kernel_to_llama4(
# clone config to avoid modifying the original
config = deepcopy(config)
if intermediate_size:
- setattr(config, "intermediate_size", intermediate_size)
+ config.intermediate_size = intermediate_size
return LigerSwiGLUMLP(config, **kwargs)
modeling_llama4.Llama4TextMLP = _liger_swiglu_mlp_wrapper
diff --git a/src/axolotl/integrations/liger/models/qwen3.py b/src/axolotl/integrations/liger/models/qwen3.py
index 1dc19eaf9..b008755da 100644
--- a/src/axolotl/integrations/liger/models/qwen3.py
+++ b/src/axolotl/integrations/liger/models/qwen3.py
@@ -43,7 +43,6 @@ def lce_forward(
Returns:
"""
- # pylint: disable=duplicate-code
output_attentions = (
output_attentions
if output_attentions is not None
@@ -113,9 +112,8 @@ def apply_liger_kernel_to_qwen3(
rms_norm: bool = False,
glu_activation: bool = False,
layer_norm: bool = False,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
) -> None:
- # pylint: disable=duplicate-code
"""
Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
@@ -130,15 +128,15 @@ def apply_liger_kernel_to_qwen3(
layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
"""
- import transformers.models.qwen3.modeling_qwen3 # noqa: F401 # pylint: disable=unused-import
+ import transformers.models.qwen3.modeling_qwen3 # noqa: F401
from liger_kernel.transformers.functional import liger_cross_entropy
from liger_kernel.transformers.layer_norm import LigerLayerNorm
from liger_kernel.transformers.rms_norm import LigerRMSNorm
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
- assert not (
- cross_entropy and fused_linear_cross_entropy
- ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ assert not (cross_entropy and fused_linear_cross_entropy), (
+ "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ )
modeling_qwen3 = sys.modules["transformers.models.qwen3.modeling_qwen3"]
diff --git a/src/axolotl/integrations/liger/models/qwen3_moe.py b/src/axolotl/integrations/liger/models/qwen3_moe.py
index 89bdc5bcc..40bee110c 100644
--- a/src/axolotl/integrations/liger/models/qwen3_moe.py
+++ b/src/axolotl/integrations/liger/models/qwen3_moe.py
@@ -45,7 +45,6 @@ def lce_forward(
Returns:
"""
- # pylint: disable=duplicate-code
output_attentions = (
output_attentions
if output_attentions is not None
@@ -135,9 +134,8 @@ def apply_liger_kernel_to_qwen3_moe(
rms_norm: bool = False,
glu_activation: bool = False,
layer_norm: bool = False,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
) -> None:
- # pylint: disable=duplicate-code
"""
Apply Liger kernels to replace original implementation in HuggingFace Llama models (2 and 3)
@@ -152,15 +150,15 @@ def apply_liger_kernel_to_qwen3_moe(
layer_norm (bool): Whether to apply Liger's LayerNorm. Default is False.
"""
- import transformers.models.qwen3_moe.modeling_qwen3_moe # noqa: F401 # pylint: disable=unused-import
+ import transformers.models.qwen3_moe.modeling_qwen3_moe # noqa: F401
from liger_kernel.transformers.functional import liger_cross_entropy
from liger_kernel.transformers.layer_norm import LigerLayerNorm
from liger_kernel.transformers.rms_norm import LigerRMSNorm
from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
- assert not (
- cross_entropy and fused_linear_cross_entropy
- ), "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ assert not (cross_entropy and fused_linear_cross_entropy), (
+ "cross_entropy and fused_linear_cross_entropy cannot both be True."
+ )
modeling_qwen3_moe = sys.modules["transformers.models.qwen3_moe.modeling_qwen3_moe"]
@@ -174,7 +172,7 @@ def apply_liger_kernel_to_qwen3_moe(
# clone config to avoid modifying the original
config = deepcopy(config)
if intermediate_size:
- setattr(config, "intermediate_size", intermediate_size)
+ config.intermediate_size = intermediate_size
return LigerSwiGLUMLP(config, **kwargs)
modeling_qwen3_moe.Qwen3MoeMLP = _liger_swiglu_mlp_wrapper
diff --git a/src/axolotl/integrations/liger/plugin.py b/src/axolotl/integrations/liger/plugin.py
new file mode 100644
index 000000000..89f7c37b7
--- /dev/null
+++ b/src/axolotl/integrations/liger/plugin.py
@@ -0,0 +1,182 @@
+"""
+Liger-Kernel Plugin for Axolotl
+"""
+
+import inspect
+import sys
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.logging import get_logger
+
+from .models.base import patch_lce_forward
+from .utils import patch_with_compile_disable
+
+LOG = get_logger(__name__)
+
+
+class LigerPlugin(BasePlugin):
+ """
+ Plugin for LIGER integraton with Axolotl.
+ """
+
+ def get_input_args(self):
+ return "axolotl.integrations.liger.LigerArgs"
+
+ def pre_model_load(self, cfg):
+ if cfg.torch_compile:
+ # torch compile will unnecessarily attempt to optimize the triton kernel unless explicitly disabled
+ import liger_kernel.ops.fused_linear_cross_entropy
+
+ patch_with_compile_disable(
+ liger_kernel.ops.fused_linear_cross_entropy,
+ "fused_linear_cross_entropy_forward",
+ )
+ patch_with_compile_disable(
+ liger_kernel.ops.fused_linear_cross_entropy,
+ "fused_linear_cross_entropy_backward",
+ )
+ from liger_kernel.transformers.cross_entropy import LigerCrossEntropyLoss
+ from liger_kernel.transformers.functional import liger_cross_entropy
+ from liger_kernel.transformers.layer_norm import LigerLayerNorm
+ from liger_kernel.transformers.monkey_patch import MODEL_TYPE_TO_APPLY_LIGER_FN
+ from liger_kernel.transformers.rms_norm import LigerRMSNorm
+ from liger_kernel.transformers.rope import liger_rotary_pos_emb
+ from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+ if cfg.liger_cross_entropy and cfg.liger_fused_linear_cross_entropy:
+ raise ValueError(
+ "Cannot have both `liger_cross_entropy` and `liger_fused_linear_cross_entropy` set."
+ )
+
+ if cfg.model_config_type in MODEL_TYPE_TO_APPLY_LIGER_FN:
+ apply_liger_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[cfg.model_config_type]
+ liger_fn_sig = inspect.signature(apply_liger_fn)
+ kwargs = {}
+ if "rope" in liger_fn_sig.parameters:
+ kwargs["rope"] = cfg.liger_rope
+ if "cross_entropy" in liger_fn_sig.parameters:
+ kwargs["cross_entropy"] = cfg.liger_cross_entropy
+ if "fused_linear_cross_entropy" in liger_fn_sig.parameters:
+ kwargs["fused_linear_cross_entropy"] = (
+ cfg.liger_fused_linear_cross_entropy
+ )
+ if "rms_norm" in liger_fn_sig.parameters:
+ kwargs["rms_norm"] = cfg.liger_rms_norm
+ if "layer_norm" in liger_fn_sig.parameters:
+ kwargs["layer_norm"] = cfg.liger_layer_norm
+ if "geglu" in liger_fn_sig.parameters:
+ kwargs["geglu"] = cfg.liger_glu_activation
+ elif "swiglu" in liger_fn_sig.parameters:
+ kwargs["swiglu"] = cfg.liger_glu_activation
+ LOG.info(f"Applying LIGER to {cfg.model_config_type} with kwargs: {kwargs}")
+ apply_liger_fn(**kwargs)
+ elif cfg.model_config_type == "jamba":
+ from transformers.models.jamba import modeling_jamba
+
+ from .models.jamba import lce_forward as jamba_lce_forward
+
+ if cfg.liger_rope:
+ modeling_jamba.apply_rotary_pos_emb = liger_rotary_pos_emb
+ if cfg.liger_rms_norm:
+ modeling_jamba.JambaRMSNorm = LigerRMSNorm
+ if cfg.liger_glu_activation:
+ modeling_jamba.JambaMLP = LigerSwiGLUMLP
+ if cfg.liger_layer_norm:
+ modeling_jamba.nn.LayerNorm = LigerLayerNorm
+ if cfg.liger_cross_entropy:
+ from transformers.loss.loss_utils import nn
+
+ nn.functional.cross_entropy = liger_cross_entropy
+ if cfg.liger_fused_linear_cross_entropy:
+ modeling_jamba.JambaForCausalLM.forward = jamba_lce_forward
+ elif cfg.model_config_type == "deepseek_v2":
+ from accelerate import init_empty_weights
+ from transformers import AutoModelForCausalLM
+
+ with init_empty_weights():
+ model = AutoModelForCausalLM.from_pretrained(
+ cfg.base_model, trust_remote_code=cfg.trust_remote_code or False
+ )
+ modeling_mod = sys.modules[model.__class__.__module__]
+
+ from .models.deepseekv2 import lce_forward as deepseekv2_lce_forward
+
+ if cfg.liger_rope:
+ # The DeepseekV2 version of RoPE is different than upstream LLaMA.
+ # See https://github.com/linkedin/Liger-Kernel/issues/129#issuecomment-2313763528
+ LOG.warning("Fused liger_rope is not supported for DeepseekV2.")
+ if cfg.liger_rms_norm:
+ modeling_mod.DeepseekV2RMSNorm = LigerRMSNorm
+ if cfg.liger_glu_activation:
+ modeling_mod.DeepseekV2MLP.forward = LigerSwiGLUMLP.forward
+ if cfg.liger_layer_norm:
+ LOG.warning("liger_layer_norm is not supported for DeepseekV2.")
+ if cfg.liger_cross_entropy:
+ # We do not patch `nn.functional.cross_entropy` for DeepseekV2 as it still uses
+ # nn.CrossEntropyLoss in the forward method.
+ modeling_mod.CrossEntropyLoss = LigerCrossEntropyLoss
+ if cfg.liger_fused_linear_cross_entropy:
+ modeling_mod.DeepseekV2ForCausalLM.forward = deepseekv2_lce_forward
+ elif cfg.model_config_type == "llama4":
+ from axolotl.integrations.liger.models.llama4 import (
+ apply_liger_kernel_to_llama4,
+ )
+
+ apply_liger_kernel_to_llama4(
+ cross_entropy=cfg.liger_cross_entropy,
+ fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+ glu_activation=cfg.liger_glu_activation,
+ rms_norm=cfg.liger_rms_norm,
+ layer_norm=cfg.liger_layer_norm,
+ )
+ elif cfg.model_config_type == "qwen3":
+ from axolotl.integrations.liger.models.qwen3 import (
+ apply_liger_kernel_to_qwen3,
+ )
+
+ apply_liger_kernel_to_qwen3(
+ cross_entropy=cfg.liger_cross_entropy,
+ fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+ glu_activation=cfg.liger_glu_activation,
+ rms_norm=cfg.liger_rms_norm,
+ layer_norm=cfg.liger_layer_norm,
+ )
+ elif cfg.model_config_type == "qwen3_moe":
+ from axolotl.integrations.liger.models.qwen3_moe import (
+ apply_liger_kernel_to_qwen3_moe,
+ )
+
+ apply_liger_kernel_to_qwen3_moe(
+ cross_entropy=cfg.liger_cross_entropy,
+ fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+ glu_activation=cfg.liger_glu_activation,
+ rms_norm=cfg.liger_rms_norm,
+ layer_norm=cfg.liger_layer_norm,
+ )
+ elif cfg.model_config_type == "granitemoe":
+ from liger_kernel.transformers import apply_liger_kernel_to_granite
+
+ apply_liger_kernel_to_granite(
+ rope=cfg.liger_rope,
+ cross_entropy=cfg.liger_cross_entropy,
+ fused_linear_cross_entropy=cfg.liger_fused_linear_cross_entropy,
+ rms_norm=cfg.liger_rms_norm,
+ swiglu=cfg.liger_glu_activation,
+ )
+ elif cfg.liger_fused_linear_cross_entropy:
+ try:
+ patch_lce_forward(cfg.model_config_type)
+ LOG.warning_once(
+ f"Applied ONLY liger_fused_linear_cross_entropy genericpatches for model type: {cfg.model_config_type}"
+ )
+ LOG.warning_once(
+ f"Liger + {cfg.model_config_type} generic FLCE support is experimental and may not work as expected."
+ )
+ except RuntimeError:
+ LOG.warning(
+ f"Unsupported model config type: {cfg.model_config_type}. Liger not applied."
+ )
+ else:
+ LOG.warning(
+ f"Unsupported model config type: {cfg.model_config_type}. Liger not applied."
+ )
diff --git a/src/axolotl/integrations/lm_eval/__init__.py b/src/axolotl/integrations/lm_eval/__init__.py
index 8db4dc634..0ab6b8697 100644
--- a/src/axolotl/integrations/lm_eval/__init__.py
+++ b/src/axolotl/integrations/lm_eval/__init__.py
@@ -7,7 +7,7 @@ import subprocess # nosec
from axolotl.integrations.base import BasePlugin
from axolotl.integrations.lm_eval.cli import build_lm_eval_command
-from .args import LMEvalArgs # pylint: disable=unused-import. # noqa: F401
+from .args import LMEvalArgs as LMEvalArgs
class LMEvalPlugin(BasePlugin):
@@ -20,7 +20,6 @@ class LMEvalPlugin(BasePlugin):
def post_train_unload(self, cfg):
if cfg.lm_eval_post_train:
- # pylint: disable=duplicate-code
for lm_eval_args in build_lm_eval_command(
cfg.lm_eval_tasks,
bfloat16=cfg.bfloat16 or cfg.bf16,
diff --git a/src/axolotl/integrations/lm_eval/cli.py b/src/axolotl/integrations/lm_eval/cli.py
index 19608e1d9..ead82dcb7 100644
--- a/src/axolotl/integrations/lm_eval/cli.py
+++ b/src/axolotl/integrations/lm_eval/cli.py
@@ -99,7 +99,6 @@ def lm_eval(config: str, cloud: Optional[str] = None):
with open(config, encoding="utf-8") as file:
cfg: DictDefault = DictDefault(yaml.safe_load(file))
- # pylint: disable=duplicate-code
for lm_eval_args in build_lm_eval_command(
cfg.lm_eval_tasks,
bfloat16=cfg.bfloat16 or cfg.bf16,
diff --git a/src/axolotl/integrations/spectrum/__init__.py b/src/axolotl/integrations/spectrum/__init__.py
index 9f66aef97..5e8f9128d 100644
--- a/src/axolotl/integrations/spectrum/__init__.py
+++ b/src/axolotl/integrations/spectrum/__init__.py
@@ -23,7 +23,7 @@ import requests
from axolotl.integrations.base import BasePlugin
from axolotl.utils.logging import get_logger
-from .args import SpectrumArgs # pylint: disable=unused-import. # noqa: F401
+from .args import SpectrumArgs as SpectrumArgs
LOG = get_logger(__name__)
@@ -46,7 +46,7 @@ def _generate_unfrozen_params_yaml(snr_data, top_fraction=0.5):
"^lm_head.weight$",
"^model.embed_tokens.weight$",
]
- for layer_type, layer_names in top_layers_by_type.items():
+ for _, layer_names in top_layers_by_type.items():
for layer_name in layer_names:
unfrozen_parameters.append(layer_name)
return unfrozen_parameters
@@ -84,7 +84,7 @@ class SpectrumPlugin(BasePlugin):
snr_data = json.load(fin)
except FileNotFoundError:
pass
- except Exception as exc: # pylint: disable=broad-exception-caught
+ except Exception as exc:
LOG.warning(f"Failed to read SNR data from {snr_path}: {exc}")
if not snr_data:
diff --git a/src/axolotl/integrations/spectrum/args.py b/src/axolotl/integrations/spectrum/args.py
index df5756038..be6ca4bfc 100644
--- a/src/axolotl/integrations/spectrum/args.py
+++ b/src/axolotl/integrations/spectrum/args.py
@@ -15,6 +15,7 @@
"""
Module for handling Spectrum input arguments.
"""
+
from typing import Optional
from pydantic import BaseModel, model_validator
diff --git a/src/axolotl/kernels/geglu.py b/src/axolotl/kernels/geglu.py
index 6acbea0d4..ee3260ebd 100644
--- a/src/axolotl/kernels/geglu.py
+++ b/src/axolotl/kernels/geglu.py
@@ -5,8 +5,6 @@ See "GLU Variants Improve Transformer" (https://arxiv.org/abs/2002.05202).
Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
"""
-# pylint: disable=invalid-name,unnecessary-lambda-assignment,duplicate-code
-
import torch
import triton
import triton.language as tl
diff --git a/src/axolotl/kernels/lora.py b/src/axolotl/kernels/lora.py
index 63c9e57bd..c3356fb90 100644
--- a/src/axolotl/kernels/lora.py
+++ b/src/axolotl/kernels/lora.py
@@ -7,13 +7,12 @@ See "LoRA: Low-Rank Adaptation of Large Language Models"
Credit to `unsloth` (https://unsloth.ai/) for inspiration for this implementation.
"""
-# pylint: disable=invalid-name
-
from typing import Callable
import torch
from bitsandbytes.functional import QuantState
from torch import nn
+from torch.distributed.tensor import DTensor
from .geglu import geglu_backward, geglu_forward
from .quantize import dequantize
@@ -25,6 +24,7 @@ def get_lora_parameters(
proj: nn.Module,
) -> tuple[
torch.Tensor,
+ torch.Tensor | None,
QuantState | None,
torch.Tensor | None,
torch.Tensor | None,
@@ -37,39 +37,54 @@ def get_lora_parameters(
proj: The projection module to extract parameters from.
Returns:
- A tuple containing the base weight matrix, quantization state, LoRA A matrix,
- LoRA B matrix, and scaling factor. States and matrices may be None if not
- available.
+ A tuple containing the base weights, quantization state, LoRA A and B weights,
+ scaling factor, and base layer bias. Quant state, weights, and bias may be
+ `None` if not available.
"""
# For DPO or disabled adapters
base_layer = proj.base_layer if hasattr(proj, "base_layer") else proj
W = base_layer.weight
+ b = base_layer.bias
if not hasattr(proj, "disable_adapters") or proj.disable_adapters or proj.merged:
quant_state = getattr(W, "quant_state", None)
- return W, quant_state, None, None, None
+ return W, b, quant_state, None, None, None
+
+ quant_state = getattr(W, "quant_state", None)
active_adapter = (
proj.active_adapters[0]
if hasattr(proj, "active_adapters")
else proj.active_adapter
)
- A = proj.lora_A[active_adapter].weight
- B = proj.lora_B[active_adapter].weight
+
+ linear_A = proj.lora_A[active_adapter]
+ linear_B = proj.lora_B[active_adapter]
+
+ # This manual unsharding is needed for FSDP2 + LoRA kernels compatibility.
+ # We fuse linear layers + LoRA adapters calculations into a single
+ # torch.autograd.Function, bypassing the registered unshard / reshard behavior.
+ # Note that we don't apply resharding later in this module (it gets messy quickly),
+ # but LoRA parameters are generally small enough that this is not an issue.
+ if isinstance(linear_A.weight, DTensor):
+ linear_A.unshard()
+ linear_B.unshard()
+
+ A = linear_A.weight
+ B = linear_B.weight
s = proj.scaling[active_adapter]
- quant_state = getattr(W, "quant_state", None)
-
- return W, quant_state, A, B, s
+ return W, b, quant_state, A, B, s
def matmul_lora(
X: torch.Tensor,
W: torch.Tensor,
- W_quant: QuantState,
- A: torch.Tensor,
- B: torch.Tensor,
- s: float,
+ b: torch.Tensor | None,
+ W_quant: QuantState | None,
+ A: torch.Tensor | None,
+ B: torch.Tensor | None,
+ s: float | None,
out: torch.Tensor | None = None,
) -> torch.Tensor:
"""
@@ -90,20 +105,22 @@ def matmul_lora(
dtype = X.dtype
W = dequantize(W.t(), W_quant)
+ reshape = False
if X.dim() == 3:
batch, seq_len, _ = X.shape
X = X.view(-1, X.shape[-1])
reshape = True
- else:
- reshape = False
out = torch.matmul(X, W, out=out)
if W_quant is not None:
del W
if A is not None:
- A, B = A.t(), B.t()
- out += (X @ A.to(dtype)) @ (s * B.to(dtype))
+ A, B = A.t().to(dtype), B.t().to(dtype) # type: ignore[union-attr]
+ out += s * X @ A @ B
+
+ if b is not None:
+ out += b
return out.view(batch, seq_len, -1) if reshape else out
@@ -117,17 +134,20 @@ class LoRA_MLP(torch.autograd.Function):
ctx,
X: torch.Tensor,
gate_weight: torch.Tensor,
- gate_quant: object | None,
+ gate_bias: torch.Tensor | None,
+ gate_quant: QuantState | None,
gate_A: torch.Tensor | None,
gate_B: torch.Tensor | None,
gate_scale: float,
up_weight: torch.Tensor,
- up_quant: object | None,
+ up_bias: torch.Tensor | None,
+ up_quant: QuantState | None,
up_A: torch.Tensor | None,
up_B: torch.Tensor | None,
up_scale: float,
down_weight: torch.Tensor,
- down_quant: object | None,
+ down_bias: torch.Tensor | None,
+ down_quant: QuantState | None,
down_A: torch.Tensor | None,
down_B: torch.Tensor | None,
down_scale: float,
@@ -142,20 +162,22 @@ class LoRA_MLP(torch.autograd.Function):
ctx: Autograd context
X: Input features
gate_weight: Gate projection weight
+ gate_bias: Gate projection bias
gate_quant: Gate quantization state
gate_A: Gate LoRA A matrix
gate_B: Gate LoRA B matrix
gate_scale: Gate LoRA scale
- up_weight: Up-projection weight
- up_quant: Up-projection quantization state
- up_A: Up-projection LoRA A matrix
- up_B: Up-projection LoRA B matrix
- up_scale: Up-projection LoRA scale
- down_weight: Down-projection weight
- down_quant: Down-projection quantization state
- down_A: Down-projection LoRA A matrix
- down_B: Down-projection LoRA B matrix
- down_scale: Down-projection LoRA scale
+ up_weight: Up projection weight
+ up_quant: Up projection quantization state
+ up_A: Up projection LoRA A matrix
+ up_B: Up projection LoRA B matrix
+ up_scale: Up projection LoRA scale
+ down_weight: Down projection weight
+ down_bias: Down projection bias
+ down_quant: Down projection quantization state
+ down_A: Down projection LoRA A matrix
+ down_B: Down projection LoRA B matrix
+ down_scale: Down projection LoRA scale
activation_fn: Forward activation function
activation_fn_backward: Backward activation function
inplace: Whether to perform operations in-place
@@ -164,15 +186,17 @@ class LoRA_MLP(torch.autograd.Function):
Output transformed by multi-layer perceptron and activation function
"""
# Compute projections
- gate = matmul_lora(X, gate_weight, gate_quant, gate_A, gate_B, gate_scale)
- up = matmul_lora(X, up_weight, up_quant, up_A, up_B, up_scale)
+ gate = matmul_lora(
+ X, gate_weight, gate_bias, gate_quant, gate_A, gate_B, gate_scale
+ )
+ up = matmul_lora(X, up_weight, up_bias, up_quant, up_A, up_B, up_scale)
# Activation
hidden = activation_fn(gate, up)
# Down projection
output = matmul_lora(
- hidden, down_weight, down_quant, down_A, down_B, down_scale
+ hidden, down_weight, down_bias, down_quant, down_A, down_B, down_scale
)
# Save for backward
@@ -195,22 +219,26 @@ class LoRA_MLP(torch.autograd.Function):
torch.Tensor | None,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
None,
None,
None,
+ None,
]:
"""
Performs backward pass computation for LoRA MLP.
@@ -222,7 +250,7 @@ class LoRA_MLP(torch.autograd.Function):
Returns:
Tuple containing gradients for all inputs from forward pass:
- Input gradient tensor (or `None`)
- - `None` for weights/quantization states
+ - `None` for weights/biases/quantization states
- LoRA A/B matrix gradients (or `None`)
- `None` for scaling factors
- `None` for activation functions and flags
@@ -265,9 +293,10 @@ class LoRA_MLP(torch.autograd.Function):
dtype = X.dtype
# Down projection
- DW = matmul_lora(
+ grad_down = matmul_lora(
grad_output,
down_weight.t(),
+ None,
down_quant,
down_B,
down_A,
@@ -275,7 +304,7 @@ class LoRA_MLP(torch.autograd.Function):
)
# Activation backward
- h, grad_gate, grad_up = ctx.activation_fn_backward(DW, gate, up)
+ h, grad_gate, grad_up = ctx.activation_fn_backward(grad_down, gate, up)
# Initialize and compute LoRA gradients
d_down_A = d_down_B = d_up_A = d_up_B = d_gate_A = d_gate_B = None
@@ -315,8 +344,8 @@ class LoRA_MLP(torch.autograd.Function):
dX += grad_up @ up_B.to(dtype).t() @ (up_scale * up_A.to(dtype).t())
# Gate projection gradients
- gate_weight = dequantize(gate_weight.t(), gate_quant)
- dX += grad_gate @ gate_weight.t()
+ gate_weight = dequantize(gate_weight, gate_quant)
+ dX += grad_gate @ gate_weight
del gate_weight
if gate_A is not None and gate_B is not None:
@@ -334,22 +363,26 @@ class LoRA_MLP(torch.autograd.Function):
dX,
None,
None,
+ None,
d_gate_A.t() if d_gate_A is not None else None,
d_gate_B.t() if d_gate_B is not None else None,
None,
None,
None,
+ None,
d_up_A.t() if d_up_A is not None else None,
d_up_B.t() if d_up_B is not None else None,
None,
None,
None,
+ None,
d_down_A.t() if d_down_A is not None else None,
d_down_B.t() if d_down_B is not None else None,
None,
None,
None,
None,
+ None,
)
@@ -364,23 +397,26 @@ def apply_lora_mlp_swiglu(self, X: torch.Tensor, inplace: bool = True) -> torch.
Returns:
Output tensor after applying LoRA-adapted MLP with SwiGLU activation
"""
- gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
- upW, upW_quant, upA, upB, upS = get_lora_parameters(self.up_proj)
- downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
+ gateW, gateb, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
+ upW, upb, upW_quant, upA, upB, upS = get_lora_parameters(self.up_proj)
+ downW, downb, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
out = LoRA_MLP.apply(
X,
gateW,
+ gateb,
gateW_quant,
gateA,
gateB,
gateS,
upW,
+ upb,
upW_quant,
upA,
upB,
upS,
downW,
+ downb,
downW_quant,
downA,
downB,
@@ -404,22 +440,25 @@ def apply_lora_mlp_geglu(self, X: torch.Tensor, inplace: bool = True) -> torch.T
Returns:
Output tensor after applying LoRA-adapted MLP with GEGLU activation
"""
- gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
- upW, upW_quant, upA, upB, upS = get_lora_parameters(self.up_proj)
- downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
+ gateW, gateb, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj)
+ upW, upb, upW_quant, upA, upB, upS = get_lora_parameters(self.up_proj)
+ downW, downb, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj)
out = LoRA_MLP.apply(
X,
gateW,
+ gateb,
gateW_quant,
gateA,
gateB,
gateS,
upW,
+ upb,
upW_quant,
upA,
upB,
upS,
downW,
+ downb,
downW_quant,
downA,
downB,
@@ -446,16 +485,19 @@ class LoRA_QKV(torch.autograd.Function):
ctx: torch.autograd.function.FunctionCtx,
X: torch.Tensor,
q_weight: torch.Tensor,
+ q_bias: torch.Tensor | None,
q_quant: QuantState | None,
q_A: torch.Tensor | None,
q_B: torch.Tensor | None,
q_scale: float,
k_weight: torch.Tensor,
+ k_bias: torch.Tensor | None,
k_quant: QuantState | None,
k_A: torch.Tensor | None,
k_B: torch.Tensor | None,
k_scale: float,
v_weight: torch.Tensor,
+ v_bias: torch.Tensor | None,
v_quant: QuantState | None,
v_A: torch.Tensor | None,
v_B: torch.Tensor | None,
@@ -469,16 +511,19 @@ class LoRA_QKV(torch.autograd.Function):
ctx: Autograd context
X: Input tensor
q_weight: Query projection weight
+ q_bias: Query projection bias
q_quant: Query quantization state
q_A: Query LoRA A matrix
q_B: Query LoRA B matrix
q_scale: Query LoRA scale
k_weight: Key projection weight
+ k_bias: Key projection bias
k_quant: Key quantization state
k_A: Key LoRA A matrix
k_B: Key LoRA B matrix
k_scale: Key LoRA scale
v_weight: Value projection weight
+ v_bias: Value projection bias
v_quant: Value quantization state
v_A: Value LoRA A matrix
v_B: Value LoRA B matrix
@@ -488,20 +533,21 @@ class LoRA_QKV(torch.autograd.Function):
Returns:
Tuple of (Query, Key, Value) projection tensors
"""
- Q = matmul_lora(X, q_weight, q_quant, q_A, q_B, q_scale)
- K = matmul_lora(X, k_weight, k_quant, k_A, k_B, k_scale)
- V = matmul_lora(X, v_weight, v_quant, v_A, v_B, v_scale)
+ Q = matmul_lora(X, q_weight, q_bias, q_quant, q_A, q_B, q_scale)
+ K = matmul_lora(X, k_weight, k_bias, k_quant, k_A, k_B, k_scale)
+ V = matmul_lora(X, v_weight, v_bias, v_quant, v_A, v_B, v_scale)
ctx.save_for_backward(X, q_A, q_B, k_A, k_B, v_A, v_B)
ctx.scales = (q_scale, k_scale, v_scale)
ctx.quants = (q_quant, k_quant, v_quant)
ctx.weights = (q_weight, k_weight, v_weight)
+ ctx.biases = (q_bias, k_bias, v_bias)
ctx.inplace = inplace
return Q, K, V
@staticmethod
- @torch_amp_custom_fwd
+ @torch_amp_custom_bwd
def backward(
ctx: torch.autograd.function.FunctionCtx,
q_grad: torch.Tensor,
@@ -511,16 +557,19 @@ class LoRA_QKV(torch.autograd.Function):
torch.Tensor,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
None,
None,
+ None,
torch.Tensor | None,
torch.Tensor | None,
None,
@@ -608,31 +657,31 @@ class LoRA_QKV(torch.autograd.Function):
# Transpose gradients if needed
if d_A_q is not None:
d_A_q = d_A_q.t()
- if d_B_q is not None:
- d_B_q = d_B_q.t()
+ d_B_q = d_B_q.t() # type: ignore[union-attr]
if d_A_k is not None:
d_A_k = d_A_k.t()
- if d_B_k is not None:
- d_B_k = d_B_k.t()
+ d_B_k = d_B_k.t() # type: ignore[union-attr]
if d_A_v is not None:
d_A_v = d_A_v.t()
- if d_B_v is not None:
- d_B_v = d_B_v.t()
+ d_B_v = d_B_v.t() # type: ignore[union-attr]
return (
grad_X.view(batch, seq_len, -1),
None,
None,
+ None,
d_A_q,
d_B_q,
None,
None,
None,
+ None,
d_A_k,
d_B_k,
None,
None,
None,
+ None,
d_A_v,
d_B_v,
None,
@@ -653,22 +702,25 @@ def apply_lora_qkv(
Returns:
Tuple of (Query, Key, Value) projection tensors
"""
- QW, QW_quant, QA, QB, QS = get_lora_parameters(self.q_proj)
- KW, KW_quant, KA, KB, KS = get_lora_parameters(self.k_proj)
- VW, VW_quant, VA, VB, VS = get_lora_parameters(self.v_proj)
+ QW, Qb, QW_quant, QA, QB, QS = get_lora_parameters(self.q_proj)
+ KW, Kb, KW_quant, KA, KB, KS = get_lora_parameters(self.k_proj)
+ VW, Vb, VW_quant, VA, VB, VS = get_lora_parameters(self.v_proj)
Q, K, V = LoRA_QKV.apply(
X,
QW,
+ Qb,
QW_quant,
QA,
QB,
QS,
KW,
+ Kb,
KW_quant,
KA,
KB,
KS,
VW,
+ Vb,
VW_quant,
VA,
VB,
@@ -688,10 +740,11 @@ class LoRA_O(torch.autograd.Function):
ctx: torch.autograd.function.FunctionCtx,
X: torch.Tensor,
W: torch.Tensor,
+ b: torch.Tensor,
W_quant: QuantState | None,
- A: torch.Tensor | None,
- B: torch.Tensor | None,
- S: float,
+ A: torch.Tensor,
+ B: torch.Tensor,
+ s: float,
) -> torch.Tensor:
"""
Forward pass for output projection with LoRA.
@@ -700,19 +753,20 @@ class LoRA_O(torch.autograd.Function):
ctx: Autograd context
X: Input tensor
W: Output projection weight
+ b: Output projection bias
W_quant: Weight quantization state
A: LoRA A matrix
B: LoRA B matrix
- S: LoRA scaling factor
+ s: LoRA scaling factor
Returns:
- Output projection tensor
+ Output projection result
"""
- XW = matmul_lora(X, W, W_quant, A, B, S)
+ XW = matmul_lora(X, W, b, W_quant, A, B, s)
ctx.custom_saved_tensors = (
W,
W_quant,
- S,
+ s,
)
ctx.save_for_backward(A, B, X)
@@ -727,8 +781,9 @@ class LoRA_O(torch.autograd.Function):
torch.Tensor,
None,
None,
- torch.Tensor | None,
- torch.Tensor | None,
+ None,
+ torch.Tensor,
+ torch.Tensor,
None,
]:
"""
@@ -741,7 +796,7 @@ class LoRA_O(torch.autograd.Function):
Returns:
Tuple containing gradients for all forward inputs
"""
- W, W_quant, S = ctx.custom_saved_tensors
+ W, W_quant, s = ctx.custom_saved_tensors
A, B, X = ctx.saved_tensors
batch, seq_len, hd = X.shape
@@ -751,17 +806,19 @@ class LoRA_O(torch.autograd.Function):
# Weight projection
dY_X = X.t() @ dY
- d_A = S * dY_X @ B
- d_B = S * A @ dY_X
+ d_A = s * dY_X @ B
+ d_B = s * A @ dY_X
# Get derivative for dX
W = dequantize(W.t(), W_quant)
dX = dY @ W.t()
del W
- dX += dY @ B.to(dtype) @ (S * A.to(dtype))
- # W, W_quant, A, B, S
- return dX.view(batch, seq_len, hd), None, None, d_A.t(), d_B.t(), None
+ A, B = A.to(dtype), B.to(dtype)
+ dX += s * dY @ B @ A
+
+ # W, b, W_quant, A, B, s
+ return dX.view(batch, seq_len, hd), None, None, None, d_A.t(), d_B.t(), None
def apply_lora_o(self, X: torch.Tensor) -> torch.Tensor:
@@ -774,7 +831,7 @@ def apply_lora_o(self, X: torch.Tensor) -> torch.Tensor:
Returns:
Transformed output tensor
"""
- OW, OW_quant, OA, OB, OS = get_lora_parameters(self.o_proj)
- output = LoRA_O.apply(X, OW, OW_quant, OA, OB, OS)
+ OW, Ob, OW_quant, OA, OB, OS = get_lora_parameters(self.o_proj)
+ output = LoRA_O.apply(X, OW, Ob, OW_quant, OA, OB, OS)
return output
diff --git a/src/axolotl/kernels/quantize.py b/src/axolotl/kernels/quantize.py
index b61603fbc..d094f2381 100644
--- a/src/axolotl/kernels/quantize.py
+++ b/src/axolotl/kernels/quantize.py
@@ -1,7 +1,5 @@
"""Dequantization utilities for `bitsandbytes` integration."""
-# pylint: disable=invalid-name,global-statement
-
import ctypes
import bitsandbytes as bnb
diff --git a/src/axolotl/kernels/swiglu.py b/src/axolotl/kernels/swiglu.py
index 43a798edc..b13bcd350 100644
--- a/src/axolotl/kernels/swiglu.py
+++ b/src/axolotl/kernels/swiglu.py
@@ -99,7 +99,6 @@ def _swiglu_bwd_kernel(
tl.store(up_ptr + offsets, grad_up, mask=mask) # grad wrt up
-# pylint: disable=unnecessary-lambda-assignment
def swiglu_forward(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
"""
SwiGLU forward pass. Computes SwiGLU activation: `x * sigmoid(x) * up`, where
@@ -128,7 +127,6 @@ def swiglu_forward(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
return out
-# pylint: disable=unnecessary-lambda-assignment
def swiglu_backward(
grad_output: torch.Tensor, gate: torch.Tensor, up: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
diff --git a/src/axolotl/loaders/__init__.py b/src/axolotl/loaders/__init__.py
index 3eef75e58..ae99bf16d 100644
--- a/src/axolotl/loaders/__init__.py
+++ b/src/axolotl/loaders/__init__.py
@@ -1,6 +1,5 @@
"""Init for axolotl.loaders module"""
-# pylint: disable=unused-import
# flake8: noqa
from .adapter import load_adapter, load_lora
diff --git a/src/axolotl/loaders/adapter.py b/src/axolotl/loaders/adapter.py
index 5517ff50a..8e8177b62 100644
--- a/src/axolotl/loaders/adapter.py
+++ b/src/axolotl/loaders/adapter.py
@@ -14,6 +14,7 @@ from peft import (
PeftConfig,
PeftMixedModel,
PeftModel,
+ TaskType,
get_peft_model,
)
from transformers import PreTrainedModel
@@ -29,14 +30,12 @@ LOG = get_logger(__name__)
def setup_quantized_meta_for_peft(model: torch.nn.Module):
"""Replaces `quant_state.to` with a dummy function to prevent PEFT from moving `quant_state` to meta device"""
- def temp_to_method(self, *args, **kwargs): # pylint: disable=unused-argument
+ def temp_to_method(self, *args, **kwargs):
return self
for param in model.parameters():
if isinstance(param, Params4bit):
- param.quant_state._orig_to = ( # pylint: disable=protected-access
- param.quant_state.to
- )
+ param.quant_state._orig_to = param.quant_state.to
param.quant_state.to = types.MethodType(temp_to_method, param.quant_state)
@@ -44,10 +43,8 @@ def setup_quantized_peft_meta_for_training(model: torch.nn.Module):
"""Replaces dummy `quant_state.to` method with the original function to allow training to continue"""
for param in model.parameters():
if isinstance(param, Params4bit) and hasattr(param.quant_state, "_orig_to"):
- param.quant_state.to = (
- param.quant_state._orig_to # pylint: disable=protected-access
- )
- param.quant_state._orig_to = None # pylint: disable=protected-access
+ param.quant_state.to = param.quant_state._orig_to
+ param.quant_state._orig_to = None
def find_all_linear_names(model):
@@ -77,6 +74,7 @@ def load_lora(
config_only: bool = False,
) -> tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]:
lora_target_modules = cfg.lora_target_modules or []
+ lora_target_parameters = cfg.lora_target_parameters or []
if cfg.lora_target_linear:
linear_names = find_all_linear_names(model)
@@ -102,18 +100,30 @@ def load_lora(
lora_config_kwargs["use_rslora"] = cfg.peft_use_rslora
if cfg.peft_layer_replication:
lora_config_kwargs["layer_replication"] = cfg.peft_layer_replication
+ if cfg.peft_trainable_token_indices:
+ lora_config_kwargs["trainable_token_indices"] = cfg.peft_trainable_token_indices
+
+ # Determine the correct PEFT task type
+ model_cls = type(model).__name__
+ if "SequenceClassification" in model_cls:
+ task_type = TaskType.SEQ_CLS
+ elif "TokenClassification" in model_cls:
+ task_type = TaskType.TOKEN_CLS
+ else:
+ task_type = TaskType.CAUSAL_LM
lora_config = LoraConfig(
r=cfg.lora_r,
lora_alpha=cfg.lora_alpha,
target_modules=lora_target_modules,
+ target_parameters=lora_target_parameters,
layers_to_transform=cfg.peft_layers_to_transform,
layers_pattern=cfg.peft_layers_pattern,
lora_dropout=cfg.lora_dropout,
fan_in_fan_out=cfg.lora_fan_in_fan_out,
modules_to_save=cfg.lora_modules_to_save if cfg.lora_modules_to_save else None,
bias="none",
- task_type="CAUSAL_LM",
+ task_type=task_type,
**lora_config_kwargs,
)
@@ -123,9 +133,9 @@ def load_lora(
rank = int(os.environ.get("LOCAL_RANK", 0))
if (
- cfg.fsdp
+ cfg.fsdp_config
and cfg.adapter
- and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
+ and cfg.fsdp_config.cpu_ram_efficient_loading
and rank != 0
):
setup_quantized_meta_for_peft(model)
@@ -153,9 +163,9 @@ def load_lora(
"Exception caught during model.print_trainable_parameters(): %s", exc
)
elif (
- cfg.fsdp
+ cfg.fsdp_config
and cfg.adapter
- and cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
+ and cfg.fsdp_config.cpu_ram_efficient_loading
and rank != 0
):
setup_quantized_peft_meta_for_training(model)
diff --git a/src/axolotl/loaders/adapters/__init__.py b/src/axolotl/loaders/adapters/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/loaders/constants.py b/src/axolotl/loaders/constants.py
index c08518dd6..4939cb28d 100644
--- a/src/axolotl/loaders/constants.py
+++ b/src/axolotl/loaders/constants.py
@@ -1,21 +1,18 @@
"""Shared constants for axolotl.loaders module"""
-from transformers import (
- Gemma3ForConditionalGeneration,
- Llama4ForConditionalGeneration,
- LlavaForConditionalGeneration,
- Mistral3ForConditionalGeneration,
- MllamaForConditionalGeneration,
- Qwen2_5_VLForConditionalGeneration,
- Qwen2VLForConditionalGeneration,
+from transformers import AutoModelForImageTextToText
+from transformers.models.auto.modeling_auto import (
+ MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
)
-MULTIMODAL_AUTO_MODEL_MAPPING = {
- "mllama": MllamaForConditionalGeneration,
- "llama4": Llama4ForConditionalGeneration,
- "llava": LlavaForConditionalGeneration,
- "qwen2_vl": Qwen2VLForConditionalGeneration,
- "qwen2_5_vl": Qwen2_5_VLForConditionalGeneration,
- "mistral3": Mistral3ForConditionalGeneration,
- "gemma3": Gemma3ForConditionalGeneration,
-}
+MULTIMODAL_AUTO_MODEL_MAPPING = dict(MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES)
+
+MULTIMODAL_AUTO_MODEL_MAPPING["lfm2-vl"] = AutoModelForImageTextToText
+
+try:
+ from transformers import VoxtralForConditionalGeneration
+
+ # transformers >4.53.2
+ MULTIMODAL_AUTO_MODEL_MAPPING["voxtral"] = VoxtralForConditionalGeneration
+except ImportError:
+ pass
diff --git a/src/axolotl/loaders/model.py b/src/axolotl/loaders/model.py
index ed1e8bbf2..d1c3052e7 100644
--- a/src/axolotl/loaders/model.py
+++ b/src/axolotl/loaders/model.py
@@ -1,5 +1,5 @@
-"""Model loader class implementation for loading, configuring, and patching various
-models.
+"""
+Model loader class implementation for loading, configuring, and patching various models.
"""
import gc
@@ -14,6 +14,7 @@ import torch
import transformers
import transformers.modeling_utils
from accelerate import init_empty_weights
+from accelerate.parallelism_config import ParallelismConfig
from peft import (
PeftConfig,
PeftMixedModel,
@@ -21,8 +22,10 @@ from peft import (
PeftModelForCausalLM,
prepare_model_for_kbit_training,
)
+from torch.distributed import DeviceMesh
from transformers import (
AutoModelForCausalLM,
+ AutoModelForImageTextToText,
AutoModelForVision2Seq,
AwqConfig,
BitsAndBytesConfig,
@@ -50,6 +53,7 @@ from axolotl.telemetry.errors import send_errors
from axolotl.utils.bench import log_gpu_memory_usage
from axolotl.utils.dict import DictDefault
from axolotl.utils.distributed import (
+ build_parallelism_config,
get_device_count,
get_device_type,
)
@@ -88,6 +92,10 @@ class ModelLoader:
`AutoModelForCausalLM`).
"""
+ use_parallel_config: bool | None = False
+ parallelism_config: ParallelismConfig | None = None
+ device_mesh: DeviceMesh | None = None
+
def __init__(
self,
cfg: DictDefault,
@@ -95,7 +103,7 @@ class ModelLoader:
*,
inference: bool = False,
reference_model: bool = False,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
"""Initializes the ModelLoader.
@@ -127,7 +135,7 @@ class ModelLoader:
# Init model config
self.model_config = load_model_config(cfg)
- self.auto_model_loader = AutoModelForCausalLM # pylint: disable=invalid-name
+ self.auto_model_loader = AutoModelForCausalLM
# Initialize the patch manager
self.patch_manager = PatchManager(
@@ -141,10 +149,15 @@ class ModelLoader:
"""Check if flash attention is installed."""
return find_spec("flash_attn") is not None
- @cached_property
- def qlora_fsdp(self):
+ @property
+ def is_fsdp_enabled(self):
+ """Property that determines if FSDP is enabled."""
+ return self.cfg.fsdp_config is not None or self.cfg.fsdp is not None
+
+ @property
+ def is_qlora_and_fsdp_enabled(self):
"""Property that determines if FSDP with QLoRA is enabled."""
- return self.cfg.fsdp and self.cfg.adapter == "qlora"
+ return self.is_fsdp_enabled and self.cfg.adapter == "qlora"
@send_errors
def load(self) -> tuple[PreTrainedModel | PeftModelForCausalLM, PeftConfig | None]:
@@ -159,6 +172,7 @@ class ModelLoader:
# Build the model
PLUGIN_MANAGER.pre_model_load(self.cfg)
+ self.patch_manager.apply_post_plugin_pre_model_load_patches()
skip_move_to_device = self._build_model()
PLUGIN_MANAGER.post_model_build(self.cfg, self.model)
@@ -179,27 +193,54 @@ class ModelLoader:
def _apply_pre_model_load_setup(self):
"""Apply patches and setup configurations before model loading."""
+ if self.use_parallel_config is not None:
+ self.use_parallel_config = (
+ self.cfg.fsdp_config
+ or (self.cfg.tensor_parallel_size and self.cfg.tensor_parallel_size > 1)
+ or (
+ self.cfg.context_parallel_size
+ and self.cfg.context_parallel_size > 1
+ )
+ )
+ if self.cfg.fsdp_config and self.cfg.fsdp_version != 2:
+ self.use_parallel_config = False
+
+ if self.use_parallel_config:
+ self._set_parallel_config()
self._set_auto_model_loader()
self._set_device_map_config()
if self.cfg.revision_of_model:
self.model_kwargs["revision"] = self.cfg.revision_of_model
+ if self.cfg.use_kernels:
+ self.model_kwargs["use_kernels"] = self.cfg.use_kernels
self._set_quantization_config()
self._set_attention_config()
+ self._check_model_requirements()
def _apply_post_model_load_setup(self):
"""Configure the model after it has been loaded."""
# Handle PeftModel if needed
if (
isinstance(self.model, (peft.PeftModel, peft.PeftModelForCausalLM))
- and not self.qlora_fsdp
+ and not self.is_qlora_and_fsdp_enabled
):
self.model = self.model.merge_and_unload()
+ self._apply_activation_checkpointing()
self._resize_token_embeddings()
self._adjust_model_config()
- self._log_memory_usage()
self._configure_embedding_dtypes()
self._configure_qat()
+ log_gpu_memory_usage(LOG, "Memory usage after model load", 0)
+
+ def _apply_activation_checkpointing(self):
+ if self.cfg.activation_offloading is True:
+ from axolotl.core.trainers.mixins.activation_checkpointing import (
+ ac_wrap_hf_model,
+ )
+
+ # ^^ importing this at the module level breaks plugins
+ ac_wrap_hf_model(self.model)
def _resize_token_embeddings(self):
"""Resize token embeddings if needed."""
@@ -253,22 +294,13 @@ class ModelLoader:
):
self.model.config.eos_token_id = self.tokenizer.eos_token_id
- def _log_memory_usage(self):
- """Log device memory usage after model load."""
- if hasattr(self.model, "device") and self.model.device.type in (
- "cuda",
- "mps",
- "npu",
- ):
- log_gpu_memory_usage(LOG, "after model load", self.model.device)
-
def _configure_embedding_dtypes(self):
"""Configure embedding module dtypes."""
# Get embedding modules
embedding_modules = get_linear_embedding_layers(self.cfg.model_config_type)
# Initial dtype conversion
- if not self.cfg.fsdp:
+ if not self.is_fsdp_enabled:
# We don't run this during FSDP because this will leave mixed and bfloat16
# dtypes in the model which FSDP doesn't like
if self.cfg.load_in_4bit and self.cfg.embeddings_skip_upcast:
@@ -280,11 +312,14 @@ class ModelLoader:
)
# Handle DeepSpeed Zero3
- if is_deepspeed_zero3_enabled():
+ if (
+ is_deepspeed_zero3_enabled()
+ or os.getenv("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3"
+ ):
self._set_z3_leaf_modules()
# Apply gradient checkpointing if needed
- needs_fa2_dtype = self.cfg.adapter or self.cfg.fsdp
+ needs_fa2_dtype = self.cfg.adapter or self.is_fsdp_enabled
if self.cfg.adapter in ["lora", "qlora"]:
needs_fa2_dtype = True
if self.cfg.gradient_checkpointing:
@@ -300,10 +335,12 @@ class ModelLoader:
# we need to convert them back to fp16/bf16 for flash-attn compatibility.
(
(needs_fa2_dtype or self.cfg.flash_attention or self.cfg.flex_attention)
- and not self.qlora_fsdp
+ and not self.is_qlora_and_fsdp_enabled
+ )
+ or (
+ # CCE requires embedding layers to be in fp16/bf16 for backward pass
+ self.cfg.cut_cross_entropy
)
- # CCE requires embedding layers to be in fp16/bf16 for backward pass
- or self.cfg.cut_cross_entropy
)
if should_convert:
@@ -359,7 +396,6 @@ class ModelLoader:
and not (self.cfg.rl and self.cfg.load_in_4bit)
and not skip_move_to_device
):
- # TODO: validate this conditional
self.model.to(f"{str(get_device_type())}:{self.cfg.local_rank}")
if get_device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) == 1:
@@ -384,6 +420,13 @@ class ModelLoader:
gc.collect()
torch.cuda.empty_cache()
+ def _set_parallel_config(self):
+ """Set parallelism configuration (DP, FSDP, TP, CP) in PartialState/Accelerator"""
+ parallelism_config, device_mesh = build_parallelism_config(self.cfg)
+ if parallelism_config:
+ self.parallelism_config = parallelism_config
+ self.device_mesh = device_mesh
+
def _set_auto_model_loader(self):
"""Set `self.auto_model_loader`. Defaults to `transformers.AutoModelForCausalLM`
(set at `__init__`). When using a multimodal model, `self.auto_model_loader`
@@ -393,6 +436,8 @@ class ModelLoader:
self.auto_model_loader = MULTIMODAL_AUTO_MODEL_MAPPING.get(
self.model_config.model_type, AutoModelForVision2Seq
)
+ if isinstance(self.auto_model_loader, str):
+ self.auto_model_loader = AutoModelForImageTextToText
def _set_device_map_config(self):
"""Setup `device_map` according to config"""
@@ -432,7 +477,17 @@ class ModelLoader:
self.model_kwargs["torch_dtype"] = self.cfg.torch_dtype
- if not is_deepspeed_zero3_enabled():
+ is_ds_zero3 = is_deepspeed_zero3_enabled()
+
+ # FSDP requires control over device placement, so don't set device_map when FSDP is enabled
+ if self.is_fsdp_enabled:
+ # For QLoRA + FSDP, we still need to set device_map to "auto" for proper initialization
+ if self.is_qlora_and_fsdp_enabled:
+ self.model_kwargs["device_map"] = {
+ "": int(os.environ.get("LOCAL_RANK", 0))
+ }
+ # For other FSDP cases, don't set device_map at all
+ elif not is_ds_zero3:
self.model_kwargs["device_map"] = device_map
cur_device = get_device_type()
@@ -454,8 +509,17 @@ class ModelLoader:
def _set_quantization_config(self):
"""Set up quantization config (bitsandbytes, awq, gptq, etc.)"""
- self.model_kwargs["load_in_8bit"] = self.cfg.load_in_8bit
- self.model_kwargs["load_in_4bit"] = self.cfg.load_in_4bit
+
+ if self.cfg.model_quantization_config == "Mxfp4Config":
+ from transformers import Mxfp4Config
+
+ mxfp4_kwargs = {}
+ if self.cfg.model_quantization_config_kwargs:
+ mxfp4_kwargs = self.cfg.model_quantization_config_kwargs
+ self.model_kwargs["quantization_config"] = Mxfp4Config(**mxfp4_kwargs)
+ else:
+ self.model_kwargs["load_in_8bit"] = self.cfg.load_in_8bit
+ self.model_kwargs["load_in_4bit"] = self.cfg.load_in_4bit
if self.cfg.gptq:
if not hasattr(self.model_config, "quantization_config"):
@@ -490,7 +554,9 @@ class ModelLoader:
self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
**self.model_config.quantization_config
)
- elif self.cfg.adapter == "qlora" and self.model_kwargs["load_in_4bit"]:
+ elif self.cfg.adapter == "qlora" and self.model_kwargs.get(
+ "load_in_4bit", False
+ ):
bnb_config = {
"load_in_4bit": True,
"llm_int8_threshold": 6.0,
@@ -501,11 +567,14 @@ class ModelLoader:
"bnb_4bit_quant_storage": torch.bfloat16,
}
if self.cfg.model_config_type in ["jamba", "qwen2_moe"] and not (
- self.cfg.deepspeed or self.cfg.fsdp
+ self.cfg.deepspeed or self.is_fsdp_enabled
):
# for some reason, this causes the loss to be off by an order of magnitude
# but deepspeed needs this still in bfloat16
bnb_config["bnb_4bit_quant_storage"] = torch.float32
+ if self.cfg.model_config_type == "falcon_h1":
+ # output projection cannot be quantized for Falcon-H1 models
+ bnb_config["llm_int8_skip_modules"] = ["out_proj"]
if self.cfg.bnb_config_kwargs:
bnb_config.update(self.cfg.bnb_config_kwargs)
@@ -513,13 +582,18 @@ class ModelLoader:
self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
**bnb_config,
)
- elif self.cfg.adapter == "lora" and self.model_kwargs["load_in_8bit"]:
+ elif self.cfg.adapter == "lora" and self.model_kwargs.get(
+ "load_in_8bit", False
+ ):
bnb_config = {
"load_in_8bit": True,
}
# Exclude mamba blocks from int8 quantization for jamba
if self.cfg.model_config_type == "jamba":
bnb_config["llm_int8_skip_modules"] = ["mamba"]
+ if self.cfg.model_config_type == "falcon_h1":
+ # output projection cannot be quantized for Falcon-H1 models
+ bnb_config["llm_int8_skip_modules"] = ["out_proj"]
self.model_kwargs["quantization_config"] = BitsAndBytesConfig(
**bnb_config,
)
@@ -531,33 +605,37 @@ class ModelLoader:
def _set_attention_config(self):
"""Sample packing uses custom FA2 patch"""
- if self.cfg.flex_attention:
+ if self.cfg.attn_implementation:
+ self.model_kwargs["attn_implementation"] = self.cfg.attn_implementation
+ elif self.cfg.flex_attention:
self.model_kwargs["attn_implementation"] = "flex_attention"
- self.model_config._attn_implementation = ( # pylint: disable=protected-access
- "flex_attention"
- )
+ self.model_config._attn_implementation = "flex_attention"
elif self.cfg.flash_attention:
if not self.cfg.sample_packing and self.cfg.s2_attention:
pass
self.model_kwargs["attn_implementation"] = "flash_attention_2"
- self.model_config._attn_implementation = ( # pylint: disable=protected-access
- "flash_attention_2"
- )
+ self.model_config._attn_implementation = "flash_attention_2"
elif self.cfg.sdp_attention:
self.model_kwargs["attn_implementation"] = "sdpa"
- self.model_config._attn_implementation = ( # pylint: disable=protected-access
- "sdpa"
- )
+ self.model_config._attn_implementation = "sdpa"
elif self.cfg.eager_attention:
self.model_kwargs["attn_implementation"] = "eager"
- self.model_config._attn_implementation = ( # pylint: disable=protected-access
- "eager"
- )
+ self.model_config._attn_implementation = "eager"
if self.cfg.low_cpu_mem_usage:
self.model_kwargs["low_cpu_mem_usage"] = True
+ def _check_model_requirements(self):
+ if self.cfg.model_config_type in ["lfm2-vl", "lfm2"]:
+ from transformers.utils.import_utils import is_causal_conv1d_available
+
+ if is_causal_conv1d_available():
+ raise ImportError(
+ "The 'causal-conv1d' package is installed but causes compatibility issues with LFM2 models. "
+ "Please uninstall it by running: `pip uninstall -y causal-conv1d`"
+ )
+
def _configure_zero3_memory_efficient_loading(
self,
) -> HfTrainerDeepSpeedConfig | None:
@@ -597,17 +675,82 @@ class ModelLoader:
return hf_ds_cfg
+ def _load_model_from_config(self, model_loader_class=None) -> PreTrainedModel:
+ """
+ Load model with random initialization using from_config.
+
+ Uses the selected loader when provided; otherwise falls back to the auto loader.
+ """
+ loader = model_loader_class or self.auto_model_loader
+ if loader in [AutoModelForCausalLM, AutoModelForVision2Seq]:
+ model = loader.from_config(
+ config=self.model_config,
+ trust_remote_code=self.cfg.trust_remote_code or False,
+ )
+ else:
+ model = loader(config=self.model_config)
+
+ return model
+
+ def _load_model_from_pretrained(self, model_loader_class=None) -> PreTrainedModel:
+ """Load model from pretrained weights."""
+ loader = model_loader_class or self.auto_model_loader
+ kwargs = {
+ "config": self.model_config,
+ "trust_remote_code": self.cfg.trust_remote_code or False,
+ **self.model_kwargs,
+ }
+ return loader.from_pretrained(self.base_model, **kwargs)
+
def _build_model(self) -> bool:
"""Load model, with load strategy depending on config."""
skip_move_to_device = False
+
+ if self.cfg.tensor_parallel_size > 1:
+ self.model_kwargs["tp_size"] = self.cfg.tensor_parallel_size
+ self.model_kwargs["tp_plan"] = "auto"
+ self.model_kwargs["device_mesh"] = self.device_mesh
+ if "device_map" in self.model_kwargs:
+ del self.model_kwargs["device_map"] # not compatible with `tp_plan`
+
+ if self.is_fsdp_enabled:
+ if self.cfg.fsdp_config.cpu_ram_efficient_loading:
+ skip_move_to_device = True
+ # Don't delete device_map for QLoRA + FSDP - it was set correctly in
+ # _set_device_map
+ if (
+ "device_map" in self.model_kwargs
+ and not self.is_qlora_and_fsdp_enabled
+ ):
+ del self.model_kwargs["device_map"]
+ elif self.is_qlora_and_fsdp_enabled:
+ skip_move_to_device = True
+
+ if (
+ self.cfg.tensor_parallel_size <= 1
+ and self.cfg.fsdp_config.cpu_ram_efficient_loading
+ and self.cfg.fsdp_version == 2
+ ):
+ # setting device_map for TP is not supported
+ local_rank = int(os.getenv("LOCAL_RANK", "0"))
+ if local_rank == 0:
+ self.model_kwargs["device_map"] = "cpu"
+ else:
+ self.model_kwargs["device_map"] = "meta"
+
if (
- self.qlora_fsdp
- and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
+ self.is_qlora_and_fsdp_enabled
+ and self.cfg.fsdp_config.cpu_ram_efficient_loading
and (
self.cfg.model_config_type == "dbrx"
or self.cfg.qlora_sharded_model_loading
)
):
+ if self.cfg.reinit_weights:
+ LOG.warning(
+ "reinit_weights is not supported with sharded quantized loading. "
+ "Loading from pretrained weights instead."
+ )
quant_storage = self.cfg.torch_dtype
quantization_config = getattr(
self.model_config, "quantization_config", None
@@ -623,41 +766,14 @@ class ModelLoader:
quantization_config=quantization_config,
)
skip_move_to_device = True
- elif (
- self.model_config.model_type in ["llama", "llama4"]
- and not self.cfg.trust_remote_code
- and not self.cfg.gptq
- ):
- # TODO: Do we need to open this up for all models?
- if self.cfg.fsdp and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading:
- skip_move_to_device = True
- if "device_map" in self.model_kwargs:
- del self.model_kwargs["device_map"]
-
- # Please don't remove underscore binding without reading the fn docstring.
- _ = self._configure_zero3_memory_efficient_loading()
-
- # Load model with random initialization if specified
- if self.cfg.random_init_weights:
- # AutoModel classes support the from_config method
- if self.auto_model_loader in [
- AutoModelForCausalLM,
- AutoModelForVision2Seq,
- ]:
- self.model = self.auto_model_loader.from_config(
- config=self.model_config,
- )
- else:
- self.model = self.auto_model_loader(config=self.model_config)
- else:
- self.model = self.auto_model_loader.from_pretrained(
- self.base_model,
- config=self.model_config,
- **self.model_kwargs,
- )
elif self.model_type == "MambaLMHeadModel":
+ if self.cfg.reinit_weights:
+ LOG.warning(
+ "reinit_weights is not supported with MambaLMHeadModel. "
+ "Loading from pretrained weights instead."
+ )
# FIXME this is janky at best and hacked together to make it work
- MambaLMHeadModel = fix_mamba_attn_for_loss() # pylint: disable=invalid-name
+ MambaLMHeadModel = fix_mamba_attn_for_loss()
self.model_kwargs["dtype"] = self.model_kwargs["torch_dtype"]
self.model_kwargs["device"] = torch.cuda.current_device()
@@ -668,55 +784,40 @@ class ModelLoader:
self.base_model,
**self.model_kwargs,
)
- elif (
- self.model_type
- and self.model_type != "AutoModelForCausalLM"
- and not self.cfg.trust_remote_code
- ):
- if self.cfg.gptq:
- self.model = self.auto_model_loader.from_pretrained(
- self.base_model,
- config=self.model_config,
- trust_remote_code=self.cfg.trust_remote_code or False,
- **self.model_kwargs,
- )
- else:
- self.model = getattr(transformers, self.model_type).from_pretrained(
- self.base_model,
- config=self.model_config,
- trust_remote_code=self.cfg.trust_remote_code or False,
- **self.model_kwargs,
- )
else:
- if self.cfg.gptq:
- self.model = self.auto_model_loader.from_pretrained(
- self.base_model,
- config=self.model_config,
- trust_remote_code=self.cfg.trust_remote_code or False,
- **self.model_kwargs,
- )
+ # Please don't remove underscore binding without reading the fn docstring
+ _ = self._configure_zero3_memory_efficient_loading()
+
+ if (
+ self.model_type
+ and self.model_type != "AutoModelForCausalLM"
+ and not self.cfg.trust_remote_code
+ and not self.cfg.gptq
+ ):
+ # Use model type from transformers
+ model_loader_class = getattr(transformers, self.model_type)
else:
- if (
- self.cfg.fsdp
- and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading
- ):
- # disabling either of these two still leads to VRAM spike before setting back down
- skip_move_to_device = True
- if "device_map" in self.model_kwargs:
- del self.model_kwargs["device_map"]
+ # Use auto model loader (handles gptq and default cases)
+ model_loader_class = self.auto_model_loader
- # Please don't remove underscore binding without reading the fn docstring.
- _ = self._configure_zero3_memory_efficient_loading()
+ if self.cfg.reinit_weights:
+ self.model = self._load_model_from_config(model_loader_class)
+ else:
+ self.model = self._load_model_from_pretrained(model_loader_class)
- self.model = self.auto_model_loader.from_pretrained(
- self.base_model,
- config=self.model_config,
- trust_remote_code=self.cfg.trust_remote_code or False,
- **self.model_kwargs,
- )
if is_deepspeed_zero3_enabled():
skip_move_to_device = True
+ if self.cfg.tensor_parallel_size > 1:
+ # workaround for upstream 4.54.0 not setting _tp_size or _device_mesh
+ # TODO(wing): remove once 4.54.1 is released
+ if self.model._tp_size != self.cfg.tensor_parallel_size:
+ self.model._tp_size = self.cfg.tensor_parallel_size
+ self.model._device_mesh = self.model_kwargs["device_mesh"]
+
+ if self.cfg.experimental_skip_move_to_device is not None:
+ skip_move_to_device = self.cfg.experimental_skip_move_to_device
+
return skip_move_to_device
def _set_z3_leaf_modules(self):
@@ -749,8 +850,8 @@ class ModelLoader:
skip_prepare_model_for_kbit_training = True
if (
- self.qlora_fsdp
- or (self.cfg.fsdp and self.cfg.fsdp_config.fsdp_cpu_ram_efficient_loading)
+ self.is_qlora_and_fsdp_enabled
+ or (self.is_fsdp_enabled and self.cfg.fsdp_config.cpu_ram_efficient_loading)
or is_deepspeed_zero3_enabled()
):
# Make sure everything is in the same dtype
@@ -772,6 +873,9 @@ class ModelLoader:
dist_dtype: torch.dtype,
before_kbit_train_or_finetune: bool,
):
+ dest = {"dtype": dist_dtype}
+ if self.cfg.lora_on_cpu:
+ dest["device"] = "cpu"
for name, module in self.model.named_modules():
if "norm" in name:
module.to(dist_dtype)
@@ -782,4 +886,4 @@ class ModelLoader:
# don't upcast lm_head for btlm
continue
if any(m in name for m in embedding_modules) and hasattr(module, "weight"):
- module.to(dist_dtype)
+ module.to(**dest)
diff --git a/src/axolotl/loaders/patch_manager.py b/src/axolotl/loaders/patch_manager.py
index 23f79d368..1e46f5c34 100644
--- a/src/axolotl/loaders/patch_manager.py
+++ b/src/axolotl/loaders/patch_manager.py
@@ -4,6 +4,7 @@ Applies pre- and post-model load patches for various fixes and optimizations.
"""
import importlib.util
+import os
from functools import cached_property
import addict
@@ -49,10 +50,12 @@ class PatchManager:
def apply_pre_model_load_patches(self):
"""Apply pre-model load patches based on config."""
+ self._apply_transformers_patches()
+ # self._apply_flex_attention_patches()
self._apply_flash_attention_patches()
+ self._apply_chunked_cross_entropy_patch()
self._apply_fsdp_patches()
self._apply_adapter_patches()
- self._apply_flex_attention_patches()
self._apply_model_specific_patches()
self._apply_fp8_patches()
self._apply_flash_attention_peft_patches()
@@ -63,6 +66,30 @@ class PatchManager:
self._patch_llama_derived_model()
self._apply_mistral_cross_entropy_patch()
self._apply_self_attention_lora_patch()
+ self._apply_fsdp2_bnb_patches()
+ self._apply_patch_deepspeed_zero3()
+ self._apply_voxtral_patches()
+ self._apply_apertus_patches()
+
+ def apply_post_plugin_pre_model_load_patches(self):
+ """Apply post plugin-pre_model_load load patches based on config."""
+ self._apply_tiled_mlp(self.cfg.model_config_type)
+
+ def _apply_transformers_patches(self):
+ from axolotl.monkeypatch.transformers.trainer_loss_calc import (
+ patch_evaluation_loop,
+ patch_maybe_log_save_evaluate,
+ )
+
+ patch_evaluation_loop()
+ patch_maybe_log_save_evaluate()
+
+ if self.cfg.context_parallel_size > 1:
+ from axolotl.monkeypatch.transformers.trainer_context_parallel import (
+ patch_prepare_context_parallel_inputs,
+ )
+
+ patch_prepare_context_parallel_inputs()
def apply_post_model_load_patches(self, model: PreTrainedModel):
"""Apply patches that require the model instance."""
@@ -78,12 +105,41 @@ class PatchManager:
patch_xformers_attn_over_fa2()
self.cfg.flash_attention = True
+ def _apply_chunked_cross_entropy_patch(self):
+ if self.cfg.chunked_cross_entropy:
+ from axolotl.monkeypatch.loss.chunked import patch_chunked_ce_loss_fn
+
+ if self.cfg.chunked_cross_entropy_num_chunks:
+ patch_chunked_ce_loss_fn(self.cfg.chunked_cross_entropy_num_chunks)
+ else:
+ patch_chunked_ce_loss_fn()
+
def _apply_fsdp_patches(self):
"""Apply patches for FSDP configurations."""
- if self.cfg.fsdp_config and str(self.cfg.fsdp_config.fsdp_version) == "2":
+ if self.cfg.context_parallel_size > 1 or (
+ self.cfg.fsdp_config and str(self.cfg.fsdp_version) == "2"
+ ):
+ from axolotl.monkeypatch.accelerate.parallelism_config import (
+ patch_parallelism_config,
+ )
+
+ patch_parallelism_config()
+ if self.cfg.fsdp_config and str(self.cfg.fsdp_version) == "2":
from axolotl.monkeypatch.accelerate.fsdp2 import patch_accelerate_fsdp2
patch_accelerate_fsdp2()
+ if self.cfg.rl:
+ from axolotl.monkeypatch.trainer.trl import patch_trl_prepare_fsdp2
+
+ patch_trl_prepare_fsdp2()
+
+ # if self.cfg.fsdp_config:
+ # # see transformers#39152
+ # from axolotl.monkeypatch.trainer_fsdp_optim import (
+ # patch_training_loop_for_fsdp,
+ # )
+ #
+ # patch_training_loop_for_fsdp()
def _apply_adapter_patches(self):
"""Apply patches for adapter configurations."""
@@ -96,13 +152,17 @@ class PatchManager:
"""Apply patches for flexible attention."""
if self.cfg.flex_attention:
from axolotl.monkeypatch.attention.flex_attn import (
- patch_flex_make_mask,
patch_flex_wrapper,
)
flex_attn_compile_kwargs = self.cfg.flex_attn_compile_kwargs or {}
patch_flex_wrapper(**flex_attn_compile_kwargs)
- patch_flex_make_mask()
+ if self.cfg.sample_packing:
+ from axolotl.core.attention.flex_block_mask import (
+ patch_create_causal_mask,
+ )
+
+ patch_create_causal_mask(self.cfg.model_config_type)
def _apply_model_specific_patches(self):
"""Apply patches specific to model architectures."""
@@ -116,6 +176,20 @@ class PatchManager:
patch_llama4_linearized_modeling()
+ if self.cfg.model_config_type == "qwen3_next" and self.cfg.sample_packing:
+ from axolotl.monkeypatch.models.qwen3_next.modeling import (
+ patch_qwen3_next_modeling_packing,
+ )
+
+ patch_qwen3_next_modeling_packing()
+
+ if self.cfg.model_config_type == "mistral3" and self.cfg.processor_type:
+ from axolotl.monkeypatch.models.mistral3.mistral_common_tokenizer import (
+ apply_mistral_tokenizer_image_patch,
+ )
+
+ apply_mistral_tokenizer_image_patch()
+
def _apply_fp8_patches(self):
"""Apply patches for FP8 support."""
if self.cfg.fp8:
@@ -123,7 +197,9 @@ class PatchManager:
patch_create_accelerate_code_for_fp8,
)
- patch_create_accelerate_code_for_fp8()
+ patch_create_accelerate_code_for_fp8(
+ self.cfg.fp8_enable_fsdp_float8_all_gather
+ )
def _apply_flash_attention_peft_patches(self):
"""Apply patches for Flash Attention with PEFT."""
@@ -136,13 +212,19 @@ class PatchManager:
def _apply_gradient_checkpointing_patches(self):
"""Apply patches for gradient checkpointing."""
- if self.cfg.gradient_checkpointing in ["unsloth", "offload"]:
+ if (
+ self.cfg.gradient_checkpointing
+ and self.cfg.activation_offloading == "legacy"
+ ):
from axolotl.monkeypatch.gradient_checkpointing import (
hf_grad_checkpoint_offload_wrapper,
)
transformers.modeling_utils.checkpoint = hf_grad_checkpoint_offload_wrapper
- if self.cfg.gradient_checkpointing == "offload_disk":
+ elif (
+ self.cfg.gradient_checkpointing
+ and self.cfg.activation_offloading == "offload_disk"
+ ):
from axolotl.monkeypatch.gradient_checkpointing import (
hf_grad_checkpoint_disk_offload_wrapper,
)
@@ -166,6 +248,17 @@ class PatchManager:
def _apply_self_attention_lora_patch(self):
"""Apply self-attention LoRA patches if configured."""
if self.cfg.lora_qkv_kernel or self.cfg.lora_o_kernel:
+ # Only patch if conditions are met
+ can_patch = (
+ self.cfg.lora_dropout == 0
+ if hasattr(self.cfg, "lora_dropout")
+ else True
+ ) # default to True if lora_dropout is not set
+
+ if not can_patch:
+ LOG.warning("Cannot patch self-attention - requires no dropout")
+ return
+
from axolotl.monkeypatch.lora_kernels import patch_self_attn_lora
patch_self_attn_lora(self.cfg)
@@ -200,6 +293,50 @@ class PatchManager:
has_remote_code=has_remote_code,
)
+ if self.cfg.sample_packing:
+ from axolotl.monkeypatch.data.batch_dataset_fetcher import (
+ apply_multipack_dataloader_patch,
+ )
+
+ LOG.info("Applying multipack dataloader patch for sample packing...")
+ apply_multipack_dataloader_patch()
+
+ def _apply_fsdp2_bnb_patches(self):
+ """Apply FSDP2 BNB patches."""
+ if (
+ self.cfg.fsdp_config
+ and str(self.cfg.fsdp_version) == "2"
+ and self.cfg.adapter == "qlora"
+ ):
+ from axolotl.monkeypatch.fsdp2_qlora import (
+ apply_init_sharded_param_patch,
+ apply_init_unsharded_param_patch,
+ )
+
+ apply_init_sharded_param_patch()
+ apply_init_unsharded_param_patch()
+
+ def _apply_tiled_mlp(self, model_type: str):
+ if self.cfg.tiled_mlp:
+ from axolotl.monkeypatch.tiled_mlp import (
+ patch_tiled_mlp,
+ )
+
+ patch_tiled_mlp(
+ model_type,
+ use_original_mlp=self.cfg.tiled_mlp_use_original_mlp,
+ cfg_num_shards=self.cfg.tiled_mlp_num_shards,
+ )
+
+ def _apply_voxtral_patches(self):
+ """Apply patches for Voxtral model."""
+ if self.cfg.model_config_type == "voxtral":
+ from axolotl.monkeypatch.models.voxtral.modeling import (
+ patch_voxtral_conditional_generation_forward,
+ )
+
+ patch_voxtral_conditional_generation_forward()
+
def _patch_attention(self):
"""Apply attention-specific patches based on model type."""
if not (self.cfg.flash_attention and hasattr(self.model_config, "model_type")):
@@ -219,6 +356,13 @@ class PatchManager:
replace_stablelm_attn_with_flash_attn(self.cfg.base_model)
+ if self.model_config.model_type in ("mistral3", "llava"):
+ from axolotl.monkeypatch.models.pixtral.modeling_flash_attention_utils import (
+ apply_patch_is_packed_sequence,
+ )
+
+ apply_patch_is_packed_sequence()
+
def _patch_loss_llama(self):
"""Patch loss functions and other optimizations for LLaMA models."""
if not self.cfg.is_llama_derived_model:
@@ -249,31 +393,21 @@ class PatchManager:
patch_self_attn_lora()
- def _patch_llama_flash_attention(self, packed=False):
+ def _patch_llama_flash_attention(self):
"""Apply Flash Attention patches for LLaMA models."""
from axolotl.monkeypatch.llama_attn_hijack_flash import (
replace_llama_attn_with_flash_attn,
)
- if packed:
- if self.cfg.device not in ["mps", "cpu"] and not self.inference:
- LOG.info("patching with flash attention for sample packing")
- replace_llama_attn_with_flash_attn(
- packed=True,
- cross_entropy=self.cfg.flash_attn_cross_entropy,
- rms_norm=self.cfg.flash_attn_rms_norm,
- )
- elif self.cfg.s2_attention:
+ if self.cfg.s2_attention:
LOG.info("patching w/ flash-enabled, shifted-sparse attention")
replace_llama_attn_with_flash_attn(
- packed=False,
cross_entropy=self.cfg.flash_attn_cross_entropy,
rms_norm=self.cfg.flash_attn_rms_norm,
use_shifted_sparse_attn=True,
)
elif self.cfg.flash_attn_cross_entropy or self.cfg.flash_attn_rms_norm:
replace_llama_attn_with_flash_attn(
- packed=False,
cross_entropy=self.cfg.flash_attn_cross_entropy,
rms_norm=self.cfg.flash_attn_rms_norm,
)
@@ -304,7 +438,7 @@ class PatchManager:
and self.cfg.sample_packing
):
if self.cfg.flash_attention:
- self._patch_llama_flash_attention(packed=self.cfg.sample_packing)
+ self._patch_llama_flash_attention()
elif self.cfg.xformers_attention:
self._patch_llama_xformers_attention()
elif self.cfg.sample_packing:
@@ -327,17 +461,12 @@ class PatchManager:
from axolotl.monkeypatch.llama_attn_hijack_flash import (
is_xformers_swiglu_available,
replace_llama_mlp_with_swiglu,
- replace_llama_qkv_with_fused,
)
if self.cfg.flash_attn_fuse_mlp and is_xformers_swiglu_available():
LOG.info("Patching with SwiGLU...")
replace_llama_mlp_with_swiglu(model)
- if self.cfg.flash_attn_fuse_qkv:
- LOG.info("Patching with fused QKV...")
- replace_llama_qkv_with_fused(model)
-
def _apply_unsloth_patches(self, model):
"""Apply unsloth optimization patches."""
if self.cfg.unsloth_lora_mlp:
@@ -365,3 +494,26 @@ class PatchManager:
from axolotl.monkeypatch.lora_kernels import apply_lora_kernel_patches
apply_lora_kernel_patches(model=model, cfg=self.cfg)
+
+ def _apply_patch_deepspeed_zero3(self):
+ try:
+ from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+
+ from axolotl.monkeypatch.deepspeed_utils import apply_deepspeed_patches
+
+ if self.cfg.activation_offloading is True and (
+ is_deepspeed_zero3_enabled()
+ or os.getenv("ACCELERATE_DEEPSPEED_ZERO_STAGE") == "3"
+ ):
+ apply_deepspeed_patches()
+ except ImportError as e:
+ LOG.warning(f"DeepSpeed patches not applied: {e}")
+
+ def _apply_apertus_patches(self):
+ """Apply patches for Apertus model."""
+ if self.cfg.model_config_type == "apertus":
+ from axolotl.monkeypatch.models.apertus.activation import (
+ patch_apertus_xielu_activation,
+ )
+
+ patch_apertus_xielu_activation()
diff --git a/src/axolotl/loaders/processor.py b/src/axolotl/loaders/processor.py
index 5d1c36618..e6fd4c0ed 100644
--- a/src/axolotl/loaders/processor.py
+++ b/src/axolotl/loaders/processor.py
@@ -23,6 +23,13 @@ def load_processor(cfg: DictDefault, tokenizer: PreTrainedTokenizerBase):
if cfg.processor_type:
processor_cls = getattr(transformers, cfg.processor_type)
+ if cfg.tokenizer_use_mistral_common:
+ from axolotl.utils.mistral import Mistral3Processor
+
+ return Mistral3Processor(
+ tokenizer=tokenizer,
+ )
+
processor = processor_cls.from_pretrained(
cfg.processor_config,
trust_remote_code=cfg.trust_remote_code or False,
diff --git a/src/axolotl/loaders/tokenizer.py b/src/axolotl/loaders/tokenizer.py
index cfc5da42c..48856116c 100644
--- a/src/axolotl/loaders/tokenizer.py
+++ b/src/axolotl/loaders/tokenizer.py
@@ -7,6 +7,7 @@ import transformers
from transformers import (
AddedToken,
AutoTokenizer,
+ PreTrainedTokenizer,
)
from axolotl.integrations.base import PluginManager
@@ -14,6 +15,7 @@ from axolotl.loaders.utils import get_linear_embedding_layers, load_model_config
from axolotl.prompt_tokenizers import LLAMA_DEFAULT_EOS_TOKEN
from axolotl.telemetry.errors import send_errors
from axolotl.utils.chat_templates import get_chat_template_from_config
+from axolotl.utils.dict import DictDefault
from axolotl.utils.distributed import (
barrier,
is_local_main_process,
@@ -49,7 +51,7 @@ def modify_tokenizer_files(
tokenizer_dir = os.path.join(output_dir, "tokenizer")
os.makedirs(tokenizer_dir, exist_ok=True)
- if is_local_main_process(): # pylint: disable=too-many-nested-blocks
+ if is_local_main_process():
# Load the tokenizer
temp_tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
@@ -72,9 +74,9 @@ def modify_tokenizer_files(
for token_id, new_value in token_id_mappings.items():
token_id_str = str(token_id)
if token_id_str in config_data["added_tokens_decoder"]:
- config_data["added_tokens_decoder"][token_id_str][
- "content"
- ] = new_value
+ config_data["added_tokens_decoder"][token_id_str]["content"] = (
+ new_value
+ )
else:
raise ValueError(
f"Token ID {token_id_str} not found in added_tokens_decoder"
@@ -119,8 +121,21 @@ def modify_tokenizer_files(
@send_errors
-def load_tokenizer(cfg):
+def load_tokenizer(cfg: DictDefault) -> PreTrainedTokenizer:
"""Load and configure the tokenizer based on the provided config."""
+
+ def _load_mistral_common_tokenizer(cfg: DictDefault):
+ """Load mistral-common tokenizer"""
+ from axolotl.utils.mistral import HFMistralTokenizer
+
+ # Load the HF-compatible wrapper around MistralTokenizer
+ tokenizer = HFMistralTokenizer.from_pretrained(cfg.tokenizer_config)
+
+ return tokenizer
+
+ if cfg.tokenizer_use_mistral_common:
+ return _load_mistral_common_tokenizer(cfg)
+
model_config = load_model_config(cfg)
tokenizer_kwargs = {}
use_fast = True # this is the default
@@ -175,7 +190,8 @@ def load_tokenizer(cfg):
tokenizer.padding_side = "left"
# Qwen base only has single token, so we need to set the special tokens
- if cfg.is_qwen_derived_model:
+ # the following check is for Qwen1 base models
+ if cfg.is_qwen_derived_model and hasattr(tokenizer, "eod_id"):
token_ids = ["bos_token_id", "eos_token_id", "pad_token_id", "unk_token_id"]
for attr_name in token_ids:
if getattr(tokenizer, attr_name) is None:
@@ -196,7 +212,7 @@ def load_tokenizer(cfg):
for k, val in special_tokens.items():
# check if new special token is not already in tokenizer and
# is adapter training to make sure lora_modules_to_save is set
- # pylint: disable=too-many-boolean-expressions
+
if (
(getattr(tokenizer, k) is None or getattr(tokenizer, k) != val)
and (len(tokenizer.encode(val, add_special_tokens=False)) > 2)
@@ -209,11 +225,12 @@ def load_tokenizer(cfg):
)
and k != "pad_token"
):
- lora_modules_to_save = ", ".join(
+ lora_modules_to_save_str = ", ".join(
[f"`{x}`" for x in lora_modules_to_save]
)
raise ValueError(
- f"Please set lora_modules_to_save to [{lora_modules_to_save}] when using an adapter and changing the special tokens."
+ f"Please set lora_modules_to_save to [{lora_modules_to_save_str}] "
+ "when using an adapter and changing the special tokens."
)
tokenizer.add_special_tokens(
@@ -259,7 +276,7 @@ def load_tokenizer(cfg):
{"additional_special_tokens": additional_special_tokens}
)
- if is_main_process(use_environ=True):
+ if is_main_process():
LOG.debug(f"EOS: {tokenizer.eos_token_id} / {tokenizer.eos_token}")
LOG.debug(f"BOS: {tokenizer.bos_token_id} / {tokenizer.bos_token}")
LOG.debug(f"PAD: {tokenizer.pad_token_id} / {tokenizer.pad_token}")
@@ -276,8 +293,13 @@ def load_tokenizer(cfg):
)
tokenizer.chat_template = chat_template_string
- else:
+ elif getattr(tokenizer, "chat_template", None) is None:
LOG.info(
"No Chat template selected. Consider adding a chat template for easier inference."
)
+
+ # make the tokenizer.pad call quieter 🤐
+ if hasattr(tokenizer, "deprecation_warnings"):
+ tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
return tokenizer
diff --git a/src/axolotl/loaders/utils.py b/src/axolotl/loaders/utils.py
index 28c935085..240e00da7 100644
--- a/src/axolotl/loaders/utils.py
+++ b/src/axolotl/loaders/utils.py
@@ -131,6 +131,17 @@ def check_model_config(cfg: DictDefault, model_config: PretrainedConfig):
f"Please include [{lora_modules_to_save_joined}] in `lora_modules_to_save`."
)
+ if (
+ cfg.tensor_parallel_size
+ and cfg.tensor_parallel_size > 1
+ and hasattr(model_config, "tie_word_embeddings")
+ and model_config.tie_word_embeddings
+ ):
+ raise ValueError(
+ "Tensor parallelism is incompatible with models configured with `tie_word_embeddings` enabled. "
+ "Please use a model without `tie_word_embeddings`, or disable tensor parallelism."
+ )
+
def load_model_config(cfg: DictDefault) -> PretrainedConfig | addict.Dict:
"""Loads and configures a model configuration from HuggingFace or local sources.
@@ -195,9 +206,11 @@ def ensure_dtype(model: PreTrainedModel, dtype: torch.dtype = torch.bfloat16):
bias_mismatch = module.bias.dtype != dtype
if weight_mismatch:
- print(f"Converting module {name}.weight: {module.weight.dtype} -> {dtype}")
+ LOG.debug(
+ f"Converting module {name}.weight: {module.weight.dtype} -> {dtype}"
+ )
if bias_mismatch:
- print(f"Converting module {name}.bias: {module.bias.dtype} -> {dtype}")
+ LOG.debug(f"Converting module {name}.bias: {module.bias.dtype} -> {dtype}")
if weight_mismatch or bias_mismatch:
module.to(dtype)
diff --git a/src/axolotl/logging_config.py b/src/axolotl/logging_config.py
index b8dc6479d..67b1d32f1 100644
--- a/src/axolotl/logging_config.py
+++ b/src/axolotl/logging_config.py
@@ -1,10 +1,7 @@
-"""
-Common logging module for axolotl
-"""
+"""Common logging module for axolotl."""
import logging
import os
-import sys
from logging import Formatter, Logger, LogRecord
from logging.config import dictConfig
from typing import Any, Dict
@@ -17,20 +14,28 @@ DEFAULT_LOG_LEVEL = "WARNING"
class AxolotlOrWarnErrorFilter(logging.Filter):
"""
- Allows ANY WARNING or higher (unless overridden by LOG_LEVEL)
- Allows axolotl.* at INFO or higher (unless overridden by AXOLOTL_LOG_LEVEL)
- Drops all other records (i.e. non-axolotl.INFO, DEBUG, etc. by default)
+ Allows ANY WARNING or higher (unless overridden by LOG_LEVEL). Allows axolotl.* at
+ INFO or higher (unless overridden by AXOLOTL_LOG_LEVEL). Drops all other records
+ (i.e. non-axolotl.INFO, DEBUG, etc. by default).
"""
def __init__(self, **kwargs):
super().__init__(**kwargs)
- self.axolotl_level = logging.getLevelNamesMapping()[
- os.getenv("AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL)
- ]
- self.other_level = logging.getLevelNamesMapping()[
- os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL)
- ]
+ axolotl_log_level = os.getenv(
+ "AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL
+ ).upper()
+ other_log_level = os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL).upper()
+
+ try:
+ # py311+ only
+ level_mapping = logging.getLevelNamesMapping()
+ self.axolotl_level = level_mapping[axolotl_log_level]
+ self.other_level = level_mapping[other_log_level]
+ except AttributeError:
+ # For py310, use getLevelName directly
+ self.axolotl_level = logging.getLevelName(axolotl_log_level)
+ self.other_level = logging.getLevelName(other_log_level)
def filter(self, record: LogRecord) -> bool:
# General filter
@@ -44,13 +49,12 @@ class AxolotlOrWarnErrorFilter(logging.Filter):
class AxolotlLogger(Logger):
- """A Logger that automatically rejects non-axolotl INFOs."""
+ """Logger that applies filtering to non-axolotl loggers."""
def __init__(self, name: str, level: int = logging.NOTSET):
super().__init__(name, level)
-
- # set global filter on the logger itself
- self.addFilter(AxolotlOrWarnErrorFilter())
+ if not name.startswith("axolotl"):
+ self.addFilter(AxolotlOrWarnErrorFilter())
class ColorfulFormatter(Formatter):
@@ -66,6 +70,7 @@ class ColorfulFormatter(Formatter):
def format(self, record):
record.rank = int(os.getenv("LOCAL_RANK", "0"))
+ record.rank_fmt = f" [RANK:{record.rank}]" if record.rank != 0 else ""
log_message = super().format(record)
return self.COLORS.get(record.levelname, "") + log_message + Fore.RESET
@@ -79,33 +84,55 @@ DEFAULT_LOGGING_CONFIG: Dict[str, Any] = {
},
"colorful": {
"()": ColorfulFormatter,
- "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d] [RANK:%(rank)d] %(message)s",
+ "format": "[%(asctime)s] [%(levelname)s] [%(name)s.%(funcName)s:%(lineno)d] [PID:%(process)d]%(rank_fmt)s %(message)s",
+ },
+ "concise": {
+ "format": "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s",
+ },
+ "concise_color": {
+ "()": ColorfulFormatter,
+ "format": "[%(asctime)s] [%(levelname)s] [%(name)s]%(rank_fmt)s %(message)s",
+ },
+ },
+ "filters": {
+ "ax_or_warn": {
+ "()": "axolotl.logging_config.AxolotlOrWarnErrorFilter",
},
},
- "filters": {},
"handlers": {
"console": {
"class": "logging.StreamHandler",
- "formatter": "simple",
- "filters": [],
- "stream": sys.stdout,
+ "formatter": "concise",
+ "filters": ["ax_or_warn"],
+ "stream": "ext://sys.stdout",
},
"color_console": {
"class": "logging.StreamHandler",
- "formatter": "colorful",
- "filters": [],
- "stream": sys.stdout,
+ "formatter": "concise_color",
+ "filters": ["ax_or_warn"],
+ "stream": "ext://sys.stdout",
+ },
+ "ax_file_only": {
+ "class": "logging.StreamHandler",
+ "level": "DEBUG",
+ "formatter": "simple",
+ "stream": "ext://axolotl.utils.tee.file_only_stream",
+ },
+ "root_file_only": {
+ "class": "logging.StreamHandler",
+ "level": "DEBUG",
+ "formatter": "simple",
+ "stream": "ext://axolotl.utils.tee.file_only_stream",
},
},
- # log level will be superseded by the AxolotlLogger
"root": {
- "handlers": ["console"],
- "level": os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL),
+ "handlers": ["console", "root_file_only"],
+ "level": os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL).upper(),
},
"loggers": {
"axolotl": {
- "handlers": ["color_console"],
- "level": os.getenv("AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL),
+ "handlers": ["color_console", "ax_file_only"],
+ "level": os.getenv("AXOLOTL_LOG_LEVEL", DEFAULT_AXOLOTL_LOG_LEVEL).upper(),
"propagate": False,
},
},
@@ -115,9 +142,15 @@ DEFAULT_LOGGING_CONFIG: Dict[str, Any] = {
def configure_logging():
"""Configure with default logging"""
init() # Initialize colorama
+
dictConfig(DEFAULT_LOGGING_CONFIG)
logging.setLoggerClass(AxolotlLogger)
- # set default `ACCELERATE_LOG_LEVEL` to `LOG_LEVEL` if available and not set
+ # Route Python warnings through logging so they reach file handlers
+ logging.captureWarnings(True)
+
+ # Set default `ACCELERATE_LOG_LEVEL` to `LOG_LEVEL` if available and not set
if "ACCELERATE_LOG_LEVEL" not in os.environ:
- os.environ["ACCELERATE_LOG_LEVEL"] = os.getenv("LOG_LEVEL", DEFAULT_LOG_LEVEL)
+ os.environ["ACCELERATE_LOG_LEVEL"] = os.getenv(
+ "LOG_LEVEL", DEFAULT_LOG_LEVEL
+ ).upper()
diff --git a/src/axolotl/models/mamba/__init__.py b/src/axolotl/models/mamba/__init__.py
index fee88e3a4..d6bb40d99 100644
--- a/src/axolotl/models/mamba/__init__.py
+++ b/src/axolotl/models/mamba/__init__.py
@@ -21,4 +21,4 @@ def fix_mamba_attn_for_loss():
from .modeling_mamba import MambaLMHeadModel as MambaLMHeadModelFixed
mixer_seq_simple.MambaLMHeadModel = MambaLMHeadModelFixed
- return mixer_seq_simple.MambaLMHeadModel # pylint: disable=invalid-name
+ return mixer_seq_simple.MambaLMHeadModel
diff --git a/src/axolotl/models/mamba/modeling_mamba.py b/src/axolotl/models/mamba/modeling_mamba.py
index 70e9c88c8..2cfe11544 100644
--- a/src/axolotl/models/mamba/modeling_mamba.py
+++ b/src/axolotl/models/mamba/modeling_mamba.py
@@ -1,4 +1,3 @@
-# pylint: skip-file
import os
from collections import namedtuple
from functools import partial
@@ -112,7 +111,7 @@ class MambaLMHeadModel(nn.Module, GenerationMixin):
self,
save_directory: Union[str, os.PathLike],
state_dict: Optional[dict] = None,
- safe_serialization: Optional[bool] = None, # pylint: disable=unused-argument
+ safe_serialization: Optional[bool] = None,
):
if state_dict is None:
state_dict = self.state_dict()
diff --git a/src/axolotl/monkeypatch/accelerate/fsdp2.py b/src/axolotl/monkeypatch/accelerate/fsdp2.py
index 955c06cbe..af6f24a63 100644
--- a/src/axolotl/monkeypatch/accelerate/fsdp2.py
+++ b/src/axolotl/monkeypatch/accelerate/fsdp2.py
@@ -2,102 +2,91 @@
monkeypatch for accelerate fsdp2 fix when modifying ordereddict during interation, and saving full state dicts
"""
+import copy
+import functools
+import os
import sys
import torch
+import torch.distributed as dist
+from torch import nn
+from axolotl.utils.bench import log_gpu_memory_usage
from axolotl.utils.logging import get_logger
LOG = get_logger(__name__)
-def fsdp2_load_full_state_dict(accelerator, model: torch.nn.Module, full_sd: dict):
+def fsdp2_load_full_state_dict(
+ _accelerator, model: torch.nn.Module, full_sd: dict, offload_to_cpu: bool = False
+):
"""
Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the
parameters from rank 0 to all other ranks. This function modifies the model in-place.
-
Args:
accelerator (`Accelerator`): The accelerator instance
model (`torch.nn.Module`):
The model to load the state dict into, expected to be on meta device or a VRAM spike can occur
full_sd (`dict`): The full state dict to load, can only be on rank 0
"""
- import torch.distributed as dist
from torch.distributed.tensor import distribute_tensor
- # Model was previously copied to meta device
+ LOG.info("Broadcasting full state dict to all ranks...")
+ import time
+
+ start_time = time.time()
+
meta_sharded_sd = model.state_dict()
sharded_sd = {}
+ for param_name, sharded_meta_param in meta_sharded_sd.items():
+ full_tensor = None
+ if _accelerator.is_main_process:
+ full_tensor = full_sd[param_name]
+ full_tensor = full_tensor.to(sharded_meta_param.dtype)
- # Rank 0 distributes the full state dict to other ranks
- def _infer_parameter_dtype(model, param_name, empty_param):
- try:
- old_param = model.get_parameter_or_buffer(param_name)
- except AttributeError:
- # Need this for LORA, as there some params are not *parameters* of sorts
- base_param_name, local_param_name = param_name.rsplit(".", 1)
- submodule = model.get_submodule(base_param_name)
- old_param = getattr(submodule, local_param_name)
-
- is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
- casting_dtype = None
- is_param_float8_e4m3fn = (
- is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn
- )
-
- if empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn:
- casting_dtype = old_param.dtype
-
- return old_param is not None and old_param.is_contiguous(), casting_dtype
-
- def _cast_and_contiguous(tensor, to_contiguous, dtype):
- if dtype is not None:
- tensor = tensor.to(dtype=dtype)
- if to_contiguous:
- tensor = tensor.contiguous()
- return tensor
-
- param_names = sorted(meta_sharded_sd.keys())
-
- for param_name in param_names:
- mesh = meta_sharded_sd[param_name].device_mesh
- if accelerator.is_main_process:
- full_param = full_sd[param_name].detach().cuda()
- dist.broadcast(full_param, src=0, group=mesh.get_group())
- sharded_tensor = distribute_tensor(
- full_param, mesh, sharded_sd[param_name].placements
- )
- to_contiguous, casting_dtype = _infer_parameter_dtype(
- model,
- param_name,
- full_param,
- )
- sharded_tensor = _cast_and_contiguous(
- sharded_tensor, to_contiguous, casting_dtype
- )
- sharded_sd[param_name] = sharded_tensor
- else:
- full_tensor = torch.empty(
- sharded_sd[param_name].size(),
- device="cuda",
- dtype=sharded_sd[param_name].dtype,
- )
- dist.broadcast(full_tensor, src=0, group=mesh.get_group())
- sharded_tensor = distribute_tensor(
- full_tensor, mesh, sharded_sd[param_name].placements
- )
- to_contiguous, casting_dtype = _infer_parameter_dtype(
- model,
- param_name,
+ if hasattr(sharded_meta_param, "device_mesh"):
+ device_mesh = sharded_meta_param.device_mesh
+ if _accelerator.is_main_process:
+ full_tensor = full_tensor.to(device_mesh.device_type)
+ else:
+ full_tensor = torch.empty(
+ sharded_meta_param.size(),
+ device=device_mesh.device_type,
+ dtype=sharded_meta_param.dtype,
+ )
+ sharded_param = distribute_tensor(
full_tensor,
+ device_mesh,
+ sharded_meta_param.placements,
+ src_data_rank=0,
)
- sharded_tensor = _cast_and_contiguous(
- sharded_tensor, to_contiguous, casting_dtype
- )
- sharded_sd[param_name] = sharded_tensor
+ else:
+ # Non-sharded parameters
+ if _accelerator.is_main_process:
+ sharded_param = full_tensor.to(torch.device("cuda"))
+ else:
+ # broadcast manually
+ sharded_param = torch.empty_like(
+ sharded_meta_param,
+ device=torch.device("cuda"),
+ dtype=sharded_meta_param.dtype,
+ )
+ dist.broadcast(sharded_param, src=0)
- # we set `assign=True` because our params are on meta device
- model.load_state_dict(sharded_sd, assign=True)
+ if offload_to_cpu:
+ sharded_param = sharded_param.cpu()
+
+ sharded_sd[param_name] = nn.Parameter(sharded_param)
+
+ del full_tensor
+ full_sd[param_name] = None
+
+ model.load_state_dict(sharded_sd, assign=True, strict=True)
+ end_time = time.time()
+ LOG.debug(
+ f"Time taken to load full state dict: {(end_time - start_time):.2f} seconds"
+ )
+ log_gpu_memory_usage(LOG, "Memory usage after broadcasting full state dict", 0)
return model
@@ -142,9 +131,9 @@ def get_state_dict(self, model, unwrap=True):
"Deepspeed TP requires deepspeed >= 0.16.4, Please update DeepSpeed via `pip install deepspeed -U`."
)
state_dict = (
- model._consolidated_16bit_state_dict() # pylint: disable=protected-access
+ model._consolidated_16bit_state_dict()
if tp_sharding
- else model._zero3_consolidated_16bit_state_dict() # pylint: disable=protected-access
+ else model._zero3_consolidated_16bit_state_dict()
)
else:
raise ValueError(
@@ -172,9 +161,11 @@ def get_state_dict(self, model, unwrap=True):
state_dict[param_name] = param.cpu()
torch.distributed.barrier()
elif self.distributed_type == DistributedType.FSDP:
- from torch.distributed.fsdp import FullStateDictConfig
- from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
- from torch.distributed.fsdp import StateDictType
+ from torch.distributed.fsdp import (
+ FullStateDictConfig,
+ FullyShardedDataParallel as FSDP,
+ StateDictType,
+ )
full_state_dict_config = FullStateDictConfig(
offload_to_cpu=True, rank0_only=True
@@ -191,17 +182,204 @@ def get_state_dict(self, model, unwrap=True):
return state_dict
-def patch_accelerate_fsdp2():
- import accelerate
- from accelerate.utils import fsdp_utils
+def _process_lora_module_for_fsdp(module, fsdp2_kwargs):
+ """Helper function to process LoRA modules for FSDP2."""
+ from torch.distributed.fsdp import fully_shard
- fsdp_utils.fsdp2_load_full_state_dict = fsdp2_load_full_state_dict
- setattr(
- sys.modules["accelerate.utils.fsdp_utils"],
- "fsdp2_load_full_state_dict",
- fsdp2_load_full_state_dict,
+ log_bias_dtype_mismatch = False
+
+ # Linear4Bit will keep it's bias term in fp32. If the weight dtype is in bf16 we are not able to
+ # wrap this. Therefore we must ensure the bias has the same dtype as the weight
+ if hasattr(module.base_layer, "bias") and module.base_layer.bias is not None:
+ if module.base_layer.weight.dtype != module.base_layer.bias.dtype:
+ log_bias_dtype_mismatch = True
+ module.base_layer.bias.data = module.base_layer.bias.data.to(
+ module.base_layer.weight.dtype
+ )
+
+ for active_adapter in module.active_adapters:
+ if module.lora_A:
+ fully_shard(module.lora_A[active_adapter], **fsdp2_kwargs)
+ if module.lora_B:
+ fully_shard(module.lora_B[active_adapter], **fsdp2_kwargs)
+ if module.lora_embedding_A:
+ fully_shard(module.lora_embedding_A[active_adapter], **fsdp2_kwargs)
+ if module.lora_embedding_B:
+ fully_shard(module.lora_embedding_B[active_adapter], **fsdp2_kwargs)
+ if module.lora_magnitude_vector:
+ fully_shard(module.lora_magnitude_vector[active_adapter], **fsdp2_kwargs)
+ return log_bias_dtype_mismatch
+
+
+def fsdp2_prepare_model(accelerator, model: torch.nn.Module) -> torch.nn.Module:
+ """Prepares the model for FSDP2 in-place. Also returns the model to avoid misuse of the original model.
+
+ Args:
+ accelerator (`Accelerator`): The accelerator instance
+ model (`torch.nn.Module`): The model to prepare
+
+ Returns:
+ `torch.nn.Module`: Prepared model
+ """
+ from accelerate.utils import get_module_children_bottom_up, is_compiled_module
+ from accelerate.utils.fsdp_utils import fsdp2_prepare_auto_wrap_policy
+ from accelerate.utils.modeling import get_non_persistent_buffers
+ from peft import PeftModel
+ from peft.tuners.lora import LoraLayer
+ from torch.distributed.fsdp import (
+ CPUOffloadPolicy,
+ FSDPModule,
+ MixedPrecisionPolicy,
+ fully_shard,
)
+ is_type_fsdp = isinstance(model, FSDPModule) or (
+ is_compiled_module(model) and isinstance(model._orig_mod, FSDPModule)
+ )
+ if is_type_fsdp:
+ return model
+
+ fsdp2_plugin = accelerator.state.fsdp_plugin
+
+ original_sd = model.state_dict()
+
+ from torch.distributed.fsdp.wrap import (
+ size_based_auto_wrap_policy,
+ transformer_auto_wrap_policy,
+ )
+
+ # We need the `auto_wrap_policy` original type to create a custom poilicy function for sharding
+ # This is because `fully_shard` doesn't support old auto wrap policies, rather we have to imitate the behaviour
+ if fsdp2_plugin.auto_wrap_policy is transformer_auto_wrap_policy:
+ pass # auto_wrap_policy_type = "transformer"
+ elif fsdp2_plugin.auto_wrap_policy is size_based_auto_wrap_policy:
+ pass # auto_wrap_policy_type = "size"
+
+ # We set `auto_wrap_policy` to `functools.partial` to avoid creating it again
+ # This is because of `apply_activation_checkpointing` which will can reuse this function
+ fsdp2_plugin.set_auto_wrap_policy(model)
+
+ if fsdp2_plugin.activation_checkpointing:
+ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+ CheckpointImpl,
+ apply_activation_checkpointing,
+ checkpoint_wrapper,
+ )
+
+ # Apply activation checkpointing before applying `fully_shard`
+ apply_activation_checkpointing(
+ model,
+ checkpoint_wrapper_fn=functools.partial(
+ checkpoint_wrapper,
+ checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+ ),
+ auto_wrap_policy=fsdp2_plugin.auto_wrap_policy,
+ )
+
+ mesh = getattr(accelerator.state, "device_mesh", None)
+
+ # Disable memory pinning if requested
+ offload_to_cpu = isinstance(fsdp2_plugin.cpu_offload, CPUOffloadPolicy)
+ if offload_to_cpu and os.environ.get("FSDP_CPU_OFFLOAD_PIN_MEMORY", "") == "false":
+ fsdp2_plugin.cpu_offload.pin_memory = False
+
+ fsdp2_kwargs = {
+ "reshard_after_forward": fsdp2_plugin.reshard_after_forward,
+ "offload_policy": fsdp2_plugin.cpu_offload,
+ # `fully_shard` doesn't accept `None` in case of `MixedPrecisionPolicy`
+ "mp_policy": fsdp2_plugin.mixed_precision_policy or MixedPrecisionPolicy(),
+ "mesh": (
+ mesh[tuple(accelerator.state.parallelism_config.fsdp_dim_names)]
+ if mesh is not None
+ else None
+ ),
+ }
+ model_has_params4bit = False
+ for _, param in model.named_parameters():
+ # this is a temporary fix whereby loading models with bnb params cannot be moved from
+ # GPU to a meta device due with FSDP2 because torch operations don't return the original class type
+ # bypassing the move to meta will still cause the VRAM spike, but at least it still will load
+ if param.__class__.__name__ == "Params4bit":
+ model_has_params4bit = True
+ break
+
+ if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit:
+ # Context: `fully_shard` moves the model to GPU if it was on CPU, however it can also be on `meta` and then it stays there even after `fully_shard`
+ # For this reason, we need to move the model to `meta` device, as then sharding happens on `meta` device
+ # If we kept the model on CPU (`cpu_ram_efficient_loading` has model be on CPU on all ranks, though non-main ranks only have `torch.emtpy`), `fully_shard` would move it to GPU
+ # Afterwards, when we call `fsdp2_load_full_state_dict`, us creating the state_dict would result into briefly having two copies of model state_dict on the GPU -> VRAM spike
+
+ # We need to keep the original non-persistent buffers, as those MAY not be in the state_dict, resulting in them staying on meta device
+ # Also, these buffers aren't getting sharded by default
+ # We get the FQNs of all non-persistent buffers, to re-register them after
+ non_persistent_buffer_fqns = get_non_persistent_buffers(
+ model, recurse=True, fqns=True
+ )
+ original_non_persistent_buffers = copy.deepcopy(
+ {k: v for k, v in model.named_buffers() if k in non_persistent_buffer_fqns}
+ )
+ # We move the model to meta device, as then sharding happens on meta device
+ model = model.to(torch.device("meta"))
+ # We need to re-tie the weights, not exactly sure why, but if we don't do this, reference to `lm_head/embed_tokens` stay hanging -> more VRAM usage
+ # We assume `transformers` models have a `tie_weights` method if they support it
+ if hasattr(model, "tie_weights"):
+ model.tie_weights()
+
+ is_peft_model = isinstance(model, PeftModel)
+
+ auto_wrap_policy = fsdp2_prepare_auto_wrap_policy(fsdp2_plugin, model)
+ log_bias_dtype_mismatch = False
+ if auto_wrap_policy is not None:
+ for module in get_module_children_bottom_up(model)[:-1]:
+ if is_peft_model and isinstance(module, LoraLayer):
+ module_log_bias_mismatch = _process_lora_module_for_fsdp(
+ module, fsdp2_kwargs
+ )
+ log_bias_dtype_mismatch |= module_log_bias_mismatch
+ if auto_wrap_policy(module) and not isinstance(module, FSDPModule):
+ fully_shard(module, **fsdp2_kwargs)
+
+ fully_shard(model, **fsdp2_kwargs)
+
+ if log_bias_dtype_mismatch:
+ LOG.warning(
+ "Bias dtype mismatch detected in LoRA base linear layer. Bias parameters have been cast to weight dtype."
+ )
+
+ if fsdp2_plugin.cpu_ram_efficient_loading:
+ fsdp2_load_full_state_dict(
+ accelerator, model, original_sd, offload_to_cpu=offload_to_cpu
+ )
+
+ if fsdp2_plugin.cpu_ram_efficient_loading and not model_has_params4bit:
+ # We re-register the buffers, as they may not be in the state_dict
+ for fqn, buffer_tensor in original_non_persistent_buffers.items():
+ buffer_tensor = buffer_tensor.to(accelerator.device)
+
+ if "." in fqn:
+ parent_fqn, local_buffer_name = fqn.rsplit(".", 1)
+ parent_module = model.get_submodule(parent_fqn)
+ else:
+ local_buffer_name = fqn
+ parent_module = model
+
+ parent_module.register_buffer(
+ local_buffer_name, buffer_tensor, persistent=False
+ )
+
+ # We need to tie the weights again, as call to `load_full_state_dict` breaks the tie
+ # Needs to be called both here and above
+ # removing this call makes the have slightly different loss
+ # removing the call above leads to extra memory usage as explained in the comment above
+ if hasattr(model, "tie_weights"):
+ model.tie_weights()
+ return model
+
+
+def patch_accelerate_fsdp2():
+ import accelerate
+
+ accelerate.accelerator.fsdp2_prepare_model = fsdp2_prepare_model
accelerate.Accelerator.get_state_dict = get_state_dict
setattr(
sys.modules["accelerate"],
diff --git a/src/axolotl/monkeypatch/accelerate/parallelism_config.py b/src/axolotl/monkeypatch/accelerate/parallelism_config.py
new file mode 100644
index 000000000..b2157fb6b
--- /dev/null
+++ b/src/axolotl/monkeypatch/accelerate/parallelism_config.py
@@ -0,0 +1,77 @@
+"""
+workaround to allow parallelism config for pure CP
+"""
+
+import os
+import warnings
+
+from accelerate import DistributedType
+
+
+def _validate_accelerator(self, accelerator):
+ _warnings = set()
+ if not accelerator.multi_device and self.total_size == 1:
+ # No distributed setup, valid parallelism config
+ return
+
+ # We need this to ensure DDP works
+ if self.total_size == 1:
+ self._set_size("dp_replicate", accelerator.num_processes)
+
+ if self.total_size != accelerator.num_processes:
+ raise ValueError(
+ f"ParallelismConfig total_size ({self.total_size}) does not match "
+ f"num_processes ({accelerator.num_processes}). Please adjust dp_replicate_size/ "
+ f"dp_shard_size/tp_size/cp_size."
+ )
+
+ # allow parallelism config when not using fsdp if using pure context parallelism
+ allow_parallelism_config = False
+
+ if (
+ self.cp_size > 1
+ and self.dp_shard_size <= 1
+ and os.environ.get("ACCELERATE_ALLOW_CP_STANDALONE", "false").lower() == "true"
+ ):
+ allow_parallelism_config = True
+
+ if (
+ self.total_size > 1
+ and not allow_parallelism_config
+ and not (accelerator.is_fsdp2 or accelerator.multi_device)
+ ):
+ raise ValueError(
+ f"ParallelismConfig is only compatible DistributedType.FSDP (version 2) or DistributedType.Multi{{Device}}, but got {accelerator.distributed_type}."
+ )
+
+ for parallelism, size in self._sizes.items():
+ if size == 1 and getattr(self, f"{parallelism}_handler", None) is not None:
+ _warnings.add(
+ f"ParallelismConfig.{parallelism}_handler is set, but {parallelism}_size is set to 1. This handler will be ignored."
+ )
+
+ if _warnings and accelerator.is_main_process:
+ warnings.warn(
+ "ParallelismConfig has the following warnings:\n" + "\n".join(_warnings),
+ UserWarning,
+ stacklevel=2,
+ )
+
+
+def patched_is_fsdp2(self) -> bool:
+ """
+ Patched version of is_fsdp2 that guards against a None fsdp_plugin.
+ """
+ # The new logic checks if fsdp_plugin exists before accessing its attributes
+ return (
+ self.distributed_type == DistributedType.FSDP
+ and self.fsdp_plugin
+ and self.fsdp_plugin.fsdp_version == 2
+ )
+
+
+def patch_parallelism_config():
+ from accelerate.accelerator import AcceleratorState, ParallelismConfig
+
+ ParallelismConfig._validate_accelerator = _validate_accelerator
+ AcceleratorState.is_fsdp2 = property(patched_is_fsdp2)
diff --git a/src/axolotl/monkeypatch/attention/flex_attn.py b/src/axolotl/monkeypatch/attention/flex_attn.py
index 3652a30b3..678f65bee 100644
--- a/src/axolotl/monkeypatch/attention/flex_attn.py
+++ b/src/axolotl/monkeypatch/attention/flex_attn.py
@@ -1,10 +1,15 @@
"""Flex attention monkey patch"""
import sys
-from typing import Optional, Tuple, Union
import torch
import transformers
+from packaging import version
+from transformers.utils.import_utils import _torch_version, is_torch_less_or_equal
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
def patch_flex_wrapper(**flex_attn_compile_kwargs):
@@ -42,162 +47,39 @@ def patch_flex_wrapper(**flex_attn_compile_kwargs):
"""
self.training = None
if not self._is_flex_compiled or training != self.training:
+ self.training = training
+ if is_torch_less_or_equal("2.5.1"):
+ self._compiled_flex_attention = torch.compile(
+ flex_attention, dynamic=False
+ )
# In PyTorch 2.6.0, there's a known issue with flex attention compilation which may
# cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs"
# see https://github.com/pytorch/pytorch/issues/146260 for training
- self.training = training
- self._compiled_flex_attention = torch.compile(
- flex_attention,
- **flex_attn_compile_kwargs,
- )
+ elif version.parse(_torch_version).base_version == "2.6.0" and training:
+ self._compiled_flex_attention = torch.compile(
+ flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"
+ )
+ # Fallback, usually the most recent torch 2.7.x+ versions
+ else:
+ LOG.info(
+ "Compiling flex attention with kwargs: %s. This may take a while...",
+ flex_attn_compile_kwargs,
+ main_process_only=True,
+ )
+ self._compiled_flex_attention = torch.compile(
+ flex_attention,
+ **flex_attn_compile_kwargs,
+ )
+ LOG.info(
+ "Flex attention compiled successfully.", main_process_only=True
+ )
+
self._is_flex_compiled = True
def __call__(self):
return self._compiled_flex_attention
transformers.integrations.flex_attention.WrappedFlexAttention = WrappedFlexAttention
- setattr(
- sys.modules["transformers.integrations.flex_attention"],
- "WrappedFlexAttention",
- WrappedFlexAttention,
- )
-
-
-def patch_flex_make_mask():
- is_torch_2_6 = torch.__version__.startswith("2.6")
-
- if not is_torch_2_6:
- return
-
- from torch.nn.attention.flex_attention import (
- _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size,
- )
- from torch.nn.attention.flex_attention import (
- BlockMask,
- )
- from torch.nn.attention.flex_attention import (
- create_block_mask as create_block_causal_mask_flex,
- )
-
- Offset = Union[torch.Tensor, int]
-
- def patched_make_flex_block_causal_mask(
- attention_mask_2d: torch.Tensor,
- attention_chunk_size: Optional[int] = None,
- query_length=None,
- key_length=None,
- offsets: Optional[Tuple[Offset, Offset]] = None,
- ) -> "BlockMask":
- """
- Create a block causal document mask for a batch of sequences, both packed and unpacked.
- Create Block causal logic and passing it into :func:`torch.nn.attention.flex_attention.create_block_mask`.
- The resultant BlockMask is a compressed representation of the full block causal
- mask. BlockMask is essential for performant computation of flex attention.
- See: https://pytorch.org/blog/flexattention/
-
- Args:
- attention_mask_2d (torch.Tensor): Attention mask for packed and padded sequences
- of shape (batch_size, total_seq_len). e.g.
-
- For unpacked sequence:
- [[1, 1, 1, 1, 0, 0, 0],
- [1, 1, 1, 1, 1, 0, 0]]
-
- For packed sequence:
- [[1, 1, 1, 2, 2, 2, 0],
- [1, 1, 2, 2, 2, 3, 3]]
-
- Returns:
- BlockMask
- """
-
- batch_size, total_seq_len = attention_mask_2d.shape
- if not key_length:
- key_length = total_seq_len
- if not query_length:
- query_length = total_seq_len
- attention_mask_2d = torch.nn.functional.pad(
- attention_mask_2d,
- value=0,
- pad=(0, abs(total_seq_len - max(key_length, flex_default_block_size))),
- )
- device = attention_mask_2d.device
- document_ids = attention_mask_2d.clone()
-
- if attention_chunk_size is not None:
- # we create an arange, then we just // by chunk size to get [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]
- chunk_idxs = (document_ids.clone().fill_(1).cumsum(-1) - 1) // (
- attention_chunk_size
- )
-
- # Instead of passing a tensor mask, flex attention requires a mask_mod function
- # that determines which elements of QK^T should be included in the attention
- # computation prior to the softmax. For sample packing, we need both the
- # logic for both causal mask and document mask. See PyTorch's official
- # blog post for more details: https://pytorch.org/blog/flexattention/#mask-mods
- def causal_mask_mod(
- batch_idx, head_idx, q_idx, kv_idx
- ): # pylint: disable=unused-argument
- """
- Defines the logic of a block causal mask by combining both a standard causal mask
- and a block diagonal document mask.
-
- See :func:`~torchtune.modules.attention_utils.create_block_causal_mask`
- for an illustration.
- """
- causal_mask = q_idx >= kv_idx # not valid when decoding
- document_mask = (
- document_ids[batch_idx, q_idx] == document_ids[batch_idx, kv_idx]
- )
- padding_mask = attention_mask_2d[batch_idx, q_idx] > 0
- final_mask = causal_mask & padding_mask & document_mask
- return final_mask
-
- def chunk_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):
- """
- Combines the chunk mask with the causal mask for chunked attention.
- """
- chunk_mask = chunk_idxs[batch_idx, q_idx] == chunk_idxs[batch_idx, kv_idx]
- causal_doc_mask = causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx)
- return chunk_mask & causal_doc_mask
-
- mask_mod_maybe_combined = (
- causal_mask_mod if attention_chunk_size is None else chunk_causal_mask_mod
- )
-
- if offsets is not None:
- q_offset = offsets[0]
- kv_offset = offsets[1]
-
- def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
- offset_q = q_idx + q_offset
- offset_kv = kv_idx + kv_offset
- return mask_mod_maybe_combined(batch_idx, head_idx, offset_q, offset_kv)
-
- else:
- mask_mod = mask_mod_maybe_combined
- return create_block_causal_mask_flex(
- mask_mod=mask_mod,
- B=batch_size,
- H=None, # attention head
- Q_LEN=query_length,
- KV_LEN=key_length,
- device=device,
- _compile=True,
- )
-
- for n in tuple(sys.modules):
- if ".modeling_" in n:
- if hasattr(sys.modules[n], "make_flex_block_causal_mask"):
- sys.modules[n].make_flex_block_causal_mask = (
- patched_make_flex_block_causal_mask
- )
- setattr(
- sys.modules[n],
- "make_flex_block_causal_mask",
- patched_make_flex_block_causal_mask,
- )
-
- transformers.integrations.flex_attention.make_flex_block_causal_mask = (
- patched_make_flex_block_causal_mask
- )
+ sys.modules[
+ "transformers.integrations.flex_attention"
+ ].WrappedFlexAttention = WrappedFlexAttention
diff --git a/src/axolotl/monkeypatch/attention/xformers.py b/src/axolotl/monkeypatch/attention/xformers.py
index 5901963f0..eca95797a 100644
--- a/src/axolotl/monkeypatch/attention/xformers.py
+++ b/src/axolotl/monkeypatch/attention/xformers.py
@@ -23,15 +23,15 @@ def xformers_attention_forward(
value: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
- dropout: float = 0.0, # pylint: disable=unused-argument
- scaling: Optional[float] = None, # pylint: disable=unused-argument
- sliding_window: Optional[int] = None, # pylint: disable=unused-argument
- softcap: Optional[float] = None, # pylint: disable=unused-argument
+ dropout: float = 0.0,
+ scaling: Optional[float] = None,
+ sliding_window: Optional[int] = None,
+ softcap: Optional[float] = None,
cu_seq_lens_q: Optional[torch.LongTensor] = None,
cu_seq_lens_k: Optional[torch.LongTensor] = None,
max_length_q: Optional[int] = None,
- max_length_k: Optional[int] = None, # pylint: disable=unused-argument
- **kwargs, # pylint: disable=unused-argument
+ max_length_k: Optional[int] = None,
+ **kwargs,
):
# Get dimensions
# query: [batch, heads, seq_len, hidden_dim]
diff --git a/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py b/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
index 589980c8b..2c5077392 100644
--- a/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/btlm_attn_hijack_flash.py
@@ -25,9 +25,7 @@ def replace_btlm_attn_with_flash_attn(model_name="cerebras/btlm-3b-8k-base"):
".configuration_btlm", ".modeling_btlm"
)
modeling_btlm = importlib.import_module(module_name)
- modeling_btlm.BTLMAttention._attn = ( # pylint: disable=protected-access
- flashattn_attn
- )
+ modeling_btlm.BTLMAttention._attn = flashattn_attn
def flashattn_attn(
@@ -35,9 +33,9 @@ def flashattn_attn(
query: torch.Tensor,
key: Optional[torch.Tensor] = None,
value: Optional[torch.Tensor] = None,
- attention_mask: Optional[torch.Tensor] = None, # pylint: disable=unused-argument
+ attention_mask: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
- position_bias: Optional[torch.Tensor] = None, # pylint: disable=unused-argument
+ position_bias: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
softmax_scale = (
1 / (key.size(-1) ** self.attn_scale_power) if self.scale_attn_weights else None
diff --git a/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py b/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
index df8d106fd..c426344a6 100644
--- a/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
+++ b/src/axolotl/monkeypatch/data/batch_dataset_fetcher.py
@@ -1,15 +1,23 @@
-"""monkey patches for the dataset fetcher to handle batches of packed indexes"""
-
-# pylint: disable=protected-access
+"""Monkey patches for the dataset fetcher to handle batches of packed indexes."""
import torch
from torch.utils.data._utils.fetch import _BaseDatasetFetcher
from torch.utils.data._utils.worker import _worker_loop
+_ORIGINAL_MAP_DATASET_FETCHER = None
+_ORIGINAL_WORKER_LOOP = None
+_IS_PATCHED = False
+
class _MapDatasetFetcher(_BaseDatasetFetcher):
+ """
+ Custom dataset fetcher that handles nested batch structures from
+ MultipackBatchSampler.
+ """
+
def fetch(self, possibly_batched_index):
if isinstance(possibly_batched_index[0], list):
+ # Handle nested structure from MultipackBatchSampler
data = [None for i in possibly_batched_index]
for i, possibly_batched_index_ in enumerate(possibly_batched_index):
if self.auto_collation:
@@ -23,6 +31,7 @@ class _MapDatasetFetcher(_BaseDatasetFetcher):
else:
data[i] = self.dataset[possibly_batched_index_]
else:
+ # Standard batch handling
if self.auto_collation:
if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
data = self.dataset.__getitems__(possibly_batched_index)
@@ -34,14 +43,54 @@ class _MapDatasetFetcher(_BaseDatasetFetcher):
def patch_fetchers():
+ """Apply patches to PyTorch's DataLoader components."""
torch.utils.data._utils.fetch._MapDatasetFetcher = _MapDatasetFetcher
torch.utils.data.dataloader._utils.fetch._MapDatasetFetcher = _MapDatasetFetcher
def patched_worker_loop(*args, **kwargs):
+ """Worker loop that ensures patches are applied in worker processes."""
patch_fetchers()
return _worker_loop(*args, **kwargs)
-torch.utils.data._utils.worker._worker_loop = patched_worker_loop
-patch_fetchers()
+def apply_multipack_dataloader_patch():
+ """
+ This patch allows DataLoader to correctly process batches that contain multiple bins
+ of packed sequences.
+ """
+ # pylint: disable=global-statement
+ global _ORIGINAL_MAP_DATASET_FETCHER, _ORIGINAL_WORKER_LOOP, _IS_PATCHED
+
+ if _IS_PATCHED:
+ return
+
+ # Store original implementations
+ _ORIGINAL_MAP_DATASET_FETCHER = torch.utils.data._utils.fetch._MapDatasetFetcher
+ _ORIGINAL_WORKER_LOOP = torch.utils.data._utils.worker._worker_loop
+
+ # Apply patches
+ patch_fetchers()
+ torch.utils.data._utils.worker._worker_loop = patched_worker_loop
+
+ _IS_PATCHED = True
+
+
+def remove_multipack_dataloader_patch():
+ """Remove the monkeypatch and restore original PyTorch DataLoader behavior."""
+ # pylint: disable=global-statement
+ global _IS_PATCHED
+
+ if not _IS_PATCHED:
+ return
+
+ if _ORIGINAL_MAP_DATASET_FETCHER:
+ torch.utils.data._utils.fetch._MapDatasetFetcher = _ORIGINAL_MAP_DATASET_FETCHER
+ torch.utils.data.dataloader._utils.fetch._MapDatasetFetcher = (
+ _ORIGINAL_MAP_DATASET_FETCHER
+ )
+
+ if _ORIGINAL_WORKER_LOOP:
+ torch.utils.data._utils.worker._worker_loop = _ORIGINAL_WORKER_LOOP
+
+ _IS_PATCHED = False
diff --git a/src/axolotl/monkeypatch/deepspeed_utils.py b/src/axolotl/monkeypatch/deepspeed_utils.py
new file mode 100644
index 000000000..d7e69e112
--- /dev/null
+++ b/src/axolotl/monkeypatch/deepspeed_utils.py
@@ -0,0 +1,67 @@
+import importlib
+import importlib.util
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_checkpoint_wrapper_setattr():
+ """
+ Patch CheckpointWrapper to properly forward DeepSpeed attributes to wrapped modules.
+
+ This fixes the issue where CheckpointWrapper doesn't forward ds_* attributes
+ (like ds_grads_remaining) to the actual wrapped module, causing DeepSpeed
+ ZeRO-3 to fail when gradient checkpointing is enabled.
+
+ This issue occurs specifically with:
+ - QLoRA + DeepSpeed ZeRO-3
+ - gradient_checkpointing: true
+ - activation_offloading: true
+
+ References:
+ - https://github.com/deepspeedai/DeepSpeed/issues/7203
+ - https://github.com/deepspeedai/DeepSpeed/blob/38d1a9eb64c9e01e32eccc50b25ba18925287441/deepspeed/runtime/zero/parameter_offload.py#L424-L458
+ - https://github.com/axolotl-ai-cloud/axolotl/pull/3102
+ """
+
+ try:
+ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+ CheckpointWrapper,
+ )
+
+ # Check if already patched
+ if hasattr(CheckpointWrapper, "_axolotl_setattr_patched"):
+ LOG.debug("CheckpointWrapper already patched")
+ return
+
+ original_setattr = CheckpointWrapper.__setattr__
+
+ def new_setattr(self, name: str, value) -> None:
+ if name.startswith("ds_") and hasattr(self, "_checkpoint_wrapped_module"):
+ setattr(self._checkpoint_wrapped_module, name, value)
+ LOG.debug(
+ f"Forwarded {name} to wrapped module {type(self._checkpoint_wrapped_module).__name__}"
+ )
+ else:
+ original_setattr(self, name, value)
+
+ CheckpointWrapper.__setattr__ = new_setattr
+ CheckpointWrapper._axolotl_setattr_patched = True
+
+ LOG.info("CheckpointWrapper patched to forward DeepSpeed attributes")
+
+ except ImportError as e:
+ LOG.debug(f"CheckpointWrapper not available: {e}")
+ except Exception as e:
+ LOG.warning(f"Failed to patch CheckpointWrapper: {e}")
+
+
+def apply_deepspeed_patches():
+ """
+ Apply DeepSpeed-related patches
+ """
+ if importlib.util.find_spec("deepspeed") is not None:
+ patch_checkpoint_wrapper_setattr()
+ else:
+ LOG.debug("DeepSpeed not available, skipping patches")
diff --git a/src/axolotl/monkeypatch/fsdp2_qlora.py b/src/axolotl/monkeypatch/fsdp2_qlora.py
new file mode 100644
index 000000000..04d0d1971
--- /dev/null
+++ b/src/axolotl/monkeypatch/fsdp2_qlora.py
@@ -0,0 +1,143 @@
+"""
+Monkeypatch to add Params4bit support to FSDP2. This enables QLoRA + FSDP2, as well as
+our LoRA / QLoRA Triton kernels to work with FSDP2.
+
+This patch modifies the _init_sharded_param method in FSDPParam to handle bitsandbytes
+Params4bit parameters.
+"""
+
+import importlib
+import inspect
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def apply_init_sharded_param_patch():
+ """Apply patch to FSDPParam._init_sharded_param to support Params4bit."""
+ from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+ # Get original source
+ original_source = inspect.getsource(FSDPParam._init_sharded_param)
+ original_source, _ = detab_code(original_source)
+
+ # Define the replacement
+ original_param_creation = """ self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))
+ self.sharded_param.requires_grad_(param.requires_grad)"""
+
+ patched_param_creation = """ import bitsandbytes as bnb
+ if isinstance(param, bnb.nn.modules.Params4bit):
+ self.sharded_param = bnb.nn.modules.Params4bit(
+ data=sharded_param,
+ requires_grad=param.requires_grad,
+ quant_state=param.quant_state,
+ blocksize=param.blocksize,
+ compress_statistics=param.compress_statistics,
+ quant_type=param.quant_type,
+ quant_storage=param.quant_storage,
+ module=param.module,
+ bnb_quantized=param.bnb_quantized,
+ )
+ self.sharded_param = self.to_sharded_dtensor(self.sharded_param)
+ else:
+ self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))
+ self.sharded_param.requires_grad_(param.requires_grad)"""
+
+ # Apply the replacement
+ if original_param_creation in original_source:
+ patched_source = original_source.replace(
+ original_param_creation, patched_param_creation
+ )
+ patched_source = patched_source.replace(
+ "def _init_sharded_param(",
+ "def patched_init_sharded_param(",
+ 1,
+ )
+
+ # Load necessary imports
+ module_name = FSDPParam.__module__
+ module = importlib.import_module(module_name)
+
+ items_to_import = []
+ for item in dir(module):
+ if item in patched_source:
+ items_to_import.append(item)
+
+ exec( # nosec B102
+ f"from {module_name} import ({', '.join(items_to_import)})",
+ globals(),
+ )
+ exec(patched_source, globals()) # nosec B102
+
+ # Replace the method
+ FSDPParam._init_sharded_param = patched_init_sharded_param
+ LOG.info("Successfully applied FSDP _init_sharded_param patch")
+ else:
+ LOG.warning("Could not find target code for _init_sharded_param patching")
+
+
+def apply_init_unsharded_param_patch():
+ """Apply patch to FSDPParam.init_unsharded_param to support Params4bit."""
+ from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
+
+ # Get original source
+ original_source = inspect.getsource(FSDPParam.init_unsharded_param)
+ original_source, _ = detab_code(original_source)
+
+ # Define the replacement
+ original_param_creation = """ self._unsharded_param = nn.Parameter(
+ unsharded_param, requires_grad=self.sharded_param.requires_grad
+ )"""
+
+ patched_param_creation = """ import bitsandbytes as bnb
+ local_tensor = self.sharded_param._local_tensor
+ if isinstance(local_tensor, bnb.nn.modules.Params4bit):
+ self._unsharded_param = bnb.nn.modules.Params4bit(
+ data=unsharded_param,
+ requires_grad=self.sharded_param.requires_grad,
+ quant_state=local_tensor.quant_state,
+ blocksize=local_tensor.blocksize,
+ compress_statistics=local_tensor.compress_statistics,
+ quant_type=local_tensor.quant_type,
+ quant_storage=local_tensor.quant_storage,
+ module=local_tensor.module,
+ bnb_quantized=local_tensor.bnb_quantized,
+ )
+ else:
+ self._unsharded_param = nn.Parameter(
+ unsharded_param, requires_grad=self.sharded_param.requires_grad
+ )"""
+
+ # Apply the replacement
+ if original_param_creation in original_source:
+ patched_source = original_source.replace(
+ original_param_creation, patched_param_creation
+ )
+ patched_source = patched_source.replace(
+ "def init_unsharded_param(",
+ "def patched_init_unsharded_param(",
+ 1,
+ )
+
+ # Load necessary imports
+ module_name = FSDPParam.__module__
+ module = importlib.import_module(module_name)
+
+ items_to_import = []
+ for item in dir(module):
+ if item in patched_source:
+ items_to_import.append(item)
+
+ exec( # nosec B102
+ f"from {module_name} import ({', '.join(items_to_import)})",
+ globals(),
+ )
+ exec(patched_source, globals()) # nosec B102
+
+ # Replace the method
+ FSDPParam.init_unsharded_param = patched_init_unsharded_param
+ LOG.info("Successfully applied FSDP init_unsharded_param patch")
+ else:
+ LOG.warning("Could not find target code for patching")
diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py b/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py
index 5d631776b..b58bbb67c 100644
--- a/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/__init__.py
@@ -5,7 +5,7 @@ from functools import partial
from packaging import version
-from axolotl.monkeypatch.gradient_checkpointing.offload_cpu import (
+from axolotl.monkeypatch.gradient_checkpointing.offload_cpu import ( # noqa: F401
CPU_Offloaded_Gradient_Checkpointer,
)
from axolotl.monkeypatch.gradient_checkpointing.offload_disk import (
@@ -25,9 +25,7 @@ else:
return False
-def hf_grad_checkpoint_offload_wrapper(
- decoder_layer, *args, use_reentrant=None
-): # pylint: disable=unused-argument
+def hf_grad_checkpoint_offload_wrapper(decoder_layer, *args, use_reentrant=None):
if uses_gc_layers(decoder_layer):
return CPU_Offloaded_Gradient_Checkpointer.apply(
decoder_layer,
@@ -44,9 +42,7 @@ def hf_grad_checkpoint_offload_wrapper(
)
-def hf_grad_checkpoint_disk_offload_wrapper(
- decoder_layer, *args, use_reentrant=None
-): # pylint: disable=unused-argument
+def hf_grad_checkpoint_disk_offload_wrapper(decoder_layer, *args, use_reentrant=None):
if uses_gc_layers(decoder_layer):
return Disco.apply(
decoder_layer,
diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
index bbb5ad40d..8d06f172d 100644
--- a/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_cpu.py
@@ -13,8 +13,17 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+
+import inspect
+
import torch
from packaging import version
+from torch.utils.checkpoint import (
+ set_device_states,
+)
+
+# support different pytorch versions
+has_device_type = "device_type" in inspect.signature(set_device_states).parameters
torch_version = version.parse(torch.__version__)
@@ -26,9 +35,7 @@ else:
torch_cuda_amp_custom_bwd = torch.amp.custom_bwd(device_type="cuda")
-class CPU_Offloaded_Gradient_Checkpointer( # pylint: disable=invalid-name
- torch.autograd.Function
-):
+class CPU_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
"""
Saves VRAM by smartly offloading to RAM.
Tiny hit to performance, since we mask the movement via non blocking calls.
@@ -57,6 +64,4 @@ class CPU_Offloaded_Gradient_Checkpointer( # pylint: disable=invalid-name
return (
None,
hidden_states.grad,
- ) + (
- None,
- ) * len(ctx.args)
+ ) + (None,) * len(ctx.args)
diff --git a/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py
index 792d3c6ef..220799fbf 100644
--- a/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py
+++ b/src/axolotl/monkeypatch/gradient_checkpointing/offload_disk.py
@@ -62,9 +62,9 @@ class DiskOffloadManager:
# Track tensor paths and their status
self.tensor_paths: deque = deque() # Ordered history of tensor paths (LIFO)
- self.file_locks: Dict[str, threading.Lock] = (
- {}
- ) # Maps file_path -> threading.Lock()
+ self.file_locks: Dict[
+ str, threading.Lock
+ ] = {} # Maps file_path -> threading.Lock()
# Maps file_path -> status ("saving", "ready", "prefetching", "loaded", "deleted")
self.file_status: Dict[str, str] = {}
@@ -236,7 +236,7 @@ class DiskOffloadManager:
self.tensor_paths.append(file_path)
# Acquire semaphore to limit concurrent save operations
- self.save_semaphore.acquire() # pylint: disable=consider-using-with
+ self.save_semaphore.acquire()
# Queue tensor for saving in background
self.save_queue.put((tensor.detach(), file_path))
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
index 70e36714c..3953cb138 100644
--- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -2,40 +2,28 @@
# copied from https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py
+import importlib.util
import warnings
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Tuple
import torch
-import torch.nn.functional as F
import transformers
from einops import rearrange
from flash_attn.bert_padding import pad_input, unpad_input
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.llama.modeling_llama import (
- LlamaAttention,
-)
-from transformers.models.llama.modeling_llama import (
- LlamaDecoderLayer as OriginalLlamaDecoderLayer,
-)
from transformers.models.llama.modeling_llama import (
LlamaMLP,
apply_rotary_pos_emb,
repeat_kv,
)
-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids, set_module_name
+from axolotl.monkeypatch.utils import set_module_name
from axolotl.utils.logging import get_logger
try:
- from flash_attn.flash_attn_interface import ( # pylint: disable=ungrouped-imports
- flash_attn_kvpacked_func,
- flash_attn_varlen_kvpacked_func,
+ from flash_attn.flash_attn_interface import (
flash_attn_varlen_qkvpacked_func,
)
except ImportError:
- from flash_attn.flash_attn_interface import (
- flash_attn_unpadded_kvpacked_func as flash_attn_varlen_kvpacked_func,
- )
from flash_attn.flash_attn_interface import (
flash_attn_unpadded_qkvpacked_func as flash_attn_varlen_qkvpacked_func,
)
@@ -45,12 +33,7 @@ LOG = get_logger(__name__)
def is_xformers_available() -> bool:
- try:
- import xformers # pylint: disable=unused-import # noqa: F401
-
- return True
- except ImportError:
- return False
+ return importlib.util.find_spec("xformers") is not None
def is_xformers_swiglu_available() -> bool:
@@ -82,19 +65,6 @@ def replace_llama_mlp_with_swiglu(model):
set_module_name(model, name, mlp)
-def replace_llama_qkv_with_fused(model):
- for name, module in model.named_modules():
- if isinstance(module, LlamaAttention):
- qkv = FusedAttention(
- module.config,
- module.q_proj,
- module.k_proj,
- module.v_proj,
- module.o_proj,
- )
- set_module_name(model, name, qkv)
-
-
def patch_fa_llama_cross_entropy():
LOG.info(
"patching transformers.loss.loss_utils.fixed_cross_entropy with flash_attn.ops.triton.cross_entropy"
@@ -109,7 +79,7 @@ def patch_fa_llama_cross_entropy():
num_items_in_batch: int = None,
ignore_index: int = -100,
**kwargs,
- ): # pylint: disable=unused-argument
+ ):
reduction = "sum" if num_items_in_batch is not None else "mean"
loss, _ = flash_attn_cross_entropy_loss(
source, target, ignore_index=ignore_index
@@ -142,28 +112,15 @@ def patch_llama_rms_norm():
def replace_llama_attn_with_flash_attn(
- packed: Optional[bool] = False,
cross_entropy: Optional[bool] = False,
rms_norm: Optional[bool] = False,
use_shifted_sparse_attn: Optional[bool] = False,
):
- transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = ( # pylint: disable=protected-access
- _prepare_decoder_attention_mask
- )
+ transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = _prepare_decoder_attention_mask
if use_shifted_sparse_attn:
transformers.models.llama.modeling_llama.LlamaAttention.forward = (
flashattn_forward_with_s2attn
)
- else:
- transformers.models.llama.modeling_llama.LlamaAttention.forward = (
- flashattn_forward
- )
-
- if packed:
- transformers.models.llama.modeling_llama.LlamaDecoderLayer = LlamaDecoderLayer
- transformers.models.llama.modeling_llama.LlamaModel.forward = (
- llama_model_forward
- )
# skip only if explicitly disabled
if cross_entropy:
@@ -174,49 +131,6 @@ def replace_llama_attn_with_flash_attn(
patch_llama_rms_norm()
-class FusedAttention(LlamaAttention):
- """
- Fused QKV Attention layer for incrementally improved training efficiency
- """
-
- def __init__(
- self,
- config,
- q: torch.nn.Linear, # pylint: disable=invalid-name
- k: torch.nn.Linear, # pylint: disable=invalid-name
- v: torch.nn.Linear, # pylint: disable=invalid-name
- o: torch.nn.Linear, # pylint: disable=invalid-name
- ):
- super().__init__(config)
- self.config = config
- self.init_device = next(iter(q.state_dict().values())).device
-
- # define equivalent fused qkv projection
- self.out_features: List[int] = [q.out_features, k.out_features, v.out_features]
- self.qkv_proj = torch.nn.Linear(
- q.in_features, sum(self.out_features), device=self.init_device, bias=False
- )
- self.o_proj = o
-
- # overwrite initialized weights with pretrained weights
- self.qkv_proj.weight.data = torch.cat(
- (q.weight.data, k.weight.data, v.weight.data), dim=0
- )
-
- def _post_training(self, model, name):
- q_proj, k_proj, v_proj = torch.split(
- self.qkv_proj.weight.data, self.out_features, dim=0
- )
-
- new_attn = LlamaAttention(self.config)
- new_attn.q_proj.weight.data = q_proj
- new_attn.k_proj.weight.data = k_proj
- new_attn.v_proj.weight.data = v_proj
- new_attn.o_proj.weight.data = self.o_proj.weight.data
-
- set_module_name(model, name, new_attn)
-
-
# Disable the transformation of the attention mask in LlamaModel as the flash attention
# requires the attention mask to be the same as the key_padding_mask
def _prepare_decoder_attention_mask(
@@ -225,7 +139,7 @@ def _prepare_decoder_attention_mask(
input_shape,
inputs_embeds,
past_key_values_length,
-): # pylint: disable=unused-argument
+):
# [bsz, seq_len]
return attention_mask
@@ -241,9 +155,9 @@ def flashattn_forward_with_s2attn(
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: bool = False,
use_cache: bool = False,
- padding_mask: Optional[torch.LongTensor] = None, # pylint: disable=unused-argument
- cu_seqlens: Optional[torch.Tensor] = None, # pylint: disable=unused-argument
- max_seqlen: Optional[torch.Tensor] = None, # pylint: disable=unused-argument
+ padding_mask: Optional[torch.LongTensor] = None,
+ cu_seqlens: Optional[torch.Tensor] = None,
+ max_seqlen: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
"""Input shape: Batch x Time x Channel
@@ -256,7 +170,8 @@ def flashattn_forward_with_s2attn(
"""
if output_attentions:
warnings.warn(
- "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+ "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead.",
+ stacklevel=2,
)
bsz, q_len, _ = hidden_states.size()
@@ -278,7 +193,6 @@ def flashattn_forward_with_s2attn(
)
# [bsz, q_len, nh, hd]
# [bsz, nh, q_len, hd]
- # pylint: disable=duplicate-code
cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
query_states, key_states = apply_rotary_pos_emb(
@@ -324,9 +238,7 @@ def flashattn_forward_with_s2attn(
.permute(0, 3, 1, 2, 4, 5)
.reshape(bsz * 2, q_len, 3, self.num_heads // 2, self.head_dim)
)
- x = rearrange( # pylint: disable=invalid-name
- qkv, "b s three h d -> b s (three h d)"
- )
+ x = rearrange(qkv, "b s three h d -> b s (three h d)")
x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)
cu_q_len_tmp = torch.arange(
0, max_s, group_size, device=key_padding_mask.device, dtype=cu_q_lens.dtype
@@ -355,576 +267,3 @@ def flashattn_forward_with_s2attn(
.reshape(bsz, q_len, nheads, self.head_dim)
)
return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, past_key_value
-
-
-def flashattn_forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.Tensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- padding_mask: Optional[torch.LongTensor] = None, # pylint: disable=unused-argument
- cu_seqlens: Optional[torch.Tensor] = None,
- max_seqlen: Optional[torch.Tensor] = None,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- """Input shape: Batch x Time x Channel
-
- attention_mask: [bsz, q_len]
- """
- # pylint: disable=duplicate-code
- bsz, q_len, _ = hidden_states.size()
-
- if not hasattr(self, "pretraining_tp"):
- self.pretraining_tp = 1
-
- if self.pretraining_tp > 1:
- key_value_slicing = (
- self.num_key_value_heads * self.head_dim
- ) // self.pretraining_tp
- query_slices = self.q_proj.weight.split(
- (self.num_heads * self.head_dim) // self.pretraining_tp, dim=0
- )
- key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
- value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
-
- query_states = [
- F.linear(hidden_states, query_slices[i]) for i in range(self.pretraining_tp)
- ]
- query_states = torch.cat(query_states, dim=-1)
-
- key_states = [
- F.linear(hidden_states, key_slices[i]) for i in range(self.pretraining_tp)
- ]
- key_states = torch.cat(key_states, dim=-1)
-
- value_states = [
- F.linear(hidden_states, value_slices[i]) for i in range(self.pretraining_tp)
- ]
- value_states = torch.cat(value_states, dim=-1)
-
- else:
- if isinstance(self, FusedAttention):
- query_states, key_states, value_states = self.qkv_proj(hidden_states).split(
- self.out_features, dim=-1
- )
- else:
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(
- bsz, q_len, self.num_heads, self.head_dim
- ).transpose(1, 2)
- key_states = key_states.view(
- bsz, q_len, self.num_key_value_heads, self.head_dim
- ).transpose(1, 2)
- value_states = value_states.view(
- bsz, q_len, self.num_key_value_heads, self.head_dim
- ).transpose(1, 2)
- # [bsz, q_len, nh, hd]
- # [bsz, nh, q_len, hd]
-
- cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
- query_states, key_states = apply_rotary_pos_emb(
- query_states, key_states, cos, sin, position_ids
- )
- # [bsz, nh, t, hd]
-
- if past_key_value is not None:
- # reuse k, v, self_attention
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
- past_key_value = (key_states, value_states) if use_cache else None
-
- # repeat k/v heads if n_kv_heads < n_heads
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- if output_attentions:
- warnings.warn(
- "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
- )
-
- #
- # flash-attn v2 start
- #
-
- if self.training:
- # during training q,k,v always have same seqlen
- assert key_states.shape == query_states.shape
- is_causal = True
- else:
- # turn off FA causal mask after first inference autoregressive iteration
- # only on first autoregressive step q,k,v have same seqlen
- is_causal = key_states.shape == query_states.shape
-
- dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
- if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
- # special handling using sample packing
- qkv = torch.stack(
- [query_states, key_states, value_states], dim=2
- ) # [bsz, nh, 3, q_len, hd]
- qkv = qkv.transpose(1, 3) # [bsz, q_len, 3, nh, hd]
- qkv = rearrange(qkv, "b s ... -> (b s) ...")
-
- output = flash_attn_varlen_qkvpacked_func(
- qkv,
- cu_seqlens,
- max_seqlen,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=True,
- )
- output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
- elif query_states.shape == key_states.shape:
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
- qkv_unpad, cu_seqlens_q, max_seqlen_q, _, output_pad_fn = generate_qkv(
- query_states,
- key_states,
- value_states,
- qkvpacked=True,
- # We have disabled _prepare_decoder_attention_mask in LlamaModel
- # the attention_mask should be the same as the key_padding_mask
- key_padding_mask=attention_mask,
- query_padding_mask=(
- attention_mask[:, -query_states.size(1) :]
- if attention_mask is not None
- else None
- ),
- )
- output_unpad = flash_attn_varlen_qkvpacked_func(
- qkv_unpad,
- cu_seqlens_q,
- max_seqlen_q,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=is_causal,
- )
- output = output_pad_fn(output_unpad)
- else:
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
- if attention_mask is None or attention_mask.all().item():
- output = flash_attn_kvpacked_func(
- query_states,
- torch.stack([key_states, value_states], 2),
- dropout_p=dropout_rate,
- causal=is_causal,
- )
- else:
- ( # pylint: disable=unbalanced-tuple-unpacking
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- _,
- _,
- output_pad_fn,
- ) = generate_qkv(
- query_states,
- key_states,
- value_states,
- kvpacked=True,
- key_padding_mask=attention_mask,
- query_padding_mask=(
- attention_mask[:, -query_states.size(1) :]
- if attention_mask is not None
- else None
- ),
- )
- if q_unpad.dtype != kv_unpad.dtype:
- kv_unpad = kv_unpad.to(q_unpad.dtype)
- output_unpad = flash_attn_varlen_kvpacked_func(
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=is_causal,
- )
- output = output_pad_fn(output_unpad)
-
- attn_output = output
- if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
- attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
-
- #
- # flash-attn v2 end
- #
-
- if self.pretraining_tp > 1:
- attn_output = attn_output.split(self.hidden_size // self.pretraining_tp, dim=2)
- o_proj_slices = self.o_proj.weight.split(
- self.hidden_size // self.pretraining_tp, dim=1
- )
- attn_output = sum(
- F.linear(attn_output[i], o_proj_slices[i])
- for i in range(self.pretraining_tp)
- )
- else:
- attn_output = self.o_proj(attn_output)
-
- return attn_output, None, past_key_value
-
-
-# based on https://github.com/Dao-AILab/flash-attention/blob/364a5b/tests/test_flash_attn.py#L38
-def generate_qkv(
- q,
- k,
- v,
- query_padding_mask=None,
- key_padding_mask=None,
- kvpacked=False,
- qkvpacked=False,
-): # pylint: disable=invalid-name,unnecessary-lambda-assignment
- """
- Arguments:
- q: (batch_size, seqlen_q, nheads, d)
- k: (batch_size, seqlen_k, nheads_k, d)
- v: (batch_size, seqlen_k, nheads_k, d)
- query_padding_mask: (batch_size, seqlen), bool
- key_padding_mask: (batch_size, seqlen), bool
- """
- assert not (kvpacked and qkvpacked)
- batch_size, seqlen_q, nheads, d = q.shape
- _, seqlen_k, nheads_k, _ = k.shape
- assert k.shape == (batch_size, seqlen_k, nheads_k, d)
- assert v.shape == (batch_size, seqlen_k, nheads_k, d)
-
- if query_padding_mask is not None:
- q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
- q, query_padding_mask
- )
-
- def output_pad_fn(output_unpad):
- return pad_input( # noqa: E731
- output_unpad, indices_q, batch_size, seqlen_q
- )
-
- else:
- q_unpad = rearrange(q, "b s h d -> (b s) h d")
- cu_seqlens_q = torch.arange(
- 0,
- (batch_size + 1) * seqlen_q,
- step=seqlen_q,
- dtype=torch.int32,
- device=q_unpad.device,
- )
- max_seqlen_q = seqlen_q
-
- def output_pad_fn(output_unpad):
- return rearrange( # noqa: E731
- output_unpad, "(b s) h d -> b s h d", b=batch_size
- )
-
- if key_padding_mask is not None:
- k_unpad, _, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
- v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
- else:
- k_unpad = rearrange(k, "b s h d -> (b s) h d")
- v_unpad = rearrange(v, "b s h d -> (b s) h d")
- cu_seqlens_k = torch.arange(
- 0,
- (batch_size + 1) * seqlen_k,
- step=seqlen_k,
- dtype=torch.int32,
- device=k_unpad.device,
- )
- max_seqlen_k = seqlen_k
-
- if qkvpacked:
- assert nheads == nheads_k
- qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
- qkv = torch.stack([q, k, v], dim=2)
- return (qkv_unpad, cu_seqlens_q, max_seqlen_q, qkv, output_pad_fn)
-
- if kvpacked:
- kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
- kv = torch.stack([k, v], dim=2)
- return (
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- q,
- kv,
- output_pad_fn,
- )
-
- return (
- q_unpad,
- k_unpad,
- v_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- q,
- k,
- v,
- output_pad_fn,
- )
-
-
-def llama_model_forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[ # pylint: disable=unused-argument
- torch.LongTensor
- ] = None,
-) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # retrieve input_ids and inputs_embeds
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError(
- "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
- )
- if input_ids is not None:
- batch_size, seq_length = input_ids.shape
- elif inputs_embeds is not None:
- batch_size, seq_length, _ = inputs_embeds.shape
- else:
- raise ValueError(
- "You have to specify either decoder_input_ids or decoder_inputs_embeds"
- )
-
- seq_length_with_past = seq_length
- past_key_values_length = 0
-
- if past_key_values is not None:
- past_key_values_length = past_key_values[0][0].shape[2]
- seq_length_with_past = seq_length_with_past + past_key_values_length
-
- cu_seqlens = None
- max_seqlen = None
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(
- past_key_values_length,
- seq_length + past_key_values_length,
- dtype=torch.long,
- device=device,
- )
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
- else:
- position_ids = position_ids.view(-1, seq_length).long()
- cu_seqlens, max_seqlen = get_cu_seqlens_from_pos_ids(position_ids)
- cu_seqlens = cu_seqlens.squeeze()
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_tokens(input_ids)
- # embed positions
- if attention_mask is None:
- attention_mask = torch.ones(
- (batch_size, seq_length_with_past),
- dtype=torch.bool,
- device=inputs_embeds.device,
- )
- padding_mask = None
- else:
- if 0 in attention_mask:
- padding_mask = attention_mask
- else:
- padding_mask = None
-
- attention_mask = (
- self._prepare_decoder_attention_mask( # pylint: disable=protected-access
- attention_mask,
- (batch_size, seq_length),
- inputs_embeds,
- past_key_values_length,
- )
- )
-
- hidden_states = inputs_embeds
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- transformers.logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- next_decoder_cache = () if use_cache else None
-
- for idx, decoder_layer in enumerate(self.layers):
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- past_key_value = past_key_values[idx] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
-
- def create_custom_forward(module):
- def custom_forward(*inputs):
- # None for past_key_value
- return module(
- *inputs,
- )
-
- return custom_forward
-
- layer_outputs = torch.utils.checkpoint.checkpoint(
- create_custom_forward(decoder_layer),
- hidden_states,
- attention_mask,
- position_ids,
- past_key_value,
- output_attentions,
- None,
- padding_mask,
- cu_seqlens,
- max_seqlen,
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- )
-
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- hidden_states = self.norm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = next_decoder_cache if use_cache else None
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
- if v is not None
- )
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- )
-
-
-class LlamaDecoderLayer(OriginalLlamaDecoderLayer):
- """
- patched version of LlamaDecoderLayer to pass through the precalculated cu_seqlens
- """
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- padding_mask: Optional[torch.LongTensor] = None,
- cu_seqlens: Optional[torch.Tensor] = None,
- max_seqlen: Optional[torch.Tensor] = None,
- ) -> Tuple[
- torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
- ]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
- cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
- """
-
- residual = hidden_states
-
- hidden_states = self.input_layernorm(hidden_states)
-
- # Self Attention
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- padding_mask=padding_mask,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- )
- hidden_states = residual + hidden_states
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.post_attention_layernorm(hidden_states)
- hidden_states = self.mlp(hidden_states)
- hidden_states = residual + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- return outputs
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py b/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
index 28223eee3..332242e2c 100644
--- a/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_xformers.py
@@ -32,10 +32,9 @@ def xformers_forward(
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: bool = False,
use_cache: bool = False,
- padding_mask: Optional[torch.LongTensor] = None, # pylint: disable=unused-argument
- **kwargs, # pylint: disable=unused-argument
+ padding_mask: Optional[torch.LongTensor] = None,
+ **kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- # pylint: disable=duplicate-code
bsz, q_len, _ = hidden_states.size()
if not hasattr(self, "pretraining_tp"):
@@ -102,7 +101,8 @@ def xformers_forward(
if output_attentions:
warnings.warn(
- "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead."
+ "Output attentions is not supported for patched `LlamaAttention`, returning `None` instead.",
+ stacklevel=2,
)
#
diff --git a/src/axolotl/monkeypatch/llama_expand_mask.py b/src/axolotl/monkeypatch/llama_expand_mask.py
index 0277c212a..5cfb7818e 100644
--- a/src/axolotl/monkeypatch/llama_expand_mask.py
+++ b/src/axolotl/monkeypatch/llama_expand_mask.py
@@ -21,6 +21,4 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]
def hijack_expand_mask():
import transformers
- transformers.models.llama.modeling_llama._expand_mask = ( # pylint: disable=protected-access
- _expand_mask
- )
+ transformers.models.llama.modeling_llama._expand_mask = _expand_mask
diff --git a/src/axolotl/monkeypatch/llama_patch_multipack.py b/src/axolotl/monkeypatch/llama_patch_multipack.py
index cfd525367..8d234881f 100644
--- a/src/axolotl/monkeypatch/llama_patch_multipack.py
+++ b/src/axolotl/monkeypatch/llama_patch_multipack.py
@@ -12,15 +12,15 @@ def hijack_llama_prepare_4d_mask():
from transformers import modeling_attn_mask_utils
from transformers.models.llama import modeling_llama
- modeling_llama._prepare_4d_causal_attention_mask_for_sdpa = ( # pylint: disable=protected-access
+ modeling_llama._prepare_4d_causal_attention_mask_for_sdpa = (
patched_prepare_4d_causal_attention_mask_for_sdpa
)
- modeling_attn_mask_utils._prepare_4d_causal_attention_mask_for_sdpa = ( # pylint: disable=protected-access
+ modeling_attn_mask_utils._prepare_4d_causal_attention_mask_for_sdpa = (
patched_prepare_4d_causal_attention_mask_for_sdpa
)
- modeling_llama._prepare_4d_causal_attention_mask = ( # pylint: disable=protected-access
+ modeling_llama._prepare_4d_causal_attention_mask = (
patched_prepare_4d_causal_attention_mask
)
- modeling_attn_mask_utils._prepare_4d_causal_attention_mask = ( # pylint: disable=protected-access
+ modeling_attn_mask_utils._prepare_4d_causal_attention_mask = (
patched_prepare_4d_causal_attention_mask
)
diff --git a/src/axolotl/monkeypatch/lora_kernels.py b/src/axolotl/monkeypatch/lora_kernels.py
index a7875eefe..e845dc6ce 100644
--- a/src/axolotl/monkeypatch/lora_kernels.py
+++ b/src/axolotl/monkeypatch/lora_kernels.py
@@ -18,6 +18,7 @@ from axolotl.kernels.lora import (
apply_lora_qkv,
)
from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
from axolotl.utils.dict import DictDefault
from axolotl.utils.logging import get_logger
@@ -29,48 +30,36 @@ QKV_PATCHES = [
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-""".lstrip(
- "\n"
- ),
+""".lstrip("\n"),
"""
query_states, key_states, value_states = self.apply_qkv(hidden_states)
query_states = query_states.view(hidden_shape).transpose(1, 2)
key_states = key_states.view(hidden_shape).transpose(1, 2)
value_states = value_states.view(hidden_shape).transpose(1, 2)
-""".lstrip(
- "\n"
- ),
+""".lstrip("\n"),
),
(
"""
query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-""".lstrip(
- "\n"
- ),
+""".lstrip("\n"),
"""
query_states, key_states, value_states = self.apply_qkv(hidden_states)
query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
key_states = self.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
value_states = value_states.view(hidden_shape).transpose(1, 2)
-""".lstrip(
- "\n"
- ),
+""".lstrip("\n"),
),
]
ORIGINAL_O_CODE = """
attn_output = self.o_proj(attn_output)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
PATCHED_O_CODE = """
attn_output = self.apply_o(attn_output)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
SUPPORTED_ACTIVATIONS = ["silu", "gelu"]
APPLY_FN_MAPPING = {
@@ -145,12 +134,30 @@ def get_attention_cls_from_config(cfg: DictDefault) -> Type[nn.Module]:
return Qwen2Attention
+ if model_type == "mllama":
+ from transformers.models.mllama.modeling_mllama import MllamaTextSelfAttention
+
+ return MllamaTextSelfAttention
+
+ if model_type == "llama4":
+ from transformers.models.llama4.modeling_llama4 import Llama4TextAttention
+
+ return Llama4TextAttention
+
+ if model_type == "mistral3":
+ from transformers.models.mistral.modeling_mistral import MistralAttention
+
+ return MistralAttention
+
+ if model_type == "gemma3_text":
+ from transformers.models.gemma3.modeling_gemma3 import Gemma3Attention
+
+ return Gemma3Attention
+
try:
# Dynamically import the module and attention class
module_path = f"transformers.models.{model_type}.modeling_{model_type}"
- model_cls_prefix = "".join(
- [part.capitalize() for part in model_type.split("_")]
- )
+ model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
module = __import__(module_path, fromlist=[f"{model_cls_prefix}Attention"])
attention_cls = getattr(module, f"{model_cls_prefix}Attention")
@@ -162,7 +169,6 @@ def get_attention_cls_from_config(cfg: DictDefault) -> Type[nn.Module]:
) from e
-# pylint: disable=protected-access
def patch_self_attn_lora(cfg: DictDefault):
"""
Given an `axolotl` config, this method patches the inferred attention class forward
@@ -189,9 +195,9 @@ def patch_self_attn_lora(cfg: DictDefault):
attention_cls._original_forward = self_attn_forward
self_attn_forward, _ = detab_code(self_attn_forward)
- assert any(
- qkv_options[0] in self_attn_forward for qkv_options in QKV_PATCHES
- ), "Original QKV code not found"
+ assert any(qkv_options[0] in self_attn_forward for qkv_options in QKV_PATCHES), (
+ "Original QKV code not found"
+ )
assert ORIGINAL_O_CODE in self_attn_forward, "Original O code not found"
for qkv_orig, qkv_patched in QKV_PATCHES:
@@ -217,16 +223,14 @@ def patch_self_attn_lora(cfg: DictDefault):
if item in self_attn_forward:
items_to_import.append(item)
- exec( # pylint: disable=exec-used # nosec B102
+ exec(
f"from {module_name} import ({', '.join(items_to_import)})",
globals(),
)
- exec(self_attn_forward, globals()) # pylint: disable=exec-used # nosec B102
+ exec(self_attn_forward, globals())
LOG.info(f"Patched attention class with LoRA optims: {attention_cls.__name__}")
- attention_cls.forward = (
- axolotl_attn_forward # pylint: disable=undefined-variable # noqa: F821
- )
+ attention_cls.forward = axolotl_attn_forward
def find_self_attn_in_layer(
@@ -263,12 +267,39 @@ def find_mlp_in_layer(
layer.feedforward.experts.gate_projs,
layer.feedforward.experts.up_projs,
layer.feedforward.experts.down_projs,
+ strict=False,
):
- yield gate_proj, up_proj, down_proj, FakeMLP(
- gate_proj, up_proj, down_proj
+ yield (
+ gate_proj,
+ up_proj,
+ down_proj,
+ FakeMLP(gate_proj, up_proj, down_proj),
)
+def get_layers(model: PeftModelForCausalLM) -> list[nn.Module]:
+ """
+ Get the layers of the model. Handles text-only and multimodal models.
+
+ Args:
+ model: A PEFT model.
+
+ Returns:
+ A list of layers.
+ """
+ pretrained_model = model.model
+
+ # check for multimodal models first
+ if hasattr(pretrained_model, "language_model"):
+ return pretrained_model.language_model.layers
+ if hasattr(pretrained_model, "model"):
+ return pretrained_model.model.layers
+
+ raise NotImplementedError(
+ f"Model type {model.config.model_type} is not supported yet. Please create an Issue."
+ )
+
+
def apply_lora_kernel_patches(
model: PeftModelForCausalLM, cfg: DictDefault
) -> PeftModelForCausalLM:
@@ -300,9 +331,9 @@ def apply_lora_kernel_patches(
# Get active LoRA adapter config
if hasattr(model, "active_adapters"):
- assert (
- len(model.active_adapters) == 1
- ), "Axolotl currently does not support LoRA Triton kernels for multiple adapters"
+ assert len(model.active_adapters) == 1, (
+ "Axolotl currently does not support LoRA Triton kernels for multiple adapters"
+ )
active_adapter = model.active_adapters[0]
else:
active_adapter = model.active_adapter
@@ -340,17 +371,7 @@ def apply_lora_kernel_patches(
if activation not in SUPPORTED_ACTIVATIONS:
raise NotImplementedError(f"Activation {activation} is not supported")
- layers = []
- # check for multimodal models first
- pretrained_model = model.model
- if hasattr(pretrained_model, "language_model"):
- layers = pretrained_model.language_model.layers
- elif hasattr(pretrained_model, "model"):
- layers = pretrained_model.model.layers
- else:
- raise NotImplementedError(
- f"Model type {model.config.model_type} is not supported yet. Please create an Issue."
- )
+ layers = get_layers(model)
# Patch each layer
for layer in layers:
@@ -368,7 +389,6 @@ def apply_lora_kernel_patches(
]
can_patch_qkv = all(
hasattr(module, "lora_A")
- and getattr(module, "base_layer", module).bias is None
and len(getattr(module, "lora_magnitude_vector", []) or []) == 0
for module in layer_modules
)
@@ -378,7 +398,8 @@ def apply_lora_kernel_patches(
self_attn.apply_qkv = types.MethodType(apply_lora_qkv, self_attn)
else:
LOG.warning_once(
- "Cannot patch some attention QKV projections - requires LoRA adapters with no bias"
+ "Cannot patch some attention QKV projections - requires LoRA "
+ "adapters and no lora_magnitude_vector (DoRA)"
)
if cfg.lora_o_kernel:
# Output patching
@@ -387,7 +408,6 @@ def apply_lora_kernel_patches(
]
can_patch_o = all(
hasattr(module, "lora_A")
- and getattr(module, "base_layer", module).bias is None
and len(getattr(module, "lora_magnitude_vector", []) or []) == 0
for module in layer_modules
)
@@ -396,14 +416,14 @@ def apply_lora_kernel_patches(
self_attn.apply_o = types.MethodType(apply_lora_o, self_attn)
else:
LOG.warning_once(
- "Cannot patch some attention output projection - requires LoRA adapters with no bias"
+ "Cannot patch some attention output projection - requires LoRA "
+ "adapters and no lora_magnitude_vector (DoRA)"
)
for gate_proj, up_proj, down_proj, mlp in find_mlp_in_layer(layer):
if cfg.lora_mlp_kernel:
# MLP patching
can_patch_mlp = all(
hasattr(proj, "lora_A")
- and getattr(proj, "base_layer", proj).bias is None
and len(getattr(proj, "lora_magnitude_vector", []) or []) == 0
for proj in (gate_proj, up_proj, down_proj)
)
@@ -413,7 +433,8 @@ def apply_lora_kernel_patches(
layer.mlp.forward = types.MethodType(apply_fn, mlp)
else:
LOG.warning_once(
- "Cannot patch some MLP layers - requires LoRA adapters with no bias"
+ "Cannot patch some MLP layers - requires LoRA adapters and no "
+ "lora_magnitude_vector (DoRA)"
)
LOG.setLevel(original_level)
diff --git a/src/axolotl/monkeypatch/loss/__init__.py b/src/axolotl/monkeypatch/loss/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/loss/chunked.py b/src/axolotl/monkeypatch/loss/chunked.py
new file mode 100644
index 000000000..26a52f898
--- /dev/null
+++ b/src/axolotl/monkeypatch/loss/chunked.py
@@ -0,0 +1,134 @@
+"""
+chunked ce loss
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+# copied and modified from torchtune.modules.loss.CEWithChunkedOutputLoss
+class CEWithChunkedOutputLoss(torch.nn.Module):
+ """
+ Cross-entropy with chunked outputs that saves memory by only upcasting one chunk at a time.
+
+ For more details, please refer to: https://github.com/pytorch/torchtune/pull/1390
+ """
+
+ def __init__(self, num_output_chunks: int = 8, ignore_index: int = -100):
+ super().__init__()
+ self.num_output_chunks = num_output_chunks
+ self.ignore_index = ignore_index
+
+ def compute_cross_entropy(
+ self,
+ logits: torch.Tensor,
+ labels: torch.Tensor,
+ normalize: bool = True,
+ ) -> torch.Tensor:
+ """
+ Upcast logits to fp32 and compute cross entropy loss.
+ """
+ return F.cross_entropy(
+ logits.float(), labels, ignore_index=self.ignore_index, reduction="sum"
+ )
+
+ def forward(
+ self, logits: List[torch.Tensor], labels: torch.Tensor, reduction="sum"
+ ) -> torch.Tensor:
+ """
+ Args:
+ logits (List[torch.Tensor]): List of chunked logits of length
+ ``self.num_output_chunks``, where each chunk has shape
+ ``(batch_size, num_tokens / num_output_chunks, vocab_size)``.
+ labels (torch.Tensor): Ground truth labels of shape ``(batch_size, num_tokens)``.
+ reduction (str): The reduction to apply to the output.
+
+ Returns:
+ torch.Tensor: Cross entropy loss of shape (1,).
+ """
+
+ total_elements = (labels != self.ignore_index).sum()
+
+ # chunk and reshape labels (bsz, num_tokens, vocab) -> [(bsz*num_tokens/num_chunks, vocab)]
+ labels = [
+ target_chunk.reshape(-1)
+ for target_chunk in labels.chunk(self.num_output_chunks, dim=1)
+ ]
+ # reshape logits [(bsz, num_tokens/num_chunks, vocab)] -> [(bsz*num_tokens/num_chunks, vocab)]
+ logits = [
+ logit_chunk.reshape(-1, logit_chunk.size(-1)) for logit_chunk in logits
+ ]
+
+ # compute one chunk at a time
+ total_loss = 0.0
+ for logits_chunk, labels_chunk in zip(logits, labels, strict=False):
+ total_loss += self.compute_cross_entropy(logits_chunk, labels_chunk)
+
+ if reduction == "sum":
+ return total_loss
+ return total_loss / total_elements
+
+
+def _build_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
+ loss_fn_ce = CEWithChunkedOutputLoss(num_output_chunks, ignore_index)
+ loss_fn_ce.compute_cross_entropy = torch.compile(
+ loss_fn_ce.compute_cross_entropy, backend="inductor"
+ )
+ return loss_fn_ce
+
+
+def get_causal_lm_loss(num_output_chunks: int = 8, ignore_index: int = -100):
+ loss_fn_ce = _build_chunked_ce_loss_fn(num_output_chunks, ignore_index)
+
+ def chunked_fix_cross_entropy(
+ source,
+ target,
+ num_items_in_batch: int = None,
+ ignore_index: int = -100,
+ **kwargs,
+ ):
+ reduction = "sum" if num_items_in_batch is not None else "mean"
+ logit_chunks = [
+ chunk for chunk in source.chunk(loss_fn_ce.num_output_chunks, dim=1)
+ ]
+ loss = loss_fn_ce(logit_chunks, target, reduction=reduction)
+ if reduction == "sum":
+ loss = loss / num_items_in_batch
+ return loss
+
+ def for_causal_lm_chunked_loss(
+ logits,
+ labels,
+ vocab_size: int = None,
+ num_items_in_batch: Optional[int] = None,
+ ignore_index: int = -100,
+ shift_labels: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ # skip the upcast to float since we handle that in the chunking loss
+ if shift_labels is None:
+ # Shift so that tokens < n predict n
+ labels = F.pad(labels, (0, 1), value=ignore_index)
+ shift_labels = labels[..., 1:].contiguous()
+
+ # Skip Flattening the tokens
+ # Enable model parallelism
+ shift_labels = shift_labels.to(logits.device)
+ loss = chunked_fix_cross_entropy(
+ logits, shift_labels, num_items_in_batch, ignore_index, **kwargs
+ )
+ return loss
+
+ return for_causal_lm_chunked_loss
+
+
+def patch_chunked_ce_loss_fn(num_output_chunks: int = 8, ignore_index: int = -100):
+ import transformers.loss.loss_utils
+
+ for_causal_lm_chunked_loss = get_causal_lm_loss(num_output_chunks, ignore_index)
+ transformers.loss.loss_utils.ForCausalLMLoss = for_causal_lm_chunked_loss
+ transformers.loss.loss_utils.LOSS_MAPPING["ForCausalLM"] = (
+ for_causal_lm_chunked_loss
+ )
diff --git a/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py b/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
index 3fc22917f..0994da91c 100644
--- a/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/mistral_attn_hijack_flash.py
@@ -1,55 +1,14 @@
"""Flash attention monkey patch for mistral model"""
-# pylint: disable=duplicate-code
-
from functools import partial
-from typing import List, Optional, Tuple, Union
-import torch
import transformers
-from einops import rearrange
-from flash_attn.bert_padding import pad_input, unpad_input
-from flash_attn.flash_attn_interface import ( # pylint: disable=ungrouped-imports
- flash_attn_kvpacked_func,
- flash_attn_varlen_kvpacked_func,
- flash_attn_varlen_qkvpacked_func,
-)
-from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.models.mistral.modeling_mistral import (
- MistralAttention as OriginalMistralAttention,
-)
-from transformers.models.mistral.modeling_mistral import (
- MistralDecoderLayer as OriginalMistralDecoderLayer,
-)
-from transformers.models.mistral.modeling_mistral import (
- apply_rotary_pos_emb,
- repeat_kv,
-)
-from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
from axolotl.utils.logging import get_logger
LOG = get_logger(__name__)
-def replace_mistral_attn_with_flash_attn(
- packed: Optional[bool] = False,
-):
- transformers.models.mistral.modeling_mistral.MistralModel._prepare_decoder_attention_mask = ( # pylint: disable=protected-access
- _prepare_decoder_attention_mask
- )
- transformers.models.mistral.modeling_mistral.MistralAttention.forward = (
- flashattn_forward
- )
- if packed:
- transformers.models.mistral.modeling_mistral.MistralDecoderLayer = (
- MistralDecoderLayer
- )
- transformers.models.mistral.modeling_mistral.MistralModel.forward = (
- mistral_model_forward
- )
-
-
def patch_mistral_cross_entropy():
from flash_attn.losses.cross_entropy import CrossEntropyLoss
@@ -57,604 +16,3 @@ def patch_mistral_cross_entropy():
transformers.models.mistral.modeling_mistral.CrossEntropyLoss = partial(
CrossEntropyLoss, inplace_backward=True
)
-
-
-@torch.jit.script
-def _make_sliding_window_causal_mask(
- bsz: int,
- tgt_len: int,
- dtype: torch.dtype,
- device: torch.device,
- past_key_values_length: int = 0,
- sliding_window: int = 4096,
-):
- """
- Make causal mask used for sliding window attention
- """
- tensor = torch.full(
- (tgt_len, tgt_len),
- fill_value=1,
- device=device,
- )
- mask = torch.tril(tensor, diagonal=0)
- # make the mask banded to account for sliding window
- # NOTE: HF implementation is wrong as of 14-10-2023 for torch.triu, needs +1
- mask = torch.triu(mask, diagonal=-sliding_window + 1)
- mask = torch.log(mask).to(dtype)
-
- if past_key_values_length > 0:
- mask = torch.cat(
- [
- torch.zeros(
- tgt_len, past_key_values_length, dtype=dtype, device=device
- ),
- mask,
- ],
- dim=-1,
- )
- return mask[None, None, :, :].expand(
- bsz, 1, tgt_len, tgt_len + past_key_values_length
- )
-
-
-# Disable the transformation of the attention mask in LlamaModel as the flash attention
-# requires the attention mask to be the same as the key_padding_mask
-def _prepare_decoder_attention_mask(
- self,
- attention_mask,
- input_shape,
- inputs_embeds,
- past_key_values_length,
- sliding_window,
-): # pylint: disable=unused-argument
- # [bsz, seq_len]
- if attention_mask is None or sliding_window is None:
- return attention_mask
-
- # NOTE: attention mask and sliding masks are only broadcastable in certain scenarios.
- # Without attention_mask.shape[0] == 1, error will trigger after eval loss but only when wandb is enabled.
- if input_shape[-1] > 1 and attention_mask.shape[0] == 1:
- sliding_window_mask = _make_sliding_window_causal_mask(
- bsz=input_shape[0],
- tgt_len=input_shape[1],
- dtype=inputs_embeds.dtype,
- device=inputs_embeds.device,
- past_key_values_length=past_key_values_length,
- sliding_window=sliding_window,
- )
- attention_mask = attention_mask + sliding_window_mask
- else:
- LOG.info("skipping sliding window mask, not broadcastable with attention mask")
-
- return attention_mask
-
-
-def flashattn_forward(
- self: OriginalMistralAttention,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: bool = False,
- use_cache: bool = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- max_seqlen: Optional[torch.Tensor] = None,
-) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
- bsz, q_len, _ = hidden_states.size()
-
- query_states = self.q_proj(hidden_states)
- key_states = self.k_proj(hidden_states)
- value_states = self.v_proj(hidden_states)
-
- query_states = query_states.view(
- bsz, q_len, self.num_heads, self.head_dim
- ).transpose(1, 2)
- key_states = key_states.view(
- bsz, q_len, self.num_key_value_heads, self.head_dim
- ).transpose(1, 2)
- value_states = value_states.view(
- bsz, q_len, self.num_key_value_heads, self.head_dim
- ).transpose(1, 2)
-
- kv_seq_len = key_states.shape[-2]
- if past_key_value is not None:
- kv_seq_len += past_key_value[0].shape[-2]
- cos, sin = self.rotary_emb(value_states, position_ids=position_ids)
- query_states, key_states = apply_rotary_pos_emb(
- query_states, key_states, cos, sin, position_ids
- )
-
- use_sliding_windows = (
- getattr(self.config, "sliding_window") is not None
- and kv_seq_len > self.config.sliding_window
- )
-
- if use_sliding_windows:
- window_size = (self.config.sliding_window, self.config.sliding_window)
- else:
- window_size = (-1, -1)
-
- if past_key_value is not None:
- # Activate slicing cache only if the config has a value `sliding_windows` attribute
- if (
- hasattr(self.config, "sliding_window")
- and kv_seq_len > self.config.sliding_window
- ):
- slicing_tokens = kv_seq_len - self.config.sliding_window
-
- past_key = past_key_value[0]
- past_value = past_key_value[1]
-
- past_key = past_key[:, :, slicing_tokens:, :].contiguous()
- past_value = past_value[:, :, slicing_tokens:, :].contiguous()
-
- if past_key.shape[-2] != self.config.sliding_window - 1:
- raise ValueError(
- f"past key much have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
- f" {past_key.shape}"
- )
-
- past_key_value = (past_key, past_value) if use_cache else None
-
- if past_key_value is not None:
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
-
- past_key_value = (key_states, value_states) if use_cache else None
-
- # repeat k/v heads if n_kv_heads < n_heads
- key_states = repeat_kv(key_states, self.num_key_value_groups)
- value_states = repeat_kv(value_states, self.num_key_value_groups)
-
- if self.training:
- # during training q,k,v always have same seqlen
- assert key_states.shape == query_states.shape
- is_causal = True
- else:
- # turn off FA causal mask after first inference autoregressive iteration
- # only on first autoregressive step q,k,v have same seqlen
- is_causal = key_states.shape == query_states.shape
-
- dropout_rate = 0.0 if not self.training else getattr(self, "attention_dropout", 0.0)
-
- if cu_seqlens is not None and max_seqlen is not None and cu_seqlens.dim() == 1:
- # special handling using sample packing
- qkv = torch.stack(
- [query_states, key_states, value_states], dim=2
- ) # [bsz, nh, 3, q_len, hd]
- qkv = qkv.transpose(1, 3) # [bsz, q_len, 3, nh, hd]
- qkv = rearrange(qkv, "b s ... -> (b s) ...")
-
- output = flash_attn_varlen_qkvpacked_func(
- qkv,
- cu_seqlens,
- max_seqlen,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=True,
- window_size=window_size,
- )
- output = rearrange(output, "(b s) ... -> b s ...", b=bsz)
- elif query_states.shape == key_states.shape:
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
- qkv_unpad, cu_seqlens_q, max_seqlen_q, _, output_pad_fn = generate_qkv(
- query_states,
- key_states,
- value_states,
- qkvpacked=True,
- # We have disabled _prepare_decoder_attention_mask in LlamaModel
- # the attention_mask should be the same as the key_padding_mask
- key_padding_mask=attention_mask,
- query_padding_mask=(
- attention_mask[:, -query_states.size(1) :]
- if attention_mask is not None
- else None
- ),
- )
- output_unpad = flash_attn_varlen_qkvpacked_func(
- qkv_unpad,
- cu_seqlens_q,
- max_seqlen_q,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=is_causal,
- window_size=window_size,
- )
- output = output_pad_fn(output_unpad)
- else:
- query_states = query_states.transpose(1, 2)
- key_states = key_states.transpose(1, 2)
- value_states = value_states.transpose(1, 2)
- if attention_mask is None or attention_mask.all().item():
- output = flash_attn_kvpacked_func(
- query_states,
- torch.stack([key_states, value_states], 2),
- dropout_p=dropout_rate,
- causal=is_causal,
- window_size=window_size,
- )
- else:
- ( # pylint: disable=unbalanced-tuple-unpacking
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- _,
- _,
- output_pad_fn,
- ) = generate_qkv(
- query_states,
- key_states,
- value_states,
- kvpacked=True,
- key_padding_mask=attention_mask,
- query_padding_mask=(
- attention_mask[:, -query_states.size(1) :]
- if attention_mask is not None
- else None
- ),
- )
- if q_unpad.dtype != kv_unpad.dtype:
- kv_unpad = kv_unpad.to(q_unpad.dtype)
- output_unpad = flash_attn_varlen_kvpacked_func(
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- dropout_p=dropout_rate,
- softmax_scale=None,
- causal=is_causal,
- window_size=window_size,
- )
- output = output_pad_fn(output_unpad)
-
- attn_output = output
- if attn_output.size() != (bsz, q_len, self.num_heads, self.head_dim):
- raise ValueError(
- f"`attn_output` should be of size {(bsz, q_len, self.num_heads, self.head_dim)}, but is"
- f" {attn_output.size()}"
- )
- attn_output = rearrange(attn_output, "b s h d -> b s (h d)")
-
- attn_output = self.o_proj(attn_output)
-
- if not output_attentions:
- attn_weights = None
-
- return attn_output, attn_weights, past_key_value
-
-
-# based on https://github.com/Dao-AILab/flash-attention/blob/364a5b/tests/test_flash_attn.py#L38
-def generate_qkv(
- q,
- k,
- v,
- query_padding_mask=None,
- key_padding_mask=None,
- kvpacked=False,
- qkvpacked=False,
-): # pylint: disable=invalid-name,unnecessary-lambda-assignment
- """
- Arguments:
- q: (batch_size, seqlen_q, nheads, d)
- k: (batch_size, seqlen_k, nheads_k, d)
- v: (batch_size, seqlen_k, nheads_k, d)
- query_padding_mask: (batch_size, seqlen), bool
- key_padding_mask: (batch_size, seqlen), bool
- """
- assert not (kvpacked and qkvpacked)
- batch_size, seqlen_q, nheads, d = q.shape
- _, seqlen_k, nheads_k, _ = k.shape
- assert k.shape == (batch_size, seqlen_k, nheads_k, d)
- assert v.shape == (batch_size, seqlen_k, nheads_k, d)
-
- if query_padding_mask is not None:
- q_unpad, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(
- q, query_padding_mask
- )
-
- def output_pad_fn(output_unpad):
- return pad_input( # noqa: E731
- output_unpad, indices_q, batch_size, seqlen_q
- )
-
- else:
- q_unpad = rearrange(q, "b s h d -> (b s) h d")
- cu_seqlens_q = torch.arange(
- 0,
- (batch_size + 1) * seqlen_q,
- step=seqlen_q,
- dtype=torch.int32,
- device=q_unpad.device,
- )
- max_seqlen_q = seqlen_q
-
- def output_pad_fn(output_unpad):
- return rearrange( # noqa: E731
- output_unpad, "(b s) h d -> b s h d", b=batch_size
- )
-
- if key_padding_mask is not None:
- k_unpad, _, cu_seqlens_k, max_seqlen_k = unpad_input(k, key_padding_mask)
- v_unpad, _, _, _ = unpad_input(v, key_padding_mask)
- else:
- k_unpad = rearrange(k, "b s h d -> (b s) h d")
- v_unpad = rearrange(v, "b s h d -> (b s) h d")
- cu_seqlens_k = torch.arange(
- 0,
- (batch_size + 1) * seqlen_k,
- step=seqlen_k,
- dtype=torch.int32,
- device=k_unpad.device,
- )
- max_seqlen_k = seqlen_k
-
- if qkvpacked:
- assert nheads == nheads_k
- qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1)
- qkv = torch.stack([q, k, v], dim=2)
- return (qkv_unpad, cu_seqlens_q, max_seqlen_q, qkv, output_pad_fn)
-
- if kvpacked:
- kv_unpad = torch.stack([k_unpad, v_unpad], dim=1)
- kv = torch.stack([k, v], dim=2)
- return (
- q_unpad,
- kv_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- q,
- kv,
- output_pad_fn,
- )
-
- return (
- q_unpad,
- k_unpad,
- v_unpad,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- q,
- k,
- v,
- output_pad_fn,
- )
-
-
-def mistral_model_forward(
- self,
- input_ids: torch.LongTensor = None,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_values: Optional[List[torch.FloatTensor]] = None,
- inputs_embeds: Optional[torch.FloatTensor] = None,
- use_cache: Optional[bool] = None,
- output_attentions: Optional[bool] = None,
- output_hidden_states: Optional[bool] = None,
- return_dict: Optional[bool] = None,
- cache_position: Optional[ # pylint: disable=unused-argument
- torch.LongTensor
- ] = None,
-) -> Union[Tuple, BaseModelOutputWithPast]:
- output_attentions = (
- output_attentions
- if output_attentions is not None
- else self.config.output_attentions
- )
- output_hidden_states = (
- output_hidden_states
- if output_hidden_states is not None
- else self.config.output_hidden_states
- )
- use_cache = use_cache if use_cache is not None else self.config.use_cache
-
- return_dict = (
- return_dict if return_dict is not None else self.config.use_return_dict
- )
-
- # retrieve input_ids and inputs_embeds
- if input_ids is not None and inputs_embeds is not None:
- raise ValueError(
- "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
- )
- if input_ids is not None:
- batch_size, seq_length = input_ids.shape
- elif inputs_embeds is not None:
- batch_size, seq_length, _ = inputs_embeds.shape
- else:
- raise ValueError(
- "You have to specify either decoder_input_ids or decoder_inputs_embeds"
- )
-
- seq_length_with_past = seq_length
- past_key_values_length = 0
-
- if past_key_values is not None:
- past_key_values_length = past_key_values[0][0].shape[2]
- seq_length_with_past = seq_length_with_past + past_key_values_length
-
- cu_seqlens = None
- max_seqlen = None
- if position_ids is None:
- device = input_ids.device if input_ids is not None else inputs_embeds.device
- position_ids = torch.arange(
- past_key_values_length,
- seq_length + past_key_values_length,
- dtype=torch.long,
- device=device,
- )
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
- else:
- position_ids = position_ids.view(-1, seq_length).long()
- cu_seqlens, max_seqlen = get_cu_seqlens_from_pos_ids(position_ids)
- cu_seqlens = cu_seqlens.squeeze()
-
- if inputs_embeds is None:
- inputs_embeds = self.embed_tokens(input_ids)
- # embed positions
- if attention_mask is None:
- attention_mask = torch.ones(
- (batch_size, seq_length_with_past),
- dtype=torch.bool,
- device=inputs_embeds.device,
- )
- attention_mask = (
- self._prepare_decoder_attention_mask( # pylint: disable=protected-access
- attention_mask,
- (batch_size, seq_length),
- inputs_embeds,
- past_key_values_length,
- sliding_window=self.config.sliding_window,
- )
- )
-
- hidden_states = inputs_embeds
-
- if self.gradient_checkpointing and self.training:
- if use_cache:
- transformers.logger.warning_once(
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
- )
- use_cache = False
-
- # decoder layers
- all_hidden_states = () if output_hidden_states else None
- all_self_attns = () if output_attentions else None
- next_decoder_cache = () if use_cache else None
-
- for idx, decoder_layer in enumerate(self.layers):
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- past_key_value = past_key_values[idx] if past_key_values is not None else None
-
- if self.gradient_checkpointing and self.training:
- layer_outputs = (
- self._gradient_checkpointing_func( # pylint: disable=protected-access
- decoder_layer.__call__,
- hidden_states,
- attention_mask,
- position_ids,
- past_key_value,
- output_attentions,
- None,
- cu_seqlens,
- max_seqlen,
- )
- )
- else:
- layer_outputs = decoder_layer(
- hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- )
-
- hidden_states = layer_outputs[0]
-
- if use_cache:
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
-
- if output_attentions:
- all_self_attns += (layer_outputs[1],)
-
- hidden_states = self.norm(hidden_states)
-
- # add hidden states from the last decoder layer
- if output_hidden_states:
- all_hidden_states += (hidden_states,)
-
- next_cache = next_decoder_cache if use_cache else None
- if not return_dict:
- return tuple(
- v
- for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
- if v is not None
- )
- return BaseModelOutputWithPast(
- last_hidden_state=hidden_states,
- past_key_values=next_cache,
- hidden_states=all_hidden_states,
- attentions=all_self_attns,
- )
-
-
-class MistralDecoderLayer(OriginalMistralDecoderLayer):
- """
- patched version of MistralDecoderLayer to pass through the precalculated cu_seqlens
- """
-
- def forward(
- self,
- hidden_states: torch.Tensor,
- attention_mask: Optional[torch.Tensor] = None,
- position_ids: Optional[torch.LongTensor] = None,
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False,
- use_cache: Optional[bool] = False,
- cu_seqlens: Optional[torch.Tensor] = None,
- max_seqlen: Optional[torch.Tensor] = None,
- ) -> Tuple[
- torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
- ]:
- """
- Args:
- hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
- attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
- `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
- output_attentions (`bool`, *optional*):
- Whether or not to return the attentions tensors of all attention layers. See `attentions` under
- returned tensors for more detail.
- use_cache (`bool`, *optional*):
- If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
- (see `past_key_values`).
- past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
- cu_seqlens (`torch.Tensor`, *optional*) cumulative sequence len when packing
- """
-
- residual = hidden_states
-
- hidden_states = self.input_layernorm(hidden_states)
-
- # Self Attention
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- position_ids=position_ids,
- past_key_value=past_key_value,
- output_attentions=output_attentions,
- use_cache=use_cache,
- cu_seqlens=cu_seqlens,
- max_seqlen=max_seqlen,
- )
- hidden_states = residual + hidden_states
-
- # Fully Connected
- residual = hidden_states
- hidden_states = self.post_attention_layernorm(hidden_states)
- hidden_states = self.mlp(hidden_states)
- hidden_states = residual + hidden_states
-
- outputs = (hidden_states,)
-
- if output_attentions:
- outputs += (self_attn_weights,)
-
- if use_cache:
- outputs += (present_key_value,)
-
- return outputs
diff --git a/src/axolotl/monkeypatch/mixtral/__init__.py b/src/axolotl/monkeypatch/mixtral/__init__.py
index 5b8054000..b353b12cf 100644
--- a/src/axolotl/monkeypatch/mixtral/__init__.py
+++ b/src/axolotl/monkeypatch/mixtral/__init__.py
@@ -31,14 +31,12 @@ def patch_mixtral_moe_forward_zero3() -> None:
topk_weight = topk_weight.to(hidden_states.dtype)
hidden_states = hidden_states.repeat_interleave(self.top_k, dim=0)
- y = torch.empty_like(hidden_states) # pylint: disable=invalid-name
+ y = torch.empty_like(hidden_states)
flat_topk_idx = topk_idx.view(-1)
for i in range(self.num_experts):
expert = self.experts[i]
y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
- y = ( # pylint: disable=invalid-name
- y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)
- ).sum(dim=1)
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
final_hidden_states = y.reshape(batch_size, sequence_length, hidden_dim)
return final_hidden_states, router_logits
diff --git a/src/axolotl/monkeypatch/models/__init__.py b/src/axolotl/monkeypatch/models/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/models/apertus/__init__.py b/src/axolotl/monkeypatch/models/apertus/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/models/apertus/activation.py b/src/axolotl/monkeypatch/models/apertus/activation.py
new file mode 100644
index 000000000..d5470aceb
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/apertus/activation.py
@@ -0,0 +1,52 @@
+"""Monkeypatch for Apertus to dtype mismatch in XIELU act"""
+
+from torch import Tensor
+
+
+def patch_apertus_xielu_activation():
+ try:
+ from transformers.activations import XIELUActivation
+ except ImportError as err:
+ raise ImportError(
+ "Cannot import XIELUActivation. "
+ "Please make sure to update your transformers version >= 4.56.1."
+ ) from err
+
+ from transformers.activations import logger
+
+ # Store the original method
+ old_fn = XIELUActivation._xielu_cuda
+
+ def _xielu_cuda_fixed(self, x: Tensor) -> Tensor:
+ """Firewall function to prevent torch.compile from seeing .item() calls"""
+ original_shape = x.shape
+ # CUDA kernel expects 3D tensors, reshape if needed
+ while x.dim() < 3:
+ x = x.unsqueeze(0)
+ if x.dim() > 3:
+ x = x.view(-1, 1, x.size(-1))
+ if original_shape != x.shape:
+ logger.warning_once(
+ "Warning: xIELU input tensor expects 3 dimensions but got (shape: %s). Reshaping to (shape: %s).",
+ original_shape,
+ x.shape,
+ )
+ result = self._xielu_cuda_obj.forward(
+ x,
+ self.alpha_p.to(x.dtype),
+ self.alpha_n.to(x.dtype),
+ # Temporary until xIELU CUDA fully implemented -> self.{beta,eps}.item()
+ self._beta_scalar,
+ self._eps_scalar,
+ self.with_vector_loads,
+ )
+ return result.view(original_shape)
+
+ # Apply the patch
+ XIELUActivation._xielu_cuda = _xielu_cuda_fixed
+
+ def unpatch():
+ """Restore the original method"""
+ XIELUActivation._xielu_cuda = old_fn
+
+ return unpatch
diff --git a/src/axolotl/monkeypatch/models/llama4/modeling.py b/src/axolotl/monkeypatch/models/llama4/modeling.py
index 4127793e7..0fc8f5699 100644
--- a/src/axolotl/monkeypatch/models/llama4/modeling.py
+++ b/src/axolotl/monkeypatch/models/llama4/modeling.py
@@ -95,18 +95,12 @@ def patch_llama4_linearized_modeling():
old_lamma_4_text_experts = modeling_llama4.Llama4TextExperts
modeling_llama4.Llama4TextExperts = Llama4TextExperts
- setattr(
- sys.modules["transformers.models.llama4"],
- "Llama4TextExperts",
- Llama4TextExperts,
- )
+ sys.modules["transformers.models.llama4"].Llama4TextExperts = Llama4TextExperts
def unpatch():
modeling_llama4.Llama4TextExperts = old_lamma_4_text_experts
- setattr(
- sys.modules["transformers.models.llama4"],
- "Llama4TextExperts",
- old_lamma_4_text_experts,
- )
+ sys.modules[
+ "transformers.models.llama4"
+ ].Llama4TextExperts = old_lamma_4_text_experts
return unpatch
diff --git a/src/axolotl/monkeypatch/models/mistral3/__init__.py b/src/axolotl/monkeypatch/models/mistral3/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py b/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py
new file mode 100644
index 000000000..9e7259a05
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/mistral3/mistral_common_tokenizer.py
@@ -0,0 +1,85 @@
+"""
+Monkeypatch to fix inefficient tensor conversion in MistralCommonTokenizer.apply_chat_template
+"""
+
+import importlib
+import inspect
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def apply_mistral_tokenizer_image_patch():
+ """Apply patch to MistralCommonTokenizer.apply_chat_template to fix image tensor conversion."""
+ from transformers.tokenization_mistral_common import MistralCommonTokenizer
+
+ # Get original source
+ original_source = inspect.getsource(MistralCommonTokenizer.apply_chat_template)
+ original_source, _ = detab_code(original_source)
+
+ # Define the replacement
+ original_tensor_conversion = (
+ " pixel_values = torch.tensor(images)"
+ )
+
+ patched_tensor_conversion = """ if isinstance(images, list) and len(images) > 0 and isinstance(images[0], np.ndarray):
+ pixel_values = torch.tensor(np.array(images))
+ else:
+ pixel_values = torch.tensor(images)"""
+
+ # Apply the replacement
+ if original_tensor_conversion in original_source:
+ patched_source = original_source.replace(
+ original_tensor_conversion, patched_tensor_conversion
+ )
+ patched_source = patched_source.replace(
+ "def apply_chat_template(",
+ "def patched_apply_chat_template(",
+ 1,
+ )
+
+ # Load necessary imports from the module
+ module_name = MistralCommonTokenizer.__module__
+ module = importlib.import_module(module_name)
+
+ # Detect what needs to be imported
+ items_to_import = []
+ for item in dir(module):
+ if item in patched_source and not item.startswith("_"):
+ items_to_import.append(item)
+
+ # Execute imports in global scope
+ if items_to_import:
+ exec( # nosec B102
+ f"from {module_name} import ({', '.join(items_to_import)})",
+ globals(),
+ )
+
+ # Also need standard imports that might be used
+ exec("import numpy as np", globals()) # nosec B102
+ exec("import torch", globals()) # nosec B102
+ exec("from typing import Union, Optional, List, Dict, Any, Callable", globals()) # nosec B102
+ exec("from pathlib import Path", globals()) # nosec B102
+
+ # Import other dependencies that might be needed
+ try:
+ exec("from transformers.utils import is_torch_available", globals()) # nosec B102
+ exec(
+ "from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, TensorType",
+ globals(),
+ ) # nosec B102
+ exec("from transformers.utils import logging", globals()) # nosec B102
+ exec("logger = logging.get_logger(__name__)", globals()) # nosec B102
+ except ImportError as e:
+ LOG.warning(f"Could not import some dependencies: {e}")
+
+ # Execute the patched source
+ exec(patched_source, globals()) # nosec B102
+
+ # Replace the method
+ MistralCommonTokenizer.apply_chat_template = patched_apply_chat_template
+ LOG.info("Successfully applied MistralCommonTokenizer tensor conversion patch")
+ else:
+ LOG.warning("Could not find target code for MistralCommonTokenizer patching")
diff --git a/src/axolotl/monkeypatch/models/pixtral/__init__.py b/src/axolotl/monkeypatch/models/pixtral/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py b/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py
new file mode 100644
index 000000000..d2b482f19
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/pixtral/modeling_flash_attention_utils.py
@@ -0,0 +1,42 @@
+"""Monkeypatch for FA utils to accept 1D position_ids from Pixtral's position_ids_in_meshgrid"""
+
+import torch
+
+
+def apply_patch_is_packed_sequence():
+ """Apply patch to FA utils to accept 1D position_ids from Pixtral's position_ids_in_meshgrid"""
+ from transformers import modeling_flash_attention_utils
+
+ def fixed_is_packed_sequence(position_ids, batch_size):
+ """
+ Check the position ids whether packed sequences are indicated or not
+ 1. Position ids exist
+ 2. Flattened sequences only are supported
+ 3. Compile-friendly `not (torch.diff(position_ids, dim=-1) >= 0).all()`, i.e. we have multiple increasing sequences
+ """
+ if position_ids is None:
+ return False
+
+ if position_ids.ndim == 1:
+ position_ids = position_ids.unsqueeze(0) # [N] -> [1, N]
+
+ increasing_position_sequences = (
+ torch.arange(position_ids.shape[1], device=position_ids.device)
+ + position_ids.min()
+ )
+ return (
+ batch_size == 1
+ and (increasing_position_sequences - position_ids).abs().sum().bool().item()
+ )
+
+ # Store original method
+ old_fn = modeling_flash_attention_utils._is_packed_sequence
+
+ # Apply the patch
+ modeling_flash_attention_utils._is_packed_sequence = fixed_is_packed_sequence
+
+ def unpatch():
+ """Restore the original method"""
+ modeling_flash_attention_utils._is_packed_sequence = old_fn
+
+ return unpatch
diff --git a/src/axolotl/monkeypatch/models/qwen3_next/__init__.py b/src/axolotl/monkeypatch/models/qwen3_next/__init__.py
new file mode 100644
index 000000000..39bcd4115
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_next/__init__.py
@@ -0,0 +1 @@
+"""Qwen3_Next model monkeypatches."""
diff --git a/src/axolotl/monkeypatch/models/qwen3_next/modeling.py b/src/axolotl/monkeypatch/models/qwen3_next/modeling.py
new file mode 100644
index 000000000..d68992d0e
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/qwen3_next/modeling.py
@@ -0,0 +1,317 @@
+"""Monkeypatch for Qwen3_Next model to pass position_ids to linear attention."""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def get_cu_seqlens(position_ids):
+ """
+ Adapted from transformers.modeling_flash_attention_utils.prepare_fa_kwargs_from_position_ids.
+
+ https://github.com/huggingface/transformers/blob/0f1b128d3359a26bd18be99c26d7f04fb3cba914/src/transformers/modeling_flash_attention_utils.py#L316
+ """
+ tensor_kwargs = {"dtype": torch.int32, "device": position_ids.device}
+
+ position_ids = position_ids.view(-1)
+ indices_q = (position_ids == 0).nonzero().view(-1)
+
+ cu_seq_lens_q = torch.cat(
+ (
+ indices_q.to(**tensor_kwargs),
+ torch.tensor(position_ids.size(), **tensor_kwargs),
+ )
+ )
+
+ return cu_seq_lens_q
+
+
+def patch_qwen3_next_decoder_layer():
+ """Patch Qwen3NextDecoderLayer to pass position_ids to linear attention."""
+ try:
+ from transformers.models.qwen3_next.modeling_qwen3_next import (
+ Qwen3NextDecoderLayer,
+ )
+ except ImportError:
+ LOG.warning("Qwen3Next model not found, skipping patch")
+ return
+
+ # Store original forward method
+ original_decoder_forward = Qwen3NextDecoderLayer.forward
+
+ def patched_decoder_forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Tuple[torch.Tensor]] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Token Mixer
+ if self.layer_type == "linear_attention":
+ hidden_states = self.linear_attn(
+ hidden_states=hidden_states,
+ cache_params=past_key_values,
+ cache_position=cache_position,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ )
+ elif self.layer_type == "full_attention":
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ cache_position=cache_position,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ # For the MoE layers, we need to unpack
+ if isinstance(hidden_states, Tuple):
+ hidden_states, _ = hidden_states
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+ # Apply the patches
+ Qwen3NextDecoderLayer.forward = patched_decoder_forward
+
+ def unpatch():
+ """Restore the original forward method"""
+ Qwen3NextDecoderLayer.forward = original_decoder_forward
+
+ return unpatch
+
+
+def patch_qwen3_next_gateddelta_layer():
+ """Patch Qwen3NextGatedDeltaNet to parse cu_seqlens and pass to chunk_gated_delta_rule"""
+ try:
+ from transformers.models.qwen3_next.modeling_qwen3_next import (
+ Qwen3NextDynamicCache,
+ Qwen3NextGatedDeltaNet,
+ apply_mask_to_padding_states,
+ )
+ except ImportError:
+ LOG.warning("Qwen3Next model not found, skipping patch")
+ return
+
+ # Store original forward method
+ original_gated_delta_net_forward = Qwen3NextGatedDeltaNet.forward
+
+ def patched_gated_delta_net_forward(
+ self,
+ hidden_states: torch.Tensor,
+ cache_params: Optional[Qwen3NextDynamicCache] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ ):
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+
+ # Set up dimensions for reshapes later
+ batch_size, seq_len, _ = hidden_states.shape
+
+ use_precomputed_states = (
+ cache_params is not None
+ and cache_params.has_previous_state
+ and seq_len == 1
+ and cache_position is not None
+ )
+
+ # getting projected states from cache if it exists
+ if cache_params is not None:
+ conv_state = cache_params.conv_states[self.layer_idx]
+ recurrent_state = cache_params.recurrent_states[self.layer_idx]
+
+ projected_states_qkvz = self.in_proj_qkvz(hidden_states)
+ projected_states_ba = self.in_proj_ba(hidden_states)
+ query, key, value, z, b, a = self.fix_query_key_value_ordering(
+ projected_states_qkvz, projected_states_ba
+ )
+ query, key, value = (
+ x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value)
+ )
+
+ mixed_qkv = torch.cat((query, key, value), dim=-1)
+ mixed_qkv = mixed_qkv.transpose(1, 2)
+
+ if use_precomputed_states:
+ # 2. Convolution sequence transformation
+ # NOTE: the conv state is updated in `causal_conv1d_update`
+ mixed_qkv = self.causal_conv1d_update(
+ mixed_qkv,
+ conv_state,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.activation,
+ )
+ else:
+ if cache_params is not None:
+ conv_state = F.pad(
+ mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)
+ )
+ cache_params.conv_states[self.layer_idx] = conv_state
+ if self.causal_conv1d_fn is not None:
+ mixed_qkv = self.causal_conv1d_fn(
+ x=mixed_qkv,
+ weight=self.conv1d.weight.squeeze(1),
+ bias=self.conv1d.bias,
+ activation=self.activation,
+ seq_idx=None,
+ )
+ else:
+ mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])
+
+ mixed_qkv = mixed_qkv.transpose(1, 2)
+ query, key, value = torch.split(
+ mixed_qkv,
+ [
+ self.key_dim,
+ self.key_dim,
+ self.value_dim,
+ ],
+ dim=-1,
+ )
+ query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim)
+ key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim)
+ value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim)
+
+ beta = b.sigmoid()
+ # If the model is loaded in fp16, without the .float() here, A might be -inf
+ g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+ if self.num_v_heads // self.num_k_heads > 1:
+ query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+ key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+
+ if not use_precomputed_states:
+ cu_seqlens = get_cu_seqlens(position_ids=position_ids)
+ core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule(
+ query,
+ key,
+ value,
+ g=g,
+ beta=beta,
+ initial_state=None,
+ output_final_state=cache_params is not None,
+ use_qk_l2norm_in_kernel=True,
+ cu_seqlens=cu_seqlens,
+ )
+
+ else:
+ core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule(
+ query,
+ key,
+ value,
+ g=g,
+ beta=beta,
+ initial_state=recurrent_state,
+ output_final_state=cache_params is not None,
+ use_qk_l2norm_in_kernel=True,
+ )
+
+ # Update cache
+ if cache_params is not None:
+ cache_params.recurrent_states[self.layer_idx] = last_recurrent_state
+
+ z_shape_og = z.shape
+ # reshape input data into 2D tensor
+ core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
+ z = z.reshape(-1, z.shape[-1])
+ core_attn_out = self.norm(core_attn_out, z)
+ core_attn_out = core_attn_out.reshape(z_shape_og)
+ core_attn_out = core_attn_out.reshape(
+ core_attn_out.shape[0], core_attn_out.shape[1], -1
+ )
+
+ output = self.out_proj(core_attn_out)
+ return output
+
+ # Apply the patches
+ Qwen3NextGatedDeltaNet.forward = patched_gated_delta_net_forward
+
+ def unpatch():
+ """Restore the original forward method"""
+ Qwen3NextGatedDeltaNet.forward = original_gated_delta_net_forward
+
+ return unpatch
+
+
+def patch_qwen3_next_imports():
+ """Patch Qwen3Next imports to use try/except instead of is_flash_linear_attention_available."""
+ try:
+ import transformers.models.qwen3_next.modeling_qwen3_next as qwen3_modeling
+ except ImportError:
+ LOG.warning("Qwen3Next model not found, skipping import patch")
+ return
+
+ # Save original values for unpatch
+ original_FusedRMSNormGated = getattr(qwen3_modeling, "FusedRMSNormGated", None)
+ original_chunk_gated_delta_rule = getattr(
+ qwen3_modeling, "chunk_gated_delta_rule", None
+ )
+ original_fused_recurrent_gated_delta_rule = getattr(
+ qwen3_modeling, "fused_recurrent_gated_delta_rule", None
+ )
+ original_is_fast_path_available = getattr(
+ qwen3_modeling, "is_fast_path_available", False
+ )
+
+ try:
+ from fla.modules import FusedRMSNormGated
+ from fla.ops.gated_delta_rule import (
+ chunk_gated_delta_rule,
+ fused_recurrent_gated_delta_rule,
+ )
+
+ qwen3_modeling.FusedRMSNormGated = FusedRMSNormGated
+ qwen3_modeling.chunk_gated_delta_rule = chunk_gated_delta_rule
+ qwen3_modeling.fused_recurrent_gated_delta_rule = (
+ fused_recurrent_gated_delta_rule
+ )
+
+ # Force is_fast_path_available to be True
+ # fla has triton kernels for causal_conv1d
+ qwen3_modeling.is_fast_path_available = True
+ except ImportError:
+ qwen3_modeling.chunk_gated_delta_rule = None
+ qwen3_modeling.fused_recurrent_gated_delta_rule = None
+ qwen3_modeling.FusedRMSNormGated = None
+
+ def unpatch():
+ """Restore the original import values"""
+ qwen3_modeling.FusedRMSNormGated = original_FusedRMSNormGated
+ qwen3_modeling.chunk_gated_delta_rule = original_chunk_gated_delta_rule
+ qwen3_modeling.fused_recurrent_gated_delta_rule = (
+ original_fused_recurrent_gated_delta_rule
+ )
+ qwen3_modeling.is_fast_path_available = original_is_fast_path_available
+
+ return unpatch
+
+
+def patch_qwen3_next_modeling_packing():
+ """Apply all Qwen3Next model patches."""
+ patch_qwen3_next_imports()
+ patch_qwen3_next_decoder_layer()
+ patch_qwen3_next_gateddelta_layer()
+
+ LOG.info("Applied Qwen3Next patch for packing")
diff --git a/src/axolotl/monkeypatch/models/voxtral/__init__.py b/src/axolotl/monkeypatch/models/voxtral/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/models/voxtral/modeling.py b/src/axolotl/monkeypatch/models/voxtral/modeling.py
new file mode 100644
index 000000000..3dd652dd8
--- /dev/null
+++ b/src/axolotl/monkeypatch/models/voxtral/modeling.py
@@ -0,0 +1,67 @@
+"""Monkeypatch for voxtral to fix leaf node and dtype mismatch"""
+
+from typing import Optional, Union
+
+import torch
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def patch_voxtral_conditional_generation_forward():
+ from transformers.models.voxtral.modeling_voxtral import (
+ VoxtralForConditionalGeneration,
+ )
+
+ # Store the original forward method
+ old_forward = VoxtralForConditionalGeneration.forward
+
+ def _forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ input_features: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[Cache] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ cache_position: Optional[torch.LongTensor] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **kwargs,
+ ) -> CausalLMOutputWithPast:
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if input_features is not None:
+ audio_embeds = self.get_audio_embeds(input_features)
+
+ # Cast audio_embeds to match inputs_embeds dtype
+ audio_embeds = audio_embeds.to(inputs_embeds.dtype)
+
+ # replace text-audio token placeholders with audio embeddings
+ audio_token_mask = input_ids == self.config.audio_token_id
+
+ inputs_embeds = inputs_embeds.clone()
+ inputs_embeds[audio_token_mask] = audio_embeds
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ labels=labels,
+ use_cache=use_cache,
+ cache_position=cache_position,
+ logits_to_keep=logits_to_keep,
+ **kwargs,
+ )
+ return outputs
+
+ # Apply the patch
+ VoxtralForConditionalGeneration.forward = _forward
+
+ def unpatch():
+ """Restore the original forward method"""
+ VoxtralForConditionalGeneration.forward = old_forward
+
+ return unpatch
diff --git a/src/axolotl/monkeypatch/multipack.py b/src/axolotl/monkeypatch/multipack.py
index 1467f9e29..48b4ea10e 100644
--- a/src/axolotl/monkeypatch/multipack.py
+++ b/src/axolotl/monkeypatch/multipack.py
@@ -11,6 +11,7 @@ from axolotl.monkeypatch.mixtral import patch_mixtral_moe_forward_zero3
from axolotl.monkeypatch.utils import get_unpad_data
SUPPORTED_MULTIPACK_MODEL_TYPES = [
+ "apertus",
"mllama_text_model",
"llama",
"llama4",
@@ -20,6 +21,7 @@ SUPPORTED_MULTIPACK_MODEL_TYPES = [
"qwen2_moe",
"qwen3",
"qwen3_moe",
+ "qwen3_next",
"falcon",
"phi",
"phi3",
@@ -35,6 +37,16 @@ SUPPORTED_MULTIPACK_MODEL_TYPES = [
"deepseek_v3",
"glm",
"glm4",
+ "smollm3",
+ "granite",
+ "granitemoe",
+ "hunyuan_v1_dense",
+ "hunyuan_v1_moe",
+ "gpt_oss",
+ "arcee",
+ "seed_oss",
+ "lfm2",
+ "lfm2_moe",
]
@@ -42,9 +54,11 @@ def patch_for_multipack(model_type, model_name=None, has_remote_code=False):
if has_remote_code:
patch_remote(model_name)
elif hasattr(transformers, "modeling_flash_attention_utils"):
- transformers.modeling_flash_attention_utils._get_unpad_data = ( # pylint: disable=protected-access
- get_unpad_data
- )
+ # sanity check in case upstream api changes on this
+ assert hasattr(
+ transformers.modeling_flash_attention_utils, "_get_unpad_data"
+ ), "transformers api changed for _get_unpad_data for flash attention"
+ transformers.modeling_flash_attention_utils._get_unpad_data = get_unpad_data
if model_type == "mixtral" and is_deepspeed_zero3_enabled():
patch_mixtral_moe_forward_zero3()
@@ -60,6 +74,4 @@ def patch_remote(model_name):
module_name = ".".join(parts)
modeling_arch = importlib.import_module(module_name)
if hasattr(modeling_arch, "_get_unpad_data"):
- modeling_arch._get_unpad_data = ( # pylint: disable=protected-access
- get_unpad_data
- )
+ modeling_arch._get_unpad_data = get_unpad_data
diff --git a/src/axolotl/monkeypatch/peft/utils.py b/src/axolotl/monkeypatch/peft/utils.py
index 0c571fbd2..d1011f5eb 100644
--- a/src/axolotl/monkeypatch/peft/utils.py
+++ b/src/axolotl/monkeypatch/peft/utils.py
@@ -49,9 +49,7 @@ def patch_peft_prep_code():
prep_code = get_peft_prep_code()
except OSError:
return
- peft.utils.other._original_create_accelerator_and_postprocess = ( # pylint: disable=protected-access
- prep_code
- )
+ peft.utils.other._original_create_accelerator_and_postprocess = prep_code
prep_code, _ = detab_code(prep_code)
if ORIGINAL_PREPARE_CODE not in prep_code:
return
@@ -68,11 +66,15 @@ def patch_peft_prep_code():
if item in prep_code:
items_to_import.append(item)
- exec( # pylint: disable=exec-used # nosec B102
+ exec(
"from peft.utils.other import (" + ", ".join(x for x in items_to_import) + ")",
globals(),
)
- exec(prep_code, globals()) # pylint: disable=exec-used # nosec B102
+ exec(prep_code, globals())
LOG.info("patching prepare_model_for_kbit_training to allow for overrides")
- peft.utils.other.prepare_model_for_kbit_training = fixed_prepare_model_for_kbit_training # pylint: disable=protected-access # pylint: disable=undefined-variable # noqa: F821
- axolotl.loaders.model.prepare_model_for_kbit_training = fixed_prepare_model_for_kbit_training # pylint: disable=protected-access # pylint: disable=undefined-variable # noqa: F821
+ peft.utils.other.prepare_model_for_kbit_training = (
+ fixed_prepare_model_for_kbit_training
+ )
+ axolotl.loaders.model.prepare_model_for_kbit_training = (
+ fixed_prepare_model_for_kbit_training
+ )
diff --git a/src/axolotl/monkeypatch/relora.py b/src/axolotl/monkeypatch/relora.py
index 5b7418e39..a01d850b3 100644
--- a/src/axolotl/monkeypatch/relora.py
+++ b/src/axolotl/monkeypatch/relora.py
@@ -6,7 +6,7 @@ import os.path
import shutil
from functools import partial
from pathlib import Path
-from typing import Dict, List, Sequence, Union
+from typing import Dict, List, Union
import bitsandbytes as bnb
import peft
@@ -14,8 +14,6 @@ import safetensors.torch as st
import torch
from huggingface_hub import snapshot_download
from torch.distributed.optim import ZeroRedundancyOptimizer
-from torch.optim.lr_scheduler import LRScheduler
-from torch.optim.optimizer import Optimizer
from transformers import (
TrainerCallback,
TrainerControl,
@@ -84,7 +82,7 @@ class ReLoRACallback(TrainerCallback):
"""Callback to merge LoRA weights into the base model and save full-weight checkpoints"""
def __init__(self, cfg: DictDefault):
- self.relora_steps = cfg.relora_steps
+ self.relora_steps = cfg.jagged_restart_steps
self.cpu_offload = cfg.relora_cpu_offload
self.quantized = cfg.load_in_4bit or cfg.load_in_8bit
self.last_full_model = cfg.base_model
@@ -93,9 +91,9 @@ class ReLoRACallback(TrainerCallback):
if not os.path.exists(self.last_full_model):
self.last_full_model = str(Path(snapshot_download(cfg.base_model)))
- assert os.path.exists(
- self.last_full_model
- ), "for ReLORA base_model must be a local path"
+ assert os.path.exists(self.last_full_model), (
+ "for ReLORA base_model must be a local path"
+ )
self.num_lora_restarts = 0
self.need_full_save = False
@@ -255,51 +253,6 @@ class ReLoRACallback(TrainerCallback):
return control
-class ReLoRAScheduler(LRScheduler):
- """Wraps another scheduler to apply per-lora-restart learning rate warmups."""
-
- def __init__(
- self,
- optimizer: Optimizer,
- inner_schedule: LRScheduler,
- relora_steps: int,
- warmup_steps: int,
- anneal_steps: int = 1,
- min_lr_scale: float = 0.001,
- ) -> None:
- self.inner_schedule = inner_schedule
- self.relora_steps = relora_steps
- self.warmup_steps = warmup_steps
- self.anneal_steps = anneal_steps
- self.min_lr_scale = min_lr_scale
- super().__init__(optimizer, inner_schedule.last_epoch)
-
- def get_lr(self) -> float:
- self.inner_schedule.last_epoch = self.last_epoch
-
- original = self.inner_schedule.get_lr()
- step = self.last_epoch
-
- if step < self.relora_steps - self.warmup_steps:
- scale = 1
- else:
- per_relora_progress = step % self.relora_steps
- if per_relora_progress < self.warmup_steps:
- cycle_t = min(1.0, (per_relora_progress) / self.warmup_steps)
- elif per_relora_progress > (self.relora_steps - self.anneal_steps):
- cycle_t = min(
- 1.0,
- (self.relora_steps - per_relora_progress) / self.anneal_steps,
- )
- else:
- cycle_t = 1
- scale = cycle_t * (1 - self.min_lr_scale) + self.min_lr_scale
-
- if isinstance(original, Sequence):
- return [lr * scale for lr in original]
- return original * scale
-
-
def sharded_paths(path: str, module_names: List[str]) -> Dict[str, str]:
model_name = "model.safetensors"
if not os.path.exists(str(Path(path) / model_name)) and not os.path.exists(
@@ -340,7 +293,6 @@ def find_lora_modules(model: peft.LoraModel) -> Dict[str, peft.tuners.lora.LoraL
key_list = [key for key, _ in model.model.named_modules() if "lora" not in key]
for key in key_list:
try:
- # pylint: disable=protected-access
_parent, target, _target_name = peft.utils._get_submodules(model.model, key)
except AttributeError:
continue
@@ -388,7 +340,7 @@ def merge_and_save(
modules = find_lora_modules(model)
if not quantized:
- for module_name, target in modules.items():
+ for _, target in modules.items():
active_adapter = target.active_adapter
if isinstance(active_adapter, list):
active_adapter = active_adapter[0]
diff --git a/src/axolotl/monkeypatch/ring_attn/__init__.py b/src/axolotl/monkeypatch/ring_attn/__init__.py
index 5833b9ce4..1c14776c9 100644
--- a/src/axolotl/monkeypatch/ring_attn/__init__.py
+++ b/src/axolotl/monkeypatch/ring_attn/__init__.py
@@ -1,22 +1,17 @@
"""Init for ring attention monkeypatch module"""
-# pylint: disable=unused-import
# flake8: noqa
from .patch import (
get_ring_attn_group,
- patch_prepare_data_loader,
- patch_prepare_device_mesh,
- register_ring_attn,
+ register_ring_attn_from_device_mesh,
set_ring_attn_group,
update_ring_attn_params,
)
__all__ = (
"get_ring_attn_group",
- "patch_prepare_data_loader",
- "patch_prepare_device_mesh",
- "register_ring_attn",
+ "register_ring_attn_from_device_mesh",
"set_ring_attn_group",
"update_ring_attn_params",
)
diff --git a/src/axolotl/monkeypatch/ring_attn/adapters/batch.py b/src/axolotl/monkeypatch/ring_attn/adapters/batch.py
index e556ba5e3..74d33ed4a 100644
--- a/src/axolotl/monkeypatch/ring_attn/adapters/batch.py
+++ b/src/axolotl/monkeypatch/ring_attn/adapters/batch.py
@@ -7,8 +7,6 @@ Our implementation closely follows the structure of that module, but we've minif
somewhat to support only the latest versions of transformers.
"""
-# pylint: disable=protected-access,cyclic-import
-
import os
from typing import Callable
@@ -18,10 +16,18 @@ import transformers
import transformers.modeling_flash_attention_utils
from ring_flash_attn import ring_flash_attn_func
from ring_flash_attn.adapters.hf_adapter import check_params
-from transformers.modeling_flash_attention_utils import (
- _flash_supports_window_size,
- is_flash_attn_greater_or_equal,
-)
+from transformers.modeling_flash_attention_utils import is_flash_attn_greater_or_equal
+
+try:
+ from transformers.modeling_flash_attention_utils import _flash_supports_window
+except ImportError:
+ try:
+ from transformers.modeling_flash_attention_utils import (
+ _flash_supports_window_size as _flash_supports_window,
+ )
+ except ImportError:
+ _flash_supports_window = True
+
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from axolotl.utils.schemas.enums import RingAttnFunc
@@ -33,7 +39,7 @@ RING_ATTN_FUNC_MAPPING = {
}
-def create_flash_attn_forward(
+def create_flash_attn_forward_varlen_llama3(
process_group: dist.ProcessGroup, ring_attn_func: RingAttnFunc
) -> Callable:
"""
@@ -51,7 +57,7 @@ def create_flash_attn_forward(
"""
# transformers 4.48+
- # pylint: disable=unused-argument
+
def _flash_attention_forward(
query_states: torch.Tensor,
key_states: torch.Tensor,
@@ -71,6 +77,7 @@ def create_flash_attn_forward(
max_length_q: int | None = None,
max_length_k: int | None = None,
target_dtype: torch.dtype | None = None,
+ attn_implementation: str | None = None,
**kwargs,
):
"""
@@ -97,6 +104,7 @@ def create_flash_attn_forward(
max_length_q: Not used in this implementation.
max_length_k: Not used in this implementation.
target_dtype: Not used in this implementation.
+ attn_implementation: Not used in this implementation.
**kwargs: Additional keyword arguments. Not used in this implementation.
Returns:
@@ -110,7 +118,7 @@ def create_flash_attn_forward(
# Handle sliding window
use_sliding_windows = (
- _flash_supports_window_size
+ _flash_supports_window
and sliding_window is not None
and key_states.shape[1] > sliding_window
)
@@ -161,7 +169,7 @@ def substitute_hf_flash_attn(
old_flash_attention_forward = (
transformers.modeling_flash_attention_utils._flash_attention_forward
)
- new_flash_attention_forward = create_flash_attn_forward(
+ new_flash_attention_forward = create_flash_attn_forward_varlen_llama3(
process_group=process_group, ring_attn_func=ring_attn_func
)
diff --git a/src/axolotl/monkeypatch/ring_attn/patch.py b/src/axolotl/monkeypatch/ring_attn/patch.py
index 7d733cfc1..e1fd10b3a 100644
--- a/src/axolotl/monkeypatch/ring_attn/patch.py
+++ b/src/axolotl/monkeypatch/ring_attn/patch.py
@@ -8,69 +8,141 @@ We also provide some patches for accelerate functions to prepare the dataloader
sequence parallelism training.
"""
-import inspect
+import os
+from typing import Optional
-import accelerate
import torch
import torch.distributed as dist
-from accelerate.logging import get_logger
+from torch.distributed import DeviceMesh
+
+try:
+ from transformers.modeling_flash_attention_utils import _flash_supports_window
+except ImportError:
+ try:
+ from transformers.modeling_flash_attention_utils import (
+ _flash_supports_window_size as _flash_supports_window,
+ )
+ except ImportError:
+ _flash_supports_window = True
from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
+from axolotl.utils.logging import get_logger
from axolotl.utils.schemas.enums import RingAttnFunc
LOG = get_logger(__name__)
-
RING_ATTN_GROUP = None
-ORIGINAL_PREPARE_DATALOADER_CODE = """ submesh_fsdp_size = 1
- submesh_dp_size = 1
- submesh_tp_size = 1
- if "tp" in torch_device_mesh.mesh_dim_names:
- submesh_tp_size = torch_device_mesh["tp"].size()
- if "dp" in torch_device_mesh.mesh_dim_names:
- submesh_dp_size = torch_device_mesh["dp"].size()
- if "fsdp" in torch_device_mesh.mesh_dim_names:
- submesh_fsdp_size = torch_device_mesh["fsdp"].size()
- process_index = process_index // submesh_tp_size"""
-
-NEW_PREPARE_DATALOADER_CODE = """ submesh_fsdp_size = 1
- submesh_dp_size = 1
- submesh_tp_size = 1
- submesh_cp_size = 1
- if "cp" in torch_device_mesh.mesh_dim_names:
- submesh_cp_size = torch_device_mesh["cp"].size()
- if "tp" in torch_device_mesh.mesh_dim_names:
- submesh_tp_size = torch_device_mesh["tp"].size()
- if "dp" in torch_device_mesh.mesh_dim_names:
- submesh_dp_size = torch_device_mesh["dp"].size()
- if "fsdp" in torch_device_mesh.mesh_dim_names:
- submesh_fsdp_size = torch_device_mesh["fsdp"].size()
- process_index = process_index // (submesh_tp_size * submesh_cp_size)"""
-
def get_ring_attn_group() -> dist.ProcessGroup:
"""Getter for ring attention group on this rank."""
if RING_ATTN_GROUP is None:
- raise RuntimeError("register_ring_attn() not yet called")
+ raise RuntimeError("register_ring_attn_from_device_mesh() not yet called")
return RING_ATTN_GROUP
def set_ring_attn_group(ring_attn_group: dist.ProcessGroup | None):
"""Setter for ring attention group on this rank."""
- global RING_ATTN_GROUP # pylint: disable=global-statement
+ global RING_ATTN_GROUP
RING_ATTN_GROUP = ring_attn_group
-def register_ring_attn(
- sequence_parallel_degree: int,
+def create_ring_flash_attention_forward(
+ process_group: dist.ProcessGroup, heads_k_stride: int
+):
+ from ring_flash_attn import llama3_flash_attn_varlen_func
+ from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS
+
+ def _flash_attention_forward_v3(
+ query_states: torch.Tensor,
+ key_states: torch.Tensor,
+ value_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ query_length: int,
+ is_causal: bool,
+ dropout: float = 0.0,
+ position_ids: Optional[torch.Tensor] = None,
+ softmax_scale: Optional[float] = None,
+ sliding_window: Optional[int] = None,
+ use_top_left_mask: bool = False,
+ softcap: Optional[float] = None,
+ deterministic: bool = None,
+ cu_seq_lens_q: Optional[torch.LongTensor] = None,
+ cu_seq_lens_k: Optional[torch.LongTensor] = None,
+ max_length_q: Optional[int] = None,
+ max_length_k: Optional[int] = None,
+ target_dtype: Optional[torch.dtype] = None,
+ attn_implementation: Optional[str] = None,
+ **kwargs,
+ ):
+ if not use_top_left_mask:
+ causal = is_causal
+ else:
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__.
+ causal = is_causal and query_length != 1
+
+ # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).
+ use_sliding_windows = (
+ _flash_supports_window
+ and sliding_window is not None
+ and key_states.shape[1] > sliding_window
+ )
+ flash_kwargs = (
+ {"window_size": (sliding_window, sliding_window)}
+ if use_sliding_windows
+ else {}
+ )
+
+ if deterministic is None:
+ deterministic = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
+ flash_kwargs["deterministic"] = deterministic
+ assert softcap is None, (
+ "llama3_flash_attn_varlen_func does not support softcap yet."
+ )
+ # flash_kwargs["softcap"] = softcap
+ flash_kwargs["group"] = process_group
+
+ # not sure why attention_mask can be not None...
+ assert causal, "only causal attention is supported yet."
+ batch_size = query_states.size(0)
+ assert batch_size == 1, "varlen data should be processed in advance."
+
+ attn_output = llama3_flash_attn_varlen_func(
+ query_states.squeeze(dim=0),
+ key_states.squeeze(dim=0),
+ value_states.squeeze(dim=0),
+ cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"],
+ cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"],
+ max_seqlen_q=DATA_PARAMS["max_seqlen_q"],
+ max_seqlen_k=DATA_PARAMS["max_seqlen_k"],
+ heads_k_stride=heads_k_stride,
+ local_k_slice=DATA_PARAMS["local_k_slice"],
+ dropout_p=dropout,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ **flash_kwargs,
+ )
+
+ attn_output = attn_output.unsqueeze(dim=0)
+
+ return attn_output
+
+ return [
+ _flash_attention_forward_v3,
+ ]
+
+
+def register_ring_attn_from_device_mesh(
+ device_mesh: "DeviceMesh",
+ context_parallel_dim: tuple[str, ...],
heads_k_stride: int | None,
ring_attn_func: RingAttnFunc | None,
):
- """Create ring attention group and substitute flash attn with ring flash attn.
+ """Create ring attention group using DeviceMesh and substitute flash attn with ring flash attn.
Args:
- sequence_parallel_degree: Sequence parallelism factor.
+ device_mesh: DeviceMesh object containing the parallelism topology.
+ context_parallel_dim: Name of the sequence parallel dimension in the device mesh.
heads_k_stride: Sequence parallelism K head stride size. Passed through to
`varlen_llama3` `ring_flash_attn` implementation.
ring_attn_func: `ring_flash_attn` ring attention implemention. If sample
@@ -78,49 +150,55 @@ def register_ring_attn(
`batch` function.
"""
rank = dist.get_rank()
- world_size = dist.get_world_size()
+
+ LOG.info(
+ f"Enabling ring attention sequence parallelism using DeviceMesh "
+ f"dimension '{context_parallel_dim}'",
+ main_process_only=True,
+ )
+
+ # Extract the sequence parallel submesh
+ try:
+ sequence_mesh = device_mesh[context_parallel_dim]
+ except (KeyError, IndexError) as e:
+ raise ValueError(
+ f"Dimension '{context_parallel_dim}' not found in device_mesh. "
+ f"Available dimensions: {device_mesh.mesh_dim_names}"
+ ) from e
+
+ # Get the process group for context parallelism
+ sequence_pg = sequence_mesh.get_group()
+ context_parallel_size = sequence_mesh.size()
if rank == 0:
LOG.info(
- "Enabling ring attention sequence parallelism: "
- f"each sequence will be processed across {sequence_parallel_degree} GPUs"
+ f"Sequence parallel degree: {context_parallel_size}, "
+ f"mesh shape: {sequence_mesh.mesh.shape}"
)
- assert sequence_parallel_degree <= world_size, (
- f"sequence_parallel_degree ({sequence_parallel_degree}) "
- f"must be less than or equal to world_size ({world_size})"
- )
- assert world_size % sequence_parallel_degree == 0, (
- f"sequence_parallel_degree ({sequence_parallel_degree}) "
- f"must evenly divide world_size ({world_size})"
- )
+ # Log which ranks are in the current process group
+ if sequence_pg != dist.GroupMember.WORLD:
+ ranks_in_group = dist.get_process_group_ranks(sequence_pg)
+ LOG.info(f"Current sequence parallel group ranks: {ranks_in_group}")
- # Assign ranks to sequence parallel groups
- group_assignments = {}
- for i in range(world_size // sequence_parallel_degree):
- ring_attn_ranks = list(
- range(
- i * sequence_parallel_degree,
- (i + 1) * sequence_parallel_degree,
- )
- )
- group = dist.new_group(ranks=ring_attn_ranks, backend="nccl")
-
- # Track which GPUs are in which groups
- for r in ring_attn_ranks:
- group_assignments[r] = i
-
- if rank in ring_attn_ranks:
- set_ring_attn_group(group)
-
- # Log the GPU group assignments
- if rank == 0:
- LOG.info(f"Sequence parallel group assignments: {group_assignments}")
+ # Set the ring attention group
+ set_ring_attn_group(sequence_pg)
if ring_attn_func is RingAttnFunc.VARLEN_LLAMA3:
- from ring_flash_attn import substitute_hf_flash_attn
+ # fmt: off
+ import ring_flash_attn.adapters.hf_adapter
- substitute_hf_flash_attn(
+ from ring_flash_attn.adapters.hf_adapter import ( # isort: skip
+ create_ring_flash_attention_forward as create_ring_flash_attention_forward_orig,
+ )
+
+ create_ring_flash_attention_forward_orig = ( # noqa: F811,F841
+ create_ring_flash_attention_forward
+ )
+ ring_flash_attn.adapters.hf_adapter.create_ring_flash_attention_forward = create_ring_flash_attention_forward
+ # fmt: on
+
+ ring_flash_attn.adapters.hf_adapter.substitute_hf_flash_attn(
process_group=get_ring_attn_group(), heads_k_stride=heads_k_stride or 1
)
elif ring_attn_func is RingAttnFunc.BATCH_RING:
@@ -147,79 +225,3 @@ def update_ring_attn_params(position_ids: torch.Tensor | None):
cu_seqlens, _ = get_cu_seqlens_from_pos_ids(position_ids)
cu_seqlens = cu_seqlens.squeeze().to(device=torch.cuda.current_device())
update_ring_flash_attn_params(cu_seqlens, get_ring_attn_group())
-
-
-def patch_prepare_data_loader():
- """Patch `accelerate.data_loader.prepare_data_loader` to respect the SP degree.
-
- Raies:
- RuntimeError: If source code to patch does not exist.
- """
- original_fn = accelerate.data_loader.prepare_data_loader
- original_source = inspect.getsource(original_fn)
-
- if ORIGINAL_PREPARE_DATALOADER_CODE not in original_source:
- raise RuntimeError(
- "SP patch failed - target snippet not found. "
- "Check accelerate's version or update the patch."
- )
-
- patched_source = original_source.replace(
- ORIGINAL_PREPARE_DATALOADER_CODE, NEW_PREPARE_DATALOADER_CODE
- )
-
- # Create a new function from the patched source
- namespace = {}
- exec( # pylint: disable=exec-used # nosec B102
- patched_source, accelerate.data_loader.__dict__, namespace
- )
- patched_function = namespace["prepare_data_loader"]
-
- accelerate.data_loader.prepare_data_loader = patched_function
- LOG.info("Patched accelerate.data_loader.prepare_data_loader for SP support")
-
-
-def patch_prepare_device_mesh(sequence_parallel_degree: int):
- """Patches the `Accelerator._prepare_device_mesh` method to create a device mesh
- that includes sequence parallelism with the specified degree.
-
- Args:
- sequence_parallel_degree (int): The degree of sequence parallelism to use.
- """
-
- def _prepare_device_mesh(self):
- """Prepare the device mesh for distributed training. The dataloader will
- determine how to load data based on the device mesh.
- """
- if self.state.torch_tp_plugin:
- return self.state.torch_tp_plugin.torch_device_mesh
- if (
- self.distributed_type == accelerate.accelerator.DistributedType.DEEPSPEED
- and hasattr(self.state, "ds_device_mesh")
- ):
- return self.state.ds_device_mesh
-
- # Create device mesh with sequence parallelism
- world_size = dist.get_world_size()
- mesh_shape = (
- world_size // sequence_parallel_degree,
- sequence_parallel_degree,
- )
- device_ids = list(range(world_size))
-
- # Note that we use "cp" instead of "sp" to match the PyTorch native "context
- # parallelism" implementation naming
- return dist.DeviceMesh(
- "cuda",
- torch.tensor(device_ids).reshape(mesh_shape),
- mesh_dim_names=("dp", "cp"),
- )
-
- # Replace the original method with our new method
- # pylint: disable=protected-access
- accelerate.accelerator.Accelerator._prepare_device_mesh = _prepare_device_mesh
-
- LOG.info(
- "Successfully patched Accelerator._prepare_device_mesh "
- f"with sequence_parallel_degree={sequence_parallel_degree}"
- )
diff --git a/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py b/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
index 85454fe2e..0fa6d6424 100644
--- a/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/stablelm_attn_hijack_flash.py
@@ -16,8 +16,8 @@
# This code is based off the following work:
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
-# pylint: disable=duplicate-code
"""PyTorch StableLM Epoch model."""
+
import importlib
import math
from typing import Optional, Tuple, Union
@@ -26,7 +26,7 @@ import torch
import torch.utils.checkpoint
from accelerate import init_empty_weights
from einops import rearrange
-from flash_attn.flash_attn_interface import ( # pylint: disable=ungrouped-imports
+from flash_attn.flash_attn_interface import (
flash_attn_varlen_qkvpacked_func,
)
from torch import nn
@@ -49,27 +49,21 @@ def replace_stablelm_attn_with_flash_attn(model_name="stabilityai/stablelm-3b-4e
".configuration_stablelm_epoch", ".modeling_stablelm_epoch"
)
modeling_stablelm = importlib.import_module(module_name)
- modeling_stablelm.Attention.forward = ( # pylint: disable=protected-access
- flashattn_attn
- )
- modeling_stablelm.StableLMEpochModel.forward = ( # pylint: disable=protected-access
- stablelm_model_forward
- )
- modeling_stablelm.DecoderLayer.forward = ( # pylint: disable=protected-access
- decoder_layer_forward
- )
+ modeling_stablelm.Attention.forward = flashattn_attn
+ modeling_stablelm.StableLMEpochModel.forward = stablelm_model_forward
+ modeling_stablelm.DecoderLayer.forward = decoder_layer_forward
def rotate_half(x: torch.Tensor):
"""Rotates half the hidden dims of the input."""
- # pylint: disable=invalid-name
+
x1, x2 = torch.chunk(x, 2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
# The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
- # pylint: disable=invalid-name
+
cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
cos = cos[position_ids].unsqueeze(1) # [batch_size, 1, seq_len, dim]
@@ -99,7 +93,7 @@ def flashattn_attn(
attention_mask: torch.FloatTensor,
position_ids: torch.LongTensor,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
- output_attentions: Optional[bool] = False, # pylint: disable=unused-argument
+ output_attentions: Optional[bool] = False,
use_cache: Optional[bool] = False,
cu_seqlens: Optional[torch.Tensor] = None,
max_seqlen: Optional[torch.Tensor] = None,
@@ -216,7 +210,6 @@ def decoder_layer_forward(
) -> Union[
Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]
]:
- # pylint: disable=duplicate-code
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
@@ -263,7 +256,6 @@ def stablelm_model_forward(
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, BaseModelOutputWithPast]:
- # pylint: disable=duplicate-code
output_attentions = (
output_attentions
if output_attentions is not None
@@ -326,13 +318,11 @@ def stablelm_model_forward(
dtype=torch.bool,
device=inputs_embeds.device,
)
- attention_mask = (
- self._prepare_decoder_attention_mask( # pylint: disable=protected-access
- attention_mask,
- (batch_size, seq_length),
- inputs_embeds,
- past_key_values_length,
- )
+ attention_mask = self._prepare_decoder_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
)
hidden_states = inputs_embeds
diff --git a/src/axolotl/monkeypatch/tiled_mlp/__init__.py b/src/axolotl/monkeypatch/tiled_mlp/__init__.py
new file mode 100644
index 000000000..4ea154991
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/__init__.py
@@ -0,0 +1,11 @@
+"""
+TiledMLP monkey patches
+"""
+
+from .patch import (
+ patch_tiled_mlp,
+)
+
+__all__ = [
+ "patch_tiled_mlp",
+]
diff --git a/src/axolotl/monkeypatch/tiled_mlp/base.py b/src/axolotl/monkeypatch/tiled_mlp/base.py
new file mode 100644
index 000000000..2c9dc8e4c
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/base.py
@@ -0,0 +1,256 @@
+"""
+TiledMLP support for DDP, FSDP, and single GPU
+"""
+
+import threading
+from typing import List
+
+import torch
+
+
+class DeepSpeedTiledMLPMoE(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx,
+ fn,
+ self,
+ x,
+ shards,
+ compute_params,
+ ) -> torch.Tensor:
+ ctx.fn = fn
+ ctx.self = self
+ ctx.shards = shards
+ ctx.compute_params = [p for p in compute_params if p.requires_grad]
+ ctx.save_for_backward(x)
+
+ x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+ with torch.no_grad():
+ output_shards = [fn(self, x_shard) for x_shard in x_shards]
+
+ ctx.is_tuple_output = isinstance(output_shards[0], tuple)
+ if isinstance(output_shards[0], tuple):
+ tuple_dim_idx = [1, 0]
+ output_unsharded = tuple(
+ torch.cat(
+ [output_shard[i] for output_shard in output_shards],
+ dim=tuple_dim_idx[i],
+ )
+ for i in range(len(output_shards[0]))
+ )
+ else:
+ output_unsharded = torch.cat(output_shards, dim=1)
+
+ return output_unsharded
+
+ @staticmethod
+ def backward(ctx, *grads) -> torch.Tensor:
+ fn = ctx.fn
+ (x,) = ctx.saved_tensors
+ self = ctx.self
+ shards = ctx.shards
+ compute_params = ctx.compute_params
+ is_tuple_output = ctx.is_tuple_output
+
+ x_requires_grad = x.requires_grad
+ x = x.detach()
+ # detach() unsets `x.requires_grad`, so restore it
+ x.requires_grad_(x_requires_grad)
+
+ incoming_grad = grads[0]
+ x_grad = torch.zeros_like(x)
+ x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+
+ shard_step = x_shards[0].numel()
+ for i, x_shard in enumerate(x_shards):
+ # Tell deepspeed not to add a new grad to its ipg bucket until the last shard is run
+ if compute_params is not None:
+ if i + 1 < shards:
+ for param in compute_params:
+ param.ds_grad_is_ready = False
+ else:
+ # last shard, can add the grad
+ for param in compute_params:
+ param.ds_grad_is_ready = True
+
+ x_shard.requires_grad_(x_requires_grad)
+
+ shard_offset = i * shard_step
+ x_shard.grad = (
+ x_grad.view(-1)
+ .narrow(0, shard_offset, x_shard.numel())
+ .view_as(x_shard)
+ )
+ incoming_grad_shard = (
+ incoming_grad.view(-1)
+ .narrow(0, shard_offset, x_shard.numel())
+ .view_as(x_shard)
+ )
+ with torch.enable_grad():
+ output = fn(self, x_shard)
+ if is_tuple_output:
+ torch.autograd.backward(output[0], incoming_grad_shard)
+ else:
+ torch.autograd.backward(output, incoming_grad_shard)
+
+ return (None, None, x_grad, None, None)
+
+
+class TiledMLP(torch.autograd.Function):
+ """
+ TiledMLP implementation using gradient hooks
+ """
+
+ @staticmethod
+ def forward(
+ ctx,
+ fn,
+ self,
+ x,
+ shards,
+ compute_params,
+ ) -> torch.Tensor:
+ ctx.fn = fn
+ ctx.self = self
+ ctx.shards = shards
+ ctx.compute_params = [p for p in compute_params if p.requires_grad]
+ ctx.save_for_backward(x)
+
+ x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+ with torch.no_grad():
+ output_shards = [fn(self, x_shard) for x_shard in x_shards]
+ ctx.is_tuple_output = isinstance(output_shards[0], tuple)
+ if isinstance(output_shards[0], tuple):
+ tuple_dim_idx = [1, 0]
+ output_unsharded = tuple(
+ torch.cat(
+ [output_shard[i] for output_shard in output_shards],
+ dim=tuple_dim_idx[i],
+ )
+ for i in range(len(output_shards[0]))
+ )
+ else:
+ output_unsharded = torch.cat(output_shards, dim=1)
+
+ return output_unsharded
+
+ @staticmethod
+ def backward(ctx, *grads) -> torch.Tensor:
+ fn = ctx.fn
+ (x,) = ctx.saved_tensors
+ self = ctx.self
+ shards = ctx.shards
+ compute_params = ctx.compute_params
+ is_tuple_output = ctx.is_tuple_output
+
+ x_requires_grad = x.requires_grad
+ x = x.detach()
+ x.requires_grad_(x_requires_grad)
+
+ incoming_grad = grads[0]
+ x_grad = torch.zeros_like(x)
+ x_shards = list(torch.chunk(x, chunks=shards, dim=1))
+
+ # Create a gradient accumulator for parameters
+ grad_accumulator = GradientAccumulator(compute_params, shards, dtype=x.dtype)
+
+ shard_step = x_shards[0].numel()
+ for i, x_shard in enumerate(x_shards):
+ x_shard.requires_grad_(x_requires_grad)
+
+ shard_offset = i * shard_step
+ x_shard.grad = (
+ x_grad.view(-1)
+ .narrow(0, shard_offset, x_shard.numel())
+ .view_as(x_shard)
+ )
+ incoming_grad_shard = (
+ incoming_grad.view(-1)
+ .narrow(0, shard_offset, x_shard.numel())
+ .view_as(x_shard)
+ )
+
+ # Install hooks for this shard
+ is_last_shard = i + 1 == shards
+ grad_accumulator.install_hooks(is_last_shard)
+
+ with torch.enable_grad():
+ output = fn(self, x_shard)
+ if is_tuple_output:
+ torch.autograd.backward(output[0], incoming_grad_shard)
+ else:
+ torch.autograd.backward(output, incoming_grad_shard)
+
+ # Clean up hooks
+ grad_accumulator.cleanup()
+ del grad_accumulator
+
+ return (None, None, x_grad, None, None)
+
+
+class GradientAccumulator:
+ """
+ Manual gradient accumulator for TiledMLP with configurable precision
+ Accumulates in specified dtype and rescales the gradient at the end
+ """
+
+ def __init__(
+ self,
+ params: List[torch.nn.Parameter],
+ total_shards: int,
+ dtype: torch.dtype | None = None,
+ ):
+ self.params = params
+ self.total_shards = total_shards
+ self.grad_accumulation_dtype = dtype or torch.float32
+ self.accumulated_grads = {}
+ self.hooks = []
+ self.lock = threading.Lock()
+ self.gradient_scale = 1.0 / total_shards
+
+ # Initialize accumulated gradients in the specified dtype
+ for param in self.params:
+ if param.grad is not None:
+ self.accumulated_grads[param] = param.grad.to(
+ self.grad_accumulation_dtype
+ )
+ param.grad = None
+ else:
+ self.accumulated_grads[param] = torch.zeros_like(
+ param, dtype=self.grad_accumulation_dtype
+ )
+
+ def install_hooks(self, is_last_shard: bool):
+ """Install gradient hooks that accumulate gradients in higher precision"""
+
+ def create_hook(param):
+ def hook(grad):
+ with self.lock:
+ grad_to_accum_dtype = grad.to(self.grad_accumulation_dtype)
+ scaled_grad = grad_to_accum_dtype * self.gradient_scale
+
+ if param in self.accumulated_grads:
+ self.accumulated_grads[param] += scaled_grad
+ else:
+ self.accumulated_grads[param] = scaled_grad.clone()
+
+ # Only assign the averaged gradient on the last shard
+ if is_last_shard:
+ param.grad = self.accumulated_grads[param].to(param.dtype)
+ return param.grad
+ return None
+
+ return hook
+
+ # Install hooks on all parameters
+ for param in self.params:
+ if param.requires_grad:
+ hook = param.register_hook(create_hook(param))
+ self.hooks.append(hook)
+
+ def cleanup(self):
+ """Remove all installed hooks"""
+ for hook in self.hooks:
+ hook.remove()
+ self.hooks.clear()
+ del self.accumulated_grads
diff --git a/src/axolotl/monkeypatch/tiled_mlp/patch.py b/src/axolotl/monkeypatch/tiled_mlp/patch.py
new file mode 100644
index 000000000..c0f89236b
--- /dev/null
+++ b/src/axolotl/monkeypatch/tiled_mlp/patch.py
@@ -0,0 +1,93 @@
+"""Monkeypatch for Tiled MLP implementation"""
+
+import math
+import os
+
+import torch
+import torch.distributed as dist
+
+from axolotl.utils.callbacks.models import get_causal_lm_model_cls_prefix
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def patch_tiled_mlp(model_type, use_original_mlp=True, cfg_num_shards=None):
+ from deepspeed.runtime.sequence_parallel.ulysses_sp import (
+ TiledMLP as DeepSpeedTiledMLP,
+ )
+
+ from axolotl.monkeypatch.tiled_mlp.base import DeepSpeedTiledMLPMoE, TiledMLP
+
+ try:
+ # Dynamically import the module and MLP class
+ module_path = f"transformers.models.{model_type}.modeling_{model_type}"
+ model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
+ module = __import__(module_path, fromlist=[f"{model_cls_prefix}MLP"])
+ mlp_cls = getattr(module, f"{model_cls_prefix}MLP")
+
+ if use_original_mlp:
+ mlp_forward = mlp_cls.forward
+ else:
+
+ def generic_mlp_forward(self_, hs):
+ return self_.down_proj(
+ self_.act_fn(self_.gate_proj(hs)) * self_.up_proj(hs)
+ )
+
+ mlp_forward = torch.compile(generic_mlp_forward)
+
+ is_distributed = int(os.environ.get("WORLD_SIZE", 1)) > 1
+
+ def tiled_mlp_forward(self, x):
+ input_shape = x.shape
+ seqlen = input_shape[-2]
+ hidden = input_shape[-1]
+ if cfg_num_shards is None:
+ num_shards = math.ceil(seqlen / hidden)
+ if is_distributed:
+ num_shards_tensor = torch.tensor(num_shards, device=x.device)
+ dist.all_reduce(num_shards_tensor, op=dist.ReduceOp.MAX)
+ num_shards = num_shards_tensor.item()
+ else:
+ num_shards = cfg_num_shards
+
+ if not self._compute_params:
+ self._compute_params = [p for p in self.parameters() if p.requires_grad]
+
+ compute_params = self._compute_params
+ if not self._tiled_mlp_dist_impl:
+ if (
+ self._compute_params
+ and any(
+ hasattr(p, "ds_id") or hasattr(p, "param_idx_in_group")
+ for p in self._compute_params
+ )
+ ) or os.environ.get("ACCELERATE_USE_DEEPSPEED", "false") == "true":
+ if model_type == "gpt_oss":
+ self._tiled_mlp_dist_impl = DeepSpeedTiledMLPMoE
+ else:
+ self._tiled_mlp_dist_impl = DeepSpeedTiledMLP
+ else:
+ self._tiled_mlp_dist_impl = TiledMLP
+
+ down_res = self._tiled_mlp_dist_impl.apply(
+ mlp_forward,
+ self,
+ x,
+ num_shards,
+ compute_params,
+ )
+ return down_res
+
+ mlp_cls.forward = tiled_mlp_forward
+ mlp_cls._compute_params = []
+ mlp_cls._tiled_mlp_dist_impl = None
+ LOG.info(
+ f"Successfully monkey-patched TiledMLP for model_type: {model_type}",
+ main_process_only=True,
+ )
+ except (ImportError, AttributeError) as e:
+ raise RuntimeError(
+ f"Could not import MLP class for model_type: {model_type}. Error: {str(e)}"
+ ) from e
diff --git a/src/axolotl/monkeypatch/trainer/lr.py b/src/axolotl/monkeypatch/trainer/lr.py
index 9afc23c46..c33674cee 100644
--- a/src/axolotl/monkeypatch/trainer/lr.py
+++ b/src/axolotl/monkeypatch/trainer/lr.py
@@ -39,4 +39,4 @@ def _get_learning_rate(self):
def patch_trainer_get_lr():
from transformers.trainer import Trainer
- Trainer._get_learning_rate = _get_learning_rate # pylint: disable=protected-access
+ Trainer._get_learning_rate = _get_learning_rate
diff --git a/src/axolotl/monkeypatch/trainer/trl.py b/src/axolotl/monkeypatch/trainer/trl.py
new file mode 100644
index 000000000..bca9f92de
--- /dev/null
+++ b/src/axolotl/monkeypatch/trainer/trl.py
@@ -0,0 +1,13 @@
+"""Monkeypatch for TRL trainer FSDP preparation."""
+
+
+def prepare_fsdp(model, accelerator):
+ from axolotl.monkeypatch.accelerate.fsdp2 import fsdp2_prepare_model
+
+ return fsdp2_prepare_model(accelerator, model)
+
+
+def patch_trl_prepare_fsdp2():
+ import trl.models.utils
+
+ trl.models.utils.prepare_fsdp = prepare_fsdp
diff --git a/src/axolotl/monkeypatch/trainer_accelerator_args.py b/src/axolotl/monkeypatch/trainer_accelerator_args.py
index 0a5b27c13..9fc6e38c6 100644
--- a/src/axolotl/monkeypatch/trainer_accelerator_args.py
+++ b/src/axolotl/monkeypatch/trainer_accelerator_args.py
@@ -18,7 +18,7 @@ ORIGINAL_TRAINER_CODE = """
PATCHED_TRAINER_CODE = """
if hasattr(self, "additional_accelerator_args"):
- additional_args = self.additional_accelerator_args(fp8=True, **args)
+ additional_args = self.additional_accelerator_args(fp8=True, enable_fsdp_float8_all_gather={enable_fsdp_float8_all_gather}, **args)
if additional_args:
args.update(additional_args)
@@ -38,23 +38,24 @@ def check_create_accelerate_code_is_patchable() -> bool:
return ORIGINAL_TRAINER_CODE in create_code
-def patch_create_accelerate_code_for_fp8():
+def patch_create_accelerate_code_for_fp8(enable_fsdp_float8_all_gather: bool):
"""
- monkeypatch create_accelerator_and_postprocess so it checks for additional kwargs
+ Monkeypatch create_accelerator_and_postprocess so it checks for additional kwargs.
"""
try:
create_code = get_create_accelerate_code()
except OSError:
return
- Trainer._original_create_accelerator_and_postprocess = ( # pylint: disable=protected-access
- create_code
- )
+ Trainer._original_create_accelerator_and_postprocess = create_code
create_code, _ = detab_code(create_code)
if ORIGINAL_TRAINER_CODE not in create_code:
return
- create_code = create_code.replace(ORIGINAL_TRAINER_CODE, PATCHED_TRAINER_CODE)
+ patched_trainer_code = PATCHED_TRAINER_CODE.format(
+ enable_fsdp_float8_all_gather=enable_fsdp_float8_all_gather
+ )
+ create_code = create_code.replace(ORIGINAL_TRAINER_CODE, patched_trainer_code)
create_code = create_code.replace(
"def create_accelerator_and_postprocess(",
"def fixed_create_accelerator_and_postprocess(",
@@ -69,12 +70,14 @@ def patch_create_accelerate_code_for_fp8():
if item in create_code:
items_to_import.append(item)
- exec( # pylint: disable=exec-used # nosec B102
+ exec(
"from transformers.trainer import ("
+ ", ".join(x for x in items_to_import)
+ ")",
globals(),
)
- exec(create_code, globals()) # pylint: disable=exec-used # nosec B102
+ exec(create_code, globals())
LOG.info("patching create_accelerator_and_postprocess to allow for overrides")
- Trainer.create_accelerator_and_postprocess = fixed_create_accelerator_and_postprocess # pylint: disable=protected-access # pylint: disable=undefined-variable # noqa: F821
+ Trainer.create_accelerator_and_postprocess = (
+ fixed_create_accelerator_and_postprocess
+ )
diff --git a/src/axolotl/monkeypatch/trainer_eval_guard.py b/src/axolotl/monkeypatch/trainer_eval_guard.py
deleted file mode 100644
index 8488a16df..000000000
--- a/src/axolotl/monkeypatch/trainer_eval_guard.py
+++ /dev/null
@@ -1,78 +0,0 @@
-"""
-fix for FSDP2 evals when using torch.compile
-"""
-
-import inspect
-
-from transformers import Trainer
-
-from axolotl.monkeypatch.utils import detab_code
-from axolotl.utils.logging import get_logger
-
-LOG = get_logger(__name__)
-
-ORIGINAL_TRAINER_CODE = """
- model.eval()
-"""
-
-PATCHED_TRAINER_CODE = """
- if hasattr(model, "eval") and callable(model.eval):
- self.model.eval()
-"""
-
-
-def get_evaluation_loop_code() -> str:
- training_loop = inspect.getsource(Trainer.evaluation_loop)
- return training_loop
-
-
-def check_evaluation_loop_is_patchable() -> bool:
- eval_loop = get_evaluation_loop_code()
- eval_loop, _ = detab_code(eval_loop)
- return ORIGINAL_TRAINER_CODE in eval_loop
-
-
-def patch_evaluation_loop_for_fsdp2():
- """
- monkeypatch for fixing the eval loop for fsdp2 with torch.compile
- """
-
- try:
- evaluation_loop = get_evaluation_loop_code()
- except OSError:
- return
- Trainer._original_evaluation_loop = ( # pylint: disable=protected-access
- evaluation_loop
- )
- evaluation_loop, _ = detab_code(evaluation_loop)
- if ORIGINAL_TRAINER_CODE not in evaluation_loop:
- return
-
- evaluation_loop = evaluation_loop.replace(
- ORIGINAL_TRAINER_CODE, PATCHED_TRAINER_CODE
- )
- evaluation_loop = evaluation_loop.replace(
- "def evaluation_loop(",
- "def _fixed_evaluation_loop(",
- 1,
- )
-
- # load imports necessary
- import transformers.trainer
-
- items_to_import = []
- for item in dir(transformers.trainer):
- if item in evaluation_loop:
- items_to_import.append(item)
-
- exec( # pylint: disable=exec-used # nosec B102
- "from transformers.trainer import ("
- + ", ".join(x for x in items_to_import)
- + ")",
- globals(),
- )
- exec(evaluation_loop, globals()) # pylint: disable=exec-used # nosec B102
- LOG.info("patching _inner_training_loop for fsdp optimizer save")
- Trainer.evaluation_loop = ( # pylint: disable=protected-access
- _fixed_evaluation_loop # pylint: disable=undefined-variable # noqa: F821
- )
diff --git a/src/axolotl/monkeypatch/trainer_fsdp_optim.py b/src/axolotl/monkeypatch/trainer_fsdp_optim.py
index 4ce5b8ecd..692f754d7 100644
--- a/src/axolotl/monkeypatch/trainer_fsdp_optim.py
+++ b/src/axolotl/monkeypatch/trainer_fsdp_optim.py
@@ -12,22 +12,18 @@ from axolotl.utils.logging import get_logger
LOG = get_logger(__name__)
ORIGINAL_TRAINER_CODE = """
-
- delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled
-
+ if delay_optimizer_creation:
+ self.optimizer = self.accelerator.prepare(self.optimizer)
"""
PATCHED_TRAINER_CODE = """
-
- delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled
-
+ if delay_optimizer_creation:
+ model = self.accelerator.prepare(self.model)
"""
def get_training_loop_code() -> str:
- training_loop = inspect.getsource(
- Trainer._inner_training_loop # pylint: disable=protected-access
- )
+ training_loop = inspect.getsource(Trainer._inner_training_loop)
return training_loop
@@ -46,9 +42,7 @@ def patch_training_loop_for_fsdp():
training_loop = get_training_loop_code()
except OSError:
return
- Trainer._original_inner_training_loop = ( # pylint: disable=protected-access
- training_loop
- )
+ Trainer._original_inner_training_loop = training_loop
training_loop, _ = detab_code(training_loop)
if ORIGINAL_TRAINER_CODE not in training_loop:
return
@@ -68,14 +62,12 @@ def patch_training_loop_for_fsdp():
if item in training_loop:
items_to_import.append(item)
- exec( # pylint: disable=exec-used # nosec B102
+ exec(
"from transformers.trainer import ("
+ ", ".join(x for x in items_to_import)
+ ")",
globals(),
)
- exec(training_loop, globals()) # pylint: disable=exec-used # nosec B102
+ exec(training_loop, globals())
LOG.info("patching _inner_training_loop for fsdp optimizer save")
- Trainer._inner_training_loop = ( # pylint: disable=protected-access
- _fixed_inner_training_loop # pylint: disable=undefined-variable # noqa: F821
- )
+ Trainer._inner_training_loop = _fixed_inner_training_loop
diff --git a/src/axolotl/monkeypatch/transformers/__init__.py b/src/axolotl/monkeypatch/transformers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/src/axolotl/monkeypatch/transformers/trainer_context_parallel.py b/src/axolotl/monkeypatch/transformers/trainer_context_parallel.py
new file mode 100644
index 000000000..74a35e83f
--- /dev/null
+++ b/src/axolotl/monkeypatch/transformers/trainer_context_parallel.py
@@ -0,0 +1,68 @@
+"""Monkey patch to allow context parallelism with FlashAttention in HF Trainer."""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+
+from transformers import Trainer
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+GUARD_PATTERN = 'if model.config._attn_implementation != "sdpa":'
+PATCHED_GUARD = (
+ 'if model.config._attn_implementation not in ("sdpa", "flash_attention_2"):'
+)
+
+
+def patch_prepare_context_parallel_inputs() -> None:
+ """Relax the SDPA-only guard when running context parallelism with FlashAttention."""
+ if getattr(Trainer, "_axolotl_prepare_context_parallel_inputs_patched", False):
+ LOG.debug("Trainer._prepare_context_parallel_inputs already patched")
+ return
+
+ try:
+ original_source = inspect.getsource(Trainer._prepare_context_parallel_inputs)
+ except OSError as exc: # pragma: no cover - occurs when source is unavailable
+ LOG.warning("Unable to patch Trainer._prepare_context_parallel_inputs: %s", exc)
+ return
+
+ if GUARD_PATTERN not in original_source:
+ LOG.warning(
+ "Expected guard not found in Trainer._prepare_context_parallel_inputs; \n"
+ "skipping FlashAttention context parallelism patch"
+ )
+ return
+
+ patched_source = original_source.replace(GUARD_PATTERN, PATCHED_GUARD)
+ patched_source, _ = detab_code(patched_source)
+ patched_source = patched_source.replace(
+ "def _prepare_context_parallel_inputs(",
+ "def axolotl_prepare_context_parallel_inputs(",
+ 1,
+ )
+
+ module_name = Trainer.__module__
+ module = importlib.import_module(module_name)
+
+ # import symbols referenced in the method so exec can succeed
+ items_to_import = []
+ for item in dir(module):
+ if item in patched_source:
+ items_to_import.append(item)
+
+ exec(f"from {module_name} import ({', '.join(items_to_import)})", globals())
+ exec(patched_source, globals())
+
+ Trainer._original_prepare_context_parallel_inputs = (
+ Trainer._prepare_context_parallel_inputs
+ )
+ Trainer._prepare_context_parallel_inputs = axolotl_prepare_context_parallel_inputs
+ Trainer._axolotl_prepare_context_parallel_inputs_source = patched_source
+ Trainer._axolotl_prepare_context_parallel_inputs_patched = True
+ LOG.debug(
+ "Patched Trainer._prepare_context_parallel_inputs for FlashAttention + CP"
+ )
diff --git a/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py b/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py
new file mode 100644
index 000000000..b8172bbe6
--- /dev/null
+++ b/src/axolotl/monkeypatch/transformers/trainer_loss_calc.py
@@ -0,0 +1,139 @@
+"""
+Module for patching transformers Trainer loss calculation to use nanmean.
+
+This is needed for context parallelism since chunks of the input sequences may be fully
+masked and return NaNs in the loss calculation.
+
+Also includes a patch for FSDP2 + torch.compile. We need to bundle this together with
+the other evaluation_loop patch because we can't patch the same code twice without
+raising an OSError.
+"""
+
+import importlib
+import inspect
+
+from transformers import Trainer
+
+from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+ORIGINAL_EVAL_CODE = {
+ "list": 'metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item()',
+ "array": 'metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()',
+}
+PATCHED_EVAL_CODE = {
+ "list": 'metrics[f"{metric_key_prefix}_loss"] = np.nanmean(np.concatenate(all_losses)).item()',
+ "array": 'metrics[f"{metric_key_prefix}_loss"] = np.nanmean(all_losses).item()',
+}
+
+ORIGINAL_MAYBE_CODE = "tr_loss_scalar = self._nested_gather(tr_loss).mean().item()"
+PATCHED_MAYBE_CODE = "tr_loss_scalar = self._nested_gather(tr_loss).nanmean().item()"
+
+
+def check_evaluation_loop_is_patchable() -> bool:
+ evaluation_loop_source = inspect.getsource(Trainer.evaluation_loop)
+ return all(value in evaluation_loop_source for value in ORIGINAL_EVAL_CODE.values())
+
+
+def patch_evaluation_loop():
+ """Patch the evaluation_loop method."""
+ # Check if already patched
+ if hasattr(Trainer, "_original_evaluation_loop"):
+ LOG.debug("Trainer.evaluation_loop already patched")
+ return
+
+ # Check if the patterns exist
+ try:
+ evaluation_loop_source = inspect.getsource(Trainer.evaluation_loop)
+ except OSError:
+ return
+ Trainer.evaluation = evaluation_loop_source
+ evaluation_loop_source, _ = detab_code(evaluation_loop_source)
+
+ # Apply the nanmean patches
+ evaluation_loop_source = evaluation_loop_source.replace(
+ ORIGINAL_EVAL_CODE["list"], PATCHED_EVAL_CODE["list"]
+ )
+ evaluation_loop_source = evaluation_loop_source.replace(
+ ORIGINAL_EVAL_CODE["array"], PATCHED_EVAL_CODE["array"]
+ )
+
+ # Rename the function to avoid conflicts
+ evaluation_loop_source = evaluation_loop_source.replace(
+ "def evaluation_loop(",
+ "def axolotl_evaluation_loop(",
+ 1,
+ )
+
+ # Get the module for necessary imports
+ module_name = Trainer.__module__
+ module = importlib.import_module(module_name)
+
+ # Import necessary items from the module
+ items_to_import = []
+ for item in dir(module):
+ if item in evaluation_loop_source:
+ items_to_import.append(item)
+
+ # Execute the imports and patched method
+ exec(
+ f"from {module_name} import ({', '.join(items_to_import)})",
+ globals(),
+ )
+ exec(evaluation_loop_source, globals())
+
+ LOG.debug("Patched Trainer.evaluation_loop with nanmean loss calculation")
+ Trainer.evaluation_loop = axolotl_evaluation_loop
+
+
+def check_maybe_log_save_evaluate_is_patchable() -> bool:
+ maybe_log_source = inspect.getsource(Trainer._maybe_log_save_evaluate)
+ return ORIGINAL_MAYBE_CODE in maybe_log_source
+
+
+def patch_maybe_log_save_evaluate():
+ """Patch the _maybe_log_save_evaluate method."""
+ # Check if already patched
+ if hasattr(Trainer, "_original_maybe_log_save_evaluate"):
+ LOG.info("Trainer._maybe_log_save_evaluate already patched")
+ return
+
+ # Check if the patterns exist
+ try:
+ maybe_log_source = inspect.getsource(Trainer._maybe_log_save_evaluate)
+ except OSError:
+ return
+ Trainer._original_maybe_log_save_evaluate = maybe_log_source
+ maybe_log_source, _ = detab_code(maybe_log_source)
+
+ # Apply the patch
+ maybe_log_source = maybe_log_source.replace(ORIGINAL_MAYBE_CODE, PATCHED_MAYBE_CODE)
+
+ # Rename the function to avoid conflicts
+ maybe_log_source = maybe_log_source.replace(
+ "def _maybe_log_save_evaluate(",
+ "def axolotl_maybe_log_save_evaluate(",
+ 1,
+ )
+
+ # Get the module for necessary imports
+ module_name = Trainer.__module__
+ module = importlib.import_module(module_name)
+
+ # Import necessary items from the module
+ items_to_import = []
+ for item in dir(module):
+ if item in maybe_log_source:
+ items_to_import.append(item)
+
+ # Execute the imports and patched method
+ exec(
+ f"from {module_name} import ({', '.join(items_to_import)})",
+ globals(),
+ )
+ exec(maybe_log_source, globals())
+
+ LOG.debug("Patched Trainer._maybe_log_save_evaluate with nanmean loss calculation")
+ Trainer._maybe_log_save_evaluate = axolotl_maybe_log_save_evaluate
diff --git a/src/axolotl/monkeypatch/unsloth_.py b/src/axolotl/monkeypatch/unsloth_.py
index 61f4eeea0..59f32c6f5 100644
--- a/src/axolotl/monkeypatch/unsloth_.py
+++ b/src/axolotl/monkeypatch/unsloth_.py
@@ -4,12 +4,12 @@ import inspect
import types
import torch
-from accelerate.logging import get_logger
from peft import PeftModelForCausalLM
from torch import nn
from transformers.models.llama.modeling_llama import LlamaFlashAttention2
from axolotl.monkeypatch.utils import detab_code
+from axolotl.utils.logging import get_logger
LOG = get_logger(__name__)
@@ -17,27 +17,19 @@ ORIGINAL_QKV_CODE = """
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
PATCHED_QKV_CODE = """
query_states, key_states, value_states = self.apply_qkv(self, hidden_states)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
ORIGINAL_O_CODE = """
attn_output = self.o_proj(attn_output)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
PATCHED_O_CODE = """
attn_output = self.apply_o(self, attn_output)
-""".lstrip(
- "\n"
-)
+""".lstrip("\n")
def original_apply_qkv(self, hidden_states):
@@ -66,13 +58,13 @@ def check_self_attn_is_patchable() -> bool:
def integrate_cross_entropy_loss_patch(model_type: str = "llama") -> None:
from unsloth.kernels.cross_entropy_loss import fast_cross_entropy_loss
- def UnslothForCausalLMLoss( # pylint: disable=invalid-name
+ def UnslothForCausalLMLoss(
logits,
labels,
- vocab_size: int, # pylint: disable=unused-argument
+ vocab_size: int,
num_items_in_batch: int = None,
- ignore_index: int = -100, # pylint: disable=unused-argument
- **kwargs, # pylint: disable=unused-argument
+ ignore_index: int = -100,
+ **kwargs,
):
# Upcast to float if we need to compute the loss to avoid potential precision issues
logits = logits.float()
@@ -93,18 +85,16 @@ def integrate_cross_entropy_loss_patch(model_type: str = "llama") -> None:
raise ValueError("Unsupported model type")
-self_attn_lora_patched = False # pylint: disable=invalid-name
+self_attn_lora_patched = False
def patch_self_attn_lora():
- global self_attn_lora_patched # pylint: disable=global-statement
+ global self_attn_lora_patched
if self_attn_lora_patched:
# prevent patching multiple times
return
self_attn_forward = get_self_attn_code()
- LlamaFlashAttention2._original_forward = ( # pylint: disable=protected-access
- self_attn_forward
- )
+ LlamaFlashAttention2._original_forward = self_attn_forward
self_attn_forward, _ = detab_code(self_attn_forward)
assert ORIGINAL_QKV_CODE in self_attn_forward, "Original qkv code not found"
assert ORIGINAL_O_CODE in self_attn_forward, "Original o code not found"
@@ -125,27 +115,25 @@ def patch_self_attn_lora():
if item in self_attn_forward:
items_to_import.append(item)
- exec( # pylint: disable=exec-used # nosec B102
+ exec(
"from transformers.models.llama.modeling_llama import ("
+ ", ".join(x for x in items_to_import)
+ ")",
globals(),
)
- exec(self_attn_forward, globals()) # pylint: disable=exec-used # nosec B102
+ exec(self_attn_forward, globals())
self_attn_lora_patched = True
LOG.info("patching unsloth attn lora")
- LlamaFlashAttention2.forward = (
- unsloth_attn_forward # pylint: disable=undefined-variable # noqa: F821
- )
+ LlamaFlashAttention2.forward = unsloth_attn_forward
def integrate_rope_embeddings():
import transformers.models.llama.modeling_llama
from unsloth.kernels.rope_embedding import fast_rope_embedding
- def apply_rotary_pos_emb( # pylint: disable=unused-argument
- q, # pylint: disable=invalid-name
- k, # pylint: disable=invalid-name
+ def apply_rotary_pos_emb(
+ q,
+ k,
cos,
sin,
position_ids=None,
diff --git a/src/axolotl/monkeypatch/xformers_/__init__.py b/src/axolotl/monkeypatch/xformers_/__init__.py
index a052ea49e..6f5b43f77 100644
--- a/src/axolotl/monkeypatch/xformers_/__init__.py
+++ b/src/axolotl/monkeypatch/xformers_/__init__.py
@@ -36,7 +36,7 @@ class FusedMLP(torch.nn.Module):
self.swiglu.w3.weight.data = down_proj.weight.data
def _post_training(self, model, name):
- w1, w2 = torch.split( # pylint: disable=invalid-name
+ w1, w2 = torch.split(
self.swiglu.w12.weight.data, self.config.intermediate_size, dim=0
)
@@ -48,5 +48,5 @@ class FusedMLP(torch.nn.Module):
set_module_name(model, name, new_mlp)
- def forward(self, x: torch.Tensor) -> torch.Tensor: # pylint: disable=invalid-name
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.swiglu(x)
diff --git a/src/axolotl/processing_strategies.py b/src/axolotl/processing_strategies.py
index ce9b6a838..07b114163 100644
--- a/src/axolotl/processing_strategies.py
+++ b/src/axolotl/processing_strategies.py
@@ -5,11 +5,15 @@ from typing import Optional
from PIL import Image, ImageOps
from PIL.Image import Resampling
-from torch import Tensor
+from torch import Tensor, zeros_like
from transformers import ProcessorMixin
from transformers.image_utils import load_image
+from transformers.models.smolvlm import SmolVLMProcessor
+from transformers.models.voxtral import VoxtralProcessor
+from axolotl.utils.dict import remove_none_values
from axolotl.utils.logging import get_logger
+from axolotl.utils.mistral.mistral3_processor import Mistral3Processor
LOG = get_logger(__name__)
@@ -137,12 +141,12 @@ class ProcessingStrategy:
image_key = key
break
- # if the image key exists, add the image to the first message
+ # if the image key exists, add the image to the first user message
if image_key is not None and processed_example[image_key] is not None:
# TODO: check if it's normal to be single image only for common datasets
# From observation, it's usually a list of single image but some datasets may have several columns for images
# Temporary solution: take the first image and suggest people convert their datasets to use multi-content Messages
- if len(processed_example[image_key]) > 0:
+ if len(processed_example[image_key]) > 1:
LOG.warning(
f"Found {len(processed_example[image_key])} images in a sample. Using the first one."
"If you are using a dataset with multiple images per sample, please convert it to use multi-content Messages."
@@ -155,9 +159,9 @@ class ProcessingStrategy:
image_value = load_image(image_value)
if self.image_size is not None:
- assert hasattr(
- image_value, "resize"
- ), "Image does not have a resize method"
+ assert hasattr(image_value, "resize"), (
+ "Image does not have a resize method"
+ )
if isinstance(self.image_size, tuple):
image_value = image_value.resize(
@@ -178,39 +182,56 @@ class ProcessingStrategy:
# Look for any image type in the first message
# some dataset have an {type: "image"} in the first message
+ msg_ind_to_add = None
ind_to_add = None
+ first_user_idx = None
- for i, content in enumerate(
- processed_example["messages"][0]["content"]
- ):
- # Usually datasets created with image columns, don't have it in the messages itself
- if content["type"] == "image" and all(
- k not in content for k in ["image", "url", "path", "base64"]
+ for msg_idx, msg_content in enumerate(processed_example["messages"]):
+ if first_user_idx is None and msg_content["role"] == "user":
+ first_user_idx = msg_idx
+ for i, content in enumerate(
+ processed_example["messages"][msg_idx]["content"]
):
- ind_to_add = i
- break
+ # Usually datasets created with image columns, don't have it in the messages itself
+ if content["type"] == "image" and all(
+ k not in content for k in ["image", "url", "path", "base64"]
+ ):
+ msg_ind_to_add = msg_idx
+ ind_to_add = i
+ break
# If an image type is found, add the image to that index
- if ind_to_add is not None:
- processed_example["messages"][0]["content"][ind_to_add][
- "image"
- ] = image_value
+ if ind_to_add is not None and msg_ind_to_add is not None:
+ processed_example["messages"][msg_ind_to_add]["content"][
+ ind_to_add
+ ]["image"] = image_value
else:
- # if no image type is found, add it to end of the first message
- processed_example["messages"][0]["content"].append(
+ # if no image type is found, add it to end of the first user message
+ if first_user_idx is None:
+ first_user_idx = 0
+ processed_example["messages"][first_user_idx]["content"].append(
{
"type": "image",
"image": image_value,
}
)
- processed_examples.append(processed_example)
+ processed_examples.append(remove_none_values(processed_example))
return processed_examples
+ def _mask_non_assistant(self, labels: Tensor) -> Tensor:
+ """
+ Mask non assistant regions to -100.
+ To be implemented per subclass.
+ """
+ return labels
+
def process_labels(self, input_ids: Tensor) -> Tensor:
labels = input_ids.clone()
+ labels = self._mask_non_assistant(labels)
+
# The labels are the input_ids, and we mask the padding tokens in the loss computation
labels[labels == self.processor.tokenizer.pad_token_id] = -100
@@ -264,6 +285,175 @@ class Gemma3ProcessingStrategy(ProcessingStrategy):
return labels
+class Gemma3nProcessingStrategy(ProcessingStrategy):
+ """Processing Strategy class for Gemma3n"""
+
+ def _mask_non_assistant(self, labels: Tensor) -> Tensor:
+ def _find_token_sequence(label, start_pos, token_sequence):
+ """Check if token_sequence appears at start_pos in label"""
+ if start_pos + len(token_sequence) > len(label):
+ return False
+ if label[start_pos] != token_sequence[0]:
+ return False
+ return (
+ label[start_pos : start_pos + len(token_sequence)].tolist()
+ == token_sequence
+ )
+
+ def _find_assistant_end(label, start_pos, assistant_end_tok, mask, i):
+ """
+ Find the end of assistant response and update mask accordingly
+
+ Returns new position to continue from and whether the end seq is found
+ """
+ k = start_pos
+ while k < len(label):
+ if not _find_token_sequence(label, k, assistant_end_tok):
+ mask[i][k] = 1
+ k += 1
+ continue
+
+ return k + len(assistant_end_tok), True
+
+ return k, False
+
+ mask = zeros_like(labels)
+
+ assistant_start_str = "\n"
- "<|user|>\n"
- f"{sample['prompt']}\n"
- "<|assistant|>\n"
+ f"<|system|>\n\n<|user|>\n{sample['prompt']}\n<|assistant|>\n"
)
answers = sorted(sample["answers"], key=lambda x: x["rank"])
data["chosen"] = answers[-1]["answer"]
diff --git a/src/axolotl/prompt_strategies/input_output.py b/src/axolotl/prompt_strategies/input_output.py
index 8be745b20..c84eecffc 100644
--- a/src/axolotl/prompt_strategies/input_output.py
+++ b/src/axolotl/prompt_strategies/input_output.py
@@ -16,7 +16,6 @@ class RawInputOutputStrategy(PromptTokenizingStrategy):
self.eos_token = self.tokenizer.eos_token
def tokenize_prompt(self, prompt):
- # pylint: disable=duplicate-code
input_ids = []
labels = []
for label, text in self.prompter.build_prompt(prompt["segments"]):
diff --git a/src/axolotl/prompt_strategies/jinja_template_analyzer.py b/src/axolotl/prompt_strategies/jinja_template_analyzer.py
index a5f89cfe5..e16a1e22b 100644
--- a/src/axolotl/prompt_strategies/jinja_template_analyzer.py
+++ b/src/axolotl/prompt_strategies/jinja_template_analyzer.py
@@ -3,6 +3,7 @@
from typing import Dict, Optional, Set, TypedDict, Union
from jinja2 import Environment, meta, nodes
+from jinja2.ext import Extension
class JinjaTemplateAnalysis(TypedDict):
@@ -27,6 +28,18 @@ class JinjaTemplateAnalysis(TypedDict):
iteration_target: Optional[Union[str, list[str]]]
+class GenerationTagIgnore(Extension):
+ """
+ Ignores the generation and endgeneration tags in Jinja templates.
+ """
+
+ tags = {"generation", "endgeneration"}
+
+ def parse(self, parser):
+ parser.stream.skip(1)
+ return nodes.Const("")
+
+
class JinjaTemplateAnalyzer:
"""
Analyzes Jinja templates to extract information about variable usage,
@@ -57,7 +70,9 @@ class JinjaTemplateAnalyzer:
"""
def __init__(self, template: str):
- self.env: Environment = Environment(autoescape=True)
+ self.env: Environment = Environment(
+ autoescape=True, extensions=[GenerationTagIgnore]
+ )
self.property_access: Dict[str, Set[str]] = {}
self.iteration_targets: Dict[str, Union[str, list[str]]] = {}
self.index_access: Dict[str, Set[Union[int, float]]] = {}
diff --git a/src/axolotl/prompt_strategies/kto/chatml.py b/src/axolotl/prompt_strategies/kto/chatml.py
index 97ae59ed5..945940f3f 100644
--- a/src/axolotl/prompt_strategies/kto/chatml.py
+++ b/src/axolotl/prompt_strategies/kto/chatml.py
@@ -2,13 +2,11 @@
KTO strategies for chatml
"""
-# pylint: disable=duplicate-code
-
def argilla(
cfg,
**kwargs,
-): # pylint: disable=possibly-unused-variable,unused-argument
+):
def transform_fn(sample):
if "system" in sample and sample["system"]:
sample["prompt"] = (
@@ -28,7 +26,7 @@ def argilla(
def argilla_chat(
cfg,
**kwargs,
-): # pylint: disable=possibly-unused-variable,unused-argument
+):
"""
for argilla/kto-mix-15k conversations
"""
@@ -43,7 +41,7 @@ def argilla_chat(
return transform_fn
-def intel(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-argument
+def intel(cfg, **kwargs):
"""
For Intel Orca KTO
ex: argilla/distilabel-intel-orca-kto
@@ -65,9 +63,7 @@ def intel(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-arg
return transform_fn
-def prompt_pairs(
- cfg, **kwargs
-): # pylint: disable=possibly-unused-variable,unused-argument
+def prompt_pairs(cfg, **kwargs):
def transform_fn(sample):
if "system" in sample and sample["system"]:
sample["prompt"] = (
@@ -84,7 +80,7 @@ def prompt_pairs(
return transform_fn
-def ultra(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-argument
+def ultra(cfg, **kwargs):
"""
for ultrafeedback binarized conversations
ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
diff --git a/src/axolotl/prompt_strategies/kto/llama3.py b/src/axolotl/prompt_strategies/kto/llama3.py
index fde3c2ed4..9061f6f5e 100644
--- a/src/axolotl/prompt_strategies/kto/llama3.py
+++ b/src/axolotl/prompt_strategies/kto/llama3.py
@@ -2,13 +2,11 @@
KTO strategies for llama-3 chat template
"""
-# pylint: disable=duplicate-code
-
def argilla(
cfg,
**kwargs,
-): # pylint: disable=possibly-unused-variable,unused-argument
+):
def transform_fn(sample):
if "system" in sample and sample["system"]:
sample["prompt"] = (
@@ -28,7 +26,7 @@ def argilla(
def argilla_chat(
cfg,
**kwargs,
-): # pylint: disable=possibly-unused-variable,unused-argument
+):
"""
for argilla/kto-mix-15k conversations
"""
@@ -43,7 +41,7 @@ def argilla_chat(
return transform_fn
-def intel(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-argument
+def intel(cfg, **kwargs):
"""
For Intel Orca KTO
ex: argilla/distilabel-intel-orca-kto
@@ -65,9 +63,7 @@ def intel(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-arg
return transform_fn
-def prompt_pairs(
- cfg, **kwargs
-): # pylint: disable=possibly-unused-variable,unused-argument
+def prompt_pairs(cfg, **kwargs):
def transform_fn(sample):
if "system" in sample and sample["system"]:
sample["prompt"] = (
@@ -84,7 +80,7 @@ def prompt_pairs(
return transform_fn
-def ultra(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-argument
+def ultra(cfg, **kwargs):
"""
for ultrafeedback binarized conversations
ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
diff --git a/src/axolotl/prompt_strategies/kto/user_defined.py b/src/axolotl/prompt_strategies/kto/user_defined.py
index 7c68a3000..e26683cde 100644
--- a/src/axolotl/prompt_strategies/kto/user_defined.py
+++ b/src/axolotl/prompt_strategies/kto/user_defined.py
@@ -2,10 +2,8 @@
User-defined KTO strategies
"""
-# pylint: disable=duplicate-code
-
-def default(cfg, dataset_idx=0, **kwargs): # pylint: disable=unused-argument
+def default(cfg, dataset_idx=0, **kwargs):
ds_cfg = cfg["datasets"][dataset_idx]["type"]
if not isinstance(ds_cfg, dict):
raise ValueError(
diff --git a/src/axolotl/prompt_strategies/llama2_chat.py b/src/axolotl/prompt_strategies/llama2_chat.py
index eef2e1d4d..9eff062ec 100644
--- a/src/axolotl/prompt_strategies/llama2_chat.py
+++ b/src/axolotl/prompt_strategies/llama2_chat.py
@@ -153,7 +153,7 @@ class LLama2ChatTokenizingStrategy(PromptTokenizingStrategy):
}
-class Llama2ChatPrompter: # pylint: disable=too-few-public-methods
+class Llama2ChatPrompter:
"""
A prompter that generates prompts for Llama2 models.
"""
@@ -190,7 +190,7 @@ class Llama2ChatPrompter: # pylint: disable=too-few-public-methods
# Skip the first one if it is not from human
source = source[1:]
- conv.messages = [] # pylint: disable=R0801
+ conv.messages = []
for j, sentence in enumerate(source):
role = roles[sentence["from"]]
assert role == conv.roles[j % 2], ALTERNATING_ASSERTION_FAILED_ROLE
diff --git a/src/axolotl/prompt_strategies/messages/__init__.py b/src/axolotl/prompt_strategies/messages/__init__.py
index cc7b84da1..2c920a568 100644
--- a/src/axolotl/prompt_strategies/messages/__init__.py
+++ b/src/axolotl/prompt_strategies/messages/__init__.py
@@ -11,7 +11,7 @@ LOG = get_logger(__name__)
def load(tokenizer, cfg, ds_cfg, processor=None):
try:
strategy = ds_cfg.get("input_transform", "chat")
- # pylint: disable=duplicate-code
+
load_fn = "load"
if strategy.split(".")[-1].startswith("load_"):
load_fn = strategy.split(".")[-1]
@@ -29,7 +29,6 @@ def load(tokenizer, cfg, ds_cfg, processor=None):
return func(tokenizer, cfg, **load_kwargs)
except ModuleNotFoundError:
return None
- except Exception as exc: # pylint: disable=broad-exception-caught
+ except Exception as exc:
LOG.error(f"Failed to load prompt strategy `{strategy}`: {str(exc)}")
raise exc
- return None
diff --git a/src/axolotl/prompt_strategies/messages/chat.py b/src/axolotl/prompt_strategies/messages/chat.py
index eaed2396a..854d25e42 100644
--- a/src/axolotl/prompt_strategies/messages/chat.py
+++ b/src/axolotl/prompt_strategies/messages/chat.py
@@ -19,7 +19,7 @@ class ChatMessageDatasetWrappingStrategy(DatasetWrappingStrategy):
processor,
message_transform=None,
formatter=None,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
"""
:param processor: tokenizer or image processor
@@ -35,7 +35,7 @@ class ChatMessageDatasetWrappingStrategy(DatasetWrappingStrategy):
dataset,
process_count: Optional[int] = None,
keep_in_memory: Optional[bool] = False,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
self.dataset = TokenizedChatDataset(
dataset,
@@ -72,9 +72,10 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
builder_kwargs["message_field_training"] = message_field_training
chat_template = ds_cfg.get("chat_template", cfg.get("chat_template", "chatml"))
- format_message = (
- lambda x: x # noqa E731 # pylint: disable=unnecessary-lambda-assignment
- )
+
+ def format_message(x):
+ return x
+
if chat_template == "chatml":
from axolotl.core.chat.format.chatml import format_message # noqa F811
if chat_template.startswith("llama3"):
diff --git a/src/axolotl/prompt_strategies/metharme.py b/src/axolotl/prompt_strategies/metharme.py
index 66da72389..35f1ef3b3 100644
--- a/src/axolotl/prompt_strategies/metharme.py
+++ b/src/axolotl/prompt_strategies/metharme.py
@@ -10,8 +10,6 @@ LOG = get_logger(__name__)
IGNORE_TOKEN_ID = -100
-# pylint: disable=duplicate-code
-
class MetharmePromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
"""
@@ -66,7 +64,7 @@ class MetharmePrompter(AlpacaPrompter):
turn_format = "{instruction}"
turn_no_input_format = "{instruction}"
- def __init__(self, *args, **kwargs): # pylint: disable=super-init-not-called
+ def __init__(self, *args, **kwargs):
pass
diff --git a/src/axolotl/prompt_strategies/orpo/chat_template.py b/src/axolotl/prompt_strategies/orpo/chat_template.py
index fdee28ea1..b655bc970 100644
--- a/src/axolotl/prompt_strategies/orpo/chat_template.py
+++ b/src/axolotl/prompt_strategies/orpo/chat_template.py
@@ -23,9 +23,7 @@ class MessageList(BaseModel):
messages: List[Message]
-def load(
- tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None, **kwargs
-): # pylint: disable=possibly-unused-variable,unused-argument
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None, **kwargs):
"""
chatml transforms for datasets with system, input, chosen, rejected
"""
@@ -219,29 +217,38 @@ class ORPOPrompter(Prompter):
for message in message_list.messages:
conversation.append(message.model_dump())
if message.role == "system":
- yield self.tokenizer.apply_chat_template(
- conversation,
- add_generation_prompt=False,
- chat_template=self.chat_template,
- tokenize=False,
- ), False
+ yield (
+ self.tokenizer.apply_chat_template(
+ conversation,
+ add_generation_prompt=False,
+ chat_template=self.chat_template,
+ tokenize=False,
+ ),
+ False,
+ )
if message.role == "user":
- yield self.tokenizer.apply_chat_template(
- conversation,
- add_generation_prompt=True,
- chat_template=self.chat_template,
- tokenize=False,
- ), False
+ yield (
+ self.tokenizer.apply_chat_template(
+ conversation,
+ add_generation_prompt=True,
+ chat_template=self.chat_template,
+ tokenize=False,
+ ),
+ False,
+ )
if message.role == "assistant":
- yield self.tokenizer.apply_chat_template(
- conversation,
- add_generation_prompt=False,
- chat_template=self.chat_template,
- tokenize=False,
- ), True
+ yield (
+ self.tokenizer.apply_chat_template(
+ conversation,
+ add_generation_prompt=False,
+ chat_template=self.chat_template,
+ tokenize=False,
+ ),
+ True,
+ )
-def argilla(cfg, **kwargs): # pylint: disable=possibly-unused-variable,unused-argument
+def argilla(cfg, **kwargs):
dataset_parser = ORPODatasetParsingStrategy()
def transform_fn(sample, tokenizer=None):
diff --git a/src/axolotl/prompt_strategies/pygmalion.py b/src/axolotl/prompt_strategies/pygmalion.py
index 51f92f397..8c53a5f27 100644
--- a/src/axolotl/prompt_strategies/pygmalion.py
+++ b/src/axolotl/prompt_strategies/pygmalion.py
@@ -69,7 +69,6 @@ class PygmalionPromptTokenizingStrategy(PromptTokenizingStrategy):
LOG.warning(f"unknown role in conversation: {role}")
res = defaultdict(lambda: [])
- # pylint: disable=duplicate-code
result, current_len = parse_tokenized_to_result(
result,
current_len,
@@ -89,7 +88,10 @@ class PygmalionPrompter:
pass
def build_prompt(
- self, source, *args, **kwargs # pylint: disable=unused-argument
+ self,
+ source,
+ *args,
+ **kwargs,
) -> Generator[Tuple[str, str], None, None]:
for msg in source:
yield msg["role"], msg["value"]
diff --git a/src/axolotl/prompt_strategies/stepwise_supervised.py b/src/axolotl/prompt_strategies/stepwise_supervised.py
index 8be7c35e3..9175126e7 100644
--- a/src/axolotl/prompt_strategies/stepwise_supervised.py
+++ b/src/axolotl/prompt_strategies/stepwise_supervised.py
@@ -66,7 +66,7 @@ class StepwiseSupervisedPromptTokenizingStrategy:
# Create step-wise labels
labels = [
[IGNORE_INDEX] * (len(completion) - 1) + [label] # type: ignore
- for completion, label in zip(completions_ids, labels)
+ for completion, label in zip(completions_ids, labels, strict=False)
]
# Join all steps
diff --git a/src/axolotl/prompt_strategies/user_defined.py b/src/axolotl/prompt_strategies/user_defined.py
index e20e80c3a..0bff514e7 100644
--- a/src/axolotl/prompt_strategies/user_defined.py
+++ b/src/axolotl/prompt_strategies/user_defined.py
@@ -83,16 +83,12 @@ def load(tokenizer, cfg, ds_cfg: Optional[UserDefinedDatasetConfig] = None):
cfg.sequence_len,
)
- setattr(
- strat,
- "parse_instruction_fields",
- partial(
- parse_instruction_fields,
- ds_cfg.field_instruction,
- ds_cfg.field_input,
- ds_cfg.field_output,
- ds_cfg.field_system,
- system_prompt,
- ),
+ strat.parse_instruction_fields = partial( # type: ignore[method-assign]
+ parse_instruction_fields,
+ ds_cfg.field_instruction,
+ ds_cfg.field_input,
+ ds_cfg.field_output,
+ ds_cfg.field_system,
+ system_prompt,
)
return strat
diff --git a/src/axolotl/prompt_tokenizers.py b/src/axolotl/prompt_tokenizers.py
index cb1a1ba4e..a7bd963f8 100644
--- a/src/axolotl/prompt_tokenizers.py
+++ b/src/axolotl/prompt_tokenizers.py
@@ -3,6 +3,7 @@
import abc
from typing import Callable, Dict, List, Optional, Tuple, Union
+from datasets import Dataset
from transformers import BatchEncoding, PreTrainedTokenizer
from axolotl.prompters import Prompter
@@ -28,6 +29,16 @@ class DatasetWrappingStrategy(abc.ABC):
Abstract class for wrapping datasets for Chat Messages
"""
+ @abc.abstractmethod
+ def wrap_dataset(
+ self,
+ dataset,
+ process_count: int | None = None,
+ keep_in_memory: bool | None = False,
+ **kwargs,
+ ) -> Dataset:
+ pass
+
class PromptTokenizingStrategy(abc.ABC):
"""
@@ -64,7 +75,7 @@ class PromptTokenizingStrategy(abc.ABC):
) -> BatchEncoding:
empty = BatchEncoding(data={"input_ids": [], "attention_mask": []})
if not prompt:
- LOG.warning("Empty text requested for tokenization.")
+ LOG.warning_once("Empty text requested for tokenization.")
return empty
result = self.tokenizer(
@@ -107,7 +118,7 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
def tokenize_prompt(self, prompt):
(
instruction,
- input, # pylint: disable=redefined-builtin
+ input,
response,
) = self.parse_instruction_fields(prompt)
user_prompt = next(
@@ -133,7 +144,10 @@ class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
return tokenized_prompt
def _build_full_prompt(
- self, instruction, input, response # pylint: disable=redefined-builtin
+ self,
+ instruction,
+ input,
+ response,
):
return next(
iter(
@@ -246,10 +260,9 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy):
raise NotImplementedError
def tokenize_prompt(self, prompt):
- # pylint: disable=duplicate-code
(
instruction,
- input, # pylint: disable=redefined-builtin
+ input,
output,
reflection,
corrected,
@@ -276,9 +289,7 @@ class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy):
return tokenized_full_prompt
- def _build_full_prompt(
- self, instruction, input, output, reflection, corrected
- ): # pylint: disable=redefined-builtin
+ def _build_full_prompt(self, instruction, input, output, reflection, corrected):
return next(
iter(
self.prompter.build_prompt(
diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py
index d29da075e..9543996f7 100644
--- a/src/axolotl/prompters.py
+++ b/src/axolotl/prompters.py
@@ -46,7 +46,6 @@ class AlpacaPrompter(Prompter):
self.match_prompt_style()
def match_prompt_style(self):
- # pylint: disable=duplicate-code
if self.prompt_style == PromptStyle.INSTRUCT.value:
self.turn_format = "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
self.turn_no_input_format = (
@@ -93,7 +92,7 @@ class AlpacaPrompter(Prompter):
def build_prompt(
self,
instruction: str,
- input: Union[None, str] = None, # pylint: disable=redefined-builtin
+ input: Union[None, str] = None,
output: Union[None, str] = None,
) -> Generator[str, None, None]:
yield self._build_result(instruction, input, output)
@@ -218,7 +217,7 @@ class ReflectAlpacaPrompter(Prompter):
def _build_result(
self,
instruction: str,
- input: Union[None, str] = None, # pylint: disable=redefined-builtin
+ input: Union[None, str] = None,
output: Union[None, str] = None,
reflection: Union[None, str] = None,
corrected: Union[None, str] = None,
@@ -242,12 +241,11 @@ class ReflectAlpacaPrompter(Prompter):
def build_prompt(
self,
instruction: str,
- input: Union[None, str] = None, # pylint: disable=redefined-builtin
+ input: Union[None, str] = None,
output: Union[None, str] = None,
reflection: Union[None, str] = None,
corrected: Union[None, str] = None,
) -> Generator[str, None, None]:
- # pylint: disable=duplicate-code
yield self._build_result(
instruction,
input,
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
index f4f4a5a91..6cbcb9aec 100644
--- a/src/axolotl/train.py
+++ b/src/axolotl/train.py
@@ -1,18 +1,23 @@
"""Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
+from __future__ import annotations
+
import importlib
import inspect
+import json
import os
+import shutil
import signal
import sys
+import typing
import weakref
+from collections import OrderedDict
from contextlib import ExitStack
from pathlib import Path
from typing import Any, Dict
import torch
import transformers.modelcard
-from accelerate.utils import save_fsdp_model
from datasets import Dataset
from huggingface_hub.errors import OfflineModeIsEnabled
from peft import PeftConfig, PeftModel
@@ -20,18 +25,12 @@ from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
from transformers.trainer import Trainer
-from axolotl.cli.art import print_axolotl_text_art
from axolotl.common.datasets import TrainDatasetMeta
from axolotl.contribs.lgpl import ( # pylint: disable = no-name-in-module
fix_untrained_tokens,
)
-from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
from axolotl.integrations.base import PluginManager
-from axolotl.loaders import (
- ModelLoader,
- load_processor,
- load_tokenizer,
-)
+from axolotl.loaders import ModelLoader, load_processor, load_tokenizer
from axolotl.telemetry.errors import send_errors
from axolotl.telemetry.manager import TelemetryManager
from axolotl.utils.ctx_managers.sequence_parallel import SequenceParallelContextManager
@@ -40,12 +39,11 @@ from axolotl.utils.distributed import cleanup_distributed
from axolotl.utils.freeze import freeze_layers_except
from axolotl.utils.logging import get_logger
from axolotl.utils.schemas.enums import RLType
+from axolotl.utils.train import determine_last_checkpoint
from axolotl.utils.trainer import setup_trainer
-try:
- from optimum.bettertransformer import BetterTransformer
-except ImportError:
- BetterTransformer = None
+if typing.TYPE_CHECKING:
+ from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
LOG = get_logger(__name__)
@@ -58,8 +56,8 @@ def setup_model_and_tokenizer(
) -> tuple[
PreTrainedModel, PreTrainedTokenizer, PeftConfig | None, ProcessorMixin | None
]:
- """
- Load the tokenizer, processor (for multimodal models), and model based on configuration.
+ """Load the tokenizer, processor (for multimodal models), and model based on
+ configuration.
Args:
cfg: Dictionary mapping `axolotl` config keys to values.
@@ -80,11 +78,8 @@ def setup_model_and_tokenizer(
if cfg.is_multimodal:
processor = load_processor(cfg, tokenizer)
- # Load the model and peft_config
- msg = "loading model"
- if cfg.adapter:
- msg += " and peft_config..."
- LOG.debug(msg)
+ # Load the model
+ LOG.debug("Loading model")
model_loader = ModelLoader(cfg, tokenizer, processor=processor)
model, peft_config = model_loader.load()
@@ -131,38 +126,15 @@ def setup_reference_model(
LOG.debug("Passing model_ref: None to RL trainer")
model_ref = None # explicit setting to None
else:
+ reference_model: bool = True
+ if cfg.rl == RLType.GRPO and cfg.trl.beta == 0:
+ reference_model = False
# load the model again for model_ref/baseline
- model_loader = ModelLoader(cfg, tokenizer, reference_model=True)
+ model_loader = ModelLoader(cfg, tokenizer, reference_model=reference_model)
model_ref, _ = model_loader.load()
return model_ref
-def determine_resume_checkpoint(cfg: DictDefault) -> str | None:
- """
- Determine the checkpoint to resume from based on configuration.
-
- Args:
- cfg: Dictionary mapping `axolotl` config keys to values.
-
- Returns:
- Path to the checkpoint to resume from, or `None` if not resuming.
- """
- if cfg.resume_from_checkpoint is None and cfg.auto_resume_from_checkpoints:
- possible_checkpoints = [
- str(cp) for cp in Path(cfg.output_dir).glob("checkpoint-*")
- ]
- if len(possible_checkpoints) > 0:
- sorted_paths = sorted(
- possible_checkpoints,
- key=lambda path: int(path.split("-")[-1]),
- )
- cfg.resume_from_checkpoint = sorted_paths[-1]
- LOG.info(
- f"Using Auto-resume functionality to start with checkpoint at {cfg.resume_from_checkpoint}"
- )
- return cfg.resume_from_checkpoint
-
-
def setup_signal_handler(
cfg: DictDefault, model: PreTrainedModel, safe_serialization: bool
):
@@ -180,8 +152,6 @@ def setup_signal_handler(
def terminate_handler(_, __, model_weakref):
if model_weakref() is not None:
_model = model_weakref()
- if cfg.flash_optimum and BetterTransformer:
- _model = BetterTransformer.reverse(_model)
_model.save_pretrained(
cfg.output_dir, safe_serialization=safe_serialization
)
@@ -218,7 +188,7 @@ def execute_training(
)
)
- if cfg.sequence_parallel_degree > 1:
+ if cfg.context_parallel_size > 1:
models = [trainer.model]
if hasattr(trainer, "ref_model") and trainer.ref_model:
models.append(trainer.ref_model)
@@ -226,16 +196,25 @@ def execute_training(
stack.enter_context(
SequenceParallelContextManager(
models=models,
- sequence_parallel_degree=cfg.sequence_parallel_degree,
+ context_parallel_size=cfg.context_parallel_size,
gradient_accumulation_steps=cfg.gradient_accumulation_steps,
ring_attn_func=cfg.ring_attn_func,
heads_k_stride=cfg.heads_k_stride,
+ gather_outputs=cfg.rl is RLType.GRPO,
+ device_mesh=trainer.accelerator.torch_device_mesh,
)
)
+ # TODO: disabling for now as not compatible with FSDP2 + torchao low bit optimizers
+ # if cfg.bf16:
+ # torch.set_default_dtype(torch.bfloat16)
+
LOG.info("Starting trainer...")
trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+ plugin_manager = PluginManager.get_instance()
+ plugin_manager.post_train(cfg, trainer.model)
+
def save_trained_model(
cfg: DictDefault,
@@ -257,54 +236,80 @@ def save_trained_model(
# Post training module hooks
for name, module in model.named_modules():
if hasattr(module, "_post_training"):
- module._post_training(model, name) # pylint: disable=protected-access
+ module._post_training(model, name)
# handle QAT
if cfg.qat:
- from axolotl.utils.quantization import convert_qat_model_for_ptq
+ from axolotl.utils.quantization import convert_qat_model
- LOG.info("Processing QAT model for saving...")
- convert_qat_model_for_ptq(
+ convert_qat_model(
model,
quantize_embedding=cfg.qat.quantize_embedding,
)
LOG.info(
- "QAT modules have been converted for PTQ. Please ensure you quantize "
- "your model weights with `axolotl quantize`."
+ "QAT usage note: please ensure you quantize your model fine-tuned using QAT by running `axolotl quantize`"
+ " with the same config which you used for training."
)
-
- # Handle FSDP state dict type
- state_dict_type = "FULL_STATE_DICT"
- if trainer.is_fsdp_enabled and str(cfg.fsdp_config.fsdp_version) != "2":
- if cfg.fsdp_final_state_dict_type:
- state_dict_type = cfg.fsdp_final_state_dict_type
- trainer.accelerator.state.fsdp_plugin.set_state_dict_type(state_dict_type)
- LOG.info(f"Set FSDP state dict type to {state_dict_type} for saving.")
-
# Handle ReLoRA early return case
- if cfg.relora_steps:
+ if cfg.relora:
if cfg.adapter == "lora" and not (cfg.load_in_4bit or cfg.load_in_8bit):
model = model.merge_and_unload()
else:
# final model weights have already been saved by `ReLoRACallback.on_train_end`
return
- if cfg.fsdp:
- # TODO: do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
- # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple
- # processes attempt to write the same file
- if (
- state_dict_type == "SHARDED_STATE_DICT"
- and cfg.fsdp_config.fsdp_state_dict_type == "SHARDED_STATE_DICT"
- ):
- save_fsdp_model(
- trainer.accelerator.state.fsdp_plugin,
- trainer.accelerator,
- trainer.model,
- cfg.output_dir,
+ if trainer.is_fsdp_enabled or cfg.fsdp_config:
+ if cfg.fsdp_config or cfg.fsdp:
+ if cfg.fsdp_config.final_state_dict_type:
+ state_dict_type = cfg.fsdp_config.final_state_dict_type
+ else:
+ state_dict_type = cfg.fsdp_config.state_dict_type
+ trainer.accelerator.state.fsdp_plugin.set_state_dict_type(state_dict_type)
+ trainer.save_model(cfg.output_dir) # only handles FULL_STATE_DICT
+ if state_dict_type == "SHARDED_STATE_DICT":
+ LOG.info(
+ "The final model was saved with a sharded state dict. Please ensure you merge "
+ "the sharded weights with `merge-sharded-fsdp-weights`."
)
- elif state_dict_type == "FULL_STATE_DICT":
- trainer.save_model(cfg.output_dir)
+ checkpoint_dir = determine_last_checkpoint(cfg, update=False)
+ if (
+ not (Path(cfg.output_dir) / "model.safetensors.index.json").exists()
+ and checkpoint_dir
+ ):
+ # import here to prevent circular import
+ from axolotl.cli.merge_sharded_fsdp_weights import merge_fsdp_weights
+
+ fsdp_dir = Path(checkpoint_dir) / "pytorch_model_fsdp_0"
+ merged_path = str(Path(cfg.output_dir) / "merged")
+ merge_fsdp_weights(
+ checkpoint_dir=str(fsdp_dir),
+ output_path=merged_path,
+ safe_serialization=True,
+ )
+ trainer.accelerator.wait_for_everyone()
+ if trainer.accelerator.is_main_process:
+ # move all files in merged_path to cfg.output_dir
+ for merged_file in Path(merged_path).iterdir():
+ if (Path(cfg.output_dir) / merged_file.name).exists():
+ (Path(cfg.output_dir) / merged_file.name).unlink()
+ shutil.move(str(merged_file), cfg.output_dir)
+ shutil.rmtree(merged_path) # remove what should be an empty dir
+ # TODO(wing):see https://github.com/huggingface/transformers/pull/40207
+ # cleanup the FSDP prefix in the model config.json
+ if trainer.accelerator.is_main_process:
+ with open(
+ Path(cfg.output_dir) / "config.json", "r", encoding="utf-8"
+ ) as config_file_io:
+ # read the model config as an OrderedDict
+ config = json.load(config_file_io, object_pairs_hook=OrderedDict)
+ config["architectures"] = [
+ name.lstrip("FSDP") for name in config["architectures"]
+ ]
+ # write the updated model config back
+ with open(
+ os.path.join(cfg.output_dir, "config.json"), "w", encoding="utf-8"
+ ) as config_file_io:
+ json.dump(config, config_file_io, indent=2)
elif cfg.deepspeed and is_deepspeed_zero3_enabled():
# Copied over from: https://github.com/huggingface/accelerate/blob/5ae611118057232f441055f7ef9ba0b0f2b8d533/docs/source/usage_guides/deepspeed.md#saving-and-loading
trainer.accelerator.wait_for_everyone()
@@ -325,9 +330,6 @@ def save_trained_model(
except FileNotFoundError:
pass
elif cfg.local_rank == 0:
- if cfg.flash_optimum and BetterTransformer:
- model = BetterTransformer.reverse(model)
-
if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
trainer.model.save_pretrained(
cfg.output_dir, safe_serialization=safe_serialization
@@ -337,9 +339,7 @@ def save_trained_model(
if hasattr(cfg, "llmcompressor") and cfg.llmcompressor:
# TODO: add integration support so this can be implemented completely within the plugin
- from axolotl.integrations.llm_compressor.utils import (
- save_compressed_model,
- )
+ from axolotl.integrations.llm_compressor.utils import save_compressed_model
save_compressed_model(
model=model,
@@ -416,7 +416,9 @@ def save_initial_configs(
# Pre-save the tokenizer and model configs
LOG.info(f"Pre-saving tokenizer to {cfg.output_dir}...")
- tokenizer.save_pretrained(str(output_dir))
+ tokenizer.save_pretrained(
+ str(Path(cfg.output_dir)), save_jinja_files=cfg.tokenizer_save_jinja_files
+ )
if hasattr(model, "config"):
LOG.info(f"Pre-saving model config to {cfg.output_dir}...")
model.config.save_pretrained(str(output_dir))
@@ -436,7 +438,7 @@ def setup_model_card(cfg: DictDefault):
badge_markdown = """[
](https://github.com/axolotl-ai-cloud/axolotl)"""
transformers.modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
- if getattr(cfg, "axolotl_config_path"):
+ if cfg.axolotl_config_path:
raw_axolotl_cfg = Path(cfg.axolotl_config_path)
version = importlib.metadata.version("axolotl")
if raw_axolotl_cfg.is_file():
@@ -487,8 +489,10 @@ def handle_untrained_tokens_fix(
)
-def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> tuple[
- HFRLTrainerBuilder | HFCausalTrainerBuilder,
+def setup_model_and_trainer(
+ cfg: DictDefault, dataset_meta: TrainDatasetMeta
+) -> tuple[
+ "HFRLTrainerBuilder" | "HFCausalTrainerBuilder",
PeftModel | PreTrainedModel,
PreTrainedTokenizer,
PeftConfig | None,
@@ -535,6 +539,20 @@ def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) ->
)
PLUGIN_MANAGER.post_trainer_create(cfg, trainer)
+ plugin_manager = PluginManager.get_instance()
+ plugin_manager.post_trainer_create(cfg, trainer)
+
+ if cfg.use_ray:
+ try:
+ import ray.train.huggingface.transformers
+
+ trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
+ except ImportError:
+ LOG.warning(
+ "The Ray integration with Hugging Face Transformers is not available. "
+ "To use Ray, install the 'ray[train]' package."
+ )
+
return (
trainer,
model,
@@ -558,8 +576,6 @@ def train(
Returns:
Tuple of (model, tokenizer) after training
"""
- print_axolotl_text_art()
-
# Setup model, tokenizer, (causal or RLHF) trainer, etc.
(
trainer,
@@ -569,12 +585,6 @@ def train(
processor,
) = setup_model_and_trainer(cfg, dataset_meta)
- # Determine if we need to resume from a checkpoint
- resume_from_checkpoint = determine_resume_checkpoint(cfg)
-
- # Configuration for saving
- safe_serialization = cfg.save_safetensors is True
-
# Handle untrained tokens if configured
safe_serialization = cfg.save_safetensors is True
train_dataset = dataset_meta.train_dataset
@@ -588,10 +598,18 @@ def train(
setup_model_card(cfg)
# Execute the training
+ resume_from_checkpoint = determine_last_checkpoint(cfg)
execute_training(cfg, trainer, resume_from_checkpoint)
+ # clear cache
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+
# Save the trained model and cleanup
save_trained_model(cfg, trainer, model, safe_serialization)
+ tokenizer.save_pretrained(
+ str(Path(cfg.output_dir)), save_jinja_files=cfg.tokenizer_save_jinja_files
+ )
create_model_card(cfg, trainer)
if not cfg.use_ray:
cleanup_distributed()
diff --git a/src/axolotl/utils/__init__.py b/src/axolotl/utils/__init__.py
index 3d0ba7c9c..7256a5700 100644
--- a/src/axolotl/utils/__init__.py
+++ b/src/axolotl/utils/__init__.py
@@ -17,7 +17,6 @@ def is_comet_available():
return importlib.util.find_spec("comet_ml") is not None
-# pylint: disable=duplicate-code
def get_pytorch_version() -> tuple[int, int, int]:
"""
Get Pytorch version as a tuple of (major, minor, patch).
@@ -45,10 +44,8 @@ def set_pytorch_cuda_alloc_conf():
)
-def patch_optimized_env():
+def get_not_null(value, default=None):
"""
- Patch environment variables to improve VRAM usage and increase download speed
+ return the value if it's not None, otherwise return the default value
"""
- if os.getenv("HF_HUB_ENABLE_HF_TRANSFER") is None:
- os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
- set_pytorch_cuda_alloc_conf()
+ return value if value is not None else default
diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
index d1e972c81..0a4594991 100644
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -1,6 +1,7 @@
"""Benchmarking and measurement utilities"""
import functools
+import logging
import torch
from transformers.utils.import_utils import is_torch_npu_available
@@ -56,16 +57,17 @@ def gpu_memory_usage(device=0):
@check_cuda_device((0.0, 0.0, 0.0))
def gpu_memory_usage_all(device=0):
- usage = torch.cuda.memory_allocated(device) / 1024.0**3
- reserved = torch.cuda.memory_reserved(device) / 1024.0**3
- smi = gpu_memory_usage_smi(device)
- return usage, reserved - usage, max(0, smi - reserved)
+ active = torch.cuda.memory_stats().get("active_bytes.all.peak", 0) / 1024.0**3
+ allocated = torch.cuda.max_memory_allocated(device) / 1024.0**3
+ reserved = torch.cuda.max_memory_reserved(device) / 1024.0**3
+ torch.cuda.reset_peak_memory_stats(device)
+ return active, allocated, reserved
def mps_memory_usage_all():
- usage = torch.mps.current_allocated_memory() / 1024.0**3
- reserved = torch.mps.driver_allocated_memory() / 1024.0**3
- return usage, reserved - usage, 0
+ active = torch.mps.current_allocated_memory() / 1024.0**3
+ allocated = torch.mps.driver_allocated_memory() / 1024.0**3
+ return active, allocated, 0
def npu_memory_usage_all(device=0):
@@ -91,21 +93,38 @@ def gpu_memory_usage_smi(device=0):
return 0.0
-def log_gpu_memory_usage(log, msg, device):
- cur_device = get_device_type()
+def get_gpu_memory_usage(device: int | torch.device = 0):
+ cur_device_type = str(get_device_type())
if torch.backends.mps.is_available():
usage, cache, misc = mps_memory_usage_all()
- elif "npu" in str(cur_device) and is_torch_npu_available():
+ elif "npu" in cur_device_type and is_torch_npu_available():
usage, cache, misc = npu_memory_usage_all(device)
- else:
+ elif "cuda" in cur_device_type and torch.cuda.is_available():
usage, cache, misc = gpu_memory_usage_all(device)
+ else:
+ return 0.0, 0.0, 0.0
+
+ return usage, cache, misc
+
+
+def log_gpu_memory_usage(
+ log: logging.Logger | logging.LoggerAdapter,
+ msg: str = "",
+ device: int | torch.device = 0,
+):
+ try:
+ active, allocated, reserved = get_gpu_memory_usage(device)
+ except ValueError:
+ # likely CPU, ignore
+ return
+ cur_device_type = str(get_device_type())
extras = []
- if cache > 0:
- extras.append(f"+{cache:.03f}GB cache")
- if misc > 0:
- extras.append(f"+{misc:.03f}GB misc")
- log.info(
- f"{str(cur_device)} memory usage {msg}: {usage:.03f}GB ({', '.join(extras)})",
+ if allocated > 0:
+ extras.append(f"+{allocated:.03f}GB allocated")
+ if reserved > 0:
+ extras.append(f"+{reserved:.03f}GB reserved")
+ msg = f"{cur_device_type} memory active:" if not msg else msg
+ log.debug(
+ f"{msg} {active:.03f}GB ({', '.join(extras)})",
stacklevel=2,
)
- return usage, cache, misc
diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py
index 8b8a77611..36370ef13 100644
--- a/src/axolotl/utils/callbacks/__init__.py
+++ b/src/axolotl/utils/callbacks/__init__.py
@@ -16,8 +16,8 @@ import pandas as pd
import torch
import torch.distributed as dist
import wandb
+import yaml
from datasets import load_dataset
-from optimum.bettertransformer import BetterTransformer
from tqdm import tqdm
from transformers import (
GenerationConfig,
@@ -27,11 +27,12 @@ from transformers import (
TrainerState,
TrainingArguments,
)
-from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
+from transformers.trainer_utils import (
+ SaveStrategy,
+)
from trl.models import unwrap_model_for_generation
from axolotl.utils import is_comet_available, is_mlflow_available
-from axolotl.utils.bench import log_gpu_memory_usage
from axolotl.utils.callbacks.perplexity import Perplexity
from axolotl.utils.distributed import (
barrier,
@@ -53,100 +54,22 @@ IGNORE_INDEX = -100
LOG = get_logger(__name__)
-class EvalFirstStepCallback(
- TrainerCallback
-): # pylint: disable=too-few-public-methods disable=unused-argument
- """
- Callback to trigger evals on the first step
- """
-
- def on_step_end(
- self,
- args: TrainingArguments,
- state: TrainerState,
- control: TrainerControl,
- **kwargs,
- ):
- if args.eval_strategy == IntervalStrategy.STEPS and state.global_step == 1:
- control.should_evaluate = True
- return control
-
-
-class SaveBetterTransformerModelCallback(
- TrainerCallback
-): # pylint: disable=too-few-public-methods
- """Callback to save the BetterTransformer wrapped model"""
-
- def on_step_end(
- self,
- args: TrainingArguments,
- state: TrainerState,
- control: TrainerControl,
- **kwargs,
- ):
- # Save
- if (
- args.save_strategy == IntervalStrategy.STEPS
- and args.save_steps > 0
- and state.global_step % args.save_steps == 0
- ):
- control.should_save = True
-
- if control.should_save:
- checkpoint_folder = os.path.join(
- args.output_dir,
- f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
- )
-
- model = BetterTransformer.reverse(kwargs["model"])
- model.save_pretrained(checkpoint_folder)
- # FIXME - need to cleanup old checkpoints
-
- # since we're saving here, we don't need the trainer loop to attempt to save too b/c
- # the trainer will raise an exception since it can't save a BetterTransformer wrapped model
- control.should_save = False
- return control
-
-
-class GPUStatsCallback(
- TrainerCallback
-): # pylint: disable=too-few-public-methods disable=unused-argument
- """Callback to track GPU utilization"""
-
- def __init__(self, cfg):
- self.cfg = cfg
- self.logged = False
-
- def on_step_end(
- self,
- args: TrainingArguments,
- state: TrainerState,
- control: TrainerControl,
- **kwargs,
- ):
- if not self.logged and state.global_step > 1:
- log_gpu_memory_usage(LOG, "while training", self.cfg.device)
- self.logged = True
- return control
-
-
class LossWatchDogCallback(TrainerCallback):
"""Callback to track loss and stop training if loss is too high"""
def __init__(self, cfg):
self.cfg = cfg
- self.logged = False
self.violations = 0
self.threshold = cfg.loss_watchdog_threshold
self.patience = cfg.loss_watchdog_patience or 3
def on_step_end(
self,
- _args: TrainingArguments,
+ args: TrainingArguments,
state: TrainerState,
control: TrainerControl,
**_kwargs,
- ):
+ ) -> TrainerControl:
if len(state.log_history) > 0 and "loss" in state.log_history[-1]:
if state.log_history[-1]["loss"] > self.threshold:
self.violations += 1
@@ -160,6 +83,21 @@ class LossWatchDogCallback(TrainerCallback):
return control
+class SaveModelOnFirstStepCallback(TrainerCallback):
+ """Callback to save the model on the first step of training if enabled"""
+
+ def on_step_end(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **_kwargs,
+ ) -> TrainerControl:
+ if state.global_step == 1:
+ control.should_save = True
+ return control
+
+
def bench_eval_callback_factory(trainer, tokenizer):
accuracy = evaluate.load("accuracy")
abcd_idx = [
@@ -263,10 +201,10 @@ def bench_eval_callback_factory(trainer, tokenizer):
def on_evaluate(
self,
args: AxolotlTrainingArguments,
- state: TrainerState, # pylint: disable=unused-argument
- control: TrainerControl, # pylint: disable=unused-argument
- metrics: Dict[str, float], # pylint: disable=unused-argument
- **kwargs, # pylint: disable=unused-argument
+ state: TrainerState,
+ control: TrainerControl,
+ metrics: Dict[str, float],
+ **kwargs,
):
data_loader = trainer.get_bench_dataloader(
bench_dataset.remove_columns(["input", "subject", "output", "name"])
@@ -296,7 +234,7 @@ def bench_eval_callback_factory(trainer, tokenizer):
# Extract results by subject.
bench_name = bench_dataset["name"]
bench_names: dict = {s: {"refs": [], "preds": []} for s in set(bench_name)}
- for s, p, r in zip(bench_name, preds, refs): # pylint: disable=invalid-name
+ for s, p, r in zip(bench_name, preds, refs, strict=False):
bench_names[s]["preds"].append(p)
bench_names[s]["refs"].append(r)
barrier()
@@ -334,9 +272,7 @@ def bench_eval_callback_factory(trainer, tokenizer):
bench_scores = []
bench_refs = []
bench_preds = []
- for (
- bench_name
- ) in combined_bench_names: # pylint: disable=consider-using-dict-items
+ for bench_name in combined_bench_names:
bench_score = accuracy.compute(
references=combined_bench_names[bench_name]["refs"],
predictions=combined_bench_names[bench_name]["preds"],
@@ -385,18 +321,18 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
else:
try:
metrics[metric] = evaluate.load(metric)
- except Exception as exc: # pylint: disable=broad-exception-caught
+ except Exception as exc:
LOG.warning(f"{metric}: {exc.args}")
return metrics
def on_evaluate(
self,
- args: AxolotlTrainingArguments, # pylint: disable=unused-argument
+ args: AxolotlTrainingArguments,
state: TrainerState,
control: TrainerControl,
- train_dataloader, # pylint: disable=unused-argument
+ train_dataloader,
eval_dataloader,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
trainer.model_wrapped.eval()
@@ -404,7 +340,6 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
self.cfg.device
) # Use this instead of trainer.model_wrapped.device as it may return cpu if fsdp offloaded
- # pylint: disable=duplicate-code
generation_config = GenerationConfig(
max_new_tokens=self.cfg.eval_max_new_tokens,
bos_token_id=tokenizer.bos_token_id,
@@ -435,9 +370,7 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
try:
# Only pass the kwargs that are in the metric's feature list
metric_kwargs = {
- k: kwargs[k]
- for k in metric._feature_names() # pylint: disable=protected-access
- if k in kwargs
+ k: kwargs[k] for k in metric._feature_names() if k in kwargs
}
if isinstance(metric, Perplexity):
@@ -449,7 +382,7 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
if "score" in metric_score
else metric_score["mean_score"]
)
- except Exception: # pylint: disable=broad-exception-caught
+ except Exception:
traceback.print_exc()
LOG.debug(
f"Failed to compute metric {metric.name} with kwargs {kwargs.keys()}"
@@ -497,6 +430,7 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
batch_input_ids,
batch_labels,
batch_pos_ids,
+ strict=False,
):
if pos_ids is None:
pos_ranges = [(0, len(input_ids_all) - 1)]
@@ -547,7 +481,7 @@ def causal_lm_bench_eval_callback_factory(trainer: Trainer, tokenizer):
prediction_all_tokens = predictions["sequences"].cpu().tolist()
prediction_without_prompt_tokens_list = []
for prompt_token_ids, prediction_tokens in zip(
- prompt_token_ids_list, prediction_all_tokens
+ prompt_token_ids_list, prediction_all_tokens, strict=False
):
prediction_without_prompt_tokens = prediction_tokens[
len(prompt_token_ids) :
@@ -585,12 +519,12 @@ def log_prediction_callback_factory(trainer: Trainer, tokenizer, logger: str):
def on_evaluate(
self,
- args: AxolotlTrainingArguments, # pylint: disable=unused-argument
+ args: AxolotlTrainingArguments,
state: TrainerState,
control: TrainerControl,
- train_dataloader, # pylint: disable=unused-argument
+ train_dataloader,
eval_dataloader,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
eval_table_size = self.cfg.eval_table_size
@@ -600,7 +534,6 @@ def log_prediction_callback_factory(trainer: Trainer, tokenizer, logger: str):
trainer.model.eval()
device = torch.device(self.cfg.device)
- # pylint: disable=duplicate-code
generation_config = GenerationConfig(
max_new_tokens=self.cfg.eval_max_new_tokens,
bos_token_id=tokenizer.bos_token_id,
@@ -668,6 +601,7 @@ def log_prediction_callback_factory(trainer: Trainer, tokenizer, logger: str):
batch_labels,
batch_pos_ids,
batch_logits,
+ strict=False,
):
if pos_ids is None:
pos_ranges = [(0, len(input_ids_all) - 1)]
@@ -721,7 +655,7 @@ def log_prediction_callback_factory(trainer: Trainer, tokenizer, logger: str):
prediction_all_tokens = predictions["sequences"].cpu().tolist()
prediction_without_prompt_tokens_list = []
for prompt_token_ids, prediction_tokens in zip(
- prompt_token_ids_list, prediction_all_tokens
+ prompt_token_ids_list, prediction_all_tokens, strict=False
):
prediction_without_prompt_tokens = prediction_tokens[
len(prompt_token_ids) :
@@ -740,7 +674,11 @@ def log_prediction_callback_factory(trainer: Trainer, tokenizer, logger: str):
prediction_text,
pred_step_text,
) in zip(
- prompt_texts, completion_texts, predicted_texts, pred_step_texts
+ prompt_texts,
+ completion_texts,
+ predicted_texts,
+ pred_step_texts,
+ strict=False,
):
table_data["id"].append(row_index)
table_data["Prompt"].append(prompt_text)
@@ -798,12 +736,12 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
def on_train_begin(
self,
- args: AxolotlTrainingArguments, # pylint: disable=unused-argument
- state: TrainerState, # pylint: disable=unused-argument
+ args: AxolotlTrainingArguments,
+ state: TrainerState,
control: TrainerControl,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
- if is_main_process():
+ if state.is_world_process_zero:
try:
# sync config to top level in run, cannot delete file right away because wandb schedules it to be synced even w/policy = 'now', so let OS delete it later.
with NamedTemporaryFile(
@@ -822,6 +760,37 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
except (FileNotFoundError, ConnectionError) as err:
LOG.warning(f"Error while saving Axolotl config to WandB: {err}")
+ try:
+ with open(self.axolotl_config_path, "r", encoding="utf-8") as f:
+ cfg = yaml.safe_load(f) or {}
+
+ chat_tpl = cfg.get("chat_template_jinja")
+ if chat_tpl:
+ with NamedTemporaryFile(
+ mode="w", delete=True, suffix=".jinja", prefix="chat_template_"
+ ) as temp_ct_file:
+ if (
+ isinstance(chat_tpl, str)
+ and os.path.exists(chat_tpl)
+ and os.path.isfile(chat_tpl)
+ ):
+ copyfile(chat_tpl, temp_ct_file.name)
+ else:
+ temp_ct_file.write(str(chat_tpl))
+ temp_ct_file.flush()
+
+ artifact = wandb.Artifact(
+ f"chat-template-{wandb.run.id}", type="jinja-template"
+ )
+ artifact.add_file(temp_ct_file.name)
+ wandb.log_artifact(artifact)
+ wandb.save(temp_ct_file.name)
+ LOG.info(
+ "The chat_template_jinja has been saved to the WandB run under files."
+ )
+ except (FileNotFoundError, ConnectionError, yaml.YAMLError) as err:
+ LOG.warning(f"Error while saving chat_template_jinja to WandB: {err}")
+
if args.deepspeed:
try:
# sync config to top level in run, cannot delete file right away because wandb schedules it to be synced even w/policy = 'now', so let OS delete it later.
@@ -860,22 +829,68 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
class GCCallback(TrainerCallback):
"""Callback to garbage collect torch cache"""
- def __init__(self, gc_steps=None):
- self.gc_steps = gc_steps
+ def __init__(self, gc_steps: int | None = -1):
+ self.gc_steps: int = gc_steps or -1
+ self.next_gc_on_begin_step: int = -1
- def on_step_end(
- self, args, state, control, **kwargs # pylint: disable=unused-argument
- ):
- if self.gc_steps > 0 and state.global_step % self.gc_steps == 0:
- torch.cuda.empty_cache()
- gc.collect()
-
- def on_epoch_end(
- self, args, state, control, **kwargs # pylint: disable=unused-argument
- ):
+ def _gc(self):
torch.cuda.empty_cache()
gc.collect()
+ def on_train_begin(
+ self,
+ args,
+ state,
+ control,
+ **kwargs,
+ ):
+ self._gc()
+
+ def on_step_begin(
+ self,
+ args,
+ state,
+ control,
+ **kwargs,
+ ):
+ if self.next_gc_on_begin_step == state.global_step or state.global_step == 0:
+ self._gc()
+
+ def on_step_end(
+ self,
+ args,
+ state,
+ control,
+ **kwargs,
+ ):
+ if control.should_evaluate:
+ # automatically GC before evals so the eval memory spike from the CEL doesn't OOM the trainer
+ self._gc()
+ # also GC on the start of the next step after the eval
+ self.next_gc_on_begin_step = state.global_step + 1
+ elif self.gc_steps > 0 and state.global_step % self.gc_steps == 0:
+ self._gc()
+ elif (
+ args.save_strategy == SaveStrategy.STEPS
+ and state.save_steps > 0
+ and state.global_step % state.save_steps == 0
+ ):
+ # gc on save steps in case anything is loaded to CPU RAM like offloaded tensors
+ self._gc()
+ elif state.global_step >= state.max_steps:
+ if args.save_strategy == SaveStrategy.STEPS:
+ # gc on save steps in case anything is loaded to CPU RAM like offloaded tensors
+ self._gc()
+
+ def on_epoch_end(
+ self,
+ args,
+ state,
+ control,
+ **kwargs,
+ ):
+ self._gc()
+
def colab_inference_post_train_callback(trainer: Trainer):
class ColabCallback(TrainerCallback):
@@ -885,16 +900,12 @@ def colab_inference_post_train_callback(trainer: Trainer):
self.gpu_name = torch.cuda.get_device_name(0)
self.cfg = cfg
- def on_train_end(
- self, args, state, control, **kwargs
- ): # pylint: disable=unused-argument
+ def on_train_end(self, args, state, control, **kwargs):
"""
handle T4 gpu, we need to convert attention to eager for inference
"""
if "Tesla T4" in self.gpu_name and self.cfg.xformers_attention:
- trainer.model.config._attn_implementation = ( # pylint: disable=protected-access
- "eager"
- )
+ trainer.model.config._attn_implementation = "eager"
trainer.model.gradient_checkpointing_disable()
trainer.model.config.use_cache = True
trainer.model.eval()
diff --git a/src/axolotl/utils/callbacks/comet_.py b/src/axolotl/utils/callbacks/comet_.py
index 7dce95145..cd3bcf70e 100644
--- a/src/axolotl/utils/callbacks/comet_.py
+++ b/src/axolotl/utils/callbacks/comet_.py
@@ -22,10 +22,10 @@ class SaveAxolotlConfigtoCometCallback(TrainerCallback):
def on_train_begin(
self,
- args: "AxolotlTrainingArguments", # pylint: disable=unused-argument
- state: TrainerState, # pylint: disable=unused-argument
+ args: "AxolotlTrainingArguments",
+ state: TrainerState,
control: TrainerControl,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
if is_main_process():
try:
diff --git a/src/axolotl/utils/callbacks/lisa.py b/src/axolotl/utils/callbacks/lisa.py
index 348cdf2da..03f189d80 100644
--- a/src/axolotl/utils/callbacks/lisa.py
+++ b/src/axolotl/utils/callbacks/lisa.py
@@ -55,9 +55,7 @@ def lisa_callback_factory(trainer: "AxolotlTrainer"):
for param in layer.parameters():
param.requires_grad = False
- def on_step_begin(
- self, args, state, control, **kwargs
- ): # pylint: disable=unused-argument
+ def on_step_begin(self, args, state, control, **kwargs):
# Check if it's time to switch active layers, including at step 0
if state.global_step % self.step_interval == 0 or state.global_step == 1:
self.switch_active_layers()
diff --git a/src/axolotl/utils/callbacks/mlflow_.py b/src/axolotl/utils/callbacks/mlflow_.py
index ac72f5e6d..30120a87d 100644
--- a/src/axolotl/utils/callbacks/mlflow_.py
+++ b/src/axolotl/utils/callbacks/mlflow_.py
@@ -23,7 +23,6 @@ def should_log_artifacts() -> bool:
class SaveAxolotlConfigtoMlflowCallback(TrainerCallback):
- # pylint: disable=duplicate-code
"""Callback to save axolotl config to mlflow"""
def __init__(self, axolotl_config_path):
@@ -31,10 +30,10 @@ class SaveAxolotlConfigtoMlflowCallback(TrainerCallback):
def on_train_begin(
self,
- args: "AxolotlTrainingArguments", # pylint: disable=unused-argument
- state: TrainerState, # pylint: disable=unused-argument
+ args: "AxolotlTrainingArguments",
+ state: TrainerState,
control: TrainerControl,
- **kwargs, # pylint: disable=unused-argument
+ **kwargs,
):
if is_main_process():
try:
diff --git a/src/axolotl/utils/callbacks/models.py b/src/axolotl/utils/callbacks/models.py
new file mode 100644
index 000000000..5a20d70d9
--- /dev/null
+++ b/src/axolotl/utils/callbacks/models.py
@@ -0,0 +1,23 @@
+"""Helper functions for model classes"""
+
+from typing import Tuple
+
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+
+
+def get_causal_lm_model_cls_prefix(model_type: str) -> Tuple[str, str]:
+ if model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+ causal_lm_cls = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[model_type]
+ causal_lm_cls_prefix = causal_lm_cls
+ for suffix in [
+ "ForCausalLM",
+ "ForConditionalGeneration",
+ "LMHeadModel",
+ "GenerationDecoder",
+ ]:
+ causal_lm_cls_prefix = causal_lm_cls_prefix.replace(suffix, "")
+ return causal_lm_cls_prefix, causal_lm_cls
+ causal_lm_cls_prefix = "".join(
+ [part.capitalize() for part in model_type.split("_")]
+ )
+ return causal_lm_cls_prefix, f"{causal_lm_cls_prefix}ForCausalLM"
diff --git a/src/axolotl/utils/callbacks/profiler.py b/src/axolotl/utils/callbacks/profiler.py
index 36604813f..2cf5e0f4f 100644
--- a/src/axolotl/utils/callbacks/profiler.py
+++ b/src/axolotl/utils/callbacks/profiler.py
@@ -19,26 +19,57 @@ class PytorchProfilerCallback(TrainerCallback):
PyTorch Profiler callback to create snapshots of GPU memory usage at specified steps.
"""
- def __init__(self, steps_to_profile: int = 5):
- self.steps_to_profile = steps_to_profile
- if self.steps_to_profile:
- torch.cuda.memory._record_memory_history( # pylint: disable=protected-access
- enabled="all"
- )
+ def __init__(self, steps_to_profile: int = 5, profiler_steps_start: int = 0):
+ # steps are 0 indexed, so to start at 0-th step, we start at beginning of first step,
+ # and finish at end of last step, so 5 steps_to_profile is steps [0, 1, 2, 3, 4]
+ self.profiler_steps_end = profiler_steps_start + steps_to_profile - 1
+ if profiler_steps_start == 0:
+ # start recording memory allocations before everything is allocated, because if we start
+ # at the beginning of step 0, we won't have any memory allocations in the traces
+ torch.cuda.memory._record_memory_history(enabled="all")
+ profiler_steps_start = -1
+ self.profiler_steps_start = profiler_steps_start
- def on_step_end( # pylint: disable=unused-argument
+ def on_step_begin(
self,
- args: TrainingArguments, # pylint: disable=unused-argument
+ args: TrainingArguments,
state: TrainerState,
- control: TrainerControl, # pylint: disable=unused-argument
- **kwargs, # pylint: disable=unused-argument
+ control: TrainerControl,
+ **kwargs,
):
- if state.global_step == self.steps_to_profile:
- snapshot = torch.cuda.memory._snapshot() # pylint: disable=protected-access
+ if state.global_step == self.profiler_steps_start:
+ torch.cuda.memory._record_memory_history(enabled="all")
+
+ def on_step_end(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **kwargs,
+ ):
+ if state.global_step == self.profiler_steps_end:
+ snapshot = torch.cuda.memory._snapshot()
with open(Path(args.output_dir) / "snapshot.pickle", "wb") as fout:
dump(snapshot, fout)
# tell CUDA to stop recording memory allocations now
- torch.cuda.memory._record_memory_history( # pylint: disable=protected-access
- enabled=None
- )
+ torch.cuda.memory._record_memory_history(enabled=None)
+
+ def on_train_end(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **kwargs,
+ ):
+ # make sure to record if we happen to have more steps than steps to profile
+ if (
+ state.global_step >= self.profiler_steps_start
+ and state.global_step < self.profiler_steps_end
+ ):
+ snapshot = torch.cuda.memory._snapshot()
+ with open(Path(args.output_dir) / "snapshot.pickle", "wb") as fout:
+ dump(snapshot, fout)
+
+ # tell CUDA to stop recording memory allocations now
+ torch.cuda.memory._record_memory_history(enabled=None)
diff --git a/src/axolotl/utils/callbacks/qat.py b/src/axolotl/utils/callbacks/qat.py
index cf4d9a937..70746d6be 100644
--- a/src/axolotl/utils/callbacks/qat.py
+++ b/src/axolotl/utils/callbacks/qat.py
@@ -38,9 +38,7 @@ class QATCallback(TrainerCallback):
def __init__(self, cfg: QATConfig):
self.cfg = cfg
- def on_step_begin(
- self, args, state, control, model, **kwargs
- ): # pylint: disable=unused-argument
+ def on_step_begin(self, args, state, control, model, **kwargs):
if self.cfg.fake_quant_after_n_steps is not None:
if state.global_step == 0:
LOG.info(f"Disabling fake quantization at step {state.global_step}")
diff --git a/src/axolotl/utils/callbacks/tokens_per_second.py b/src/axolotl/utils/callbacks/tokens_per_second.py
new file mode 100644
index 000000000..ead129240
--- /dev/null
+++ b/src/axolotl/utils/callbacks/tokens_per_second.py
@@ -0,0 +1,64 @@
+"""A callback for calculating tokens per second during training."""
+
+import time
+
+import torch
+from transformers import (
+ TrainerCallback,
+ TrainerControl,
+ TrainerState,
+ TrainingArguments,
+)
+
+
+class TokensPerSecondCallback(TrainerCallback):
+ """
+ A callback to measure and log tokens per second during training.
+ """
+
+ def __init__(self, tensor_parallel_size, context_parallel_size):
+ super().__init__()
+ self.step_time = 0.0
+ self.start_time = 0.0
+ self.non_data_parallel_size = 1
+ if tensor_parallel_size is not None:
+ self.non_data_parallel_size *= tensor_parallel_size
+ if context_parallel_size is not None:
+ self.non_data_parallel_size *= context_parallel_size
+
+ def on_step_begin(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **kwargs,
+ ): # pylint: disable=unused-argument
+ self.start_time = time.perf_counter()
+ state.last_tokens_per_second = torch.zeros(1)
+
+ def on_step_end(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ **kwargs,
+ ): # pylint: disable=unused-argument
+ if hasattr(state, "num_tokens"):
+ step_time = time.perf_counter() - self.start_time
+ num_tokens_per_device = state.num_tokens.clone()
+ # non data parallel groups have duplicated tokens, so we avoid double-counting
+ num_tokens_per_device = num_tokens_per_device / self.non_data_parallel_size
+ state.last_tokens_per_second = num_tokens_per_device / step_time
+
+ def on_log(
+ self,
+ args: TrainingArguments,
+ state: TrainerState,
+ control: TrainerControl,
+ logs=None,
+ **kwargs,
+ ): # pylint: disable=unused-argument
+ # after logging, clear the running metrics
+ if hasattr(state, "last_tokens_per_second"):
+ state.last_tokens_per_second.zero_()
+ state.num_tokens = torch.zeros(1)
diff --git a/src/axolotl/utils/chat_templates.py b/src/axolotl/utils/chat_templates.py
deleted file mode 100644
index bf496d2c5..000000000
--- a/src/axolotl/utils/chat_templates.py
+++ /dev/null
@@ -1,147 +0,0 @@
-"""
-This module provides functionality for selecting chat templates based on user choices.
-These templates are used for formatting messages in a conversation.
-"""
-
-from typing import TYPE_CHECKING, Any, Dict, Optional
-
-from axolotl.utils.logging import get_logger
-
-if TYPE_CHECKING:
- from transformers import PreTrainedTokenizerBase
-
-LOG = get_logger("axolotl.utils.chat_templates")
-
-_JINJA_TEMPALTE_CHOICE = "jinja"
-_DEFAULT_TEMPLATE_CHOICE = "tokenizer_default"
-_DEFAULT_FALLBACK_CHATML_TEMPLATE_CHOICE_PREFIX = "tokenizer_default_fallback_"
-
-_CHAT_TEMPLATES = {
- "alpaca": "{{ bos_token }}{% for message in messages %}{% if message['role'] == 'system' and loop.first %}{{ message['content'] }}{% elif message['role'] == 'user' %}{{ '### Instruction:\n' + message['content'] }}{% elif message['role'] == 'assistant' %}{{ '### Response:\n' + message['content'] + eos_token }}{% endif %}{% if not loop.last %}{{ '\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '\n\n### Response:\n' }}{% endif %}",
- "mistral_v1": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ ' [INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}", # Mistral 7B V1, Mistral 7B V2, Mixtral 8x7B V1...
- "mistral_v2v3": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + '[/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}", # V3: Mistral 7B V3, Small, Large...
- "mistral_v3_tekken": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST]' + message['content'] + '[/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}", # V3-Tekken: Nemo, Pixtral...
- "mistral_v7_tekken": "{%- set today = strftime_now(\"%Y-%m-%d\") %}\n{%- set default_system_message = \"You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup headquartered in Paris.\\nYour knowledge base was last updated on 2023-10-01. The current date is \" + today + \".\\n\\nWhen you're not sure about some information, you say that you don't have the information and don't make up anything.\\nIf the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. \\\"What are some good restaurants around me?\\\" => \\\"Where are you?\\\" or \\\"When is the next flight to Tokyo\\\" => \\\"Where do you travel from?\\\")\" %}\n\n{{- bos_token }}\n\n{%- if messages[0]['role'] == 'system' %}\n {%- if messages[0]['content'] is string %}\n {%- set system_message = messages[0]['content'] %}\n {%- else %}\n {%- set system_message = messages[0]['content'][0]['text'] %}\n {%- endif %}\n {%- set loop_messages = messages[1:] %}\n{%- else %}\n {%- set system_message = default_system_message %}\n {%- set loop_messages = messages %}\n{%- endif %}\n{{- '[SYSTEM_PROMPT]' + system_message + '[/SYSTEM_PROMPT]' }}\n\n{%- for message in loop_messages %}\n {%- if message['role'] == 'user' %}\n {%- if message['content'] is string %}\n {{- '[INST]' + message['content'] + '[/INST]' }}\n {%- else %}\n {{- '[INST]' }}\n {%- for block in message['content'] %}\n {%- if block['type'] == 'text' %}\n {{- block['text'] }}\n {%- elif block['type'] in ['image', 'image_url'] %}\n {{- '[IMG]' }}\n {%- else %}\n {{- raise_exception('Only text and image blocks are supported in message content!') }}\n {%- endif %}\n {%- endfor %}\n {{- '[/INST]' }}\n {%- endif %}\n {%- elif message['role'] == 'system' %}\n {%- if message['content'] is string %}\n {{- '[SYSTEM_PROMPT]' + message['content'] + '[/SYSTEM_PROMPT]' }}\n {%- else %}\n {{- '[SYSTEM_PROMPT]' + message['content'][0]['text'] + '[/SYSTEM_PROMPT]' }}\n {%- endif %}\n {%- elif message['role'] == 'assistant' %}\n {%- if message['content'] is string %}\n {{- message['content'] + eos_token }}\n {%- else %}\n {{- message['content'][0]['text'] + eos_token }}\n {%- endif %}\n {%- else %}\n {{- raise_exception('Only user, system and assistant roles are supported!') }}\n {%- endif %}\n{%- endfor %}",
- "chatml": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
- "gemma": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '