How to Reduce Memory Usage by Pruning Intermediate Activation Tensors During LoRA Fine-tuning of Llama2 Model?

zbtrs: <p>Hello everyone,</p>
<p>I’m currently working on fine-tuning the Llama2 model using LoRA and am trying to reduce the memory usage by pruning intermediate activation tensors. I noticed that the parameters of the base model have <code>require_grad=False</code>, and only the newly added <code>loraA</code> and <code>loraB</code> layers on each <code>q_proj</code> and <code>v_proj</code> have <code>require_grad=True</code>.</p>
<p>Here is the original code for the <code>LlamaMLP</code> class:</p>
<pre><code class="lang-auto">class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        if self.config.pretraining_tp &gt; 1:
            slice = self.intermediate_size // self.config.pretraining_tp
            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
            up_proj_slices = self.up_proj.weight.split(slice, dim=0)
            down_proj_slices = self.down_proj.weight.split(slice, dim=1)

            gate_proj = torch.cat(
                [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
            )
            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)

            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
            down_proj = [
                F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
            ]
            down_proj = sum(down_proj)
        else:
            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

        return down_proj
</code></pre>
<p>To reduce memory usage, I modified the code to detach, set to <code>None</code> , and <code>torch.cuda.empty_cache()</code> for the intermediate tensors that are not needed in the backward pass. Here is my modified code:</p>
<pre><code class="lang-auto">class LlamaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        if self.config.pretraining_tp &gt; 1:
            slice = self.intermediate_size // self.config.pretraining_tp
            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
            up_proj_slices = self.up_proj.weight.split(slice, dim=0)
            down_proj_slices = self.down_proj.weight.split(slice, dim=1)

            gate_proj = torch.cat(
                [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
            )
            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)

            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
            down_proj = [
                F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
            ]
            down_proj = sum(down_proj)
        else:
            up_proj = self.up_proj(x)
            print(f"up_proj memory: {up_proj.element_size() * up_proj.nelement() / (1024 ** 3)} GB")
    
            gate_proj = self.gate_proj(x)
            print(f"gate_proj memory: {gate_proj.element_size() * gate_proj.nelement() / (1024 ** 3)} GB")
            print_memory_usage('sb1')
            x.detach()
            x = None
            torch.cuda.empty_cache()
            print_memory_usage('sb2')
    
            act_fn = self.act_fn(gate_proj)
            print(f"act_fn memory: {act_fn.element_size() * act_fn.nelement() / (1024 ** 3)} GB")
            print_memory_usage('sb3')
            gate_proj.detach()
            gate_proj = None
            torch.cuda.empty_cache()
            print_memory_usage('sb4')
            mul = act_fn * up_proj
            print(f"mul memory: {mul.element_size() * mul.nelement() / (1024 ** 3)} GB")
            print_memory_usage('sb5')
            act_fn.detach()
            act_fn = None
            up_proj.detach()
            up_proj = None
            torch.cuda.empty_cache()
            print_memory_usage('sb6')
            down_proj = self.down_proj(mul)
            print(f"down_proj memory: {down_proj.element_size() * down_proj.nelement() / (1024 ** 3)} GB")
            mul.detach()
            mul = None

        return down_proj
</code></pre>
<p>However, I observed that <code>torch.cuda.memory_allocated()</code> still increases and the increase in memory usage corresponds to the size of these intermediate tensors. Could anyone help me understand why this is happening and how to effectively reduce the memory usage during the forward pass?</p>
<p>Thank you in advance!</p>