wehub-resource-sync
569 lines
18 KiB
Python
569 lines
18 KiB
Python
# flake8: noqa
|
|
# isort: skip_file
|
|
|
|
# __pytorch_save_start__
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.optim import Adam
|
|
|
|
import ray.train.torch
|
|
from ray import train
|
|
from ray.train import Checkpoint, ScalingConfig
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
def train_func(config):
|
|
n = 100
|
|
# create a toy dataset
|
|
# data : X - dim = (n, 4)
|
|
# target : Y - dim = (n, 1)
|
|
X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
|
|
Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
|
|
# toy neural network : 1-layer
|
|
# Wrap the model in DDP
|
|
model = ray.train.torch.prepare_model(nn.Linear(4, 1))
|
|
criterion = nn.MSELoss()
|
|
|
|
optimizer = Adam(model.parameters(), lr=3e-4)
|
|
for epoch in range(config["num_epochs"]):
|
|
y = model.forward(X)
|
|
loss = criterion(y, Y)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
metrics = {"loss": loss.item()}
|
|
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
checkpoint = None
|
|
|
|
should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
|
|
# In standard DDP training, where the model is the same across all ranks,
|
|
# only the global rank 0 worker needs to save and report the checkpoint
|
|
if train.get_context().get_world_rank() == 0 and should_checkpoint:
|
|
torch.save(
|
|
model.module.state_dict(), # NOTE: Unwrap the model.
|
|
os.path.join(temp_checkpoint_dir, "model.pt"),
|
|
)
|
|
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
|
|
|
|
train.report(metrics, checkpoint=checkpoint)
|
|
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
train_loop_config={"num_epochs": 5},
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
)
|
|
result = trainer.fit()
|
|
# __pytorch_save_end__
|
|
|
|
# __pytorch_restore_start__
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.optim import Adam
|
|
|
|
import ray.train.torch
|
|
from ray import train
|
|
from ray.train import Checkpoint, ScalingConfig
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
def train_func(config):
|
|
n = 100
|
|
# create a toy dataset
|
|
# data : X - dim = (n, 4)
|
|
# target : Y - dim = (n, 1)
|
|
X = torch.Tensor(np.random.normal(0, 1, size=(n, 4)))
|
|
Y = torch.Tensor(np.random.uniform(0, 1, size=(n, 1)))
|
|
# toy neural network : 1-layer
|
|
model = nn.Linear(4, 1)
|
|
optimizer = Adam(model.parameters(), lr=3e-4)
|
|
criterion = nn.MSELoss()
|
|
|
|
# Wrap the model in DDP and move it to GPU.
|
|
model = ray.train.torch.prepare_model(model)
|
|
|
|
# ====== Resume training state from the checkpoint. ======
|
|
start_epoch = 0
|
|
checkpoint = train.get_checkpoint()
|
|
if checkpoint:
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
model_state_dict = torch.load(
|
|
os.path.join(checkpoint_dir, "model.pt"),
|
|
# map_location=..., # Load onto a different device if needed.
|
|
)
|
|
model.module.load_state_dict(model_state_dict)
|
|
optimizer.load_state_dict(
|
|
torch.load(os.path.join(checkpoint_dir, "optimizer.pt"))
|
|
)
|
|
start_epoch = (
|
|
torch.load(os.path.join(checkpoint_dir, "extra_state.pt"))["epoch"] + 1
|
|
)
|
|
# ========================================================
|
|
|
|
for epoch in range(start_epoch, config["num_epochs"]):
|
|
y = model.forward(X)
|
|
loss = criterion(y, Y)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
metrics = {"loss": loss.item()}
|
|
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
checkpoint = None
|
|
|
|
should_checkpoint = epoch % config.get("checkpoint_freq", 1) == 0
|
|
# In standard DDP training, where the model is the same across all ranks,
|
|
# only the global rank 0 worker needs to save and report the checkpoint
|
|
if train.get_context().get_world_rank() == 0 and should_checkpoint:
|
|
# === Make sure to save all state needed for resuming training ===
|
|
torch.save(
|
|
model.module.state_dict(), # NOTE: Unwrap the model.
|
|
os.path.join(temp_checkpoint_dir, "model.pt"),
|
|
)
|
|
torch.save(
|
|
optimizer.state_dict(),
|
|
os.path.join(temp_checkpoint_dir, "optimizer.pt"),
|
|
)
|
|
torch.save(
|
|
{"epoch": epoch},
|
|
os.path.join(temp_checkpoint_dir, "extra_state.pt"),
|
|
)
|
|
# ================================================================
|
|
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
|
|
|
|
train.report(metrics, checkpoint=checkpoint)
|
|
|
|
if epoch == 1:
|
|
raise RuntimeError("Intentional error to showcase restoration!")
|
|
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
train_loop_config={"num_epochs": 5},
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
run_config=train.RunConfig(failure_config=train.FailureConfig(max_failures=1)),
|
|
)
|
|
result = trainer.fit()
|
|
# __pytorch_restore_end__
|
|
|
|
# __checkpoint_from_single_worker_start__
|
|
import tempfile
|
|
|
|
from ray import train
|
|
|
|
|
|
def train_fn(config):
|
|
...
|
|
|
|
metrics = {...}
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
checkpoint = None
|
|
|
|
# Only the global rank 0 worker saves and reports the checkpoint
|
|
if train.get_context().get_world_rank() == 0:
|
|
... # Save checkpoint to temp_checkpoint_dir
|
|
|
|
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
|
|
|
|
train.report(metrics, checkpoint=checkpoint)
|
|
# __checkpoint_from_single_worker_end__
|
|
|
|
|
|
# __lightning_save_example_start__
|
|
import lightning.pytorch as pl
|
|
|
|
from ray import train
|
|
from ray.train.lightning import RayTrainReportCallback
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
class MyLightningModule(pl.LightningModule):
|
|
# ...
|
|
|
|
def on_validation_epoch_end(self):
|
|
...
|
|
mean_acc = calculate_accuracy()
|
|
self.log("mean_accuracy", mean_acc, sync_dist=True)
|
|
|
|
|
|
def train_func():
|
|
...
|
|
model = MyLightningModule(...)
|
|
datamodule = MyLightningDataModule(...)
|
|
|
|
trainer = pl.Trainer(
|
|
# ...
|
|
callbacks=[RayTrainReportCallback()]
|
|
)
|
|
trainer.fit(model, datamodule=datamodule)
|
|
|
|
|
|
ray_trainer = TorchTrainer(
|
|
train_func,
|
|
scaling_config=train.ScalingConfig(num_workers=2),
|
|
run_config=train.RunConfig(
|
|
checkpoint_config=train.CheckpointConfig(
|
|
num_to_keep=2,
|
|
checkpoint_score_attribute="mean_accuracy",
|
|
checkpoint_score_order="max",
|
|
),
|
|
),
|
|
)
|
|
# __lightning_save_example_end__
|
|
|
|
|
|
# __lightning_custom_save_example_start__
|
|
import os
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from lightning.pytorch.callbacks import Callback
|
|
|
|
import ray
|
|
import ray.train
|
|
from ray.train import Checkpoint
|
|
|
|
|
|
class CustomRayTrainReportCallback(Callback):
|
|
def on_train_epoch_end(self, trainer, pl_module):
|
|
should_checkpoint = trainer.current_epoch % 3 == 0
|
|
|
|
with TemporaryDirectory() as tmpdir:
|
|
# Fetch metrics from `self.log(..)` in the LightningModule
|
|
metrics = trainer.callback_metrics
|
|
metrics = {k: v.item() for k, v in metrics.items()}
|
|
|
|
# Add customized metrics
|
|
metrics["epoch"] = trainer.current_epoch
|
|
metrics["custom_metric"] = 123
|
|
|
|
checkpoint = None
|
|
global_rank = ray.train.get_context().get_world_rank() == 0
|
|
if global_rank == 0 and should_checkpoint:
|
|
# Save model checkpoint file to tmpdir
|
|
ckpt_path = os.path.join(tmpdir, "ckpt.pt")
|
|
trainer.save_checkpoint(ckpt_path, weights_only=False)
|
|
|
|
checkpoint = Checkpoint.from_directory(tmpdir)
|
|
|
|
# Report to train session
|
|
ray.train.report(metrics=metrics, checkpoint=checkpoint)
|
|
# __lightning_custom_save_example_end__
|
|
|
|
# __lightning_restore_example_start__
|
|
import os
|
|
|
|
from ray import train
|
|
from ray.train import Checkpoint
|
|
from ray.train.torch import TorchTrainer
|
|
from ray.train.lightning import RayTrainReportCallback
|
|
|
|
|
|
def train_func():
|
|
model = MyLightningModule(...)
|
|
datamodule = MyLightningDataModule(...)
|
|
trainer = pl.Trainer(..., callbacks=[RayTrainReportCallback()])
|
|
|
|
checkpoint = train.get_checkpoint()
|
|
if checkpoint:
|
|
with checkpoint.as_directory() as ckpt_dir:
|
|
ckpt_path = os.path.join(ckpt_dir, RayTrainReportCallback.CHECKPOINT_NAME)
|
|
trainer.fit(model, datamodule=datamodule, ckpt_path=ckpt_path)
|
|
else:
|
|
trainer.fit(model, datamodule=datamodule)
|
|
|
|
|
|
ray_trainer = TorchTrainer(
|
|
train_func,
|
|
scaling_config=train.ScalingConfig(num_workers=2),
|
|
run_config=train.RunConfig(
|
|
checkpoint_config=train.CheckpointConfig(num_to_keep=2),
|
|
),
|
|
)
|
|
# __lightning_restore_example_end__
|
|
|
|
|
|
# __transformers_save_example_start__
|
|
from transformers import TrainingArguments
|
|
|
|
from ray import train
|
|
from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
def train_func(config):
|
|
...
|
|
|
|
# Configure logging, saving, evaluation strategies as usual.
|
|
args = TrainingArguments(
|
|
...,
|
|
eval_strategy="epoch",
|
|
save_strategy="epoch",
|
|
logging_strategy="step",
|
|
)
|
|
|
|
trainer = transformers.Trainer(args, ...)
|
|
|
|
# Add a report callback to transformers Trainer
|
|
# =============================================
|
|
trainer.add_callback(RayTrainReportCallback())
|
|
trainer = prepare_trainer(trainer)
|
|
|
|
trainer.train()
|
|
|
|
|
|
ray_trainer = TorchTrainer(
|
|
train_func,
|
|
run_config=train.RunConfig(
|
|
checkpoint_config=train.CheckpointConfig(
|
|
num_to_keep=3,
|
|
checkpoint_score_attribute="eval_loss", # The monitoring metric
|
|
checkpoint_score_order="min",
|
|
)
|
|
),
|
|
)
|
|
# __transformers_save_example_end__
|
|
|
|
|
|
# __transformers_custom_save_example_start__
|
|
from ray import train
|
|
|
|
from transformers.trainer_callback import TrainerCallback
|
|
|
|
|
|
class MyTrainReportCallback(TrainerCallback):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.metrics = {}
|
|
|
|
def on_log(self, args, state, control, model=None, logs=None, **kwargs):
|
|
"""Log is called on evaluation step and logging step."""
|
|
self.metrics.update(logs)
|
|
|
|
def on_save(self, args, state, control, **kwargs):
|
|
"""Event called after a checkpoint save."""
|
|
|
|
checkpoint = None
|
|
if train.get_context().get_world_rank() == 0:
|
|
# Build a Ray Train Checkpoint from the latest checkpoint
|
|
checkpoint_path = transformers.trainer.get_last_checkpoint(args.output_dir)
|
|
checkpoint = Checkpoint.from_directory(checkpoint_path)
|
|
|
|
# Report to Ray Train with up-to-date metrics
|
|
ray.train.report(metrics=self.metrics, checkpoint=checkpoint)
|
|
|
|
# Clear the metrics buffer
|
|
self.metrics = {}
|
|
# __transformers_custom_save_example_end__
|
|
|
|
|
|
# __distributed_checkpointing_start__
|
|
from ray import train
|
|
from ray.train import Checkpoint
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
|
|
def train_func(config):
|
|
...
|
|
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
rank = train.get_context().get_world_rank()
|
|
torch.save(
|
|
...,
|
|
os.path.join(temp_checkpoint_dir, f"model-rank={rank}.pt"),
|
|
)
|
|
checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
|
|
|
|
train.report(metrics, checkpoint=checkpoint)
|
|
|
|
|
|
trainer = TorchTrainer(
|
|
train_func,
|
|
scaling_config=train.ScalingConfig(num_workers=2),
|
|
run_config=train.RunConfig(storage_path="s3://bucket/"),
|
|
)
|
|
# The checkpoint in cloud storage will contain: model-rank=0.pt, model-rank=1.pt
|
|
# __distributed_checkpointing_end__
|
|
|
|
# __inspect_checkpoint_example_start__
|
|
from pathlib import Path
|
|
|
|
from ray.train import Checkpoint
|
|
|
|
# For demonstration, create a locally available directory with a `model.pt` file.
|
|
example_checkpoint_dir = Path("/tmp/test-checkpoint")
|
|
example_checkpoint_dir.mkdir()
|
|
example_checkpoint_dir.joinpath("model.pt").touch()
|
|
|
|
# Create the checkpoint, which is a reference to the directory.
|
|
checkpoint = Checkpoint.from_directory(example_checkpoint_dir)
|
|
|
|
# Inspect the checkpoint's contents with either `as_directory` or `to_directory`:
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
assert Path(checkpoint_dir).joinpath("model.pt").exists()
|
|
|
|
checkpoint_dir = checkpoint.to_directory()
|
|
assert Path(checkpoint_dir).joinpath("model.pt").exists()
|
|
# __inspect_checkpoint_example_end__
|
|
|
|
# __inspect_transformers_checkpoint_example_start__
|
|
# After training finished
|
|
checkpoint = result.checkpoint
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
hf_checkpoint_path = f"{checkpoint_dir}/checkpoint/"
|
|
# __inspect_transformers_checkpoint_example_end__
|
|
|
|
# __inspect_lightning_checkpoint_example_start__
|
|
# After training finished
|
|
checkpoint = result.checkpoint
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
lightning_checkpoint_path = f"{checkpoint_dir}/checkpoint.ckpt"
|
|
# __inspect_lightning_checkpoint_example_end__
|
|
|
|
# __checkpoint_upload_mode_sync_start__
|
|
def train_fn(config):
|
|
...
|
|
metrics = {...}
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
... # Save checkpoint to tmpdir
|
|
checkpoint = Checkpoint.from_directory(tmpdir)
|
|
train.report(
|
|
metrics,
|
|
checkpoint=checkpoint,
|
|
checkpoint_upload_mode=train.CheckpointUploadMode.SYNC,
|
|
)
|
|
# __checkpoint_upload_mode_sync_end__
|
|
|
|
# __checkpoint_upload_mode_async_start__
|
|
def train_fn(config):
|
|
...
|
|
metrics = {...}
|
|
tmpdir = tempfile.mkdtemp()
|
|
... # Save checkpoint to tmpdir
|
|
checkpoint = Checkpoint.from_directory(tmpdir)
|
|
train.report(
|
|
metrics,
|
|
checkpoint=checkpoint,
|
|
checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
|
|
)
|
|
# __checkpoint_upload_mode_async_end__
|
|
|
|
# __checkpoint_upload_mode_no_upload_start__
|
|
from s3torchconnector.dcp import S3StorageWriter
|
|
from torch.distributed.checkpoint.state_dict_saver import save
|
|
from torch.distributed.checkpoint.state_dict import get_state_dict
|
|
|
|
|
|
def train_fn(config):
|
|
...
|
|
for epoch in range(config["num_epochs"]):
|
|
# Directly upload checkpoint to s3 with Torch
|
|
model, optimizer = ...
|
|
storage_context = ray.train.get_context().get_storage()
|
|
checkpoint_path = (
|
|
f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
|
|
)
|
|
storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
|
|
model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
|
|
save(
|
|
{"model": model_dict, "opt": opt_dict},
|
|
storage_writer=storage_writer,
|
|
)
|
|
|
|
# Report that checkpoint to Ray Train
|
|
metrics = {...}
|
|
checkpoint = Checkpoint(checkpoint_path)
|
|
train.report(
|
|
metrics,
|
|
checkpoint=checkpoint,
|
|
checkpoint_upload_mode=train.CheckpointUploadMode.NO_UPLOAD,
|
|
)
|
|
# __checkpoint_upload_mode_no_upload_end__
|
|
|
|
|
|
# __checkpoint_upload_fn_start__
|
|
from torch.distributed.checkpoint.state_dict_saver import async_save
|
|
from s3torchconnector.dcp import S3StorageWriter
|
|
from torch.distributed.checkpoint.state_dict import get_state_dict
|
|
|
|
from ray import train
|
|
from ray.train import Checkpoint
|
|
|
|
|
|
def train_fn(config):
|
|
...
|
|
for epoch in config["num_epochs"]:
|
|
# Start async checkpoint upload to s3 with Torch
|
|
model, optimizer = ...
|
|
storage_context = train.get_context().get_storage()
|
|
checkpoint_path = (
|
|
f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
|
|
)
|
|
storage_writer = S3StorageWriter(region="us-west-2", path=checkpoint_path)
|
|
model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
|
|
ckpt_ref = async_save(
|
|
{"model": model_dict, "opt": opt_dict},
|
|
storage_writer=storage_writer,
|
|
)
|
|
|
|
def wait_async_save(checkpoint, checkpoint_dir_name):
|
|
# This function waits for checkpoint to be finalized before returning it as is
|
|
ckpt_ref.result()
|
|
return checkpoint
|
|
|
|
# Ray Train kicks off a thread that waits for the async checkpoint upload to complete
|
|
# before reporting the checkpoint
|
|
metrics = {...}
|
|
checkpoint = Checkpoint(checkpoint_path)
|
|
train.report(
|
|
metrics=metrics,
|
|
checkpoint=checkpoint,
|
|
checkpoint_upload_mode=train.CheckpointUploadMode.ASYNC,
|
|
checkpoint_upload_fn=wait_async_save,
|
|
# As uploading into the experiment directory then don't delete the checkpoint after upload is complete
|
|
delete_local_checkpoint_after_upload=False,
|
|
)
|
|
|
|
trainer = TorchTrainer(
|
|
train_fn,
|
|
train_loop_config={"num_epochs": 3},
|
|
scaling_config=train.ScalingConfig(num_workers=2, use_gpu=True),
|
|
# we need a cpu backend for async_save and a gpu backend for training
|
|
torch_config=train.torch.TorchConfig(backend="cpu:gloo,cuda:nccl"),
|
|
run_config=train.RunConfig(storage_path="s3://bucket/")
|
|
)
|
|
# __checkpoint_upload_fn_end__
|
|
|
|
# __get_all_reported_checkpoints_example_start__
|
|
import ray.train
|
|
from ray.train import CheckpointConsistencyMode
|
|
|
|
def train_fn():
|
|
for epoch in range(2):
|
|
metrics = {"train/loss": 0.1}
|
|
checkpoint = ...
|
|
ray.train.report(
|
|
metrics,
|
|
checkpoint=checkpoint,
|
|
validation=...,
|
|
)
|
|
|
|
# Get committed checkpoints which may still have ongoing validations.
|
|
committed_checkpoints = ray.train.get_all_reported_checkpoints(
|
|
consistency_mode=CheckpointConsistencyMode.COMMITTED)
|
|
|
|
# Wait for all pending validations to finish to access reported checkpoints
|
|
# with validation metrics attached.
|
|
validated_checkpoints = ray.train.get_all_reported_checkpoints()
|
|
...
|
|
# __get_all_reported_checkpoints_example_end__ |