Memory leak during replacement of modules (Mixtral, Scatter MoE)
Author: casper-hansen
Pull request: no

### System Info

transformers 4.38.2, python 3.10, torch 2.2.0, cuda 12.3, 4x A6000

### Who can help?

@SunMarc @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

This is referencing code that is committed in an [axolotl branch](https://github.com/OpenAccess-AI-Collective/axolotl/blob/scatter_moe/examples/mistral/mixtral_fused.py). Go to the link to see more details, we are adopting [Scatter MoE](https://github.com/shawntan/scattermoe/) for optimized training.

Step 1: Load in the Mixtral model

```python
import gc
import torch
from tqdm import tqdm
from axolotl.monkeypatch.moe.moe import SparseMoeBlock
from transformers import AutoTokenizer, TextStreamer
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock, MixtralForCausalLM, MixtralConfig

model_path = "mistralai/Mixtral-8x7B-Instruct-v0.1"

# Load model
config = MixtralConfig.from_pretrained(model_path, max_position_embeddings=2048, use_cache=False)
model = MixtralForCausalLM.from_pretrained(
    model_path,
    config=config,
    device_map="auto",
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16
)
```

Step 2: Define optimized modules to replace the sparse MoE for faster training/inference

```python
import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseMoeBlock(nn.Module):
    def __init__(self, experts, gate, hidden_dim, ffn_dim, num_experts, top_k):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.ffn_dim = ffn_dim
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = gate
        self.experts = FusedExperts(
            experts=experts,
            hidden_dim=hidden_dim,
            ffn_dim=ffn_dim,
            num_experts=num_experts,
            top_k=top_k,
            activation=experts[0].act_fn
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)

        # router_logits: (batch * sequence_length, n_experts)
        router_logits = self.gate(hidden_states)
        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)

        # we cast back to the input dtype
        routing_weights = routing_weights.to(hidden_states.dtype)

        # Fused expert forward
        final_hidden_states = self.experts(hidden_states, routing_weights, selected_experts)

        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
        return final_hidden_states, router_logits


class FusedExperts(nn.Module):
    def __init__(
        self,
        experts: nn.ModuleList =None,
        hidden_dim=128,
        ffn_dim=512,
        num_experts=8,
        top_k=2,
        activation=nn.SiLU(),
    ):
        """
        This implements fused experts that are compatible with Mixtral.
        MLP of type Gated-Linear Unit, typically with a SiLU activation function.
        """
        super(FusedExperts, self).__init__()

        device = experts[0].w1.weight.device
        self.num_experts = num_experts
        self.hidden_dim = hidden_dim
        self.ffn_dim = ffn_dim
        self.experts = ParallelExperts(num_experts, hidden_dim, 2 * ffn_dim, device=device)
        self.output_experts = ParallelExperts(num_experts, ffn_dim, hidden_dim, device=device)
        self.top_k = min(top_k, self.num_experts)
        self.activation = activation

        with torch.no_grad():
            for i in range(len(experts)):
                self.experts.weight.data[i].copy_(
                    torch.cat(
                        [experts[i].w1.weight.detach(), experts[i].w3.weight.detach()],
                        dim=0
                    )
                )
                self.output_experts.weight.data[i].copy_(
                    experts[i].w2.weight.detach()
                )

    def forward(
        self, x: torch.Tensor, routing_weights: torch.Tensor, selected_experts: torch.Tensor
    ):
        x_shape = x.size()
        x = x.view(-1, x_shape[-1])
        with torch.no_grad():
            sorted_expert_idxs, sorted_scattered_idxs = ops.flatten_and_sort(
                selected_experts
            )
            padded_block_idxs, expert_offsets = ops.padded_block_indices(
                sorted_expert_idxs, self.num_experts
            )

        h, gates = self.experts(
            x,
            self.top_k,
            sorted_expert_idxs,
            sorted_scattered_idxs,
            padded_block_idxs,
            expert_offsets,
            grouped_out=True,
        ).chunk(2, dim=-1)
        h = self.activation(gates) * h
        y = self.output_experts(
            h,
            1,
            sorted_expert_idxs,
            sorted_scattered_idxs,
            padded_block_idxs,
            expert_offsets,
            grouped_in=True,
            gates=routing_weights,
        )
        y = y.view(*x_shape[:-1], y.size(-1))
        return y
```

Step 3: Attempt to replace the `block_sparse_moe` with optimized version.

```python

modules = {k:v for k,v in model.named_modules() if isinstance(v, MixtralSparseMoeBlock)}

with tqdm(modules.items(), desc="scatter moe") as pbar:
    for i, (name, module) in enumerate(pbar):
        smoe = SparseMoeBlock(
            experts=module.experts,
            gate=module.gate,
            hidden_dim=module.hidden_dim,
            ffn_dim=module.ffn_dim,
            num_experts=module.num_experts,
            top_k=module.top_k,
        )
        old_module = model.model.layers[i].block_sparse_moe
        setattr(model.model.layers[i], "block_sparse_moe", smoe)
        del old_module
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.empty_cache()
```

Result: We run into OOM issues because we cannot free the memory of the old module due to a memory leak. Things I have tried:
- deleting individual w1, w2, w3 in each expert during the init of FusedExperts
- various ways to run del/gc.collect()/empty_cache()/cpu()


### Expected behavior

The expected behavior is to replace the original modules with optimized modules to achieve a speedup within the Huggingface ecosystem without running into OOM issues due to not being able to free memory.

Comments:
Hey @casper-hansen 
Can you try this approach instead: https://github.com/huggingface/transformers/blob/d9fa13ce62711c5c4840d1faf0473bfb9237a508/src/transformers/integrations/bitsandbytes.py#L191 ? not sure why your approach leads to OOM though, maybe the other way would be to init the layer on the meta device from pytorch with `with torch.device("meta"):` context manager and re-assign the weights later after the module replacement. You can
1- get the state dict of the model first
2- load your new model on the meta device and perform module replacement
3- load the state dict into the meta model with `load_state_dict(xxx, assign=True, strict=True)`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.