*hpage = ERR_PTR(-ENOMEM);
                return;
        }
+       *hpage = new_page;
        count_vm_event(THP_COLLAPSE_ALLOC);
 #endif
 
-       if (unlikely(mem_cgroup_newpage_charge(new_page, mm, GFP_KERNEL))) {
-#ifdef CONFIG_NUMA
-               put_page(new_page);
-#endif
+       if (unlikely(mem_cgroup_newpage_charge(new_page, mm, GFP_KERNEL)))
                return;
-       }
 
        /*
         * Prevent all access to pagetables with the exception of
        prepare_pmd_huge_pte(pgtable, mm);
        spin_unlock(&mm->page_table_lock);
 
-#ifndef CONFIG_NUMA
        *hpage = NULL;
-#endif
+
        khugepaged_pages_collapsed++;
 out_up_write:
        up_write(&mm->mmap_sem);
 
 out:
        mem_cgroup_uncharge_page(new_page);
-#ifdef CONFIG_NUMA
-       put_page(new_page);
-#endif
        goto out_up_write;
 }
 
        barrier(); /* write khugepaged_pages_to_scan to local stack */
 
        while (progress < pages) {
-               cond_resched();
-
 #ifndef CONFIG_NUMA
                if (!hpage)
                        hpage = khugepaged_alloc_hugepage(&wait);
                                break;
                        wait = false;
                        khugepaged_alloc_sleep();
+               } else if (hpage) {
+                       put_page(hpage);
+                       hpage = NULL;
                }
 #endif
+               cond_resched();
 
                if (unlikely(kthread_should_stop() || freezing(current)))
                        break;