diff --git a/src/storage/cpu_device_storage.h b/src/storage/cpu_device_storage.h
index 0053392c67c9..ae4a2bcb9331 100644
--- a/src/storage/cpu_device_storage.h
+++ b/src/storage/cpu_device_storage.h
@@ -40,13 +40,12 @@ class CPUDeviceStorage {
  public:
   /*!
    * \brief Aligned allocation on CPU.
-   * \param size Size to allocate.
-   * \return Pointer to the storage.
+   * \param handle Handle struct.
    */
   inline static void* Alloc(Storage::Handle* handle);
   /*!
    * \brief Deallocation.
-   * \param ptr Pointer to deallocate.
+   * \param handle Handle struct.
    */
   inline static void Free(Storage::Handle handle);
 
diff --git a/src/storage/gpu_device_storage.h b/src/storage/gpu_device_storage.h
index 72b1fb3559fb..ac1160f63431 100644
--- a/src/storage/gpu_device_storage.h
+++ b/src/storage/gpu_device_storage.h
@@ -43,13 +43,12 @@ class GPUDeviceStorage {
  public:
   /*!
    * \brief Allocation.
-   * \param size Size to allocate.
-   * \return Pointer to the storage.
+   * \param handle Handle struct.
    */
   inline static void* Alloc(Storage::Handle* handle);
   /*!
    * \brief Deallocation.
-   * \param ptr Pointer to deallocate.
+   * \param handle Handle struct.
    */
   inline static void Free(Storage::Handle handle);
 };  // class GPUDeviceStorage
diff --git a/src/storage/pinned_memory_storage.h b/src/storage/pinned_memory_storage.h
index d84cd0c96554..61820bfe914a 100644
--- a/src/storage/pinned_memory_storage.h
+++ b/src/storage/pinned_memory_storage.h
@@ -38,14 +38,13 @@ class PinnedMemoryStorage {
  public:
   /*!
    * \brief Allocation.
-   * \param size Size to allocate.
-   * \return Pointer to the storage.
+   * \param handle Handle struct.
    */
   inline static void* Alloc(Storage::Handle* handle);
 
   /*!
    * \brief Deallocation.
-   * \param ptr Pointer to deallocate.
+   * \param handle Handle struct.
    */
   inline static void Free(Storage::Handle handle);
 };
diff --git a/src/storage/storage.cc b/src/storage/storage.cc
index 90d7f0dda068..7484e699d388 100644
--- a/src/storage/storage.cc
+++ b/src/storage/storage.cc
@@ -127,7 +127,8 @@ void StorageImpl::Alloc(Storage::Handle* handle) {
 }
 
 void StorageImpl::Free(Storage::Handle handle) {
-  // Do nothing if dtpr is nullptr.
+  // Do nothing if dtpr is nullptr because the handle may have already
+  // been freed or have not been allocated memory yet.
   if (handle.dptr == nullptr) return;
 
   const Context &ctx = handle.ctx;
@@ -143,7 +144,8 @@ void StorageImpl::Free(Storage::Handle handle) {
 }
 
 void StorageImpl::DirectFree(Storage::Handle handle) {
-  // Do nothing if dtpr is nullptr.
+  // Do nothing if dtpr is nullptr because the handle may have already
+  // been freed or have not been allocated memory yet.
   if (handle.dptr == nullptr) return;
 
   const Context &ctx = handle.ctx;
diff --git a/src/storage/storage_manager.h b/src/storage/storage_manager.h
index 15a2c7ecffcb..d17dc91dc2fc 100644
--- a/src/storage/storage_manager.h
+++ b/src/storage/storage_manager.h
@@ -39,20 +39,17 @@ class StorageManager {
  public:
   /*!
    * \brief Allocation.
-   * \param size Size to allocate.
-   * \return Pointer to the storage.
+   * \param handle Handle struct.
    */
   virtual void Alloc(Storage::Handle* handle) = 0;
   /*!
    * \brief Deallocation.
-   * \param ptr Pointer to deallocate.
-   * \param size Size of the storage.
+   * \param handle Handle struct.
    */
   virtual void Free(Storage::Handle handle) = 0;
   /*!
-   * \brief Direct de-allocation.
-   * \param ptr Pointer to deallocate.
-   * \param size Size of the storage.
+   * \brief Direct deallocation.
+   * \param handle Handle struct.
    */
   virtual void DirectFree(Storage::Handle handle) = 0;
   /*!