Object Storage Service:Storing and accessing checkpoints in OSS buckets

Pré-requisitos

O OSS Connector for AI/ML precisa estar instalado e configurado. Para obter mais informações, consulte Instale OSS Connector for AI/ML e Configure OSS Connector for AI/ML.

OssCheckpoint

OssCheckpoint é adequado para cenários que envolvem leitura e gravação de resultados durante o processo de treinamento de dados.

O exemplo a seguir mostra como ler e gravar checkpoints usando OssCheckpoint.

import torch
from osstorchconnector import OssCheckpoint

ENDPOINT = "endpoint"
CRED_PATH = "/root/.alibabacloud/credentials"
CONFIG_PATH = "/etc/oss-connector/config.json"

#  Create a checkpoint by using OssCheckpoint
checkpoint = OssCheckpoint(endpoint=ENDPOINT, cred_path=CRED_PATH, config_path=CONFIG_PATH)

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

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

Tipos de dados

Os objetos de checkpoint criados com OssCheckpoint disponibilizam operações comuns de E/S. Para obter mais informações, consulte data types in OSS Connector for AI/ML.

Parâmetros

A tabela a seguir descreve os parâmetros que precisam ser configurados ao usar OssCheckpoint.

Distributed checkpoint (DCP)

O OSS Connector for AI/ML oferece suporte ao recurso PyTorch Distributed Checkpoint (DCP) a partir da versão V1.2.3. Use OssDCPFileSystem para armazenar e ler checkpoints distribuídos diretamente no OSS.

O exemplo a seguir mostra como usar OssDCPFileSystem para salve e carregar um checkpoint distribuído.

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)
# Use 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

O OSS Connector for AI/ML oferece suporte ao formato safetensors a partir da versão V1.2.0rc6. Use OssSafetensor para armazenar e ler arquivos safetensors diretamente no OSS.

O exemplo a seguir mostra como usar OssSafetensor para salve e carregar arquivos 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 as a safetensors file to 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 a safetensor file from OSS.
loaded_tensors = sfts.load_file(OSS_URI, device="cpu")

# Or load tensors by using 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 key.
        tensor = f.get_tensor(key)