diff --git a/CHANGELOG.md b/CHANGELOG.md index b0d05e954..4c6febd77 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] ### Added +- Uses the per thread default stream for cublas +- Uses the strided batched gemm cublas call when possible for the batchedGemm. +- In the general batchedGemm case, reduces the number of memcpy calls from 3 to 1. +- Rounds the width of input batches to a multiple of 8 when the GPU backend is being used. This is to enable better use of tensorcores on Volta architectures and newer. +- Places NVIDIA notices to some files - Local/global sharding with MPI training via `--sharding local` - fp16 support for factors. - Correct training with fp16 via `--fp16`. diff --git a/LICENSE.md b/LICENSE.md index 878908214..269f4085b 100644 --- a/LICENSE.md +++ b/LICENSE.md @@ -3,6 +3,8 @@ MIT License Copyright (c) 2016 Marcin Junczys-Dowmunt, the University of Edinburgh, Adam Mickiewicz University +Copyright (c) 2020 NVIDIA Corporation + Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights diff --git a/src/data/corpus.cpp b/src/data/corpus.cpp index e8ce850b6..d0b6a8a6d 100755 --- a/src/data/corpus.cpp +++ b/src/data/corpus.cpp @@ -257,25 +257,33 @@ CorpusBase::batch_ptr Corpus::toBatch(const std::vector& batchVector) { std::vector sentenceIds; - std::vector maxDims; // @TODO: What's this? widths? maxLengths? + std::vector maxWordsInBatchSentence; for(auto& ex : batchVector) { // @TODO: rename 'ex' to 'sample' or 'sentenceTuple' - if(maxDims.size() < ex.size()) - maxDims.resize(ex.size(), 0); + if(maxWordsInBatchSentence.size() < ex.size()) + maxWordsInBatchSentence.resize(ex.size(), 0); for(size_t i = 0; i < ex.size(); ++i) { - if(ex[i].size() > (size_t)maxDims[i]) - maxDims[i] = (int)ex[i].size(); + if(ex[i].size() > (size_t)maxWordsInBatchSentence[i]) + maxWordsInBatchSentence[i] = (int)ex[i].size(); } sentenceIds.push_back(ex.getId()); } + + // When running on GPU, we want the batchWidth to be a multiple of 8 for better tensorcore usage + if(options_->get("cpu-threads") == 0) { + constexpr int roundingFactor = 8; + for(size_t j = 0; j < maxWordsInBatchSentence.size(); ++j) { + maxWordsInBatchSentence[j] = roundingFactor * ((maxWordsInBatchSentence[j] + roundingFactor - 1) / roundingFactor); + } + } std::vector> subBatches; - for(size_t j = 0; j < maxDims.size(); ++j) { - subBatches.emplace_back(New(batchSize, maxDims[j], vocabs_[j])); + for(size_t j = 0; j < maxWordsInBatchSentence.size(); ++j) { + subBatches.emplace_back(New(batchSize, maxWordsInBatchSentence[j], vocabs_[j])); } - std::vector words(maxDims.size(), 0); + std::vector words(maxWordsInBatchSentence.size(), 0); for(size_t b = 0; b < batchSize; ++b) { // loop over batch entries - for(size_t j = 0; j < maxDims.size(); ++j) { // loop over streams + for(size_t j = 0; j < maxWordsInBatchSentence.size(); ++j) { // loop over streams auto subBatch = subBatches[j]; for(size_t s = 0; s < batchVector[b][j].size(); ++s) { // loop over word positions subBatch->data()[subBatch->locate(/*batchIdx=*/b, /*wordPos=*/s)/*s * batchSize + b*/] = batchVector[b][j][s]; @@ -285,7 +293,7 @@ CorpusBase::batch_ptr Corpus::toBatch(const std::vector& batchVector) { } } - for(size_t j = 0; j < maxDims.size(); ++j) + for(size_t j = 0; j < maxWordsInBatchSentence.size(); ++j) subBatches[j]->setWords(words[j]); auto batch = batch_ptr(new batch_type(subBatches)); diff --git a/src/layers/generic.cpp b/src/layers/generic.cpp index d44f40206..a788f1ec3 100755 --- a/src/layers/generic.cpp +++ b/src/layers/generic.cpp @@ -91,9 +91,10 @@ namespace marian { } // if selIdx are given, then we must reshuffle accordingly - if (!hypIndices.empty()) // use the same function that shuffles decoder state - sel = rnn::State::select(sel, hypIndices, (int)beamSize, /*isBatchMajor=*/false); - + if (!hypIndices.empty()) { // use the same function that shuffles decoder state + auto indices = graph()->indices(hypIndices); + sel = rnn::State::select(sel, indices, (int)beamSize, /*isBatchMajor=*/false); + } return sel; } diff --git a/src/models/states.h b/src/models/states.h index c2f9ee05a..9ae2c8a82 100755 --- a/src/models/states.h +++ b/src/models/states.h @@ -30,7 +30,8 @@ class EncoderState { // Sub-select active batch entries from encoder context and context mask Ptr select(const std::vector& batchIndices) { // [batchIndex] indices of active batch entries // Dimension -2 is OK for both, RNN and Transformer models as the encoder context in Transformer gets transposed to the same dimension layout - return New(index_select(context_, -2, batchIndices), index_select(mask_, -2, batchIndices), batch_); + auto indices = context_->graph()->indices(batchIndices); + return New(index_select(context_, -2, indices), index_select(mask_, -2, indices), batch_); } }; diff --git a/src/rnn/types.h b/src/rnn/types.h index 47424b759..882a3939e 100644 --- a/src/rnn/types.h +++ b/src/rnn/types.h @@ -1,5 +1,7 @@ #pragma once +#include "common/definitions.h" +#include "common/shape.h" #include "marian.h" #include @@ -12,7 +14,7 @@ struct State { Expr output; Expr cell; - State select(const std::vector& selIdx, // [beamIndex * activeBatchSize + batchIndex] + State select(Expr selIdx, // [beamIndex * activeBatchSize + batchIndex] int beamSize, bool isBatchMajor) const { return{ select(output, selIdx, beamSize, isBatchMajor), select(cell, selIdx, beamSize, isBatchMajor) }; @@ -20,15 +22,14 @@ struct State { // this function is also called by Logits static Expr select(Expr sel, // [beamSize, dimTime, dimBatch, dimDepth] or [beamSize, dimBatch, dimTime, dimDepth] (dimTime = 1 for RNN) - const std::vector& selIdx, // [beamIndex * activeBatchSize + batchIndex] + Expr selIdx, // [beamIndex * activeBatchSize + batchIndex] int beamSize, bool isBatchMajor) { if (!sel) return sel; // keep nullptr untouched sel = atleast_4d(sel); - - int dimBatch = (int)selIdx.size() / beamSize; + int dimBatch =(int) selIdx->shape().elements()/beamSize; int dimDepth = sel->shape()[-1]; int dimTime = isBatchMajor ? sel->shape()[-2] : sel->shape()[-3]; @@ -83,8 +84,30 @@ class States { States select(const std::vector& selIdx, // [beamIndex * activeBatchSize + batchIndex] int beamSize, bool isBatchMajor) const { States selected; + Expr indices; + + // We need to check if either a states's cell or output fields are non-null. In this case, we need + // to select rows from at least one of the tensors. If only some exprs are non-null, the call to + // select will handle this for us by returning a null expr naturally. + for (auto& state : states_) { + if (state.cell) { + indices = state.cell->graph()->indices(selIdx); + break; + } + + if (state.output) { + indices = state.output->graph()->indices(selIdx); + break; + } + } + + // If indices is null here, then all of the state.cell and state.output entries are null. Therefore, + // select will ignore the null indices expr and simply return a null pointer which is the expected + // behavior + + // GPU OPT: Implement kernel to batch these on GPU for(auto& state : states_) - selected.push_back(state.select(selIdx, beamSize, isBatchMajor)); + selected.push_back(state.select(indices, beamSize, isBatchMajor)); return selected; } diff --git a/src/tensors/gpu/backend.h b/src/tensors/gpu/backend.h index 75cc604da..651d211f3 100644 --- a/src/tensors/gpu/backend.h +++ b/src/tensors/gpu/backend.h @@ -52,6 +52,7 @@ class Backend : public marian::Backend { if(!cublasHandle_) { // lazy initialization here to avoid memory usage when unused setDevice(); cublasCreate(&cublasHandle_); + cublasSetStream(cublasHandle_, cudaStreamPerThread); } return cublasHandle_; } @@ -60,6 +61,7 @@ class Backend : public marian::Backend { if(!cusparseHandle_) { // lazy initialization here to avoid memory usage when unused setDevice(); cusparseCreate(&cusparseHandle_); + cusparseSetStream(cusparseHandle_, cudaStreamPerThread); } return cusparseHandle_; } diff --git a/src/tensors/gpu/prod.cpp b/src/tensors/gpu/prod.cpp index bf7e5512c..b30a29a6c 100755 --- a/src/tensors/gpu/prod.cpp +++ b/src/tensors/gpu/prod.cpp @@ -1,4 +1,3 @@ - #ifdef _MSC_VER #pragma warning(disable: 4505) // warning C4505: '__float2half_rz': unreferenced local function has been removed (missing 'static inline') #endif @@ -133,6 +132,39 @@ struct TypedGemm { // specializati Barray, ldb, beta, Carray, ldc, batchCount)); } + + static void stridedBatchedGemm(cublasHandle_t handle, + CudaCompute computeCapability, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const float *alpha, + const float *A, int lda, int strideA, + const float *B, int ldb, int strideB, + const float *beta, + float *C, int ldc, int strideC, + int batchCount) { + // double #if and if unfortunately required to safeguard against compilation error + // with CUDA 8.0 and runtime error with CUDA >9.0 on GPUs with compute capability under 5 + #if CUDA_VERSION > 9000 + // query math mode and set algorithm accordingly + auto algorithm = tensorOpsEnabled(handle) ? CUBLAS_GEMM_DEFAULT_TENSOR_OP : CUBLAS_GEMM_DEFAULT; + if(computeCapability.major >= 5) + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, transa, transb, + m, n, k, alpha, + (void const*)A, CUDA_R_32F, lda, strideA, + (void const*)B, CUDA_R_32F, ldb, strideB, beta, + (void*)C, CUDA_R_32F, ldc, strideC, batchCount, + CUDA_R_32F, algorithm)); + #endif + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, transa, transb, + m, n, k, alpha, + A, lda, strideA, + B, ldb, strideB, + beta, + C, ldc, strideC, + batchCount)); + } }; #if COMPILE_FP16 @@ -181,6 +213,29 @@ struct TypedGemm { // specialization (void**)Carray, CUDA_R_16F, ldc, batchCount, CUDA_R_16F, algorithm)); } + + static void stridedBatchedGemm(cublasHandle_t handle, + CudaCompute computeCapability, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const half *alpha, + const half *A, int lda, int strideA, + const half *B, int ldb, int strideB, + const half *beta, + half *C, int ldc, int strideC, + int batchCount) { + ABORT_IF(computeCapability.major < 6, "Compute capability {} below 6 should not happen for FP16", computeCapability.major); + // query math mode and set algorithm accordingly + auto algorithm = tensorOpsEnabled(handle) ? CUBLAS_GEMM_DEFAULT_TENSOR_OP : CUBLAS_GEMM_DEFAULT; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, transa, transb, + m, n, k, alpha, + (void const*)A, CUDA_R_16F, lda, strideA, + (void const*)B, CUDA_R_16F, ldb, strideB, beta, + (void*)C, CUDA_R_16F, ldc, strideC, batchCount, + CUDA_R_16F, algorithm)); + } + }; template <> @@ -228,6 +283,28 @@ struct TypedGemm { // specializatio (void**)Carray, CUDA_R_16F, ldc, batchCount, CUDA_R_32F, algorithm)); // use 32-bit compute type for accumulation } + + static void stridedBatchedGemm(cublasHandle_t handle, + CudaCompute computeCapability, + cublasOperation_t transa, + cublasOperation_t transb, + int m, int n, int k, + const float *alpha, + const half *A, int lda, int strideA, + const half *B, int ldb, int strideB, + const float *beta, + half *C, int ldc, int strideC, + int batchCount) { + ABORT_IF(computeCapability.major < 6, "Compute capability {} below 6 should not happen for FP16", computeCapability.major); + // query math mode and set algorithm accordingly + auto algorithm = tensorOpsEnabled(handle) ? CUBLAS_GEMM_DEFAULT_TENSOR_OP : CUBLAS_GEMM_DEFAULT; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, transa, transb, + m, n, k, alpha, + (void const*)A, CUDA_R_16F, lda, strideA, + (void const*)B, CUDA_R_16F, ldb, strideB, beta, + (void*)C, CUDA_R_16F, ldc, strideC, batchCount, + CUDA_R_32F, algorithm)); + } }; #endif @@ -323,8 +400,8 @@ void ProdBatchedTyped(marian::Tensor C, CUDA_CHECK(cudaSetDevice((int)C->getDeviceId().no)); ComputeType alpha = scalar; - int batchA = A->shape().elements() / (A->shape()[-1] * A->shape()[-2]); - int batchB = B->shape().elements() / (B->shape()[-1] * B->shape()[-2]); + int batchDimA = A->shape().elements() / (A->shape()[-1] * A->shape()[-2]); + int batchDimB = B->shape().elements() / (B->shape()[-1] * B->shape()[-2]); int m = A->shape()[-2]; int k = A->shape()[-1]; @@ -350,46 +427,57 @@ void ProdBatchedTyped(marian::Tensor C, auto cublasHandle = backend->getCublasHandle(); auto compute = backend->getCudaComputeCapability(); - auto strideA = batchA == 1 ? 0 : m * k; - auto strideB = batchB == 1 ? 0 : n * k; + auto strideA = batchDimA == 1 ? 0 : m * k; + auto strideB = batchDimB == 1 ? 0 : n * k; auto strideC = n * m; - auto batchC = std::max(batchA, batchB); - std::vector aptr; - std::vector bptr; - std::vector cptr; + if(batchDimA == batchDimB) { + setTensorMode(cublasHandle); + TypedGemm::stridedBatchedGemm(cublasHandle, compute, + opB, opA, + n, m, k, + &alpha, + B->data(), ldb, strideB, + A->data(), lda, strideA, + &beta, + C->data(), ldc, strideC, + batchDimA); + unsetTensorMode(cublasHandle); + } else { + auto batchDimC = std::max(batchDimA, batchDimB); + size_t size = 3*batchDimC; + std::vector ptrs(size); + auto aStart = 0; + auto bStart = batchDimC; + auto cStart = bStart + batchDimC; + + for(int i = 0; i < batchDimC; i++) { + ptrs[aStart + i] = A->data() + (i % batchDimA) * strideA; + ptrs[bStart + i] = B->data() + (i % batchDimB) * strideB; + ptrs[cStart + i] = C->data() + i * strideC; + } - for(int i = 0; i < batchC; i++) { - aptr.push_back(A->data() + (i % batchA) * strideA); - bptr.push_back(B->data() + (i % batchB) * strideB); - cptr.push_back(C->data() + i * strideC); + // auto fails here from weird reason + IPtr mp_ptrs = allocator->alloc(size); + ElementType** dest = mp_ptrs->data(); + cudaStream_t cublasStream = 0; + CUBLAS_CHECK(cublasGetStream(cublasHandle, &cublasStream)); + CUDA_CHECK(cudaMemcpyAsync(dest, ptrs.data(), size * sizeof(ElementType*), cudaMemcpyHostToDevice, cublasStream)); + + setTensorMode(cublasHandle); + TypedGemm::batchedGemm(cublasHandle, compute, + opB, opA, + n, m, k, + &alpha, + mp_ptrs->data() + bStart, ldb, + mp_ptrs->data() + aStart, lda, + &beta, + mp_ptrs->data() + cStart, ldc, + batchDimC); + unsetTensorMode(cublasHandle); + + allocator->free(mp_ptrs); } - - // auto fails here from weird reason - IPtr mp_aptr = allocator->alloc(aptr.size()); - CudaCopy(aptr.data(), aptr.data() + aptr.size(), mp_aptr->data()); - - IPtr mp_bptr = allocator->alloc(bptr.size()); - CudaCopy(bptr.data(), bptr.data() + bptr.size(), mp_bptr->data()); - - IPtr mp_cptr = allocator->alloc(cptr.size()); - CudaCopy(cptr.data(), cptr.data() + cptr.size(), mp_cptr->data()); - - setTensorMode(cublasHandle); - TypedGemm::batchedGemm(cublasHandle, compute, - opB, opA, - n, m, k, - &alpha, - mp_bptr->data(), ldb, - mp_aptr->data(), lda, - &beta, - mp_cptr->data(), ldc, - batchC); - unsetTensorMode(cublasHandle); - - allocator->free(mp_aptr); - allocator->free(mp_bptr); - allocator->free(mp_cptr); } // @TODO: add version with compute type for completeness diff --git a/src/tests/units/operator_tests.cpp b/src/tests/units/operator_tests.cpp index 27ccf1396..1988a5eef 100644 --- a/src/tests/units/operator_tests.cpp +++ b/src/tests/units/operator_tests.cpp @@ -1,7 +1,3 @@ -/* All or part of this file was contributed by NVIDIA under license: - * Copyright (C) 2020 NVIDIA Corporation - * SPDX-License-Identifier: MIT - */ #include "catch.hpp" #include "graph/expression_graph.h" #include "graph/expression_operators.h"