GPUDirect Storage prototype tutorial

Author: mikaylagawarecki

.jenkins/validate_tutorials_built.py

@@ -31,6 +31,7 @@
     "prototype_source/vmap_recipe",
     "prototype_source/torchscript_freezing",
     "prototype_source/nestedtensor",
+    "prototype_source/gpu_direct_storage", # requires specific filesystem + GPUDirect Storage to be set up
     "recipes_source/recipes/saving_and_loading_models_for_inference",
     "recipes_source/recipes/saving_multiple_models_in_one_file",
     "recipes_source/recipes/tensorboard_with_pytorch",

prototype_source/gpu_direct_storage.py

@@ -0,0 +1,124 @@
+"""
+(prototype) Using GPUDirect Storage
+====================================
+
+GPUDirect Storage enabes a direct data path for direct memeory access transfers
+between GPU memory and storage, avoiding a bounce buffer through the CPU.
+
+In version ``2.7``, we introduced some prototype APIs to ``torch.cuda.gds`` that serve as thin wrappers around
+the `cuFile APIs <https://docs.nvidia.com/gpudirect-storage/api-reference-guide/index.html#cufile-io-api>`_
+that can be used with ``torch.Tensor``.
+
+In this tutorial, we will demonstrate how to use the ``torch.cuda.gds`` APIs in conjunction with
+checkpoints generated by ``torch.save`` and ``torch.load`` on local filesystem. 
+
+.. grid:: 2
+
+    .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn
+       :class-card: card-prerequisites
+
+       * Understand how to use the ``torch.cuda.gds`` APIs in conjunction with
+         checkpoints generated by ``torch.save`` and ``torch.load`` on local filesystem
+    
+    .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites
+       :class-card: card-prerequisites
+
+       * PyTorch v.2.7.0 or later
+       * GPUDirect Storage must be installed per
+         `the documentation <https://docs.nvidia.com/gpudirect-storage/troubleshooting-guide/contents.html>`_
+       * Ensure that the filesystem that you are saving/loading to supports GPUDirect Storage.
+"""
+
+################################################################################
+# Using GPUDirect Storage with ``torch.save`` and ``torch.load``
+# =============================================================
+# GPUDirect Storage requires a storage alignment of 4KB. One can toggle this using
+# ``torch.utils.serialization.config.save.storage_alignment`` to toggle this
+
+import torch
+from torch.utils.serialization import config as serialization_config
+
+serialization_config.save.storage_alignment = 4096
+
+################################################################################
+# Given a state dictionary of tensors that are on the GPU, one can use the ``torch.serialization.skip_data`` context
+# manager to save a checkpoint that contains all relevant metadata except the storage bytes. For each ``torch.Storage``
+# in the state dictionary, space will be reserved within the checkpoint for the storage bytes.
+
+import torch.nn as nn
+
+m = nn.Linear(5, 10, device='cuda')
+sd = m.state_dict()
+
+with torch.serialization.skip_data():
+    torch.save(sd, "checkpoint.pt")
+
+################################################################################
+# We can get the offsets that each storage should be written to within the checkpoint by loading under
+# a ``FakeTensorMode``. A FakeTensor is a tensor that has metadata (e.g. sizes, strides, dtype, device)
+# information about the tensor but does not have any storage bytes. The following snippet will not materialize
+# any data but which will tag each ``FakeTensor`` with the offset within the checkpoint that
+# corresponds to the tensor.
+
+import os
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+with FakeTensorMode() as mode:
+    fake_sd = torch.load("checkpoint.pt")
+
+for k, v in fake_sd.items():
+    print(f"key={k}, offset={v.untyped_storage()._checkpoint_offset}")
+
+f = torch.cuda.gds.GdsFile("checkpoint.pt", os.O_RDWR)
+
+for k, v in sd.items():
+    offset = fake_sd[k].untyped_storage()._checkpoint_offset
+    f.save_storage(v.untyped_storage(), offset)
+
+################################################################################
+# We verify correctness of the saved checkpoint by ``torch.load`` and comparing.
+
+sd_loaded = torch.load("checkpoint.pt")
+for k, v in sd_loaded.items():
+    assert torch.equal(v, sd[k])
+
+################################################################################
+# The loading flow is the inverse, we can ``torch.load`` under the ``torch.serialization.skip_data`` context
+# manager to load everything except the storage bytes. This means that any tensors in the checkpoint will be
+# created but their storages will be empty (i.e. the tensors will be created via ``torch.empty``). If the
+# tensors to be loaded to are persistent, one can use the ``torch.cuda.gds.gds_register_buffer`` API to register
+# the storages as gds buffers.
+
+with torch.serialization.skip_data():
+    sd_loaded = torch.load("checkpoint.pt")
+
+################################################################################
+# We once again use the ``FakeTensorMode`` to get the checkpoint offsets and
+# ascertain that the loaded checkpoint is the same as the saved checkpoint.
+
+for k, v in sd_loaded.items():
+    assert not torch.equal(v, sd[k])
+    offset = fake_sd[k].untyped_storage()._checkpoint_offset
+    f.load_storage(v.untyped_storage(), offset)
+    assert torch.equal(v, sd[k])
+
+del f
+
+
+################################################################################
+# Buffer Registration
+# ===================
+# We also provide ``torch.cuda.gds.gds_register_buffer`` to register the
+# tensor storages as GPUDirect Storage buffers. See `here
+# <https://docs.nvidia.com/gpudirect-storage/best-practices-guide/index.html#cufile-bufregister-fileread-filewrite>`_
+# for when one should do this.
+
+for v in sd.values():
+    torch.cuda.gds.gds_register_buffer(v.untyped_storage())
+
+# Summary
+# =======
+#
+# In this tutorial we have demonstrated how to use the prototype ``torch.cuda.gds`` APIs
+# in conjunction with ``torch.save`` and ``torch.load`` on local filesystem. Do
+# file in issue in the PyTorch GitHub repo if you have any feedback.

prototype_source/prototype_index.rst

@@ -247,6 +247,14 @@ Prototype features are not available as part of binary distributions like PyPI o
    :link: ../prototype/python_extension_autoload.html
    :tags: Extending-PyTorch, Frontend-APIs
 
+.. GPUDirect Storage
+.. customcarditem::
+   :header: (prototype) Using GPUDirect Storage
+   :card_description: Learn how to use GPUDirect Storage in PyTorch.
+   :image: ../_static/img/thumbnails/cropped/generic-pytorch-logo.png
+   :link: ../prototype/gpudirect_storage.html
+   :tags: GPUDirect-Storage
+
 .. End of tutorial card section
 
 .. raw:: html