Description

This PR includes the following changes:

1. Added support for resuming training from a checkpoint for algorithms other than PPO.
2. Cleaned up parts of the VLA code using pre-commit.

Usage details:

1. You should activate save_checkpoint to save the optimizer states and models:

# Execute deepspeed command
deepspeed \
     --master_port ${MASTER_PORT} \
     --module align_anything.trainers.text_image_to_text.dpo \
     ...
     --save_checkpoint True \
     --epochs 2

We set the default value of save_checkpoint as True because we don't want new users to train once again since they have no idea of it.

2. Then you can restart the training by:

MODEL_NAME_OR_PATH="/PATH/TO/YOUR/CKPT/slice_100" # for example

# Execute deepspeed command
deepspeed \
     --master_port ${MASTER_PORT} \
     --module align_anything.trainers.text_image_to_text.dpo \
     --model_name_or_path ${MODEL_NAME_OR_PATH} \
     ...
     --load_checkpoint True \
     --epochs 2

Motivation and Context

resolve #150

Types of changes

What types of changes does your code introduce? Put an x in all the boxes that apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide. (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly. (required for a bug fix or a new feature)
• I have updated the documentation accordingly.