Hi Yi Wang,
Thanks for your reply. I edit the code according to your suggestions like this (del pixel_values, del x_rref):
class TransformerShard1(nn.Module):
    def __init__(self, device, config,num_layers):
        super().__init__()
        # self.config = config
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x).to_tuple()[0]
        end = time.time()
        return x

    @torch.no_grad()
    def forward(self, pixel_values=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = pixel_values.to_here().to(self.device)
        del pixel_values
        gc.collect()
        with self._lock:
            x = self.forward_kernel(x)
        print("Shard1 finish its microbatch ")
        return x.cpu()

class TransformerShard2(nn.Module):
    def __init__(self,device, config, num_layers):
        super().__init__()
        self.model = ViTForImageClassification.from_pretrained(config)
        self.model.vit.encoder.layer = nn.Sequential(*[self.model.vit.encoder.layer[i] for i in range(num_layers, 2*num_layers)])
        self.model.vit= nn.Sequential(*list(self.model.vit.children())[1:-1])
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        self._lock = threading.Lock()
        self.device = device

    def forward_kernel(self, x):
        x = self.model(x)[0]
        return x

    @torch.no_grad()
    def forward(self, x_rref=None,
        head_mask=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None):
        x = x_rref.to_here().to(self.device)
        del x_rref
        gc.collect()
        with self._lock:
            x = self.forward_kernel(x)
        print("Shard1 finish its microbatch ")
        return x.cpu()

The overall memory usage is reduced (for the microbatch is 256, the peak memory usage is reduced from ~800MB to ~500MB). However, after each forward pass, the memory still keeps growing. It even does not be reduced after finishing one batch, which is very strange.