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

feat: move to uv first (#3545)

Browse Source 
* feat: move to uv first

* fix: update doc to uv first

* fix: merge dev/tests into uv pyproject

* fix: update docker docs to match current config

* fix: migrate examples to readme

* fix: add llmcompressor to conflict

* feat: rec uv sync with lockfile for dev/ci

* fix: update docker docs to clarify how to use uv images

* chore: docs

* fix: use system python, no venv

* fix: set backend cpu

* fix: only set for installing pytorch step

* fix: remove unsloth kernel and installs

* fix: remove U in tests

* fix: set backend in deps too

* chore: test

* chore: comments

* fix: attempt to lock torch

* fix: workaround torch cuda and not upgraded

* fix: forgot to push

* fix: missed source

* fix: nightly upstream loralinear config

* fix: nightly phi3 long rope not work

* fix: forgot commit

* fix: test phi3 template change

* fix: no more requirements

* fix: carry over changes from new requirements to pyproject

* chore: remove lockfile per discussion

* fix: set match-runtime

* fix: remove unneeded hf hub buildtime

* fix: duplicate cache delete on nightly

* fix: torchvision being overridden

* fix: migrate to uv images

* fix: leftover from merge

* fix: simplify base readme

* fix: update assertion message to be clearer

* chore: docs

* fix: change fallback for cicd script

* fix: match against main exactly

* fix: peft 0.19.1 change

* fix: e2e test

* fix: ci

* fix: e2e test
This commit is contained in:
NanoCode012
2026-04-21 21:16:03 +07:00
committed by GitHub
parent 323da791eb
commit 9de5b76336
58 changed files with 496 additions and 1520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
tests/e2e/integrations/test_scattermoe_lora_olmoe.py
View File

@@ -54,24 +54,8 @@ except (ImportError, ModuleNotFoundError):
)
def peft_lora_to_scattermoe(peft_A, peft_B, num_experts, rank):
peft_B_em = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
K_inter, N_hidden = peft_B.shape[0], peft_A.shape[1]
smoe_A = torch.zeros(
rank * num_experts,
K_inter,
device=peft_A.device,
dtype=peft_A.dtype,
)
smoe_B = torch.zeros(
N_hidden,
rank * num_experts,
device=peft_A.device,
dtype=peft_A.dtype,
)
for e in range(num_experts):
s = e * rank
smoe_A[s : s + rank, :] = peft_B_em[:, s : s + rank].T
smoe_B[:, s : s + rank] = peft_A[s : s + rank, :].T
smoe_A = peft_A
smoe_B = peft_lora_B_to_scattermoe(peft_B, num_experts, rank)
return smoe_A, smoe_B
def _unwrap_experts_lora(experts_module):
@@ -322,12 +306,14 @@ class TestLoRABLayoutConversion:
hidden, inter = 32, 16
scaling = 2.0
peft_A = torch.randn(E * r, hidden)
peft_B = torch.randn(inter, E * r)
# peft >=0.19.1 for down_proj [E, hidden, inter]:
# swaps in/out, lora_A [r*E, inter], lora_B [hidden, r*E]
peft_A = torch.randn(E * r, inter)
peft_B = torch.randn(hidden, E * r)
A_r = peft_A.reshape(E, r, hidden)
B_r = peft_B.reshape(inter, r, E)
delta_peft = torch.einsum("o r e, e r i -> e i o", B_r, A_r) * scaling
A_r = peft_A.reshape(E, r, inter)
B_r = peft_B.reshape(hidden, r, E)
delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
smoe_A, smoe_B = peft_lora_to_scattermoe(peft_A, peft_B, E, r)
for e in range(E):
@@ -339,31 +325,22 @@ class TestLoRABLayoutConversion:
)
def test_gate_up_proj_conversion(self):
"""Verify gate_up_proj LoRA conversion with non-square dims (Qwen3-like).
"""Verify gate_up_proj LoRA conversion with non-square dims.
gate_up_proj param: [E, 2*inter, hidden].
peft: in_features=2*inter, out_features=hidden.
peft lora_A: [r*E, 2*inter], lora_B: [hidden, r*E].
scattermoe W = param.T = [E, hidden, 2*inter], K=hidden, N=2*inter.
peft swaps in/out for 3D: lora_A [r*E, hidden], lora_B [2*inter, r*E].
scattermoe needs: lora_A [r*E, K=hidden], lora_B [N=2*inter, r*E].
Uses non-square dims (hidden=32 != 2*inter=24) to catch A<->B swap bugs.
"""
E, r = 4, 2
hidden, inter = 32, 12 # 2*inter=24 != hidden=32
scaling = 2.0
# peft assigns: in_features=2*inter, out_features=hidden
peft_A = torch.randn(E * r, 2 * inter) # [r*E, in_features=2*inter]
peft_B = torch.randn(hidden, E * r) # [out_features=hidden, r*E]
peft_A = torch.randn(E * r, hidden) # [r*E, in=hidden]
peft_B = torch.randn(2 * inter, E * r) # [out=2*inter, r*E]
# peft delta via einsum: "o r e, e r i -> e i o"
A_r = peft_A.reshape(E, r, 2 * inter)
B_r = peft_B.reshape(hidden, r, E)
delta_peft = torch.einsum("o r e, e r i -> e i o", B_r, A_r) * scaling
# delta_peft[e] has shape [in_features, out_features] = [2*inter, hidden]
# = param[e] shape [2*inter, hidden]
A_r = peft_A.reshape(E, r, hidden)
B_r = peft_B.reshape(2 * inter, r, E)
delta_peft = torch.einsum("o r e, e r i -> e o i", B_r, A_r) * scaling
smoe_A, smoe_B = peft_gate_up_lora_to_scattermoe(peft_A, peft_B, E, r)
# smoe_A should be [r*E, K=hidden], smoe_B should be [N=2*inter, r*E]
@@ -421,23 +398,21 @@ class TestPeftLoRAWeightExtraction:
r,
)
# gate_up_proj [E, 2*inter, hidden]
# peft: in_features=2*inter (dim 1), out_features=hidden (dim 2)
# gate_up_proj [E, 2*inter, hidden] — peft swaps in/out for 3D
assert trainable[
"base_model.model.moe.experts.base_layer.lora_A.default.weight"
].shape == (E * r, 2 * config.intermediate_size)
assert trainable[
"base_model.model.moe.experts.base_layer.lora_B.default.weight"
].shape == (config.hidden_size, E * r)
# down_proj [E, hidden, inter]
# peft: in_features=hidden (dim 1), out_features=inter (dim 2)
assert trainable[
"base_model.model.moe.experts.lora_A.default.weight"
].shape == (E * r, config.hidden_size)
assert trainable[
"base_model.model.moe.experts.base_layer.lora_B.default.weight"
].shape == (2 * config.intermediate_size, E * r)
# down_proj [E, hidden, inter] — peft swaps in/out for 3D
assert trainable[
"base_model.model.moe.experts.lora_A.default.weight"
].shape == (E * r, config.intermediate_size)
assert trainable[
"base_model.model.moe.experts.lora_B.default.weight"
].shape == (config.intermediate_size, E * r)
].shape == (config.hidden_size, E * r)
@requires_cuda
def test_peft_forward_runs(self):
@@ -488,8 +463,7 @@ class TestPeftLoRAWeightExtraction:
assert gup_lora is not None, "gate_up_proj LoRA not detected"
assert down_lora is not None, "down_proj LoRA not detected"
# Check shapes (after peft->scattermoe conversion with A<->B swap)
# gate_up_proj W = param.T = [E, hidden, 2*inter], K=hidden, N=2*inter
# gate_up_proj: K=hidden, N=2*inter
E, r = config.num_experts, 4
gup_A, gup_B, gup_s = gup_lora
assert gup_A.shape == (E * r, config.hidden_size), (
@@ -501,7 +475,7 @@ class TestPeftLoRAWeightExtraction:
f"{(2 * config.intermediate_size, E * r)}, got {gup_B.shape}"
)
# down_proj W = param.T = [E, inter, hidden], K=inter, N=hidden
# down_proj: K=inter, N=hidden
down_A, down_B, down_s = down_lora
assert down_A.shape == (E * r, config.intermediate_size), (
f"down_proj smoe_A: expected [r*E, K=inter]={(E * r, config.intermediate_size)}, "

21
tests/e2e/patched/test_unsloth_integration.py
View File

@@ -1,21 +0,0 @@
"""Test module for checking whether the integration of Unsloth with Hugging Face Transformers is working as expected."""
import unittest
import pytest
@pytest.mark.skip(
reason="Unsloth integration will be broken going into latest transformers"
)
class TestUnslothIntegration(unittest.TestCase):
"""Unsloth monkeypatch integration tests."""
def test_is_self_attn_patchable(self):
from axolotl.monkeypatch.unsloth_ import check_self_attn_is_patchable
# ensures the current version of transformers has loss code that matches our patching code
self.assertTrue(
check_self_attn_is_patchable(),
"HF transformers self attention code has changed and isn't patchable",
)

184
tests/e2e/patched/test_unsloth_qlora.py
View File

@@ -1,184 +0,0 @@
"""
e2e tests for unsloth qlora
"""
import pytest
from axolotl.common.datasets import load_datasets
from axolotl.train import train
from axolotl.utils.config import normalize_config, validate_config
from axolotl.utils.dict import DictDefault
from ..utils import check_model_output_exists, check_tensorboard
@pytest.mark.skip(
reason="Unsloth integration will be broken going into latest transformers"
)
class TestUnslothQLoRA:
"""
Test class for Unsloth QLoRA Llama models
"""
@pytest.mark.parametrize(
"sample_packing",
[True, False],
)
def test_unsloth_llama_qlora_fa2(self, temp_dir, sample_packing):
cfg = DictDefault(
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
"sequence_len": 1024,
"sample_packing": sample_packing,
"flash_attention": True,
"unsloth_lora_mlp": True,
"unsloth_lora_qkv": True,
"unsloth_lora_o": True,
"load_in_4bit": True,
"adapter": "qlora",
"lora_r": 16,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_target_linear": True,
"val_set_size": 0.05,
"special_tokens": {
"pad_token": "<|endoftext|>",
},
"datasets": [
{
"path": "mhenrichsen/alpaca_2k_test",
"type": "alpaca",
},
],
"num_epochs": 1,
"max_steps": 5,
"save_steps": 10,
"micro_batch_size": 4,
"gradient_accumulation_steps": 2,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"use_tensorboard": True,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)
normalize_config(cfg)
dataset_meta = load_datasets(cfg=cfg)
train(cfg=cfg, dataset_meta=dataset_meta)
check_model_output_exists(temp_dir, cfg)
check_tensorboard(
temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
)
def test_unsloth_llama_qlora_unpacked(self, temp_dir):
cfg = DictDefault(
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
"sequence_len": 1024,
"unsloth_lora_mlp": True,
"unsloth_lora_qkv": True,
"unsloth_lora_o": True,
"sample_packing": False,
"load_in_4bit": True,
"adapter": "qlora",
"lora_r": 16,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_target_linear": True,
"val_set_size": 0.05,
"special_tokens": {
"pad_token": "<|endoftext|>",
},
"datasets": [
{
"path": "mhenrichsen/alpaca_2k_test",
"type": "alpaca",
},
],
"num_epochs": 1,
"max_steps": 5,
"save_steps": 10,
"micro_batch_size": 4,
"gradient_accumulation_steps": 2,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"use_tensorboard": True,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)
normalize_config(cfg)
dataset_meta = load_datasets(cfg=cfg)
train(cfg=cfg, dataset_meta=dataset_meta)
check_model_output_exists(temp_dir, cfg)
check_tensorboard(
temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
)
@pytest.mark.parametrize(
"sdp_attention",
[True, False],
)
def test_unsloth_llama_qlora_unpacked_no_fa2_fp16(self, temp_dir, sdp_attention):
cfg = DictDefault(
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
"sequence_len": 1024,
"unsloth_lora_mlp": True,
"unsloth_lora_qkv": True,
"unsloth_lora_o": True,
"sample_packing": False,
"load_in_4bit": True,
"adapter": "qlora",
"lora_r": 16,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_target_linear": True,
"val_set_size": 0.05,
"special_tokens": {
"pad_token": "<|endoftext|>",
},
"datasets": [
{
"path": "mhenrichsen/alpaca_2k_test",
"type": "alpaca",
},
],
"num_epochs": 1,
"max_steps": 5,
"save_steps": 10,
"micro_batch_size": 4,
"gradient_accumulation_steps": 2,
"sdp_attention": sdp_attention,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"use_tensorboard": True,
"fp16": True,
"save_first_step": False,
}
)
cfg = validate_config(cfg)
normalize_config(cfg)
dataset_meta = load_datasets(cfg=cfg)
train(cfg=cfg, dataset_meta=dataset_meta)
check_model_output_exists(temp_dir, cfg)
check_tensorboard(
temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss (%s) is too high"
)