@eternalNight,

Thank you for submitting a great PR! I believe this will bring a significant improvement across many use cases.

I’m happy to merge this change, but I’d like your feedback on one point:
Right now, the approach is to pass the downcasted dtype to allgather. As an alternative, would it be possible to identify a downcast-taking parameter and place allgather after the downcast?

The advantage of this approach is that it keeps the C++ side (the more primitive components in DeepCompile) simpler. That said, there may be trade-offs compared to your current design.

I’d really appreciate your thoughts on this.

@eternalNight,

Thank you for submitting a great PR! I believe this will bring a significant improvement across many use cases.

I’m happy to merge this change, but I’d like your feedback on one point: Right now, the approach is to pass the downcasted dtype to allgather. As an alternative, would it be possible to identify a downcast-taking parameter and place allgather after the downcast?

The advantage of this approach is that it keeps the C++ side (the more primitive components in DeepCompile) simpler. That said, there may be trade-offs compared to your current design.

Adding allgather after the param downcast is possibly feasible. The major differences from this version are:

1. We don't need the additional dtype argument in allgather_param().
2. We still need the dtype argument in executor->allgatherParam because the output_buf needs that dtype to determine its size.
3. Similar to prefetch_params_fused() and executor->prefetchParamsFused(). But additionally we need to move the downcast ops ahead of the prefetch one.
4. Type-casting is unnecessary in launchAllGather().
5. Either the downcast op should run on the ag stream (can we move an existing op to a separate stream?), or we need an additional event to ensure that allgather does not start till the downcast completes.

On the C++ side, 1-4 reduces complexity (a bit) while 5 increases complexity. So I'm not sure if either approach is really simpler than the other.

I’d really appreciate your thoughts on this.

BTW, I found today that, with the inductor backend, loss of that openvla-like test became NaN after a few iterations, and that was so even without torch autocast. The same setup works fine with the eager backend.

I suspect that it's related to tensor lifecycle, but need some more time to find out.

@eternalNight Thanks for sharing your insights! I was considering whether it might be useful to implement various features as a graph modification, but I think your approach makes sense for this PR.

BTW, I found today that, with the inductor backend, loss of that openvla-like test became NaN after a few iterations, and that was so even without torch autocast. The same setup works fine with the eager backend.

Did you encounter this issue without the changes in this PR? I’ve run into similar problems while I was writing the code of DeepCompile. In most cases, Triton code generated by Inductor releases buffer tensors earlier than expected, causing the tensor to be destroyed.
I tested this with several models, including LLaMA and Mixtral, but it’s possible I missed a pattern that shows up in a different architecture.

BTW, I found today that, with the inductor backend, loss of that openvla-like test became NaN after a few iterations, and that was so even without torch autocast. The same setup works fine with the eager backend.

Did you encounter this issue without the changes in this PR? I’ve run into similar problems while I was writing the code of DeepCompile. In most cases, Triton code generated by Inductor releases buffer tensors earlier than expected, causing the tensor to be destroyed. I tested this with several models, including LLaMA and Mixtral, but it’s possible I missed a pattern that shows up in a different architecture.

Yes, the issue exists before this PR and without torch autocast.

@tohtana I think I have found (at least one of) the cause of that loss NaN issue.

Torch op schema can specify if the output tensor aliases the storage of an input one. Take t() as an example: aten::t(Tensor(a) self) -> Tensor(a).

Inductor-generated code explicitly drop a tensor when it's no longer needed. So for allgather_param, it looks like below (note the del statements):

buf0 = torch.ops.dc.allgather_param.default(primals_53, 139956546525584, 49)
del primals_53
buf1 = buf0
del buf0

Now for release_param(), it is declared as release_param(Tensor a, int graph_id, int id, int n_users) -> Tensor, but the returned tensor is just a, not a copy of a. With the del that follows, a can be freed early. Changing the schema to release_param(Tensor(a) a, int graph_id, int id, int n_users) -> Tensor(a) makes the situation better but only a bit: losses are back to normal only when I commented out fast_free_schedule in zero3_compile.py. There probably is yet another cause.

Anyway, this is a different issue from what this PR is aimed at. I'll open a separate issue for this.

Update: That "another cause" is the previous workaround, i.e., cloning the input tensor in release_param. After cleaning that up, the loss NaN issue disappears. Refer to #7597 for the full code change. This PR has been tested to be compatible with that one.

Let's merge this after #7597 is merged. Please let me know when you are okay for merge (I think this PR should have some conflicts with #7597)

@tohtana Rebased and conflicts resolved.