fix: match no_sync with pytorch

[k223kim]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

This PR fixes #1670; it adds a test that shows how the current thunder's no_sync context manager already behaves like torch. If accepted, I will have a follow up PR that removes sync_grads entirely and update the docs.

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[Borda]

can we have it rather as an argument ot we really need this loop?

[k223kim]

Hmm, we do not need this loop. Would it make a difference if we use arguments?

[crcrpar]

from the perspective of what we test, there's no differences but making use_no_sync argument would make test logs way better. It clearly tells us which setting fails or not. I'd have use_no_sync as an argument as @Borda suggests.

[t-vi]

I think we should split the test.
That said, I would expect the no-sync test to roughly do the following:

for j in range(2):
    for i in range(2):
        with no_sync():
            fw
            loss
            bw
        grab grad.clones for comparison  # these would not be synced, even outside the context manager.
    fw
    loss
    bw  # sync happens only(!) here and below
    grab grad.clones for comparison  # these are synced
    zero_grad
    fw
    loss
    bw  # sync happens only(!) here and above
    grab grad.clones for comparison  # these are synced
    zero_grad

and then the comparison is between the gradients in order between plain PT DDP and Thunder DDP.

[t-vi]

Note that the grabbing of the grads happens outside the no_sync.

[k223kim]

@t-vi I think this should fail if thunder and torch are different