gradient accumulation

[tonyjohnchen]

Adding new feature gradient accumulation to only update weight for every x steps.

Example command without using gradient accumulation:

python3 MaxText/train.py MaxText/configs/base.yml   base_output_directory=${MAXTEXT_OUTPUT_PATH} run_name=${RUN_NAME}    enable_checkpointing=false async_checkpointing=false    per_device_batch_size=1    skip_first_n_steps_for_profiler=5 steps=30    dataset_type=synthetic    profiler=xplane

Example command with using gradient accumulation:

python3 MaxText/train.py MaxText/configs/base.yml   base_output_directory=${MAXTEXT_OUTPUT_PATH} run_name=${RUN_NAME}    enable_checkpointing=false async_checkpointing=false    per_device_batch_size=1    skip_first_n_steps_for_profiler=5 steps=30    dataset_type=synthetic    profiler=xplane gradient_accumulation_steps=10

Result1
Result2

[anfals]

Double

[gobbleturk]

Have you tested this gives the same loss as a large batch? E.g. you could run locally on just a v4-8 with per_device_batch=1, gradient_accumulation_steps=5 for 100 steps and per_device_batch_size=5 gradient_accumulation_steps=1 for 100 steps (should give near identical loss)

[RissyRan]

Agree with above test to verify the result.

Also trying to learn more here. After some research, it seems we should do:

1. Accumulated Gradients += Gradients (from current accumulation step)
2. Averaged Gradients = Accumulated Gradients / Number of accumulation steps (instead of doing average all the way, or calculate gradients at the last microbatch only)