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

checkpoint model on first step callback (#2906)

Browse Source 
* checkpoint model on first step callback

* remove debug

* add test cases; update existing tests not to save on first step

* move test out of solo

* delete

* default to False

* typo
This commit is contained in:
Dan Saunders
2025-07-15 15:00:48 -04:00
committed by GitHub
parent d320ef6199
commit 10ba1622f7
146 changed files with 419 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/e2e/patched/test_4d_multipack_llama.py
View File

@@ -55,6 +55,7 @@ class Test4dMultipackLlama(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"fp16": True,
"save_first_step": False,
}
)
cfg = validate_config(cfg)
@@ -102,6 +103,7 @@ class Test4dMultipackLlama(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"fp16": True,
"save_first_step": False,
}
)
cfg = validate_config(cfg)

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

@@ -69,6 +69,7 @@ class TestActivationCheckpointing:
"bf16": True,
"save_safetensors": True,
"gradient_checkpointing": gradient_checkpointing,
"save_first_step": False,
}
)

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

@@ -62,6 +62,7 @@ class TestFAXentropyLlama:
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"use_tensorboard": True,
"save_first_step": False,
}
)
if is_torch_bf16_gpu_available():

2
tests/e2e/patched/test_falcon_samplepack.py
View File

@@ -58,6 +58,7 @@ class TestFalconPatched(unittest.TestCase):
"save_steps": 10,
"eval_steps": 10,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)
@@ -99,6 +100,7 @@ class TestFalconPatched(unittest.TestCase):
"save_steps": 10,
"eval_steps": 10,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)

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

@@ -61,6 +61,7 @@ class TestFAFlattening:
"optimizer": "adamw_8bit",
"lr_scheduler": "cosine",
"use_tensorboard": True,
"save_first_step": False,
}
)
if is_torch_bf16_gpu_available():

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

@@ -53,6 +53,7 @@ class TestFusedLlama(unittest.TestCase):
"max_steps": 10,
"save_steps": 5,
"eval_steps": 5,
"save_first_step": False,
}
)
if is_torch_bf16_gpu_available():

2
tests/e2e/patched/test_llama_s2_attention.py
View File

@@ -58,6 +58,7 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
"save_steps": 5,
"eval_steps": 5,
"bf16": "auto",
"save_first_step": False,
}
)
@@ -100,6 +101,7 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
"save_steps": 5,
"eval_steps": 5,
"bf16": "auto",
"save_first_step": False,
}
)

2
tests/e2e/patched/test_lora_llama_multipack.py
View File

@@ -55,6 +55,7 @@ class TestLoraLlama(unittest.TestCase):
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_first_step": False,
}
)
if is_torch_bf16_gpu_available():
@@ -108,6 +109,7 @@ class TestLoraLlama(unittest.TestCase):
"learning_rate": 0.00001,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"save_first_step": False,
}
)
cfg = validate_config(cfg)

2
tests/e2e/patched/test_mistral_samplepack.py
View File

@@ -56,6 +56,7 @@ class TestMistral(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)
@@ -97,6 +98,7 @@ class TestMistral(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)

2
tests/e2e/patched/test_mixtral_samplepack.py
View File

@@ -52,6 +52,7 @@ class TestMixtral(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)
@@ -90,6 +91,7 @@ class TestMixtral(unittest.TestCase):
"save_steps": 3,
"eval_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)
cfg = validate_config(cfg)

2
tests/e2e/patched/test_model_patches.py
View File

@@ -45,6 +45,7 @@ class TestModelPatches(unittest.TestCase):
"max_steps": 20,
"save_steps": 10,
"eval_steps": 10,
"save_first_step": False,
}
)
cfg = validate_config(cfg)
@@ -78,6 +79,7 @@ class TestModelPatches(unittest.TestCase):
"max_steps": 20,
"save_steps": 10,
"eval_steps": 10,
"save_first_step": False,
}
)
cfg = validate_config(cfg)

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

@@ -49,6 +49,7 @@ class TestLlamaPeftEmbeddings:
"bf16": "auto",
"save_safetensors": True,
"embeddings_skip_upcast": True,
"save_first_step": False,
}
)

2
tests/e2e/patched/test_phi_multipack.py
View File

@@ -54,6 +54,7 @@ class TestPhiMultipack(unittest.TestCase):
"eval_steps": 3,
"save_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)
@@ -105,6 +106,7 @@ class TestPhiMultipack(unittest.TestCase):
"eval_steps": 3,
"save_steps": 4,
"bf16": "auto",
"save_first_step": False,
}
)

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

@@ -58,6 +58,7 @@ class TestResumeLlama:
"max_steps": 15,
"use_tensorboard": True,
"save_safetensors": True,
"save_first_step": False,
}
)
if is_torch_bf16_gpu_available():

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

@@ -47,6 +47,7 @@ def fixture_cfg():
"special_tokens": {
"pad_token": "<|endoftext|>",
},
"save_first_step": False,
}
)

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

@@ -62,6 +62,7 @@ class TestUnslothQLoRA:
"lr_scheduler": "cosine",
"use_tensorboard": True,
"bf16": "auto",
"save_first_step": False,
}
)
@@ -112,6 +113,7 @@ class TestUnslothQLoRA:
"lr_scheduler": "cosine",
"use_tensorboard": True,
"bf16": "auto",
"save_first_step": False,
}
)
@@ -167,6 +169,7 @@ class TestUnslothQLoRA:
"lr_scheduler": "cosine",
"use_tensorboard": True,
"fp16": True,
"save_first_step": False,
}
)