ppe.compile Enable forward only and custom decompositions
#740
Conversation
emcastillo
requested review from
asi1024,
kmaehashi and
linshokaku
as code owners
emcastillo
force-pushed
the
no-backprop
branch
from
85aeb01 to
0a8f2e6
Compare
| else f"input_{i - len(self._parameter_names)}" | ||
| ) | ||
|
|
||
| # Remove the outputs |
There was a problem hiding this comment.
I want some detailed comments about the 'output_node' that is being deleted here. I believe it's an intermediate variable for gradient calculation, but is it sufficient to just delete one?
There was a problem hiding this comment.
added comments explaining it!
| node for node in fwd_graph.graph.nodes if node.op == "output" | ||
| ][0] | ||
| outputs = pytree.tree_flatten(output_node.args)[0] | ||
| fwd_graph.graph.erase_node(output_node) |
There was a problem hiding this comment.
It would be great if we could verify via tests whether or not the intended node has been deleted. However, if it proves difficult, I don't think there's a need to force its addition.
There was a problem hiding this comment.
yeah! thats a great idea! thanks!!
There was a problem hiding this comment.
added a check in the test!
emcastillo
force-pushed
the
no-backprop
branch
from
dd4ae46 to
b60e7b5
Compare
emcastillo
force-pushed
the
no-backprop
branch
from
cba2d6a to
b710645
Compare
|
/test |
| for n, b in module.named_buffers(): | ||
| parameters_and_buffers.append(b) | ||
| names.append(n) |
There was a problem hiding this comment.
| for i_node in primal_inputs: | ||
| bwd_outs.append( | ||
| bwd_graph.call_function( | ||
| torch.ones, (i_node.meta.get("tensor_meta").shape,) |
There was a problem hiding this comment.
nit:
It seems more natural to output zeros.
|
/test |
In some cases we want to generate a forward only graph using the aot lightweight tracing.
Also it may be necessary to avoid using default decompositions to prevent some operators to be decomposed and losing important information.