🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/117799

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit ead91ac with merge base 5c17f66 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Does this part of the optim state dict load use multiple streams (which is why memory is not freed immediately when there are no more references)?

No multiple streams. But even if the temporary tensors are not referenced in Python, it can still being used, iiuc. For example, the clone() one will keep the source tensor (and its memory) alive until the operation is done. Or do you think this should not happen?

I was wondering if you have tried del-ing the source that is being cloned and seeing if that frees it at the time of del.

Suppose at first we have this:

a = torch.empty((3,), device="cuda")
b = a.clone()
c = torch.empty((3,), device="cuda")

c cannot reuse the memory of a since a is still alive (due to the Python reference).

Now, suppose we del a:

a = torch.empty((3,), device="cuda")
b = a.clone()
del a  # <--- add this
c = torch.empty((3,), device="cuda")

c can reuse the memory of a. It does not need to wait until the GPU copy kernel from clone() finishes because any subsequent GPU kernel using the memory for c will be sequentially after the GPU copy kernel.

I did not see any special handling of clone that would record a cudaEvent and free the memory later (like in the case of multiple streams and recordStream). I would be curious if del-ing did not free the memory instantly if there is only a single stream.

@awgu that makes sense. However, the tensor should already be deleted as it has no reference after the inner util function. So I don't think del makes any difference in our case. I would do more tests to verify.

@awgu I can confirm that del alone is not enough. I'll check if multiple streams are being used.

@awgu The root cause is DTensor.to_local() has asynchronous communication by default. So I believe we need to call synchronize() in such a case.

@awgu The root cause is DTensor.to_local() has asynchronous communication by default. So I believe we need to call synchronize() in such a case.

Thanks for looking into this!

@wanchaol Is there any way to specify to to_local() to use synchronous collectives? We want to avoid recordStream, and IIUC, we only avoid it for synchronous collectives. (cc: @kwen2501).

If we can run synchronous collectives and avoid recordStream, then I think we can avoid the CPU sync from synchronize() and instead just rely on the PG NCCL calling current_stream.wait_stream(nccl_stream). Avoiding the CPU sync here could prevent some CPU boundedness (though I have not looked at any profiles).

@awgu This is optimizer state_dict code path. What's the down side of using torch.cuda.synchronize()?

@fegin The downside is just performance from synchronizing the CPU. If there is not much performance downside, then using synchronize sounds good to me!

(I do not have a good mental model of the optimizer state dict performance, so maybe this is not really an issue. However, in general, I do think that considering the performance of state dict makes sense.)

@awgu Yup, agree that performance is important. Currently, synchronize() is probably not the main bottleneck. The main bottleneck is the allgather which is done in a per-parameter way without batching. However, it is hard to batching the allgather for optimizer state dict load.

This becomes a bigger issue than just a performance issue. If we don't call torch.cuda.synchronize(), the subsequent clone() will not work -- the value will not be correctly cloned.

cc., @wanchaol @awgu @wz337 @LucasLLC

Can we use stream sync instead of device sync?

if only device sync works and stream sync doesn't, that means there is something wrong with DTensor's to_local() -- it should sync its communication work back to the "current stream" (or else provide a Work handle for user to sync at some point)

If (1) proper stream dependency is maintained by DTensor's to_local() and (2) torch.clone observes stream properly, you don't even need to call stream sync.

@kwen2501 _optim_utils.py calls to DTensor.to_local() to gather the tensor. It does not know which stream DTensor.to_local() uses. So I don't think I'm able to use stream wait.

If to_local()'s API does not have flags like async_op=True|False, it must always sync back to main stream, so that you don't need to figure out which stream to wait. Said plainly, to_local() must always call work.wait() inside or expose work out.

in the more general sense of a user using DTensor.to_local, i don't know if it makes sense that operation has to be synchronous. Wouldn't it be reasonable to expect a to_local call to return a new 'ACT' that represents ongoing tensor work?

Wouldn't it be reasonable to expect a to_local call to return a new 'ACT' that represents ongoing tensor work?

The reasonableness depends on how much tendency users have to write
x.to_local()
vesus
x = x.to_local()

In a world that tensor.to(...) is prevailing, I'd say the tendency is non-negligible...

@kwen2501 tensor.to(...) is not inplace though. Only nn.Module.to is in-place.

@awgu You are right. Thanks for the correction!

I have not found the root cause. Since _gather_state_dict() is used not only by FSDP but also PP-FSDP, I decide to change the behavior of _gather_state_dict() to always call wait() before returning the tensor. The new PR is #118197.