Make maxtext_xpk_runner support other trainers

[lukebaumann]

At the moment, the command is hard coded to execute the MaxText.train module. If the module to execute could be configured in the WorkloadConfig, that would allow for additional trainers being developed to make use of the benchmark runner.

I am running into this while wanting to test a new trainer MaxText.elastic_train in PR #1310 but there are also other trainers being developed and merged and I think configuring maxtext_xpk_runner will be useful moving forward.