Support run trainer locally

[zpcore]

Support to run train.py locally:

tp dbrun --use-hf torchprime/hf_models/train.py # huggingface run

Or

tp dbrun  torchprime/torch_xla_models/train.py 

This makes it easier to debug model code.

Currently it is still using the torch_xla from the base docker image. In the future, we may want to give the option to run with local build torch_xla.

[tengyifei]

I usually run a model locally like this:

python3 torchprime/torch_xla_models/train.py model=llama-3-8b mesh.fsdp=8 profile_step=3 max_steps=50

Is this insufficient for your use case?

[zpcore]

I usually run a model locally like this:

python3 torchprime/torch_xla_models/train.py model=llama-3-8b mesh.fsdp=8 profile_step=3 max_steps=50

Is this insufficient for your use case?

Yes, this should also work. The thing is that I saw many permission issues when run train.py directly. This just helps to run in a container.

[tengyifei]

This introduces a magical constant. I think it's simpler if we add a "push=True" function argument, and have the other file call this with push=False