Add auto training param estimation

[gkumbhat]

Description

Currently the training happens with a fixed set of parameter, but if these parameters + model pair is not suitable for resource considerations, then we would blow out with OOMs. This feature request is to add an auto mode, that would try to configure the best possible configuration and do a dry run to make sure things don't go out of memory with those configuration.