Call optimizer.zero_grad() when profiling

[mostafaelhoushi]

Didn't test locally... let's see if CI merges and I can check locally as well tomorrow

[mostafaelhoushi]

Error in CI... need to debug that

[jacobkahn]

Discussed offline -- from a correctness perspective, we need to clear accumulated gradients on each iteration, so let's get this passing in CI 👍

[mostafaelhoushi]

So far in this PR, I was calling optim.zero_grad() inside the model wrapper before creating fx graph:

This caused things to mess up and it requires more debugging. Also, I was surprised that the fx trace already had an instruction to zero out tensors (assuming these are gradient tensors). In the screenshot below, the LHS is the trace of testSimpleTrainSchedule (the test case that is failing in CI) before and after this PR:

Another option is to call optim.zero_grad() outside the wrapper and during profiling. This will require some refactoring. I will probably go through this route.

Having said that, the instructions that seem to be introduced by optim.zero_grad():

    zeros = torch.ops.aten.zeros.default([1], device = device(type='cpu'), pin_memory = False)
    empty = torch.ops.aten.empty.memory_format([280], dtype = torch.uint8, device = device(type='cpu'))

are problematic when creating or modifying graphs. This is because the instructions do not take any tensor as input, and the output of the instructions are not used by any other op. Hence, the ILP solver might randomly move them to any other location in the graph, while we want to make sure to call them before the forward pass of the model.