diff --git a/src/Tests.cpp b/src/Tests.cpp
index bf2a317..d4749f6 100644
--- a/src/Tests.cpp
+++ b/src/Tests.cpp
@@ -4935,9 +4935,12 @@ static void TestBudget()
 {
     wprintf(L"Testing budget...\n");
 
-    static const VkDeviceSize BUF_SIZE = 100ull * 1024 * 1024;
+    static const VkDeviceSize BUF_SIZE = 10ull * 1024 * 1024;
     static const uint32_t BUF_COUNT = 4;
 
+    const VkPhysicalDeviceMemoryProperties* memProps = {};
+    vmaGetMemoryProperties(g_hAllocator, &memProps);
+
     for(uint32_t testIndex = 0; testIndex < 2; ++testIndex)
     {
         vmaSetCurrentFrameIndex(g_hAllocator, ++g_FrameIndex);
@@ -4945,8 +4948,10 @@ static void TestBudget()
         VmaBudget budgetBeg[VK_MAX_MEMORY_HEAPS] = {};
         vmaGetBudget(g_hAllocator, budgetBeg);
 
-        for(uint32_t i = 0; i < VK_MAX_MEMORY_HEAPS; ++i)
+        for(uint32_t i = 0; i < memProps->memoryHeapCount; ++i)
         {
+            TEST(budgetBeg[i].budget > 0);
+            TEST(budgetBeg[i].budget <= memProps->memoryHeaps[i].size);
             TEST(budgetBeg[i].allocationBytes <= budgetBeg[i].blockBytes);
         }
 
@@ -4994,7 +4999,7 @@ static void TestBudget()
         vmaGetBudget(g_hAllocator, budgetEnd);
 
         // CHECK
-        for(uint32_t i = 0; i < VK_MAX_MEMORY_HEAPS; ++i)
+        for(uint32_t i = 0; i < memProps->memoryHeapCount; ++i)
         {
             TEST(budgetEnd[i].allocationBytes <= budgetEnd[i].blockBytes);
             if(i == heapIndex)
diff --git a/src/VulkanSample.cpp b/src/VulkanSample.cpp
index 39961ec..ccab048 100644
--- a/src/VulkanSample.cpp
+++ b/src/VulkanSample.cpp
@@ -1163,7 +1163,8 @@ void SetAllocatorCreateInfo(VmaAllocatorCreateInfo& outInfo)
         outInfo.flags |= VMA_ALLOCATOR_CREATE_KHR_BIND_MEMORY2_BIT;
     }
 #if !defined(VMA_MEMORY_BUDGET) || VMA_MEMORY_BUDGET == 1
-    if(VK_EXT_memory_budget_enabled && VK_KHR_get_physical_device_properties2_enabled)
+    if(VK_EXT_memory_budget_enabled && (
+        GetVulkanApiVersion() >= VK_API_VERSION_1_1 || VK_KHR_get_physical_device_properties2_enabled))
     {
         outInfo.flags |= VMA_ALLOCATOR_CREATE_EXT_MEMORY_BUDGET_BIT;
     }
diff --git a/src/vk_mem_alloc.h b/src/vk_mem_alloc.h
index 7d5621f..6f5294b 100644
--- a/src/vk_mem_alloc.h
+++ b/src/vk_mem_alloc.h
@@ -16723,6 +16723,20 @@ void VmaAllocator_T::UpdateVulkanBudget()
             m_Budget.m_VulkanUsage[heapIndex] = budgetProps.heapUsage[heapIndex];
             m_Budget.m_VulkanBudget[heapIndex] = budgetProps.heapBudget[heapIndex];
             m_Budget.m_BlockBytesAtBudgetFetch[heapIndex] = m_Budget.m_BlockBytes[heapIndex].load();
+
+            // Some bugged drivers return the budget incorrectly, e.g. 0 or much bigger than heap size.
+            if(m_Budget.m_VulkanBudget[heapIndex] == 0)
+            {
+                m_Budget.m_VulkanBudget[heapIndex] = m_MemProps.memoryHeaps[heapIndex].size * 8 / 10; // 80% heuristics.
+            }
+            else if(m_Budget.m_VulkanBudget[heapIndex] > m_MemProps.memoryHeaps[heapIndex].size)
+            {
+                m_Budget.m_VulkanBudget[heapIndex] = m_MemProps.memoryHeaps[heapIndex].size;
+            }
+            if(m_Budget.m_VulkanUsage[heapIndex] == 0 && m_Budget.m_BlockBytesAtBudgetFetch[heapIndex] > 0)
+            {
+                m_Budget.m_VulkanUsage[heapIndex] = m_Budget.m_BlockBytesAtBudgetFetch[heapIndex];
+            }
         }
         m_Budget.m_OperationsSinceBudgetFetch = 0;
     }