support bert/gpt2 and downstream model with new api and dtensor
support bert/gpt2 and downstream model with new api and dtensor