Best practices question for an op passing data forward -> backward

[DickJC123]

What's the recommended approach for an operator that has a computed value from Forward() that is helpful in the Backward() calculation, given that the computed value is not an I/O of the operator?

I'm familiar with the long standing approach of declaring a 'hidden output', but is that still the recommended approach?

@samskalicky @MoisesHer @szha @ptrendx @leezu

[szha]

FStatefulCompute is designed for this purpose:
https://github.com/apache/incubator-mxnet/blob/527573ec2b9b2696ffcafd1570cd94e2187f4c32/src/operator/rnn.cu#L35
https://github.com/apache/incubator-mxnet/blob/4a8da9ec62e8cadd7df6ad5e9ba305b777104068/src/operator/rnn-inl.h#L1523

[DickJC123]

I can think of two choices when using the FStatefulCompute approach:

1. Dynamically Alloc() the fwd->bwd Tensor on the Forward(), then Free() is on the Backward(), or
2. Alloc() the fwd->bwd Tensor once and hold onto it forever.

Approach 1 is not CUDA Graphs compatible, since the captured fwd->bwd Tensor is no longer valid when the graph is replayed.
Approach 2 holds onto the memory past Backward(), so past when it's actually needed.

The 'hidden output' approach does not have these problems. It is CUDA Graphs compatible, assuming a static graph memory layout, plus the fwd->bwd Tensor is up for reallocation past the Backward() call.

[DickJC123]

Another point to consider is how the operator behaves in an inference-only graph. Seems like with the stateful op approach, one can easily react to is_train == false, and not allocate the fwd->bwd Tensor if it's not needed otherwise. I'm not sure if the 'hidden output' approach can avoid allocating the space for that.