Busy wait so that non local rank zeros don't timeout while local rank zero downloads a monolithic checkpoint

[dakinggg]

What does this PR do?

Before the dist.barrier() in download_checkpoint, busy wait for local rank 0 to reach that part of the code. This prevents a timeout on the barrier when it takes local rank 0 a long time to download the checkpoint.

Manual test passed

What issue(s) does this change relate to?

Closes CO-1881

Before submitting

• Have you read the contributor guidelines?
• 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)

[eracah]

LGTM and passes manual test