Update base for Update on "[WIP] zero bubble"

To run zb test: 
`python test_runner.py ./test_out --test pp_zb`

internal mast run:
`
torchx run mast.py:train --additional_folders /home/howardhuang/local/torchtitan --twtask_bootstrap_script run_torchtitan.sh --h "grandteton" --nodes 8 train_configs/debug_model_3d_mast.toml
`
  


[ghstack-poisoned]