Fix GPU-CPU tensor manipulation. Small performance boost #178
Conversation
|
cc @masahi |
|
I think the second optimization is not really helping since the output of |
| @@ -75,13 +75,13 @@ def _is_safe_to_sample(prob_like): | |||
| logits = torch.from_dlpack(logits) | |||
| num_seq = len(sampling_params) | |||
|
|
|||
| mask_random = torch.tensor( | |||
| mask_random_dvc = torch.tensor( | |||
There was a problem hiding this comment.
dvc is a strange post fix. Just use gpu or simply mask_random (no post fix).
There was a problem hiding this comment.
Hello @masahi! I guess that potentially it can work on cpu in the future. I thought to rename it to _gpu, but it will confuse somebody in the future due to my guess. What do you think about it?
There was a problem hiding this comment.
Logits here is supposed to be always on GPU (you can add an assert). So there is no confusion.
| [p.sampling_type == SamplingType.RANDOM for p in sampling_params], | ||
| dtype=torch.bool | ||
| ) | ||
| mask_greedy_cpu = torch.logical_not(mask_random_cpu) |
There was a problem hiding this comment.
When you create a GPU mask, PT will first create a CPU mask and do cudaMemcpy under the hood. So you can create a CPU mask once, and create a GPU mask by copying the CPU mask explicitly. So you can avoid creating a CPU mask twice.
| ) | ||
| mask_greedy = torch.logical_not(mask_random) | ||
| mask_greedy_cpu = torch.logical_not(mask_random_cpu) | ||
| if logits.device == torch.device("cpu"): |
There was a problem hiding this comment.
It costs nothing but we support case with run topology on CPU. As I know @elvin-n run singlebatch on CPU after small fixes
* ios downloader * use dist as optional provided dir * Update iOS app to new reload api --------- Co-authored-by: Yaxing Cai <[email protected]>
There are two fixes:
logitsfrom GPU (potential performance reduction). After fix there are two pairs of masks on CPU andlogitsdevice side, moreover CPU masks are created only if it needs.res_random = torch.multinomial(probs, 1, True).cpu().numpy()[:, 0]byres_random = torch.multinomial(probs, 1, True)[:, 0].cpu().numpy(). As I understand in the first case we copy the slice from old numpy to new one after copying full tensor from GPU to CPU, but in the second case we get slice view (without memory copying) and copying from GPU to CPU only sliced tensor (not full)Results of benchmark:
Note: Fluctuation of time from run to run is big enough (~1-2%) therefore several runs were performed and averaged, possibly measurements should be more careful