Add demo script for single-node async checkpointing

[awonak]

Fixes b/369874146

Example demo outputs:

$ python demo/checkpointing/train.py --project=gcs-tess --bucket=awonak-us-west1 --embedding-dim=10 --num-epochs=3 
Epoch 1/3, Train Loss: 449945.0625
Epoch 2/3, Train Loss: 409629.6875
Epoch 3/3, Train Loss: 380471.9375
Average checkpoint save time: 0.8570 seconds (stdev 0.1997)
Total run time: 1m53s

$ python demo/checkpointing/train.py --project=gcs-tess --bucket=awonak-us-west1 --embedding-dim=10 --num-epochs=3 --use-async
Epoch 1/3, Train Loss: 456512.0000
Epoch 2/3, Train Loss: 413212.9375
Epoch 3/3, Train Loss: 382256.8125
Average checkpoint save time: 0.0009 seconds (stdev 0.0009)
Total run time: 1m36s

Note: I will update the main README with single node async details once either this or #190 gets merged.

• Tests pass
• Appropriate changes to documentation are included in the PR