tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

bug-fix: only apply patches when CUDA is available (#3561)

Browse Source 
* bug-fix: only apply patches when CUDA is available

This will otherwise crash when performing operations with CUDA_VISIBLE_DEVICES=, such as LoRA merging on CPU.

This patch only patches the Qwen 3.5 model, since that's the only one I've tested. This patch should most likely check torch.cuda for all other models as well. One limitation here is that I'm assuming the user runs CUDA, but that assumption is not restricted to this patch so it is probably fine.

* include patch_qwen3_next_modeling_packing, patch_qwen3_5_moe_modeling_packing, and patch_qwen3_5_vlm_flash_attention in cuda guard
This commit is contained in:
kallewoof
2026-04-01 08:05:15 +09:00
committed by GitHub
parent a81feabbd9
commit a4c94416eb
1 changed files with 35 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
src/axolotl/loaders/patch_manager.py
View File

@@ -8,6 +8,7 @@ import os
from functools import cached_property from functools import cached_property
import addict import addict
import torch
import transformers import transformers
from transformers import PretrainedConfig, PreTrainedModel from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_flash_attention_utils import is_flash_attn_available from transformers.modeling_flash_attention_utils import is_flash_attn_available
@@ -258,6 +259,32 @@ class PatchManager:
patch_llama4_linearized_modeling() patch_llama4_linearized_modeling()
if self.cfg.model_config_type == "kimi_linear":
from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
patch_kimi_model,
)
patch_kimi_model()
if self.cfg.model_config_type == "nemotron_h":
if self.cfg.sample_packing:
from transformers.models.nemotron_h.modeling_nemotron_h import (
NemotronHPreTrainedModel,
)
from axolotl.monkeypatch.models.nemotron_h.modeling import (
patch_nemotron_h_modeling_packing,
)
patch_nemotron_h_modeling_packing()
# supports_gradient_checkpointing is only enabled after
# patch_nemotron_h_modeling_packing() installs the GC-compatible
# NemotronHBlock.forward. Without the patch, upstream marks this
# False because the original block forward is not GC-safe.
NemotronHPreTrainedModel.supports_gradient_checkpointing = True
# Patches requiring CUDA
if torch.cuda.is_available():
if self.cfg.model_config_type == "qwen3_next" and self.cfg.sample_packing: if self.cfg.model_config_type == "qwen3_next" and self.cfg.sample_packing:
from axolotl.monkeypatch.models.qwen3_next.modeling import ( from axolotl.monkeypatch.models.qwen3_next.modeling import (
patch_qwen3_next_modeling_packing, patch_qwen3_next_modeling_packing,
@@ -290,30 +317,6 @@ class PatchManager:
patch_qwen3_5_vlm_flash_attention() patch_qwen3_5_vlm_flash_attention()
if self.cfg.model_config_type == "kimi_linear":
from axolotl.monkeypatch.models.kimi_linear.patch_kimi_linear import (
patch_kimi_model,
)
patch_kimi_model()
if self.cfg.model_config_type == "nemotron_h":
if self.cfg.sample_packing:
from transformers.models.nemotron_h.modeling_nemotron_h import (
NemotronHPreTrainedModel,
)
from axolotl.monkeypatch.models.nemotron_h.modeling import (
patch_nemotron_h_modeling_packing,
)
patch_nemotron_h_modeling_packing()
# supports_gradient_checkpointing is only enabled after
# patch_nemotron_h_modeling_packing() installs the GC-compatible
# NemotronHBlock.forward. Without the patch, upstream marks this
# False because the original block forward is not GC-safe.
NemotronHPreTrainedModel.supports_gradient_checkpointing = True
@staticmethod @staticmethod
def _fix_nemotron_h_conversion_mapping(): def _fix_nemotron_h_conversion_mapping():
"""Remove the spurious embedding→embeddings WeightRenaming from the """Remove the spurious embedding→embeddings WeightRenaming from the