Support AMP mode for pure TP runs

[slabhs-aws]

Hey,
At the moment, TP-only runs are forced to not be using AMP here. Can we enable AMP for TP-only as well?

[tianyu-l]

See context here #2213.

This feature is not well tested / verified.

Curious, why do you need to use AMP with TP?

[slabhs-aws]

Thanks for the pointer! Well during model bringup, we start small including TP-only strategy, and wish to conduct AMP runs for accuracy validation