perf: run k2 CTC loss on the GPU instead of the CPU - #77
Open
lumpidu wants to merge 1 commit into
Open
Conversation
CTCLossWithLabelPriors.to() left k2_device hard-set to "cpu", so every step moved the full log_probs tensor from the GPU to the CPU and ran the k2 CTC graph there. That transfer stalls the step and leaves the GPU idle for most of the alignment stage. Select the device based on what k2 was built with: use the model device when k2 has CUDA support, and fall back to the CPU otherwise. to() runs during setup of every stage, including stages that never touch the alignment loss, so the k2 import is guarded; missing k2 keeps the harmless cpu default and only the alignment stage, which genuinely needs k2, fails later. The supervision segments stay on the CPU at both DenseFsaVec call sites because k2 requires that argument there regardless of where log_probs lives. The computed loss is identical to the CPU path; only the placement changes.
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
1 participant
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.
CTCLossWithLabelPriors.to()leftk2_devicehard-set to "cpu", so every step moved the fulllog_probstensor from the GPU to the CPU and ran the k2 CTC graph there.That transfer stalls the step and leaves the GPU idle for most of the alignment stage.
Select the device based on what k2 was built with: use the model device when k2 has CUDA support, and fall back to the CPU otherwise. Operation
to()runs during setup of every stage, including stages that never touch the alignment loss, so the k2 import is guardedA missing k2 keeps the harmless cpu default and only the alignment stage, which anyways needs k2, fails later. The supervision segments stay on the CPU at both
DenseFsaVeccall sites because k2 requires that argument there regardless of wherelog_probslives.Computed loss is identical to the CPU path, only placement changes.