[Bugfix]: DeepseekR1 model load fails with weights tied error

[cennn]

Fix: 12541
This PR resolves the bug causing the DeepseekR1 model to fail during loading due to weight-tied errors.

The implemented changes are as follows:
1 Make MLACommonImpl inherit from Module for easy use of _process_weights_after_loading during parameter decompression.
2 Set tie_weights=False in maybe_offload_to_cpu to prevent weight-tying issues.
3 Add gc.collect to _process_weights_after_loading to tackle potential memory leaks that may lead to "cuda out of memory" errors.

It should be noted that due to the offload operation and the use of gc.collect, the model loading time is approximately 2 minutes longer. The startup command utilized is llm = LLM("deepseek-ai/DeepSeek-R1", tensor_parallel_size = 8, cpu_offload_gb = 50, trust_remote_code = True). (on 8*H100 80G).

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mgoin]

I think this is a good call although calling gc.collect after every single module is very conservative. I would like to not increase the loading time of weights by so much. Maybe just using torch.cuda.empty_cache would be faster? FYI @LucasWilkinson

[youkaichao]

this line should be the key to fix the bug

[youkaichao]

I think this is a good call although calling gc.collect after every single module is very conservative. I would like to not increase the loading time of weights by so much. Maybe just using torch.cuda.empty_cache would be faster? FYI @LucasWilkinson

@mgoin I think you need gc.collect() to gc python variables, otherwise calling torch.cuda.empty_cache will not work.

I suspect torch.cuda.empty_cache would be even slower as it calls cuda malloc functions.

[youkaichao]

ideally we can run this just for every layer. but we don't have the concept of layer here.

[youkaichao]

@mgoin I think ultimately we need to use the tensor subclass approach for cpu offloading, then we don't need to manually move cpu tensors back and forth. that would require #12158 though.

[BigCousin-z]

Is there any plan to implement the cpu_offload_gb technology in vLLM for Mixture of Experts (MoE) models? Specifically, can it support the configuration of offloading strategies based on usage scenarios or specified layers?

[BigCousin-z]

@cennn @youkaichao

[mgoin]

@cennn @youkaichao Maybe we could compromise by making a utility that doesn't run GC for every single module, but instead does it every N calls. What do you think of something like this for now?

class GarbageCollector:
    def __init__(self, collection_frequency: int):
        self.counter = 0
        self.collection_frequency = collection_frequency
    
    def maybe_collect(self):
        self.counter += 1
        if self.counter >= self.collection_frequency:
            self.counter = 0
            gc.collect()

def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
                               target_device: torch.device) -> None:
    collector = GarbageCollector(collection_frequency=5)
    
    for _, module in model.named_modules():
        quant_method = getattr(module, "quant_method", None)
        if isinstance(quant_method, QuantizeMethodBase):
            with device_loading_context(module, target_device):
                quant_method.process_weights_after_loading(module)
            collector.maybe_collect()

    for _, module in model.named_modules():
        if isinstance(module, Attention) and \
            hasattr(module.impl, "process_weights_after_loading"):
            with device_loading_context(module.impl, target_device):
                module.impl.process_weights_after_loading(model_config.dtype)
            collector.maybe_collect()