Layer activation discrepancy when running DeepSeekV3 after swapping `nn.Linear` with custom `FakeLinear` with disk offloading
Author: balaabhijit
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.49.0
- Platform: Linux-5.15.0-122-generic-x86_64-with-glibc2.35
- Python version: 3.11.10
- Huggingface_hub version: 0.27.0
- Safetensors version: 0.4.5
- Accelerate version: 1.4.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Using `accelerate` for big model inference and offloading layers to disk
- Using GPU in script?: Yes 8x80GB (A100) GPUs
- GPU type: NVIDIA A100 80GB PCIe

### Who can help?

@SunMarc @ArthurZucker 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```python

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
# This code just registers the DeepSeekV3 with HF autoclass
from deepseekv3.hf_utils import register_with_auto_class
register_with_auto_class()


def get_causal_model_and_tokenizer(model_name, torch_dtype, model_kwargs={}, tokenizer_kwargs={}):
    tokenizer = AutoTokenizer.from_pretrained(model_name, **tokenizer_kwargs)
    config = AutoConfig.from_pretrained(model_name)
    with init_empty_weights():
        model = AutoModelForCausalLM.from_config(
            config,
            attn_implementation=model_kwargs.pop("attn_implementation", "eager"),
        )
    model = load_checkpoint_and_dispatch(
            model=model,
            checkpoint=model_name,
            dtype=torch_dtype,
            no_split_module_classes=model._no_split_modules,
            **model_kwargs,
        )
    return model, tokenizer

# I have 8 A100 GPUs with 80GB memory in each
# I offload the rest to disk. I have also attached my `hf_device_map`
model, tokenizer = get_causal_model_and_tokenizer(
    model_name = "/data/custom_model_downloads/DeepSeek-V3-bf16",
    torch_dtype=torch.bfloat16,
    model_kwargs= {
        "attn_implementation": "eager",
        "offload_folder": "/data/model_offload_cache/deepseekv3_671b",
        "offload_buffers": True,
        "device_map": "auto",
        "max_memory":
            {
            0: 60000000000,
            1: 60000000000,
            2: 60000000000,
            3: 60000000000,
            4: 60000000000,
            5: 60000000000,
            6: 60000000000,
            7: 60000000000,
            },
    }
)

print(model)

model = model.eval()


print(model.device)


def forward_hook(name):
    def tmp(layer, inp, out):
        print(f"Done with layer {name}")
    return tmp

def forward_prehook(name):
    def tmp(layer, inp):
        print(f"Starting layer {name}")
    return tmp


forward_hook_handles = []
forward_prehook_handles = []
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        forward_hook_handles.append(module.register_forward_hook(forward_hook(name)))
        forward_prehook_handles.append(module.register_forward_pre_hook(forward_prehook(name)))


text = "Hello I am not sure what is happening can you help me debug this?"
tokens = tokenizer(text, return_tensors="pt")

# Model before swapping. Works as intended. This prints the entry and exit from the `nn.Linear` using the hooks
output = model.generate(tokens.input_ids, max_new_tokens=1, use_cache=False)


for hook in forward_hook_handles:
    hook.remove()

for hook in forward_prehook_handles:
    hook.remove()


class FakeLinear(torch.nn.Linear):
    """
    FakeLinear: A fake linear layer that can be used to replace a real linear layer.
    It is exactly same as `nn.Linear` as the forward() is inherited from `nn.Linear`
    """
    @classmethod
    @torch.no_grad()
    def from_float(
        cls,
        mod,
    ) -> "FakeLinear":
        with torch.device("meta"):
            super_kwargs = {
                "in_features": mod.in_features,
                "out_features": mod.out_features,
                "bias": False,
            }
            new_mod = cls(**super_kwargs)
        new_mod.weight = mod.weight
        new_mod.bias = mod.bias
        return new_mod


from tqdm import tqdm

def _replace_with_custom_fn_if_matches_filter(
    model,
    replacement_fn,
    filter_fn,
    pbar: tqdm,
    cur_fqn="",
) -> None:
    """
    For each `child` in `model`, replaces it with `replacement_fn(child)`
    if `filter_fn(child)` is `True`
    """
    name_to_child = dict(model.named_children())
    for name, child in name_to_child.items():
        pbar.update(1)
        if cur_fqn == "":
            new_fqn = name
        else:
            new_fqn = f"{cur_fqn}.{name}"
        if filter_fn(child, new_fqn):
            new_child = replacement_fn(child)
            setattr(model, name, new_child)
        else:
            _replace_with_custom_fn_if_matches_filter(
                child, replacement_fn, filter_fn, pbar, new_fqn
            )

def swap_linear_(model: torch.nn.Module):
    _replace_with_custom_fn_if_matches_filter(
        model=model,
        replacement_fn=lambda mod: FakeLinear.from_float(
            mod=mod,
        ),
        filter_fn=lambda mod, fqn: type(mod) == torch.nn.Linear,
        pbar=tqdm(
            desc="Swapping linear layers in model...",
        ),
    )


swap_linear_(model)

# You will see that all the `nn.Linear` is replaced with `FakeLinear` as expected
print(model)


forward_hook_handles = []
forward_prehook_handles = []
for name, module in model.named_modules():
    if type(module) == FakeLinear:
        forward_hook_handles.append(module.register_forward_hook(forward_hook(name)))
        forward_prehook_handles.append(module.register_forward_pre_hook(forward_prehook(name)))


model.hf_device_map

# The Activated layers are not same as before. Infact a lot of layers are not even activated.
# They are same for layers in the GPUs but for the layers on the disk (example layers 19 and above) only 0-7 are activated. For layers 18 there is a discrepancy in the layers numbers getting activated. But whereas before the swap almost all the MOEs are activated in each layer
output = model.generate(tokens.input_ids, max_new_tokens=1, use_cache=False)
```

I have attached 3 files:
- [hf_device_map.txt](https://github.com/user-attachments/files/18857452/hf_device_map.txt)
- [after_swapping.txt](https://github.com/user-attachments/files/18857450/after_swapping.txt): Layer activation numbers before swapping
- [before_swapping.txt](https://github.com/user-attachments/files/18857453/before_swapping.txt): Layer activation numbers after swapping

### Expected behavior

I expect both of the files to have the same layer activation numbers. The numbers match exactly until Layer 17 (The last layer which resides in the GPU). The problem starts after layer 18 (First layer in disk). In Layer 18 there are a few missing layers and from layer 19 on only 0-7 experts are activated.

Comments:
I don't think we can debug this one, unfortunately! When you're doing very unusual custom tasks like this, we really can't reproduce and dig into the bug for you, because we can't tell if the bug is simply an issue in your code or not without spending a lot of time on it.

If you really believe there's a bug here, we'll need a much simpler reproducer script, ideally as short as possible with as simple a model as possible, and a clearer idea of where exactly the bug in `Transformers` is.

Closing for now, but feel free to reopen in future if you can identify a reproducible bug here!