From 68e35585adb8e5b39b4ee7f78d3bbf55146d733a Mon Sep 17 00:00:00 2001
From: janEbert <janpublicebert@posteo.net>
Date: Fri, 25 Oct 2024 19:12:21 +0200
Subject: [PATCH] Rename distributed checkpointing import
`dist_checkpoint` -> `dcp`
The DCP name is more canonical and it makes sense to make users aware of
this for improving their terminology.
---
pytorch-fsdp-example/main.py | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/pytorch-fsdp-example/main.py b/pytorch-fsdp-example/main.py
index 5843373..991e6fc 100644
--- a/pytorch-fsdp-example/main.py
+++ b/pytorch-fsdp-example/main.py
@@ -4,7 +4,7 @@ import os
import time
import torch
-from torch.distributed import checkpoint as dist_checkpoint
+from torch.distributed import checkpoint as dcp
from torch.distributed import fsdp
import torchvision
@@ -123,9 +123,9 @@ def save_model(model, save_dir):
state_dict_config,
):
cp_state_dict = {'model': model.state_dict()}
- dist_checkpoint.save_state_dict(
+ dcp.save_state_dict(
cp_state_dict,
- dist_checkpoint.FileSystemWriter(save_dir),
+ dcp.FileSystemWriter(save_dir),
)
@@ -140,9 +140,9 @@ def load_model(model, load_dir):
state_dict_config,
):
cp_state_dict = {'model': model.state_dict()}
- dist_checkpoint.load_state_dict(
+ dcp.load_state_dict(
cp_state_dict,
- dist_checkpoint.FileSystemReader(load_dir),
+ dcp.FileSystemReader(load_dir),
)
model.load_state_dict(cp_state_dict['model'])
--
GitLab