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

use smaller pretrained models for ci

Browse Source
This commit is contained in:
Wing Lian
2026-04-23 13:51:01 +00:00
parent 1bf65c500e
commit 431888c1de
24 changed files with 614 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/e2e/multigpu/test_dist_muon_fsdp2.py
View File

@@ -37,12 +37,26 @@ def verify_training_success(temp_dir):
event_file = os.path.join(tb_log_path, event_files[0])
reader = SummaryReader(event_file)
df = reader.scalars
train_loss_df = df[df.tag == "train/train_loss"]
train_loss_df = df[df.tag == "train/loss"]
if len(train_loss_df) == 0:
train_loss_df = df[df.tag == "train/train_loss"]
if len(train_loss_df) > 0:
final_loss = train_loss_df.value.values[-1]
values = train_loss_df.value.values
final_loss = values[-1]
assert not torch.isnan(torch.tensor(final_loss)), (
f"Training loss is NaN: {final_loss}"
)
if len(values) >= 2:
initial_loss = float(values[0])
assert float(final_loss) <= initial_loss * 1.10, (
f"Training loss regressed: initial={initial_loss:.4f}, "
f"final={final_loss:.4f} — likely silent bug (e.g. "
"bad label masking) pushed loss scale up."
)
assert float(final_loss) <= 10.0, (
f"Final loss {final_loss:.4f} above sanity bound 10.0 "
"— absolute scale wrong."
)
class TestDistMuon:
@@ -52,7 +66,7 @@ class TestDistMuon:
def test_fft_sft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -109,7 +123,7 @@ class TestDistMuon:
def test_lora_sft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [

26
tests/e2e/multigpu/test_fsdp1.py
View File

@@ -38,12 +38,26 @@ def verify_training_success(temp_dir):
event_file = os.path.join(tb_log_path, event_files[0])
reader = SummaryReader(event_file)
df = reader.scalars
train_loss_df = df[df.tag == "train/train_loss"]
train_loss_df = df[df.tag == "train/loss"]
if len(train_loss_df) == 0:
train_loss_df = df[df.tag == "train/train_loss"]
if len(train_loss_df) > 0:
final_loss = train_loss_df.value.values[-1]
values = train_loss_df.value.values
final_loss = values[-1]
assert not torch.isnan(torch.tensor(final_loss)), (
f"Training loss is NaN: {final_loss}"
)
if len(values) >= 2:
initial_loss = float(values[0])
assert float(final_loss) <= initial_loss * 1.10, (
f"Training loss regressed: initial={initial_loss:.4f}, "
f"final={final_loss:.4f} — likely silent bug (e.g. "
"bad label masking) pushed loss scale up."
)
assert float(final_loss) <= 10.0, (
f"Final loss {final_loss:.4f} above sanity bound 10.0 "
"— absolute scale wrong."
)
class TestFSDP1:
@@ -56,7 +70,7 @@ class TestFSDP1:
def test_fft_sft(self, temp_dir, fsdp_cpu_ram_efficient_loading):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -126,7 +140,7 @@ class TestFSDP1:
def test_lora_sft(self, temp_dir, adapter_config):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -190,7 +204,7 @@ class TestFSDP1:
def test_dpo_fft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"rl": "dpo",
@@ -262,7 +276,7 @@ class TestFSDP1:
def test_dpo_lora(self, temp_dir, adapter_config):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"load_in_4bit": adapter_config["load_in_4bit"],
"rl": "dpo",
"chat_template": "chatml",

32
tests/e2e/multigpu/test_fsdp2.py
View File

@@ -38,12 +38,26 @@ def verify_training_success(temp_dir):
event_file = os.path.join(tb_log_path, event_files[0])
reader = SummaryReader(event_file)
df = reader.scalars
train_loss_df = df[df.tag == "train/train_loss"]
train_loss_df = df[df.tag == "train/loss"]
if len(train_loss_df) == 0:
train_loss_df = df[df.tag == "train/train_loss"]
if len(train_loss_df) > 0:
final_loss = train_loss_df.value.values[-1]
values = train_loss_df.value.values
final_loss = values[-1]
assert not torch.isnan(torch.tensor(final_loss)), (
f"Training loss is NaN: {final_loss}"
)
if len(values) >= 2:
initial_loss = float(values[0])
assert float(final_loss) <= initial_loss * 1.10, (
f"Training loss regressed: initial={initial_loss:.4f}, "
f"final={final_loss:.4f} — likely silent bug (e.g. "
"bad label masking) pushed loss scale up."
)
assert float(final_loss) <= 10.0, (
f"Final loss {final_loss:.4f} above sanity bound 10.0 "
"— absolute scale wrong."
)
class TestFSDP2:
@@ -57,7 +71,7 @@ class TestFSDP2:
def test_fft_sft(self, temp_dir, fsdp_cpu_ram_efficient_loading):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -114,7 +128,7 @@ class TestFSDP2:
def test_lora_sft(self, temp_dir, peft_use_dora):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -180,7 +194,7 @@ class TestFSDP2:
def test_lora_sft_kernels(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -243,7 +257,7 @@ class TestFSDP2:
def test_qlora_sft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -305,7 +319,7 @@ class TestFSDP2:
def test_qlora_sft_kernels(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -370,7 +384,7 @@ class TestFSDP2:
def test_dpo_fft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"rl": "dpo",
@@ -428,7 +442,7 @@ class TestFSDP2:
def test_dpo_lora(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"rl": "dpo",
"chat_template": "chatml",

2
tests/e2e/multigpu/test_fsdp2_lora_kernels.py
View File

@@ -40,7 +40,7 @@ def _run_training(temp_dir, cfg):
def _base_lora_fsdp2_config(temp_dir, **overrides):
"""Base config for LoRA + FSDP2 + kernel tests."""
cfg = {
"base_model": "Qwen/Qwen3-0.6B",
"base_model": "axolotl-ai-co/tiny-qwen3-129m",
"sequence_len": 512,
"val_set_size": 0.0,
"datasets": [

8
tests/e2e/multigpu/test_tp.py
View File

@@ -8,7 +8,7 @@ from accelerate.test_utils import execute_subprocess_async, get_torch_dist_uniqu
from axolotl.utils.dict import DictDefault
from tests.e2e.utils import check_tensorboard, require_torch_2_7_0
from tests.e2e.utils import check_tensorboard_loss_decreased, require_torch_2_7_0
class TestTensorParallel:
@@ -21,7 +21,7 @@ class TestTensorParallel:
def test_fft_sft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "Qwen/Qwen2.5-0.5B",
"base_model": "axolotl-ai-co/tiny-qwen2-129m",
"sequence_len": 2048,
"val_set_size": 0.01,
"datasets": [
@@ -63,6 +63,6 @@ class TestTensorParallel:
]
)
check_tensorboard(
temp_dir + "/runs", "train/train_loss", 1.0, "Train Loss (%s) is too high"
check_tensorboard_loss_decreased(
temp_dir + "/runs", max_initial=5.0, max_final=4.7
)