Skip to content

Fix device mismatch in VibeVoiceASR.encode_speech under device_map=auto - #441

Open
Ammar Hashmi (ammar-iitm) wants to merge 1 commit into
microsoft:mainfrom
ammar-iitm:fix/asr-speech-masks-device-mismatch
Open

Ammar Hashmi (ammar-iitm) wants to merge 1 commit into
microsoft:mainfrom
ammar-iitm:fix/asr-speech-masks-device-mismatch

Conversation

@ammar-iitm

Copy link
Copy Markdown

Summary

  • Fixes Tensor device mismatch error when using multi-GPU with device_map=auto #240: running microsoft/VibeVoice-ASR with device_map="auto" across multiple GPUs crashes with:
    RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:6)
    
    at vibevoice/modular/modeling_vibevoice_asr.py:335, inside encode_speech:
    combined_features = acoustic_features[speech_masks] + semantic_features[speech_masks]
  • Root cause: under device_map="auto", self.model.acoustic_connector/semantic_connector can be dispatched (via accelerate) to a different GPU than wherever speech_masks was created, and PyTorch requires a boolean index tensor to be on CPU or on the same device as the tensor it indexes.
  • This is the same class of multi-GPU device mismatch already handled elsewhere in this exact file — forward() does shift_labels = shift_labels.to(shift_logits.device) a few dozen lines below for the same reason. This PR applies the identical, already-established pattern to speech_masks.

Change

  • speech_masks = speech_masks.to(acoustic_features.device) right before the boolean-indexed combine, in VibeVoiceASRForConditionalGeneration.encode_speech.

Test plan

  • python3 -m py_compile vibevoice/modular/modeling_vibevoice_asr.py
  • Verified the change is a safe no-op for the common single-device case: with speech_masks/acoustic_features/semantic_features all on the same device, the output is bit-identical (torch.equal) to the pre-change code path.
  • I don't have multi-GPU hardware available to reproduce the exact device_map="auto" cross-GPU scenario from the issue end-to-end — would appreciate confirmation from the reporter or a maintainer with multi-GPU access that this resolves it.

🤖 Generated with Claude Code

https://claude.ai/code/session_016fM2LafALmBK16KnFWjsp7

With device_map="auto" across multiple GPUs, acoustic_connector/
semantic_connector can be dispatched to a different device than
speech_masks, since it isn't moved alongside the model layers. Boolean-
indexing requires the mask to be on the same device as the indexed
tensor (or CPU), so this raised:

  RuntimeError: indices should be either on cpu or on the same device
  as the indexed tensor (cuda:6)

at acoustic_features[speech_masks] (reported in microsoft#240 with a full
multi-GPU stack trace pointing at this exact line).

Move speech_masks to acoustic_features.device right before indexing,
mirroring the existing shift_labels.to(shift_logits.device) pattern
already used elsewhere in this same file for the same class of
multi-GPU device mismatch.

Fixes microsoft#240

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016fM2LafALmBK16KnFWjsp7
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Tensor device mismatch error when using multi-GPU with device_map=auto

1 participant