在OSS中儲存和訪問檢查點

本文為您介紹如何使用OssCheckpoint直接從OSS中讀寫檢查點（模型訓練過程中儲存的特定時間點的模型狀態）。

前提條件

已安裝並配置OSS Connector for AI/ML。具體操作，請參見安裝OSS Connector for AI/ML和配置OSS Connector for AI/ML。

OssCheckpoint

OssCheckpoint適用於資料訓練過程中對訓練結果進行讀寫需求的情境。

以下樣本展示了如何使用OssCheckpoint來進行Checkpoint的讀取和寫入。

import torch
from osstorchconnector import OssCheckpoint

ENDPOINT = "http://oss-cn-beijing-internal.aliyuncs.com"
CRED_PATH = "/root/.alibabacloud/credentials"
CONFIG_PATH = "/etc/oss-connector/config.json"

#  使用OssCheckpoint建立checkpoint
checkpoint = OssCheckpoint(endpoint=ENDPOINT, cred_path=CRED_PATH, config_path=CONFIG_PATH)

# 讀 checkpoint
CHECKPOINT_READ_URI = "oss://checkpoint/epoch.0"
with checkpoint.reader(CHECKPOINT_READ_URI) as reader:
   state_dict = torch.load(reader)

# 寫 checkpoint
CHECKPOINT_WRITE_URI = "oss://checkpoint/epoch.1"
with checkpoint.writer(CHECKPOINT_WRITE_URI) as writer:
   torch.save(state_dict, writer)

資料類型

通過OssCheckpoint建立的checkpoint對象實現了常用的IO介面。更多資訊，請參見OSS Connector for AI/ML中的資料類型。

參數配置

使用OssCheckpoint時需要進行相應配置，具體配置項說明請參見下表。

參數名	參數類型	是否必選	說明
endpoint	string	是	OSS對外服務的訪問網域名稱。更多資訊，請參見地區和Endpoint。
cred_path	string	是	鑒權檔案預設路徑為`/root/.alibabacloud/credentials`，更多資訊請參見配置訪問憑證。
config_path	string	是	OSS Connector設定檔預設路徑為`/etc/oss-connector/config.json`，更多資訊請參見配置OSS Connector。

分布式檢查點（DCP）

OSS Connector for AI/ML從V1.2.3起支援PyTorch分布式檢查點（DCP）功能。使用OssDCPFileSystem可以直接在OSS上儲存和讀取分布式檢查點。

以下樣本展示了如何使用OssDCPFileSystem來儲存和讀取分布式檢查點。

import torchvision
import torch.distributed.checkpoint as DCP
from osstorchconnector import OssDCPFileSystem
import torch

ENDPOINT = "http://oss-cn-beijing-internal.aliyuncs.com"
CONFIG_PATH = "/etc/oss-connector/config.json"
CRED_PATH = "/root/.alibabacloud/credentials"
OSS_URI = "oss://ossconnectorbucket/dcp-checkpoint-resnet18"

model = torchvision.models.resnet18()

# write to OSS
fs = OssDCPFileSystem(endpoint=ENDPOINT, cred_path=CRED_PATH, config_path=CONFIG_PATH)
oss_storage_writer = fs.writer(OSS_URI)
# DCP.save or DCP.async_save
checkpoint_future = DCP.async_save(
    state_dict=model.state_dict(),
    storage_writer=oss_storage_writer,
)
checkpoint_future.result()


# load from OSS
loaded_state_dict = {
    key: torch.zeros_like(value) for key, value in model.state_dict().items()
}
oss_storage_reader = fs.reader(OSS_URI)
DCP.load(
    loaded_state_dict,
    storage_reader=oss_storage_reader,
)

Safetensors

OSS Connector for AI/ML從V1.2.0rc6起支援safetensors格式。使用OssSafetensor可以直接在OSS上儲存和讀取safetensors檔案。

以下樣本展示了如何使用OssSafetensor來儲存和讀取safetensors檔案。

import torch
from osstorchconnector import OssSafetensor

ENDPOINT = "http://oss-cn-beijing-internal.aliyuncs.com"
CONFIG_PATH = "/etc/oss-connector/config.json"
CRED_PATH = "/root/.alibabacloud/credentials"
OSS_URI = "oss://ossconnectorbucket/safetensors/model.safetensors"

sfts = OssSafetensor(endpoint=ENDPOINT, cred_path=CRED_PATH, config_path=CONFIG_PATH)

# save tensors to safetensor file on OSS
tensors = {"embedding": torch.rand((512, 1024)), "attention": torch.rand((256, 256))}
metadata = {"a": "a", "b": "b"}
sfts.save_file(tensors, OSS_URI, metadata)

# load safetensor file from OSS
loaded_tensors = sfts.load_file(OSS_URI, device="cpu")

# or load tensors by safe_open
with sfts.safe_open(OSS_URI, device ="cpu") as f:
    metadata = f.metadata() # get metadata
    for key in f.keys(): # read tensors by keys
        tensor = f.get_tensor(key)