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

sample generation, tests fixes

Browse Source
This commit is contained in:
Dan Saunders
2025-08-18 18:25:04 +00:00
parent 8569675b26
commit 556a69118f
9 changed files with 585 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
tests/e2e/test_diffusion.py
View File

@@ -1,6 +1,4 @@
"""
E2E smoke test for diffusion training plugin
"""
"""E2E smoke test for diffusion training plugin."""
from axolotl.common.datasets import load_datasets
from axolotl.train import train
@@ -11,13 +9,12 @@ from tests.e2e.utils import check_model_output_exists
class TestDiffusion:
"""
Test case for diffusion training plugin
"""
"""Test case for diffusion training plugin."""
def test_diffusion_smoke_test(self, temp_dir):
"""
Smoke test for diffusion training to ensure the plugin loads and trains without error.
Smoke test for diffusion training to ensure the plugin loads and trains without
error.
"""
cfg = DictDefault(
{
@@ -36,7 +33,7 @@ class TestDiffusion:
},
],
"num_epochs": 1,
"max_steps": 3, # Very short for smoke test
"max_steps": 3,
"micro_batch_size": 1,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
@@ -48,33 +45,23 @@ class TestDiffusion:
"save_first_step": False,
"logging_steps": 1,
"eval_steps": 3,
"plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
# Diffusion-specific config
"diffusion_mask_token_id": 32000,
"plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
"diffusion_mask_token_id": 16,
"diffusion_eps": 1e-3,
"diffusion_importance_weighting": False,
}
)
# Normalize and validate config
cfg = normalize_config(cfg)
cfg = validate_config(cfg)
normalize_config(cfg)
dataset_meta = load_datasets(cfg=cfg)
# Load datasets to ensure they work with diffusion training
datasets_meta = load_datasets(cfg=cfg, cli_args=DictDefault({}))
assert datasets_meta.train_dataset is not None
assert len(datasets_meta.train_dataset) > 0
# Run training
train(cfg=cfg, cli_args=DictDefault({}), dataset_meta=datasets_meta)
# Check that model was saved
check_model_output_exists(cfg)
train(cfg=cfg, dataset_meta=dataset_meta)
check_model_output_exists(temp_dir, cfg)
def test_diffusion_sft_labels(self, temp_dir):
"""
Test that diffusion training properly handles SFT data with labels.
"""
"""Test that diffusion training properly handles SFT data with labels."""
cfg = DictDefault(
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
@@ -92,7 +79,7 @@ class TestDiffusion:
},
],
"num_epochs": 1,
"max_steps": 2, # Very short for smoke test
"max_steps": 3,
"micro_batch_size": 1,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
@@ -104,35 +91,29 @@ class TestDiffusion:
"save_first_step": False,
"logging_steps": 1,
"eval_steps": 2,
"plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
# Diffusion-specific config
"diffusion_mask_token_id": 32000,
"plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
"diffusion_mask_token_id": 16,
"diffusion_eps": 1e-3,
"diffusion_importance_weighting": True, # Test importance weighting
"diffusion_importance_weighting": True,
# Ensure we have proper SFT labels
"train_on_inputs": False, # This ensures prompt tokens get -100 labels
"train_on_inputs": False,
}
)
# Normalize and validate config
cfg = normalize_config(cfg)
cfg = validate_config(cfg)
normalize_config(cfg)
dataset_meta = load_datasets(cfg=cfg)
# Load datasets
datasets_meta = load_datasets(cfg=cfg, cli_args=DictDefault({}))
# Verify that the dataset has labels
sample = datasets_meta.train_dataset[0]
sample = dataset_meta.train_dataset[0]
assert "labels" in sample, "SFT dataset should have labels"
# Check that some labels are -100 (prompt tokens)
labels = sample["labels"]
if hasattr(labels, "tolist"):
labels = labels.tolist()
assert -100 in labels, "SFT dataset should have -100 labels for prompt tokens"
# Run training
train(cfg=cfg, cli_args=DictDefault({}), dataset_meta=datasets_meta)
# Check that model was saved
check_model_output_exists(cfg)
train(cfg=cfg, dataset_meta=dataset_meta)
check_model_output_exists(temp_dir, cfg)