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

more train steps

Browse Source
This commit is contained in:
Wing Lian
2026-04-25 02:17:48 +00:00
parent f77408a3d0
commit 1b9520cc8b
8 changed files with 105 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -32,8 +32,8 @@ def verify_training_success(temp_dir):
check_tensorboard_loss_decreased(
temp_dir + "/runs",
initial_window=3,
final_window=3,
initial_window=10,
final_window=10,
max_initial=5.0,
max_final=4.7,
)
@@ -61,7 +61,8 @@ class TestFSDP2:
},
],
"num_epochs": 1,
"max_steps": 20,
"max_steps": 80,
"warmup_steps": 5,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
@@ -79,6 +80,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
"bf16": True,
}
)
@@ -121,14 +125,15 @@ class TestFSDP2:
"adapter": "lora",
"lora_r": 8,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_dropout": 0.0,
"lora_target_linear": True,
"num_epochs": 1,
"max_steps": 20,
"max_steps": 80,
"warmup_steps": 5,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 2e-4,
"learning_rate": 1e-3,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
@@ -142,6 +147,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
"bf16": True,
# explicitly disable LORA kernels, as they may be auto-enabled
"lora_mlp_kernel": False,
@@ -188,11 +196,12 @@ class TestFSDP2:
"lora_alpha": 16,
"lora_target_linear": True,
"num_epochs": 1,
"max_steps": 20,
"max_steps": 80,
"warmup_steps": 5,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 2e-4,
"learning_rate": 1e-3,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
@@ -206,6 +215,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
"bf16": True,
"lora_mlp_kernel": True,
"lora_qkv_kernel": True,
@@ -250,14 +262,15 @@ class TestFSDP2:
"adapter": "qlora",
"lora_r": 8,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_dropout": 0.0,
"lora_target_linear": True,
"num_epochs": 1,
"max_steps": 20,
"max_steps": 80,
"warmup_steps": 5,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 2e-4,
"learning_rate": 1e-3,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
@@ -271,6 +284,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
"bf16": True,
}
)
@@ -314,12 +330,12 @@ class TestFSDP2:
"lora_alpha": 16,
"lora_target_linear": True,
"num_epochs": 1,
"max_steps": 30,
"warmup_steps": 3,
"max_steps": 80,
"warmup_steps": 5,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 2e-4,
"learning_rate": 1e-3,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
@@ -333,6 +349,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
"bf16": True,
"lora_mlp_kernel": True,
"lora_qkv_kernel": True,
@@ -395,6 +414,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
}
)
@@ -443,7 +465,7 @@ class TestFSDP2:
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 2e-4,
"learning_rate": 1e-3,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
@@ -457,6 +479,9 @@ class TestFSDP2:
"reshard_after_forward": True,
},
"use_tensorboard": True,
"seed": 42,
"sample_packing": True,
"pad_to_sequence_len": True,
}
)