Fix device mismatch in VibeVoiceASR.encode_speech under device_map=auto - #441
Open
Ammar Hashmi (ammar-iitm) wants to merge 1 commit into
Open
Ammar Hashmi (ammar-iitm) wants to merge 1 commit into
Ammar Hashmi (ammar-iitm) wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
microsoft/VibeVoice-ASRwithdevice_map="auto"across multiple GPUs crashes with:vibevoice/modular/modeling_vibevoice_asr.py:335, insideencode_speech:device_map="auto",self.model.acoustic_connector/semantic_connectorcan be dispatched (viaaccelerate) to a different GPU than whereverspeech_maskswas created, and PyTorch requires a boolean index tensor to be on CPU or on the same device as the tensor it indexes.forward()doesshift_labels = shift_labels.to(shift_logits.device)a few dozen lines below for the same reason. This PR applies the identical, already-established pattern tospeech_masks.Change
speech_masks = speech_masks.to(acoustic_features.device)right before the boolean-indexed combine, inVibeVoiceASRForConditionalGeneration.encode_speech.Test plan
python3 -m py_compile vibevoice/modular/modeling_vibevoice_asr.pyspeech_masks/acoustic_features/semantic_featuresall on the same device, the output is bit-identical (torch.equal) to the pre-change code path.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