Uneven GPU memory usage and CUDA OOM during multi-GPU inference with Qwen3-VL-30B-A3B-Thinking

[Xqle]

When performing inference on a ~3-minute video using the Qwen3-VL-30B-A3B-Thinking model from transformers in 8*3090(24 GB) and setting fps=1, a CUDA Out-of-Memory (OOM) error occurs. 

  Traceback (most recent call last):
      File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 94, in <module>
        main()
      File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 78, in main
        generated_ids = model.generate(**inputs, max_new_tokens=args.max_tokens)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
        return func(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/generation/utils.py", line 2564, in generate
        result = decoding_method(
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/generation/utils.py", line 2784, in _sample
        outputs = self(**model_inputs, return_dict=True)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/accelerate/hooks.py", line 175, in new_forward
        output = module._old_forward(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
        outputs = func(self, *args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 1601, in forward
        outputs = self.model(
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
        outputs = func(self, *args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 1389, in forward
        outputs = self.language_model(
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/utils/generic.py", line 1064, in wrapper
        outputs = func(self, *args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 962, in forward
        layer_outputs = decoder_layer(
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_layers.py", line 94, in __call__
        return super().__call__(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/accelerate/hooks.py", line 175, in new_forward
        output = module._old_forward(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
        return func(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 391, in forward
        hidden_states = self.mlp(hidden_states)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/accelerate/hooks.py", line 175, in new_forward
        output = module._old_forward(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 151, in forward
        routed_out = self.experts(hidden_states, router_weights, router_indices)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
        return self._call_impl(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
        return forward_call(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/accelerate/hooks.py", line 175, in new_forward
        output = module._old_forward(*args, **kwargs)
      File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py", line 120, in forward
        next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj)
      torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 5.06 GiB. GPU 0 has a total capacity of 23.68 GiB of which 4.40 GiB is free. Including non-PyTorch memory, this process has 19.27 GiB memory in use. Of the allocated memory 18.67 GiB is allocated by PyTorch, and 304.64 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

However, GPU monitoring shows that only one GPU has extremely high memory usage, while the others remain underutilized (less than 50% memory used).

After investigation, the memory surge appears to happen inside Qwen3VLMoeTextExperts.forward: 

Is there a known workaround for this, or will future releases include optimization or DTensor support for the MoE expert modules to improve memory distribution during multi-GPU inference?

[Rocketknight1]

cc @zucchini-nlp @ArthurZucker - maybe we should rewrite some ops when the hidden state can get very large like this, so it can be sharded! I have some capacity now while people are reviewing response schemas, so I can take this if you're busy

[AvinashDwivedi]

@Rocketknight1 Can I work on this issue? I'm new in open source but I can finish this.

[Rocketknight1]

@AvinashDwivedi let's wait until we know what we want to do first!

[ArthurZucker]

You are probably using device_map="auto" no? I think you should use tensor paralleL!

[Xqle]

@ArthurZucker
Yes, I was using device_map="auto".

Now I’m trying to enable tensor parallelism by simply replacing it with tp_plan="auto", but I encountered the following issue.

When using python infer_qwen3vl_transformers.py, I got the error:

Traceback (most recent call last):
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/integrations/tensor_parallel.py", line 57, in initialize_tensor_parallelism
    rank = int(os.environ["RANK"])
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/os.py", line 680, in __getitem__
    raise KeyError(key) from None
KeyError: 'RANK'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 195, in <module>
    main()
  File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 138, in main
    model = AutoModelForVision2Seq.from_pretrained(
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/auto/modeling_auto.py", line 2289, in from_pretrained
    return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 604, in from_pretrained
    return model_class.from_pretrained(
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 277, in _wrapper
    return func(*args, **kwargs)
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4702, in from_pretrained
    tp_plan, device_map, device_mesh, tp_size = initialize_tensor_parallelism(tp_plan, tp_size=tp_size)
  File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/integrations/tensor_parallel.py", line 74, in initialize_tensor_parallelism
    raise OSError(
OSError: We tried to initialize torch.distributed for you, but it failed. Make sure you init torch distributed in your script to use `tp_plan='auto'`.

When using torchrun --nproc_per_node=8 infer_qwen3vl_transformers.py, I ran into

  warnings.warn(
Loading checkpoint shards:  38%|███████████████████████████████████████████████████████                                                                                        | 5/13 [00:08<00:14,  1.76s/it]
[rank0]: Traceback (most recent call last):
[rank0]:   File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 195, in <module>
[rank0]:     main()
[rank0]:   File "/data01/xuqile/code/Inference/transformers/infer_qwen3vl_transformers.py", line 138, in main
[rank0]:     model = AutoModelForVision2Seq.from_pretrained(
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/auto/modeling_auto.py", line 2289, in from_pretrained
[rank0]:     return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 604, in from_pretrained
[rank0]:     return model_class.from_pretrained(
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 277, in _wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5051, in from_pretrained
[rank0]:     ) = cls._load_pretrained_model(
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 5471, in _load_pretrained_model
[rank0]:     _error_msgs, disk_offload_index = load_shard_file(args)
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 847, in load_shard_file
[rank0]:     disk_offload_index = _load_state_dict_into_meta_model(
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/modeling_utils.py", line 721, in _load_state_dict_into_meta_model
[rank0]:     shard_and_distribute_module(
[rank0]:   File "/data01/xuqile/miniforge3/envs/qwen3vl/lib/python3.10/site-packages/transformers/integrations/tensor_parallel.py", line 1050, in shard_and_distribute_module
[rank0]:     param = param[:].to(param_casting_dtype)
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 384.00 MiB. GPU 0 has a total capacity of 23.68 GiB of which 54.69 MiB is free. Including non-PyTorch memory, this process has 23.62 GiB memory in use. Of the allocated memory 23.18 GiB is allocated by PyTorch, and 204.32 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
[rank0]:[W1016 18:31:10.386689977 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

Could you please share the correct way to enable tensor parallelism in Transformers?

Here’s my current code: https://gist.github.com/Xqle/d0519e19e3f4b68978b9806b8500f27c