Mohit Khatwani

Tested on GKE with following command ``` mkdir -p dataset; gcsfuse --implicit-dirs ${DATASET_BUCKET} dataset; python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=gs://mohitkhatwani-maxtext-logs/ dataset_path=dataset/ ```

Cleanup docker image after each action run

1

# Description Cleanup docker image after each action run so that cleanup is performed on every runner rather than the best available runner which Github chooses to run CI on....

Skip checkpointing at step=0

2

# Description - Skip checkpointing at step=0 - add abs for max numerical diff log in forward_pass_checker # Tests Integration tests # Checklist Before submitting this PR, please make sure...

# Description Setting param_scan_axis=0 helps in improving perf and reducing memory required by optimizer state. FIXES: #1382 # Tests Since this is touching all the checkpoint creation files, I'm going...

Add checkpoint conversion instructions to GRPO docs

1

# Description - add environment variable section - add checkpoint conversion instructions If the change fixes a bug or a Github issue, please include a link, e.g.,: FIXES: b/123456 FIXES:...