rename cuda_mem_limit and hip_mem_limit to gpu_mem_limit for both CUDA EP and ROCm EP (#7226)

With this change, differentiating CUDA EP and ROCm EP is not needed in training script when mem_limit option needs to be set.

Co-authored-by: Weixing Zhang <wezhan@microsoft.com>
This commit is contained in:
Weixing Zhang 2021-04-05 09:04:04 -07:00 committed by GitHub
parent 68b12a6179
commit 74ee24cf7f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 60 additions and 64 deletions

View file

@ -10,7 +10,7 @@ extern "C" {
/**
* \param device_id hip device id, starts from zero.
*/
ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_ROCM, _In_ OrtSessionOptions* options, int device_id, size_t hip_mem_limit);
ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_ROCM, _In_ OrtSessionOptions* options, int device_id, size_t gpu_mem_limit);
#ifdef __cplusplus
}

View file

@ -268,7 +268,7 @@ typedef enum OrtCudnnConvAlgoSearch {
typedef struct OrtCUDAProviderOptions {
int device_id; // cuda device with id=0 as default device.
OrtCudnnConvAlgoSearch cudnn_conv_algo_search; // cudnn conv algo search option
size_t cuda_mem_limit; // default cuda memory limitation to maximum finite value of size_t.
size_t gpu_mem_limit; // default cuda memory limitation to maximum finite value of size_t.
int arena_extend_strategy; // default area extend strategy to KNextPowerOfTwo.
int do_copy_in_default_stream;
int has_user_compute_stream;
@ -281,7 +281,7 @@ typedef struct OrtCUDAProviderOptions {
typedef struct OrtROCMProviderOptions {
int device_id; // hip device with id=0 as default device.
int miopen_conv_exhaustive_search; // miopen conv algo exhaustive search option
size_t hip_mem_limit; // default hip memory limitation to maximum finite value of size_t.
size_t gpu_mem_limit; // default hip memory limitation to maximum finite value of size_t.
int arena_extend_strategy; // default area extend strategy to KNextPowerOfTwo.
} OrtROCMProviderOptions;

View file

@ -59,7 +59,7 @@ ONNX_OPERATOR_KERNEL_EX(
} // namespace cuda
AllocatorPtr CUDAExecutionProvider::CreateCudaAllocator(OrtDevice::DeviceId device_id, size_t cuda_mem_limit, ArenaExtendStrategy arena_extend_strategy,
AllocatorPtr CUDAExecutionProvider::CreateCudaAllocator(OrtDevice::DeviceId device_id, size_t gpu_mem_limit, ArenaExtendStrategy arena_extend_strategy,
CUDAExecutionProviderExternalAllocatorInfo external_allocator_info) {
if (external_allocator_info.UseExternalAllocator()) {
AllocatorCreationInfo default_memory_info(
@ -78,7 +78,7 @@ AllocatorPtr CUDAExecutionProvider::CreateCudaAllocator(OrtDevice::DeviceId devi
},
device_id,
true,
{cuda_mem_limit,
{gpu_mem_limit,
static_cast<int>(arena_extend_strategy),
-1, -1});
@ -87,7 +87,7 @@ AllocatorPtr CUDAExecutionProvider::CreateCudaAllocator(OrtDevice::DeviceId devi
}
}
CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t cuda_mem_limit,
CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t gpu_mem_limit,
ArenaExtendStrategy arena_extend_strategy, CUDAExecutionProviderExternalAllocatorInfo external_allocator_info) {
CUDA_CALL_THROW(cudaSetDevice(device_id));
stream_ = stream;
@ -99,7 +99,7 @@ CUDAExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId de
CUDNN_CALL_THROW(cudnnSetStream(cudnn_handle_, stream));
// CUDA malloc/free is expensive so always use an arena
allocator_ = CreateCudaAllocator(device_id, cuda_mem_limit, arena_extend_strategy, external_allocator_info);
allocator_ = CreateCudaAllocator(device_id, gpu_mem_limit, arena_extend_strategy, external_allocator_info);
}
CUDAExecutionProvider::PerThreadContext::~PerThreadContext() {
@ -198,7 +198,7 @@ CUDAExecutionProvider::PerThreadContext& CUDAExecutionProvider::GetPerThreadCont
// get or create a context
if (context_state_.retired_context_pool.empty()) {
context = std::make_shared<PerThreadContext>(info_.device_id, static_cast<cudaStream_t>(GetComputeStream()), info_.cuda_mem_limit, info_.arena_extend_strategy, info_.external_allocator_info);
context = std::make_shared<PerThreadContext>(info_.device_id, static_cast<cudaStream_t>(GetComputeStream()), info_.gpu_mem_limit, info_.arena_extend_strategy, info_.external_allocator_info);
} else {
context = context_state_.retired_context_pool.back();
context_state_.retired_context_pool.pop_back();
@ -2046,7 +2046,7 @@ void CUDAExecutionProvider::RegisterAllocator(std::shared_ptr<AllocatorManager>
// Used to allocate CUDA device memory
auto cuda_alloc = allocator_manager->GetAllocator(info_.device_id, OrtMemTypeDefault);
if (nullptr == cuda_alloc) {
cuda_alloc = CreateCudaAllocator(info_.device_id, info_.cuda_mem_limit, info_.arena_extend_strategy, info_.external_allocator_info);
cuda_alloc = CreateCudaAllocator(info_.device_id, info_.gpu_mem_limit, info_.arena_extend_strategy, info_.external_allocator_info);
allocator_manager->InsertAllocator(cuda_alloc);
}
TryInsertAllocator(cuda_alloc);

View file

@ -80,7 +80,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
}
void RegisterAllocator(std::shared_ptr<AllocatorManager> allocator_manager) override;
static AllocatorPtr CreateCudaAllocator(OrtDevice::DeviceId device_id, size_t cuda_mem_limit, ArenaExtendStrategy arena_extend_strategy,
static AllocatorPtr CreateCudaAllocator(OrtDevice::DeviceId device_id, size_t gpu_mem_limit, ArenaExtendStrategy arena_extend_strategy,
CUDAExecutionProviderExternalAllocatorInfo external_alloc_info);
private:
@ -99,7 +99,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
class PerThreadContext final {
public:
PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t cuda_mem_limit, ArenaExtendStrategy arena_extend_strategy,
PerThreadContext(OrtDevice::DeviceId device_id, cudaStream_t stream, size_t gpu_mem_limit, ArenaExtendStrategy arena_extend_strategy,
CUDAExecutionProviderExternalAllocatorInfo external_alloc_info);
~PerThreadContext();

View file

@ -12,7 +12,7 @@ namespace onnxruntime {
namespace cuda {
namespace provider_option_names {
constexpr const char* kDeviceId = "device_id";
constexpr const char* kMemLimit = "cuda_mem_limit";
constexpr const char* kMemLimit = "gpu_mem_limit";
constexpr const char* kArenaExtendStrategy = "arena_extend_strategy";
constexpr const char* kCudnnConvAlgoSearch = "cudnn_conv_algo_search";
constexpr const char* kDoCopyInDefaultStream = "do_copy_in_default_stream";
@ -70,7 +70,7 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P
free = reinterpret_cast<void*>(address);
return Status::OK();
})
.AddAssignmentToReference(cuda::provider_option_names::kMemLimit, info.cuda_mem_limit)
.AddAssignmentToReference(cuda::provider_option_names::kMemLimit, info.gpu_mem_limit)
.AddAssignmentToEnumReference(
cuda::provider_option_names::kArenaExtendStrategy,
arena_extend_strategy_mapping, info.arena_extend_strategy)
@ -88,7 +88,7 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P
ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const CUDAExecutionProviderInfo& info) {
const ProviderOptions options{
{cuda::provider_option_names::kDeviceId, MakeStringWithClassicLocale(info.device_id)},
{cuda::provider_option_names::kMemLimit, MakeStringWithClassicLocale(info.cuda_mem_limit)},
{cuda::provider_option_names::kMemLimit, MakeStringWithClassicLocale(info.gpu_mem_limit)},
{cuda::provider_option_names::kcudaExternalAlloc, MakeStringWithClassicLocale(reinterpret_cast<size_t>(info.external_allocator_info.alloc))},
{cuda::provider_option_names::kcudaExternalFree, MakeStringWithClassicLocale(reinterpret_cast<size_t>(info.external_allocator_info.free))},
{cuda::provider_option_names::kArenaExtendStrategy,

View file

@ -33,7 +33,7 @@ struct CUDAExecutionProviderExternalAllocatorInfo {
struct CUDAExecutionProviderInfo {
OrtDevice::DeviceId device_id{0};
size_t cuda_mem_limit{std::numeric_limits<size_t>::max()};
size_t gpu_mem_limit{std::numeric_limits<size_t>::max()};
ArenaExtendStrategy arena_extend_strategy{ArenaExtendStrategy::kNextPowerOfTwo};
OrtCudnnConvAlgoSearch cudnn_conv_algo_search{OrtCudnnConvAlgoSearch::EXHAUSTIVE};
bool do_copy_in_default_stream{true};

View file

@ -53,7 +53,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_CUDA,
_In_ OrtSessionOptions* options, _In_ const OrtCUDAProviderOptions* cuda_options) {
CUDAExecutionProviderInfo info{};
info.device_id = gsl::narrow<OrtDevice::DeviceId>(cuda_options->device_id);
info.cuda_mem_limit = cuda_options->cuda_mem_limit;
info.gpu_mem_limit = cuda_options->gpu_mem_limit;
info.arena_extend_strategy = static_cast<onnxruntime::ArenaExtendStrategy>(cuda_options->arena_extend_strategy);
info.cudnn_conv_algo_search = cuda_options->cudnn_conv_algo_search;
info.do_copy_in_default_stream = cuda_options->do_copy_in_default_stream;

View file

@ -58,7 +58,7 @@ ONNX_OPERATOR_KERNEL_EX(
} // namespace rocm
ROCMExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, hipStream_t stream, size_t hip_mem_limit, ArenaExtendStrategy arena_extend_strategy) {
ROCMExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId device_id, hipStream_t stream, size_t gpu_mem_limit, ArenaExtendStrategy arena_extend_strategy) {
HIP_CALL_THROW(hipSetDevice(device_id));
stream_ = stream;
@ -73,7 +73,7 @@ ROCMExecutionProvider::PerThreadContext::PerThreadContext(OrtDevice::DeviceId de
},
device_id,
true,
{hip_mem_limit,
{gpu_mem_limit,
static_cast<int>(arena_extend_strategy),
-1, -1});
@ -128,7 +128,7 @@ ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& in
},
info_.device_id,
true,
{info_.hip_mem_limit,
{info_.gpu_mem_limit,
static_cast<int>(info_.arena_extend_strategy),
-1, -1});
@ -206,7 +206,7 @@ ROCMExecutionProvider::PerThreadContext& ROCMExecutionProvider::GetPerThreadCont
// get or create a context
if (context_state_.retired_context_pool.empty()) {
context = std::make_shared<PerThreadContext>(info_.device_id, static_cast<hipStream_t>(GetComputeStream()), info_.hip_mem_limit, info_.arena_extend_strategy);
context = std::make_shared<PerThreadContext>(info_.device_id, static_cast<hipStream_t>(GetComputeStream()), info_.gpu_mem_limit, info_.arena_extend_strategy);
} else {
context = context_state_.retired_context_pool.back();
context_state_.retired_context_pool.pop_back();

View file

@ -105,7 +105,7 @@ class ROCMExecutionProvider : public IExecutionProvider {
class PerThreadContext final {
public:
PerThreadContext(OrtDevice::DeviceId device_id, hipStream_t stream, size_t hip_mem_limit, ArenaExtendStrategy arena_extend_strategy);
PerThreadContext(OrtDevice::DeviceId device_id, hipStream_t stream, size_t gpu_mem_limit, ArenaExtendStrategy arena_extend_strategy);
~PerThreadContext();
rocblas_handle RocblasHandle() const {

View file

@ -10,7 +10,7 @@ namespace onnxruntime {
namespace rocm {
namespace provider_option_names {
constexpr const char* kDeviceId = "device_id";
constexpr const char* kMemLimit = "hip_mem_limit";
constexpr const char* kMemLimit = "gpu_mem_limit";
constexpr const char* kArenaExtendStrategy = "arena_extend_strategy";
constexpr const char* kConvExhaustiveSearch = "conv_exhaustive_search";
} // namespace provider_option_names
@ -30,7 +30,7 @@ ROCMExecutionProviderInfo ROCMExecutionProviderInfo::FromProviderOptions(const P
ProviderOptionsParser{}
// TODO validate info.device_id
.AddAssignmentToReference(rocm::provider_option_names::kDeviceId, info.device_id)
.AddAssignmentToReference(rocm::provider_option_names::kMemLimit, info.hip_mem_limit)
.AddAssignmentToReference(rocm::provider_option_names::kMemLimit, info.gpu_mem_limit)
.AddAssignmentToReference(rocm::provider_option_names::kConvExhaustiveSearch, info.miopen_conv_exhaustive_search)
.AddAssignmentToEnumReference(
rocm::provider_option_names::kArenaExtendStrategy,
@ -43,7 +43,7 @@ ROCMExecutionProviderInfo ROCMExecutionProviderInfo::FromProviderOptions(const P
ProviderOptions ROCMExecutionProviderInfo::ToProviderOptions(const ROCMExecutionProviderInfo& info) {
const ProviderOptions options{
{rocm::provider_option_names::kDeviceId, MakeStringWithClassicLocale(info.device_id)},
{rocm::provider_option_names::kMemLimit, MakeStringWithClassicLocale(info.hip_mem_limit)},
{rocm::provider_option_names::kMemLimit, MakeStringWithClassicLocale(info.gpu_mem_limit)},
{rocm::provider_option_names::kConvExhaustiveSearch, MakeStringWithClassicLocale(info.miopen_conv_exhaustive_search)},
{rocm::provider_option_names::kArenaExtendStrategy,
EnumToName(arena_extend_strategy_mapping, info.arena_extend_strategy)},

View file

@ -13,7 +13,7 @@ namespace onnxruntime {
// Information needed to construct HIP execution providers.
struct ROCMExecutionProviderInfo {
OrtDevice::DeviceId device_id{0};
size_t hip_mem_limit{std::numeric_limits<size_t>::max()};
size_t gpu_mem_limit{std::numeric_limits<size_t>::max()};
ArenaExtendStrategy arena_extend_strategy{ArenaExtendStrategy::kNextPowerOfTwo};
bool miopen_conv_exhaustive_search{false};
bool do_copy_in_default_stream{true};

View file

@ -52,7 +52,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_ROCM,
_In_ OrtSessionOptions* options, _In_ const OrtROCMProviderOptions* rocm_options) {
ROCMExecutionProviderInfo info{};
info.device_id = gsl::narrow<OrtDevice::DeviceId>(rocm_options->device_id);
info.hip_mem_limit = rocm_options->hip_mem_limit;
info.gpu_mem_limit = rocm_options->gpu_mem_limit;
info.arena_extend_strategy = static_cast<onnxruntime::ArenaExtendStrategy>(rocm_options->arena_extend_strategy);
info.miopen_conv_exhaustive_search = rocm_options->miopen_conv_exhaustive_search;

View file

@ -153,7 +153,7 @@ onnxruntime::CUDAExecutionProviderExternalAllocatorInfo external_allocator_info{
// TODO remove deprecated global config
OrtDevice::DeviceId cuda_device_id = 0;
// TODO remove deprecated global config
size_t cuda_mem_limit = std::numeric_limits<size_t>::max();
size_t gpu_mem_limit = std::numeric_limits<size_t>::max();
// TODO remove deprecated global config
onnxruntime::ArenaExtendStrategy arena_extend_strategy = onnxruntime::ArenaExtendStrategy::kNextPowerOfTwo;
#endif
@ -458,7 +458,7 @@ static AllocatorPtr GetCudaAllocator(OrtDevice::DeviceId id) {
static std::unordered_map<OrtDevice::DeviceId, AllocatorPtr> id_to_allocator_map;
if (id_to_allocator_map.find(id) == id_to_allocator_map.end()) {
id_to_allocator_map.insert({id, CUDAExecutionProvider::CreateCudaAllocator(id, cuda_mem_limit, arena_extend_strategy, external_allocator_info)});
id_to_allocator_map.insert({id, CUDAExecutionProvider::CreateCudaAllocator(id, gpu_mem_limit, arena_extend_strategy, external_allocator_info)});
}
return id_to_allocator_map[id];
@ -564,7 +564,7 @@ static void RegisterExecutionProviders(InferenceSession* sess, const std::vector
: [&]() {
CUDAExecutionProviderInfo info{};
info.device_id = cuda_device_id;
info.cuda_mem_limit = cuda_mem_limit;
info.gpu_mem_limit = gpu_mem_limit;
info.arena_extend_strategy = arena_extend_strategy;
info.cudnn_conv_algo_search = cudnn_conv_algo_search;
info.do_copy_in_default_stream = do_copy_in_default_stream;
@ -588,7 +588,7 @@ static void RegisterExecutionProviders(InferenceSession* sess, const std::vector
: [&]() {
ROCMExecutionProviderInfo info{};
info.device_id = cuda_device_id;
info.hip_mem_limit = cuda_mem_limit;
info.gpu_mem_limit = gpu_mem_limit;
info.arena_extend_strategy = arena_extend_strategy;
return info;
}();
@ -884,7 +884,7 @@ void addGlobalMethods(py::module& m, Environment& env) {
[&]() {
CUDAExecutionProviderInfo info{};
info.device_id = cuda_device_id;
info.cuda_mem_limit = cuda_mem_limit;
info.gpu_mem_limit = gpu_mem_limit;
info.arena_extend_strategy = arena_extend_strategy;
info.cudnn_conv_algo_search = cudnn_conv_algo_search;
info.do_copy_in_default_stream = do_copy_in_default_stream;
@ -897,7 +897,7 @@ void addGlobalMethods(py::module& m, Environment& env) {
[&]() {
ROCMExecutionProviderInfo info{};
info.device_id = cuda_device_id;
info.hip_mem_limit = cuda_mem_limit;
info.gpu_mem_limit = gpu_mem_limit;
info.arena_extend_strategy = arena_extend_strategy;
return info;
}()),
@ -989,8 +989,8 @@ void addGlobalMethods(py::module& m, Environment& env) {
m.def("set_cuda_mem_limit", [](const int64_t limit) {
LogDeprecationWarning(
"set_cuda_mem_limit",
"CUDA execution provider option \"cuda_mem_limit\", ROCM execution provider option \"hip_mem_limit\"");
cuda_mem_limit = gsl::narrow<size_t>(limit);
"CUDA execution provider option \"gpu_mem_limit\", ROCM execution provider option \"gpu_mem_limit\"");
gpu_mem_limit = gsl::narrow<size_t>(limit);
});
// TODO remove deprecated global config
m.def("set_arena_extend_strategy", [](const onnxruntime::ArenaExtendStrategy strategy) {

View file

@ -82,22 +82,22 @@ class TestInferenceSession(unittest.TestCase):
sess = onnxrt.InferenceSession(get_name("mul_1.onnx"))
self.assertIn('CUDAExecutionProvider', sess.get_providers())
# test get/set of "cuda_mem_limit" configuration.
# test get/set of "gpu_mem_limit" configuration.
options = sess.get_provider_options()
self.assertIn('CUDAExecutionProvider', options)
option = options['CUDAExecutionProvider']
self.assertIn('cuda_mem_limit', option)
ori_mem_limit = option['cuda_mem_limit']
self.assertIn('gpu_mem_limit', option)
ori_mem_limit = option['gpu_mem_limit']
new_mem_limit = int(ori_mem_limit) // 2
option['cuda_mem_limit'] = new_mem_limit
option['gpu_mem_limit'] = new_mem_limit
sess.set_providers(['CUDAExecutionProvider'], [option])
options = sess.get_provider_options()
self.assertEqual(options['CUDAExecutionProvider']['cuda_mem_limit'], str(new_mem_limit))
self.assertEqual(options['CUDAExecutionProvider']['gpu_mem_limit'], str(new_mem_limit))
option['cuda_mem_limit'] = ori_mem_limit
option['gpu_mem_limit'] = ori_mem_limit
sess.set_providers(['CUDAExecutionProvider'], [option])
options = sess.get_provider_options()
self.assertEqual(options['CUDAExecutionProvider']['cuda_mem_limit'], ori_mem_limit)
self.assertEqual(options['CUDAExecutionProvider']['gpu_mem_limit'], ori_mem_limit)
def test_get_and_set_option_with_values(option_name, option_values):
provider_options = sess.get_provider_options()
@ -138,15 +138,15 @@ class TestInferenceSession(unittest.TestCase):
with self.assertRaises(RuntimeError):
sess.set_providers(['CUDAExecutionProvider'], [option])
option['cuda_mem_limit'] = -1024
option['gpu_mem_limit'] = -1024
with self.assertRaises(RuntimeError):
sess.set_providers(['CUDAExecutionProvider'], [option])
option['cuda_mem_limit'] = 1024.1024
option['gpu_mem_limit'] = 1024.1024
with self.assertRaises(RuntimeError):
sess.set_providers(['CUDAExecutionProvider'], [option])
option['cuda_mem_limit'] = 'wrong_value'
option['gpu_mem_limit'] = 'wrong_value'
with self.assertRaises(RuntimeError):
sess.set_providers(['CUDAExecutionProvider'], [option])

View file

@ -64,7 +64,7 @@ struct BertParameters : public TrainingRunner::Parameters {
float initial_lr_phase2;
size_t num_train_steps_phase2;
float warmup_ratio_phase2;
float cuda_mem_limit_in_gb = -1;
float gpu_mem_limit_in_gb = -1;
bool debug_break = false;
PathString train_data_dir_phase2;
PathString test_data_dir_phase2;
@ -181,7 +181,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
cxxopts::value<int64_t>()->default_value("0"))
("ratio_min", "Lamb min ratio parameter", cxxopts::value<float>()->default_value("0.05"))
("ratio_max", "Lamb max ratio parameter", cxxopts::value<float>()->default_value("5.0"))
("cuda_mem_limit_in_gb", "Max cuda memory ort can use, in GB", cxxopts::value<float>()->default_value("-1.0"))
("gpu_mem_limit_in_gb", "Max cuda memory ort can use, in GB", cxxopts::value<float>()->default_value("-1.0"))
("data_parallel_size", "Data parallel group size.", cxxopts::value<int>()->default_value("1"))
("horizontal_parallel_size", "Horizontal model parallel group size.", cxxopts::value<int>()->default_value("1"))
("pipeline_parallel_size", "Number of pipeline stages.", cxxopts::value<int>()->default_value("1"))
@ -238,7 +238,7 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
}
params.lr_params.warmup_ratio = ratio;
params.cuda_mem_limit_in_gb = flags["cuda_mem_limit_in_gb"].as<float>();
params.gpu_mem_limit_in_gb = flags["gpu_mem_limit_in_gb"].as<float>();
float ratio_phase2 = flags["warmup_ratio_phase2"].as<float>();
if (ratio_phase2 > 1.f || ratio_phase2 < 0.f) {
@ -601,8 +601,8 @@ void setup_training_params(BertParameters& params) {
{
CUDAExecutionProviderInfo info{};
info.device_id = gsl::narrow<OrtDevice::DeviceId>(MPIContext::GetInstance().GetLocalRank());
if (params.cuda_mem_limit_in_gb > 0) {
info.cuda_mem_limit = gsl::narrow<size_t>(params.cuda_mem_limit_in_gb * 1024 * 1024 * 1024);
if (params.gpu_mem_limit_in_gb > 0) {
info.gpu_mem_limit = gsl::narrow<size_t>(params.gpu_mem_limit_in_gb * 1024 * 1024 * 1024);
}
info.cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::EXHAUSTIVE;

View file

@ -24,7 +24,7 @@
namespace onnxruntime {
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_CUDA(OrtDevice::DeviceId device_id,
OrtCudnnConvAlgoSearch cudnn_conv_algo_search = OrtCudnnConvAlgoSearch::EXHAUSTIVE,
size_t cuda_mem_limit = std::numeric_limits<size_t>::max(),
size_t gpu_mem_limit = std::numeric_limits<size_t>::max(),
onnxruntime::ArenaExtendStrategy arena_extend_strategy = ArenaExtendStrategy::kNextPowerOfTwo,
bool do_copy_in_default_stream = true);
std::shared_ptr<IExecutionProviderFactory> CreateExecutionProviderFactory_Dnnl(int use_arena);

View file

@ -120,8 +120,8 @@ class TrainingRunner {
VectorString histogram_names;
VectorString norm_names;
//Default value is -1.0f. When cuda_mem_limit_in_gb < 0, ORT can use all cuda memory available.
float cuda_mem_limit_in_gb = -1.0f;
//Default value is -1.0f. When gpu_mem_limit_in_gb < 0, ORT can use all cuda memory available.
float gpu_mem_limit_in_gb = -1.0f;
bool EnableTensorboard() const {
return !is_perf_test && !log_dir.empty() && MPIContext::GetInstance().GetWorldRank() == 0;

View file

@ -689,21 +689,17 @@ class ORTTrainer(object):
providers[providers.index(provider_name)] = (provider_name, provider_options[provider_name])
#default: using cuda
elif 'cuda' in self.options.device.id.lower():
cuda_ep_options = {"device_id": _utils.get_device_index(self.options.device.id)}
cuda_ep_name = ("ROCMExecutionProvider" if self.is_rocm_pytorch else "CUDAExecutionProvider")
gpu_ep_options = {"device_id": _utils.get_device_index(self.options.device.id)}
gpu_ep_name = ("ROCMExecutionProvider" if self.is_rocm_pytorch else "CUDAExecutionProvider")
if self.options.device.mem_limit > 0:
if not self.is_rocm_pytorch:
cuda_ep_options["cuda_mem_limit"] = self.options.device.mem_limit
else:
warnings.warn("Ignoring 'mem_limit' for {}".format(cuda_ep_name))
gpu_ep_options["gpu_mem_limit"] = self.options.device.mem_limit
if cuda_ep_name not in providers:
if gpu_ep_name not in providers:
raise RuntimeError(
"ORTTrainer options specify a CUDA device but the {} provider is unavailable.".format(
cuda_ep_name))
providers[providers.index(cuda_ep_name)] = (cuda_ep_name, cuda_ep_options)
providers[providers.index(gpu_ep_name)] = (gpu_ep_name, gpu_ep_options)
return providers

View file

@ -434,11 +434,11 @@ void build_allreduce_graph(Graph& graph, AllreduceGraphConfigVector& config,
std::unique_ptr<IExecutionProvider> create_cuda_execution_provider() {
CUDAExecutionProviderInfo info;
OrtDevice::DeviceId device_id = static_cast<OrtDevice::DeviceId>(training::MPIContext::GetInstance().GetLocalRank());
size_t cuda_mem_limit = std::numeric_limits<size_t>::max();
cuda_mem_limit = static_cast<size_t>(1 * 1024 * 1024 * 1024);
size_t gpu_mem_limit = std::numeric_limits<size_t>::max();
gpu_mem_limit = static_cast<size_t>(1 * 1024 * 1024 * 1024);
info.device_id = device_id;
info.cuda_mem_limit = cuda_mem_limit;
info.gpu_mem_limit = gpu_mem_limit;
info.arena_extend_strategy = ArenaExtendStrategy::kNextPowerOfTwo;
return onnxruntime::make_unique<CUDAExecutionProvider>(info);
}