Speed up Axlearn CI

[soundway]

Make a few minor tweaks to make Axlearn CI faster:

• Use async mode to run different parts of the test
• Use worksteal to more evenly distribute the workload

This should cut CI time by roughly half.

[ruomingp]

Thanks! Will defer to @markblee to approve.

[madrob]

We're using cimg/base:stable which is two years old and a deprecated tag. Looking at cimg/base:current which is more recent, that still only includes bash 5.1.16 - is that new enough for this feature?

[soundway]

Yes, it seems like it. I just realized the wait feature I'm relying on requires Bash 5.1, not 5.2. So it's all good. I just updated the comment.

This change shouldn't affect this repo much other than having shorter test time due to having worksteal per each of the test shard. Given that pytype check in this repo seems to take the longest time, I recommend that we make build-and-test-job and pre-commit run in parallel, and maybe add more use machines with more CPU resources for the pre-commit check to speed things up.

[soundway]

Ahh never mind, looks like I can just make those changes in .circleci/config.yml.