Adding a feature that will stop the training/eval process after reaching some max_steps #428
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| epoch_times = [] | ||
| checkpoint_times = [] | ||
| results = {} | ||
| best_val_loss = float("inf") | ||
| total_train_steps = 0 | ||
| for epoch in range(train_config.num_epochs): |
There was a problem hiding this comment.
@wukaixingxp I wonder if we should remove some of the redundancy, maybe something like?
max_steps_reached = False # Flag to indicate max training steps reached
for epoch in range(train_config.num_epochs):
if max_steps_reached:
break # Stop starting a new epoch if max steps reached
epoch_start_time = time.perf_counter()
with MemoryTrace():
model.train()
total_loss = 0.0
total_length = len(train_dataloader) // gradient_accumulation_steps
pbar = tqdm(colour="blue", desc=f"Training Epoch: {epoch+1}", total=total_length, dynamic_ncols=True)
for step, batch in enumerate(train_dataloader):
total_train_steps += 1
# Check if max training steps reached
if train_config.max_train_step > 0 and total_train_steps > train_config.max_train_step:
print(f"Max training steps reached at step {total_train_steps}. Stopping training.")
max_steps_reached = True
break # Exit batch loop
if max_steps_reached and not train_config.enable_fsdp or local_rank == 0:
print("Training stopped after reaching the maximum number of steps.")
There was a problem hiding this comment.
I changed my code accordingly, thanks for the feedback.
HamidShojanazeri
approved these changes
Apr 9, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adding a feature that will stop the training/eval process after reaching some max_steps. By default, this feature is disabled and the max_train_step and max_eval_step are set to 0. User can pass arguments to set max_train_step or max_eval_step
to a positive integer so that the training or eval process will be stop when reaching the max_step limit. This PR will introduce two more arguments in the config file and thus the testing functions test_gradient_accumulation and test_save_to_json has been modified.
test log with max_train_step=6, max_eval_step=2