support DiT ImageNet dataset training

[zhtmike]

What does this PR do?

• Support DiT training with ImageNet dataset format, using --imagenet_format flag
• Fix the loss scale print in callback
• clean code

Fixes # (issue)

Adds # (feature)

Before submitting

• [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• [ ] Did you read the contributor guideline?
• [ ] Did you make sure to update the documentation with your changes? E.g. record bug fixes or new features in What's New. Here are the documentation guidelines
• [ ] Did you build and run the code without any errors?
• [ ] Did you report the running environment (NPU type/MS version) and performance in the doc? (better record it for data loading, model inference, or training tasks)
• [ ] Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

@xxx

[wtomin]

It looks great! Can you also provide a bash script to launch a four-card training, because I notice DiT torch repo is using a global batch size 256 for ImageNet training.

Besides, please revise the readme file for the ImageNet training. I think users are more interested in this rather than a toy dataset training.

[zhtmike]

It looks great! Can you also provide a bash script to launch a four-card training, because I notice DiT torch repo is using a global batch size 256 for ImageNet training.

Besides, please revise the readme file for the ImageNet training. I think users are more interested in this rather than a toy dataset training.

Done. Add a 4P training script and update the README