[fx] added activation checkpoint codegen

[FrankLeeeee]

In the previous PR #1349 , we annotated the nodes if they are activation-checkpointed. In this PR, we utilize these annotations and generate model forward code with activation checkpoint. The new CodeGen inherits from the PyTorch CodeGen, the code change can be found by looking for the following multi-line comment.

#########################################
# Modified for activation checkpointing #
#########################################

A unit test is added in this PR as well. Below are the different codes generated with/without activation checkpoint.

# without activation checkpoint
def forward(self, x):
    mlp1_linear1 = self.mlp1.linear1(x)
    mlp1_linear1_1 = self.mlp1.linear1(x)
    mlp2_linear1 = self.mlp2.linear1(x)
    mlp2_linear1_1 = self.mlp2.linear1(x);  x = None
    add = mlp1_linear1 + mlp1_linear1_1;  mlp1_linear1 = mlp1_linear1_1 = None
    add_1 = add + mlp2_linear1;  add = mlp2_linear1 = None
    add_2 = add_1 + mlp2_linear1_1;  add_1 = mlp2_linear1_1 = None
    return add_2

# with activation checkpoint
def forward(self, x):
    def checkpoint_0(x):
        mlp1_linear1 = self.mlp1.linear1(x)
        mlp1_linear1_1 = self.mlp1.linear1(x)
        return mlp1_linear1, mlp1_linear1_1
    mlp1_linear1, mlp1_linear1_1 = torch.utils.checkpoint.checkpoint(checkpoint_0, x)
    def checkpoint_1(x):
        mlp2_linear1 = self.mlp2.linear1(x)
        mlp2_linear1_1 = self.mlp2.linear1(x);  x = None
        return mlp2_linear1, mlp2_linear1_1
    mlp2_linear1, mlp2_linear1_1 = torch.utils.checkpoint.checkpoint(checkpoint_1, x)
    add = mlp1_linear1 + mlp1_linear1_1;  mlp1_linear1 = mlp1_linear1_1 = None
    add_1 = add + mlp2_linear1;  add = mlp2_linear1 = None
    add_2 = add_1 + mlp2_linear1_1;  add_1 = mlp2_linear1_1 = None
    return add_2

As the codegen is only provided in torch 1.12 and the CI is based on torch 1.11, the test is skipped in pytest but the full local test log is shown below.