Speed up CI tests :)

[KuuCi]

What does this PR do?

As our stack grows, we need to make sure that our CI tests is efficient and fast! (referencing https://github.com/mosaicml/composer/actions/runs/12059801726/job/33629067025 for test speed comparisons)
Old 1 GPU test:

New 1 GPU test:

Old 2 GPU test:

New 2 GPU test:

Annecdotally seeing around a 2-2:30 min speed up on 1 gpu tests and a 30 sec speed up on 2 gpu test

What issue(s) does this change relate to?

N/A

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[KuuCi]

my debt has been paid @mvpatel2000

[mvpatel2000]

yung king