saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[DML EP] Add frequent upload heap flushing (#15960)

Browse source 
This reduces peak nonlocal memory consumption when uploading large
weights for big models (e.g. LLMs), while at the same time trying to
keep the GPU as busy as possible. This change could be more
sophisticated, but at this stage it is the most minimal and least risky
change required to support LLMs.
This commit is contained in:
Patrice Vignola 2023-05-16 22:35:38 -07:00 committed by GitHub
parent 270c09a37f
commit 0ff915eba8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 18 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
View file

@ -210,6 +210,8 @@ namespace Dml
m_cpuOutputAllocator = std::make_shared<CPUAllocator>(OrtMemType::OrtMemTypeCPUOutput);
CreateDmlKernelRegistry(&m_kernelRegistry, &m_internalRegInfoMap);
m_lastUploadFlushTime = std::chrono::steady_clock::now();
}
HRESULT __stdcall ExecutionProviderImpl::GetD3DDevice(_COM_Outptr_ ID3D12Device** d3dDevice) const noexcept
@ -447,6 +449,7 @@ namespace Dml
const auto dstState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS; // GPU resources are always kept in UAV state
m_uploadHeap->BeginUploadToGpu(dstData, dstOffset, dstState, AsByteSpan(srcData, dataSizeInBytes));
FlushUploadsIfReady();
}
else if (!src->IsCpuData() && dst->IsCpuData())
{
@ -566,12 +569,23 @@ namespace Dml
assert(!m_closed);
m_uploadHeap->BeginUploadToGpu(dstData, 0, D3D12_RESOURCE_STATE_UNORDERED_ACCESS, AsByteSpan(srcData, static_cast<size_t>(srcDataSize)));
FlushUploadsIfReady();
return S_OK;
}
ORT_CATCH_RETURN
}
void ExecutionProviderImpl::FlushUploadsIfReady() const
{
// Periodically flush uploads to make sure the GPU is not idle for too long
if (std::chrono::steady_clock::now() - m_lastUploadFlushTime > m_batchFlushInterval)
{
Flush();
m_lastUploadFlushTime = std::chrono::steady_clock::now();
}
}
uint32_t ExecutionProviderImpl::GetSupportedDeviceDataTypeMask() const
{
// The DML provider registers all supported kernels up-front regardless of actual device capability,

4
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h
View file

@ -180,6 +180,8 @@ namespace Dml
uint32_t supportedDeviceDataTypeMask // Each bit corresponds to each DML_TENSOR_DATA_TYPE.
) const;
void FlushUploadsIfReady() const;
ComPtr<ID3D12Device> m_d3d12Device;
ComPtr<IDMLDevice> m_dmlDevice;
bool m_isMcdmDevice = false;
@ -195,6 +197,8 @@ namespace Dml
std::shared_ptr<const Windows::AI::MachineLearning::Adapter::InternalRegistrationInfoMap> m_internalRegInfoMap;
mutable uint64_t m_partitionKernelPrefixVal = 0;
bool m_closed = false;
mutable std::chrono::time_point<std::chrono::steady_clock> m_lastUploadFlushTime;
static constexpr std::chrono::milliseconds m_batchFlushInterval = std::chrono::milliseconds(10);
};
class DataTransfer : public onnxruntime::IDataTransfer