Some questions about operator reordering

[aonrdbit]

I'm sorry to bother you, I have a confusion while reading the paper.

I think the paper is aimed at saving memory during training rather than inference. So Figure 3 in the paper, which is the picture below, does it represent the forward propagation of the execution model? If so, then I think all the tensors needs to be saved because they are needed during the back propagation process. Why do we only need to save the input and output of a certain operator?

Looking forward to your reply, thank you very much!

[mostafaelhoushi]

Hi @aonrdbit
Thanks for your interest in our paper.
Please find below my comments/answers to your question:

I think the paper is aimed at saving memory during training rather than inference.

Operator reordering and memory fragmentation should work for both inference and training. However, operator reordering is expected to have less improvement in inference as in inference there are fewer options to reorder nodes. Models that have parallel paths like what we show in Figure 3 could benefit from operator reordering.

So Figure 3 in the paper, which is the picture below, does it represent the forward propagation of the execution model?

Figure 3 represents a simple inference model (no training).

If so, then I think all the tensors needs to be saved because they are needed during the back propagation process. Why do we only need to save the input and output of a certain operator?

For the sake of completeness, let's answer this question not in the context of Figure 3 (which was inference) but in the context in training. Let's assume we have a training model. We need to save memory for:

• input and output activations during forward pass
• during backward pass, we use those input and output activations to calculate gradients (when we calculate the gradients of an op and all of its neighbours have consumed the activations to calculate their gradients, we no longer need its input and output activations)
• and we use gradients to update weights (when we update the weights of an op, and its neighbours have consumed its gradient, then we no longer need to store the gradients)
  Hence, operator reordering utilizes these facts to change the order of operations during training. It is usually the weight update step that we have flexibility to reorder during training.

[aonrdbit]

Thanks for your answer @mostafaelhoushi. Please check if my idea is correct.

For training:
If operator reordering is not used, then all operators should be back propagated to obtain gradients first, and then the model should be updated. On the contrary, if operator reordering is used, then an operator will immediately update its parameters after getting the gradient, and then release the gradient memory. Gradient calculations and parameter updates are parallelized.

If so, then I don't think the peak memory footprint has been reduced. In my experiments, model training would get peak memory usage at the end of the forward propagation, and then the memory would be gradually released as the back propagation progresses. Operator reordering optimizes the memory release process during back propagation, but does not affect the peak memory usage throughout the entire training process.

I am still confused. Is there any other way to reorder operators in training? I look forward to your reply. Thank you very much.

[mostafaelhoushi]

Hi @aonrdbit
Sorry for the delay.

On the contrary, if operator reordering is used, then an operator will immediately update its parameters after getting the gradient

When I printed the model graph after reordering, I didn't necessarily see that.
I suggest to go to this test, and uncomment the print statements in it to see the order of operators with and without reordering:

MODel_opt/tests/fx_optimizer_test.py

Line 292 in 6ac4582

def testSimpleTrainSchedule(self):

To run the test:

python -m unittest tests.fx_optimizer_test.FXOptimizerTest.testSimpleTrainSchedule

Also, can you take a quick look at this tutorial that I just uploaded. Hope it might help explaining:
https://github.com/facebookresearch/MODel_opt/blob/main/tutorial/tutorial.ipynb

[aonrdbit]

Hi, @mostafaelhoushi
Thank you for your reply and tutorial. I understand the meaning in your tutorial.

Tutorial is the behavior in the model inference stage, and different execution orders can cause different peak memory usage. However, during the model training phase, forward propagation requires storing the output values of each operator, so different execution orders have no impact on peak memory usage.

I'm looking forward to your reply. Thank you very much.

[mostafaelhoushi]

Hi @aonrdbit Sorry for the delay
Yes during training we need to store output values in forward pass, but we don't need to store them till the end of the backward pass. Our algorithm tells PyTorch when it no longer needs a specific output value to deallocate it and store memory.

Can you try this to see how the graph changes for training a model:

I suggest to go to this test, and uncomment the print statements in it to see the order of operators with and without reordering:

MODel_opt/tests/fx_optimizer_test.py

Line 292 in 6ac4582

def testSimpleTrainSchedule(self):

To run the test:

python -m unittest tests.fx_optimizer_test.FXOptimizerTest.testSimpleTrainSchedule