diff --git a/.jenkins/validate_tutorials_built.py b/.jenkins/validate_tutorials_built.py index 3ed1e0c028..181b107102 100644 --- a/.jenkins/validate_tutorials_built.py +++ b/.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", diff --git a/prototype_source/gpu_direct_storage.py b/prototype_source/gpu_direct_storage.py new file mode 100644 index 0000000000..b02ee02261 --- /dev/null +++ b/prototype_source/gpu_direct_storage.py @@ -0,0 +1,132 @@ +""" +(prototype) Accelerating ``torch.save`` and ``torch.load`` with GPUDirect Storage +================================================================================= + +GPUDirect Storage enables a direct data path for direct memory access transfers +between GPU memory and storage, avoiding a bounce buffer through the CPU. + +In version **2.7**, we introduced new prototype APIs to ``torch.cuda.gds`` that serve as thin wrappers around +the `cuFile APIs `_ +that can be used with ``torch.Tensor`` to achieve improved I/O performance. + +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 `_ + * 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. You can toggle this by using +# ``torch.utils.serialization.config.save.storage_alignment``: + +import torch +from torch.utils.serialization import config as serialization_config + +serialization_config.save.storage_alignment = 4096 + +################################################################################ +# The steps involved in the process are as follows: +# * Write the checkpoint file without any actual data. This reserves the space on disk. +# * Read the offsets for the storage associated with each tensor in the checkpoint using ``FakeTensor``. +# * Use ``GDSFile`` to write the appropriate data at these offsets. +# +# 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 (such as sizes, strides, dtype, device) +# information about the tensor but does not have any storage bytes. The following snippet will not materialize +# any data but will tag each ``FakeTensor`` with the offset within the checkpoint that +# corresponds to the tensor. +# +# If you are continuously saving the same state dictionary during training, you +# would only need to obtain the offsets once and the same offsets can be re-used. Similarly if tensor is going to +# be saved or loaded to repeatedly you can use the ``torch.cuda.gds.gds_register_buffer`` which wraps +# ``cuFileBufRegister`` to register the storages as GDS buffers. +# +# Note that ``torch.cuda.gds.GdsFile.save_storage`` binds to the synchronous ``cuFileWrite`` API, +# so no synchronization is needed afterwards. + + +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 + # save_storage is a wrapper around `cuFileWrite` + 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: you can use ``torch.load`` with 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 (as if the tensors were created via ``torch.empty``). + +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. +# +# Similar to ``torch.cuda.gds.GdsFile.save_storage``, ``torch.cuda.gds.GdsFile.load_storage`` +# binds to the synchronous ``cuFileRead`` API, so no synchronization is needed afterwards. + +for k, v in sd_loaded.items(): + assert not torch.equal(v, sd[k]) + offset = fake_sd[k].untyped_storage()._checkpoint_offset + # load_storage is a wrapper around `cuFileRead` + f.load_storage(v.untyped_storage(), offset) + +for k, v in sd_loaded.items(): + assert torch.equal(v, sd[k]) + +del f +########################################################## +# Conclusion +# ========== +# +# 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. Please +# file an issue in the PyTorch GitHub repo if you have any feedback. diff --git a/prototype_source/prototype_index.rst b/prototype_source/prototype_index.rst index a0f7706c61..4c9a9655e5 100644 --- a/prototype_source/prototype_index.rst +++ b/prototype_source/prototype_index.rst @@ -261,6 +261,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