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

transformers v5 upgrade (#3272)

Browse Source 
* Prepare for transformers v5 upgrade

* fix hf cli

* update for hf hub changes

* fix tokenizer apply_chat_template args

* remap include_tokens_per_second

* fix tps

* handle migration for warmup

* use latest hf hub

* Fix scan -> ls

* fix import

* fix for renaming of mistral common tokenizer -> backend

* update for fixed tokenziation for llama

* Skip phi35 tests for now

* remove mistral patch fixed upstream in huggingface/transformers#41439

* use namespacing for patch

* don't rely on sdist for e2e tests for now

* run modal ci without waiting too

* Fix dep for ci

* fix imports

* Fix fp8 check

* fsdp2 fixes

* fix version handling

* update fsdp version tests for new v5 behavior

* Fail multigpu tests after 3 failures

* skip known v5 broken tests for now and cleanup

* bump deps

* unmark skipped test

* re-enable test_fsdp_qlora_prequant_packed test

* increase multigpu ci timeout

* skip broken gemma3 test

* reduce timout back to original 120min now that the hanging test is skipped

* fix for un-necessary collator for pretraining with bsz=1

* fix: safe_serialization deprecated in transformers v5 rc01 (#3318)

* torch_dtype deprecated

* load model in float32 for consistency with tests

* revert some test fixtures back

* use hf cache ls instead of scan

* don't strip fsdp_version

more fdsp_Version fixes for v5
fix version in fsdp_config
fix aliasing
fix fsdp_version check
check fsdp_version is 2 in both places

* Transformers v5 rc2 (#3347)

* bump dep

* use latest fbgemm, grab model config as part of fixture, un-skip test

* import AutoConfig

* don't need more problematic autoconfig when specifying config.json manually

* add fixtures for argilla ultrafeedback datasets

* download phi4-reasoning

* fix arg

* update tests for phi fast tokenizer changes

* use explicit model types for gemma3

---------

Co-authored-by: Wing Lian <wing@axolotl.ai>

* fix: AutoModelForVision2Seq -> AutoModelForImageTextToText

* chore: remove duplicate

* fix: attempt fix gemma3 text mode

* chore: lint

* ga release of v5

* need property setter for name_or_path for mistral tokenizer

* vllm not compatible with transformers v5

* setter for chat_template w mistral too

---------

Co-authored-by: NanoCode012 <nano@axolotl.ai>
Co-authored-by: salman <salman.mohammadi@outlook.com>
This commit is contained in:
Wing Lian
2026-01-27 17:08:24 -05:00
committed by GitHub
parent a531e9d946
commit fc4e37920b
74 changed files with 262 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/e2e/integrations/test_cut_cross_entropy.py
View File

@@ -10,7 +10,7 @@ from axolotl.utils import get_pytorch_version
from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
from axolotl.utils.dict import DictDefault
from ..utils import check_model_output_exists
from tests.e2e.utils import check_model_output_exists
@pytest.fixture()
@@ -39,7 +39,6 @@ def min_cfg(temp_dir):
"optimizer": "adamw_torch_fused",
"output_dir": temp_dir,
"lr_scheduler": "cosine",
"save_safetensors": True,
"max_steps": 10,
"bf16": "auto",
"save_first_step": False,
@@ -92,7 +91,6 @@ class TestCutCrossEntropyIntegration:
"optimizer": "adamw_torch_fused",
"output_dir": temp_dir,
"lr_scheduler": "cosine",
"save_safetensors": True,
"max_steps": 10,
"bf16": "auto",
"save_first_step": False,

1
tests/e2e/integrations/test_fp8.py
View File

@@ -48,7 +48,6 @@ class FP8IntegrationTestCase:
"sample_packing": True,
"fp8": True,
"torch_compile": True,
"save_safetensors": True,
"save_first_step": False,
}
)

2
tests/e2e/integrations/test_hooks.py
View File

@@ -11,7 +11,7 @@ from axolotl.train import train
from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
from axolotl.utils.dict import DictDefault
from ..utils import check_model_output_exists
from tests.e2e.utils import check_model_output_exists
class LogHooksPlugin(BasePlugin):

1
tests/e2e/integrations/test_kd.py
View File

@@ -65,7 +65,6 @@ def min_cfg(temp_dir):
},
"max_steps": 5,
"output_dir": temp_dir,
"save_safetensors": True,
"use_tensorboard": True,
"save_first_step": False,
}

2
tests/e2e/integrations/test_liger.py
View File

@@ -48,7 +48,6 @@ class LigerIntegrationTestCase:
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"max_steps": 5,
"save_first_step": False,
@@ -99,7 +98,6 @@ class LigerIntegrationTestCase:
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"max_steps": 5,
"save_first_step": False,

1
tests/e2e/integrations/test_llm_compressor.py
View File

@@ -57,7 +57,6 @@ class TestLLMCompressorIntegration:
"learning_rate": 1e-5,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"max_steps": 5,
"llmcompressor": {

3
tests/e2e/multigpu/solo/test_grpo.py
View File

@@ -220,7 +220,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
"learning_rate": 0.0001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,
@@ -315,7 +314,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
"learning_rate": 0.0001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,
@@ -408,7 +406,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
"learning_rate": 0.0001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,

5
tests/e2e/multigpu/test_fp8_fsdp2.py
View File

@@ -11,7 +11,7 @@ from transformers.testing_utils import get_torch_dist_unique_port
from axolotl.utils.dict import DictDefault
from tests.e2e.utils import most_recent_subdir, require_hopper, require_torch_2_7_0
from tests.e2e.utils import most_recent_subdir, require_torch_2_7_0, supports_fp8
AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent
@@ -49,7 +49,7 @@ class TestFP8FSDP2:
"""Test class for FP8 mixed precision with FSDP2 functionality."""
@require_torch_2_7_0
@require_hopper
@supports_fp8
def test_fp8_fsdp2_smoke(self, temp_dir):
"""Smoke test for 2-GPU FP8 + torch.compile + FSDP2 training"""
cfg = DictDefault(
@@ -94,7 +94,6 @@ class TestFP8FSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"save_safetensors": True,
"save_first_step": False,
}
)

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

@@ -244,6 +244,7 @@ class TestFSDP1:
verify_training_success(temp_dir)
@pytest.mark.skip("broken in transformers v5")
@pytest.mark.parametrize(
"adapter_config",
[

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

@@ -150,6 +150,10 @@ class TestFSDP2:
},
"use_tensorboard": True,
"bf16": True,
# explicitly disable LORA kernels, as they may be auto-enabled
"lora_mlp_kernel": False,
"lora_qkv_kernel": False,
"lora_o_kernel": False,
}
)

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

@@ -23,6 +23,7 @@ def download_model():
snapshot_download("axolotl-mirrors/gemma-3-4b-pt", repo_type="model")
@pytest.mark.skip(reason="FIXME")
class TestMultiGPUGemma3:
"""
Test case for Gemma3 models using LoRA
@@ -32,6 +33,7 @@ class TestMultiGPUGemma3:
cfg = DictDefault(
{
"base_model": "axolotl-mirrors/gemma-3-4b-pt",
"unfrozen_parameters": ["model.language_model.*", "lm_head"],
"sequence_len": 2048,
"ddp_find_unused_parameters": True,
"sample_packing": True,

1
tests/e2e/multigpu/test_llama.py
View File

@@ -901,7 +901,6 @@ class TestMultiGPULlama:
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
# "deepspeed": str(AXOLOTL_ROOT / "deepspeed_configs/zero1.json"),
"use_tensorboard": True,
"save_first_step": False,

1
tests/e2e/patched/test_activation_checkpointing.py
View File

@@ -66,7 +66,6 @@ class TestActivationCheckpointing:
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"gradient_checkpointing": gradient_checkpointing,
"save_first_step": False,
"dataset_num_proc": 4,

1
tests/e2e/patched/test_peft_embeddings.py
View File

@@ -46,7 +46,6 @@ class TestLlamaPeftEmbeddings:
"flash_attention": True,
"sample_packing": False,
"bf16": "auto",
"save_safetensors": True,
"embeddings_skip_upcast": True,
"save_first_step": False,
}

1
tests/e2e/patched/test_resume.py
View File

@@ -58,7 +58,6 @@ class TestResumeLlama:
"save_total_limit": 5,
"max_steps": 15,
"use_tensorboard": True,
"save_safetensors": True,
"save_first_step": False,
"include_tkps": True,
}

1
tests/e2e/solo/test_relora_llama.py
View File

@@ -63,7 +63,6 @@ class TestReLoraLlama(unittest.TestCase):
"learning_rate": 0.00001,
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"save_safetensors": True,
"use_tensorboard": True,
"save_first_step": False,
}

1
tests/e2e/test_activation_offloading.py
View File

@@ -57,7 +57,6 @@ class TestActivationOffloading:
"flash_attention": True,
"sample_packing": True,
"bf16": "auto",
"save_safetensors": True,
"gradient_checkpointing": True,
"activation_offloading": True,
"save_first_step": False,

2
tests/e2e/test_deepseekv3.py
View File

@@ -64,7 +64,6 @@ class TestDeepseekV3:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}
@@ -113,7 +112,6 @@ class TestDeepseekV3:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}

2
tests/e2e/test_diffusion.py
View File

@@ -41,7 +41,6 @@ class TestDiffusion:
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
"logging_steps": 1,
"eval_steps": 3,
@@ -97,7 +96,6 @@ class TestDiffusion:
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
"logging_steps": 1,
"eval_steps": 2,

2
tests/e2e/test_embeddings_lr.py
View File

@@ -44,7 +44,6 @@ class TestEmbeddingsLrScale(unittest.TestCase):
"optimizer": "adamw_torch_fused",
"embedding_lr_scale": 0.5,
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,
@@ -89,7 +88,6 @@ class TestEmbeddingsLrScale(unittest.TestCase):
"optimizer": "adamw_torch_fused",
"embedding_lr": 0.000005,
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,

2
tests/e2e/test_gemma2.py
View File

@@ -61,7 +61,6 @@ class TestGemma2:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
}
)
@@ -111,7 +110,6 @@ class TestGemma2:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
}
)

2
tests/e2e/test_gemma3_text.py
View File

@@ -60,7 +60,6 @@ class TestGemma3Text:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}
@@ -110,7 +109,6 @@ class TestGemma3Text:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}

4
tests/e2e/test_llama.py
View File

@@ -43,7 +43,6 @@ class TestLlama:
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
}
)
@@ -90,7 +89,6 @@ class TestLlama:
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
}
)
@@ -134,7 +132,6 @@ class TestLlama:
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
}
)
@@ -174,7 +171,6 @@ class TestLlama:
"sample_packing": False,
"batch_flattening": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
}
)

1
tests/e2e/test_llama_pretrain.py
View File

@@ -49,7 +49,6 @@ class TestPretrainLlama:
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,

2
tests/e2e/test_llama_vision.py
View File

@@ -51,7 +51,6 @@ class TestLlamaVision(unittest.TestCase):
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}
@@ -97,7 +96,6 @@ class TestLlamaVision(unittest.TestCase):
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}

1
tests/e2e/test_mamba.py
View File

@@ -49,7 +49,6 @@ class TestMamba(unittest.TestCase):
"max_steps": 20,
"save_steps": 10,
"eval_steps": None,
"save_safetensors": False,
"save_first_step": False,
}
)

1
tests/e2e/test_optimizers.py
View File

@@ -224,7 +224,6 @@ class TestCustomOptimizers(unittest.TestCase):
"learning_rate": 0.00001,
"optimizer": "schedule_free_adamw",
"lr_scheduler": "constant",
"save_safetensors": True,
"max_steps": 10,
"save_first_step": False,
}

1
tests/e2e/test_qat.py
View File

@@ -54,7 +54,6 @@ class TestQATLlama:
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_safetensors": True,
"bf16": True,
"save_first_step": False,
}

2
tests/e2e/test_save_first_step.py
View File

@@ -46,7 +46,6 @@ class TestSaveFirstStepCallback(unittest.TestCase):
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": True,
}
)
@@ -86,7 +85,6 @@ class TestSaveFirstStepCallback(unittest.TestCase):
"flash_attention": True,
"sample_packing": True,
"bf16": True,
"save_safetensors": True,
"save_first_step": False,
}
)

1
tests/e2e/test_streaming.py
View File

@@ -50,7 +50,6 @@ class TestStreamingDatasets:
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_safetensors": True,
"bf16": "auto",
"use_tensorboard": True,
"save_first_step": False,

26
tests/e2e/utils.py
View File

@@ -167,6 +167,13 @@ def require_hopper(test_case):
return unittest.skipUnless(is_hopper(), "test requires h100/hopper GPU")(test_case)
def supports_fp8(test_case):
compute_capability = torch.cuda.get_device_capability()
return unittest.skipUnless(
compute_capability >= (9, 0), "test requires h100 or newer GPU"
)(test_case)
def check_tensorboard(
temp_run_dir: str,
tag: str,
@@ -193,21 +200,10 @@ def check_model_output_exists(temp_dir: str, cfg: DictDefault) -> None:
"""
helper function to check if a model output file exists after training
checks based on adapter or not and if safetensors saves are enabled or not
checks based on adapter or not (always safetensors in Transformers V5)
"""
if cfg.save_safetensors:
if not cfg.adapter:
assert (Path(temp_dir) / "model.safetensors").exists()
else:
assert (Path(temp_dir) / "adapter_model.safetensors").exists()
if not cfg.adapter:
assert (Path(temp_dir) / "model.safetensors").exists()
else:
# check for both, b/c in trl, it often defaults to saving safetensors
if not cfg.adapter:
assert (Path(temp_dir) / "pytorch_model.bin").exists() or (
Path(temp_dir) / "model.safetensors"
).exists()
else:
assert (Path(temp_dir) / "adapter_model.bin").exists() or (
Path(temp_dir) / "adapter_model.safetensors"
).exists()
assert (Path(temp_dir) / "adapter_model.safetensors").exists()