chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,443 @@
|
||||
import collections
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pprint
|
||||
import time
|
||||
import tempfile
|
||||
from typing import Dict, Optional
|
||||
|
||||
import ray.train
|
||||
from ray.data._internal.stats import Timer
|
||||
import torch
|
||||
from logger_utils import ContextLoggerAdapter
|
||||
|
||||
from benchmark_factory import BenchmarkFactory
|
||||
|
||||
logger = ContextLoggerAdapter(logging.getLogger(__name__))
|
||||
|
||||
|
||||
class TrainLoopRunner:
|
||||
"""Generic runner that sets up the training loop scaffolding.
|
||||
|
||||
Collects perf metrics and handles periodic checkpointing and validation.
|
||||
"""
|
||||
|
||||
def __init__(self, factory: BenchmarkFactory):
|
||||
self.factory = factory
|
||||
self.benchmark_config = factory.benchmark_config
|
||||
|
||||
self._setup()
|
||||
|
||||
# Training progress state.
|
||||
self._train_batch_idx: int = 0
|
||||
self._train_epoch_idx: int = 0
|
||||
self._global_rows_processed_this_epoch: int = 0
|
||||
|
||||
# Performance metrics
|
||||
self._metrics = collections.defaultdict(lambda: Timer())
|
||||
|
||||
checkpoint = ray.train.get_checkpoint()
|
||||
if checkpoint:
|
||||
self._restore_from_checkpoint(checkpoint)
|
||||
|
||||
# Methods for subclasses to implement.
|
||||
def _setup(self):
|
||||
"""Subclasses should override this to setup the model, optimizer, etc.
|
||||
The attributes initialized in this method should only be used in the
|
||||
other overridden methods."""
|
||||
pass
|
||||
|
||||
def _cleanup(self):
|
||||
"""Subclasses can override this to cleanup any resources."""
|
||||
pass
|
||||
|
||||
def _train_step(self, train_dataloader):
|
||||
"""Subclasses should override this to implement the training step.
|
||||
A training step represents a single forward and backward pass on a batch of data.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _validate_step(self, val_dataloader):
|
||||
"""Subclasses should override this to implement the validation step.
|
||||
A validation step represents a single forward pass on a batch of data."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _save_training_state(self, local_dir: str):
|
||||
"""Subclasses should override this to save the training state.
|
||||
This should reference the model and optimizer state initialized
|
||||
in the `_setup` method."""
|
||||
pass
|
||||
|
||||
def _load_training_state(self, local_dir: str):
|
||||
"""Subclasses should override this to load the training state.
|
||||
This should reference the model and optimizer state initialized
|
||||
in the `_setup` method."""
|
||||
pass
|
||||
|
||||
def _restore_from_checkpoint(self, checkpoint: ray.train.Checkpoint):
|
||||
logger.info(
|
||||
f"Restoring from checkpoint: {checkpoint} for worker "
|
||||
f"{ray.train.get_context().get_world_rank()}"
|
||||
)
|
||||
with tempfile.TemporaryDirectory(
|
||||
dir="/mnt/local_storage"
|
||||
) as temp_checkpoint_dir:
|
||||
download_start = time.perf_counter()
|
||||
checkpoint.to_directory(temp_checkpoint_dir)
|
||||
download_time = time.perf_counter() - download_start
|
||||
|
||||
load_start = time.perf_counter()
|
||||
self._load_checkpoint(temp_checkpoint_dir)
|
||||
load_time = time.perf_counter() - load_start
|
||||
|
||||
self._metrics["checkpoint/download"].add(download_time)
|
||||
self._metrics["checkpoint/load"].add(load_time)
|
||||
|
||||
def _wrap_dataloader(self, dataloader, train: bool = True):
|
||||
dataloader_iter = iter(dataloader)
|
||||
|
||||
prefix = "train" if train else "validation"
|
||||
|
||||
def dataloader_with_timers():
|
||||
try:
|
||||
with self._metrics[f"{prefix}/iter_first_batch"].timer():
|
||||
batch = next(dataloader_iter)
|
||||
if train:
|
||||
self._train_batch_idx += 1
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
while True:
|
||||
yield batch
|
||||
|
||||
try:
|
||||
with self._metrics[f"{prefix}/iter_batch"].timer():
|
||||
batch = next(dataloader_iter)
|
||||
if train:
|
||||
self._train_batch_idx += 1
|
||||
except StopIteration:
|
||||
return
|
||||
|
||||
return dataloader_with_timers()
|
||||
|
||||
@property
|
||||
def _num_batches_to_skip(self) -> int:
|
||||
"""Calculate the number of batches to skip based on the number of rows already processed in this epoch."""
|
||||
|
||||
global_batch_size = (
|
||||
self.benchmark_config.dataloader_config.train_batch_size
|
||||
* ray.train.get_context().get_world_size()
|
||||
)
|
||||
|
||||
return self._global_rows_processed_this_epoch // global_batch_size
|
||||
|
||||
def _train_epoch(self):
|
||||
"""Subclasses can override the entrire `_train_epoch` method for more training
|
||||
logic customization."""
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
logger.info(f"Training starting @ epoch={self._train_epoch_idx}")
|
||||
|
||||
train_dataloader = self.factory.get_train_dataloader()
|
||||
train_dataloader = self._wrap_dataloader(train_dataloader, train=True)
|
||||
|
||||
# Skip through batches if we restored to a middle of the epoch.
|
||||
# TODO: Compare this baseline to the data checkpointing approach once we have it.
|
||||
if self._num_batches_to_skip:
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
logger.info(f"Skipping {self._num_batches_to_skip} batches...")
|
||||
|
||||
# Zero before the skip loop drives the wrapper, which would
|
||||
# otherwise double-count against the value restored from the
|
||||
# checkpoint. After the skip, _train_batch_idx is rebuilt to
|
||||
# _num_batches_to_skip — matching the restored value.
|
||||
self._train_batch_idx = 0
|
||||
|
||||
for _ in range(self._num_batches_to_skip):
|
||||
with self._metrics["train/iter_skip_batch"].timer():
|
||||
next(train_dataloader)
|
||||
|
||||
for batch in train_dataloader:
|
||||
with self._metrics["train/step"].timer():
|
||||
if not self.benchmark_config.skip_train_step:
|
||||
self._train_step(batch)
|
||||
if self.benchmark_config.train_step_sleep_s > 0:
|
||||
time.sleep(self.benchmark_config.train_step_sleep_s)
|
||||
|
||||
# TODO: This is slightly off if the last batch is a partial batch (if drop_last=False)
|
||||
global_batch_size = (
|
||||
self.benchmark_config.dataloader_config.train_batch_size
|
||||
* ray.train.get_context().get_world_size()
|
||||
)
|
||||
self._metrics["train/rows_processed"].add(global_batch_size)
|
||||
|
||||
self._global_rows_processed_this_epoch += global_batch_size
|
||||
|
||||
if self._should_checkpoint_during_epoch():
|
||||
self._checkpoint()
|
||||
|
||||
if self._should_validate_during_epoch():
|
||||
validation_metrics = self._validate()
|
||||
self._checkpoint(validation_metrics)
|
||||
|
||||
if self._should_log_metrics():
|
||||
logger.info(pprint.pformat(self.get_metrics(), indent=2))
|
||||
|
||||
if (
|
||||
self.benchmark_config.max_train_batches > 0
|
||||
and self._train_batch_idx >= self.benchmark_config.max_train_batches
|
||||
):
|
||||
break
|
||||
|
||||
self._train_epoch_idx += 1
|
||||
self._train_batch_idx = 0
|
||||
self._global_rows_processed_this_epoch = 0
|
||||
|
||||
def _validate_epoch(self) -> Dict[str, float]:
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
logger.info(
|
||||
f"Validation starting @ epoch={self._train_epoch_idx}, "
|
||||
f"batch={self._train_batch_idx}"
|
||||
)
|
||||
|
||||
val_dataloader = self.factory.get_val_dataloader()
|
||||
val_dataloader = self._wrap_dataloader(val_dataloader, train=False)
|
||||
|
||||
total_loss = torch.tensor(0.0).to(ray.train.torch.get_device())
|
||||
num_rows = 0
|
||||
|
||||
for batch in val_dataloader:
|
||||
with self._metrics["validation/step"].timer():
|
||||
if not self.benchmark_config.skip_validation_step:
|
||||
total_loss += self._validate_step(batch)
|
||||
|
||||
num_rows += self.benchmark_config.dataloader_config.validation_batch_size
|
||||
self._metrics["validation/rows_processed"].add(
|
||||
self.benchmark_config.dataloader_config.validation_batch_size
|
||||
)
|
||||
assert num_rows > 0, "Validation dataset yielded no batches."
|
||||
|
||||
return {"validation/loss": total_loss.item() / num_rows}
|
||||
|
||||
def _should_checkpoint_during_epoch(self) -> bool:
|
||||
"""Handles the checkpoint_every_n_steps logic."""
|
||||
return (
|
||||
self.benchmark_config.checkpoint_every_n_steps > 0
|
||||
and self._train_batch_idx % self.benchmark_config.checkpoint_every_n_steps
|
||||
== 0
|
||||
)
|
||||
|
||||
def _should_validate_during_epoch(self) -> bool:
|
||||
"""Handles the validate_every_n_steps logic."""
|
||||
return (
|
||||
self.benchmark_config.validate_every_n_steps > 0
|
||||
and self._train_batch_idx % self.benchmark_config.validate_every_n_steps
|
||||
== 0
|
||||
)
|
||||
|
||||
def _should_log_metrics(self) -> bool:
|
||||
"""Handles the log_metrics_every_n_steps logic."""
|
||||
return (
|
||||
self.benchmark_config.log_metrics_every_n_steps > 0
|
||||
and self._train_batch_idx % self.benchmark_config.log_metrics_every_n_steps
|
||||
== 0
|
||||
)
|
||||
|
||||
def _validate(self) -> Dict[str, float]:
|
||||
with self._metrics["validation/epoch"].timer():
|
||||
validation_metrics = self._validate_epoch()
|
||||
return validation_metrics
|
||||
|
||||
def _checkpoint(self, metrics: Optional[Dict[str, float]] = None):
|
||||
with tempfile.TemporaryDirectory(
|
||||
dir="/mnt/local_storage"
|
||||
) as temp_checkpoint_dir:
|
||||
with self._metrics["checkpoint/save"].timer():
|
||||
self._save_checkpoint(temp_checkpoint_dir)
|
||||
|
||||
with self._metrics["checkpoint/report"].timer():
|
||||
self._report_checkpoint(
|
||||
metrics=metrics or {},
|
||||
checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
|
||||
)
|
||||
|
||||
def _load_checkpoint(self, local_dir: str):
|
||||
self._load_training_state(local_dir)
|
||||
|
||||
run_state = torch.load(os.path.join(local_dir, "run_state.pt"))
|
||||
self._train_epoch_idx = run_state["epoch"]
|
||||
self._train_batch_idx = run_state["batch_idx"]
|
||||
self._global_rows_processed_this_epoch = run_state[
|
||||
"global_rows_processed_this_epoch"
|
||||
]
|
||||
|
||||
with open(os.path.join(local_dir, "metrics.json"), "r") as f:
|
||||
metrics_json = json.load(f)
|
||||
|
||||
for k, v in metrics_json.items():
|
||||
self._metrics[k].from_dict(v)
|
||||
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
logger.info(
|
||||
f"Restored to epoch={self._train_epoch_idx}, "
|
||||
f"train_batch_idx={self._train_batch_idx} from checkpoint: "
|
||||
f"{ray.train.get_checkpoint()}"
|
||||
)
|
||||
|
||||
def _save_checkpoint(self, local_dir: str):
|
||||
logger.info(
|
||||
f"Saving checkpoint @ epoch={self._train_epoch_idx}, "
|
||||
f"train_batch_idx={self._train_batch_idx}"
|
||||
)
|
||||
|
||||
self._save_training_state(local_dir)
|
||||
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
run_state = {
|
||||
"epoch": self._train_epoch_idx,
|
||||
"batch_idx": self._train_batch_idx,
|
||||
"global_rows_processed_this_epoch": self._global_rows_processed_this_epoch,
|
||||
}
|
||||
torch.save(run_state, os.path.join(local_dir, "run_state.pt"))
|
||||
|
||||
metrics_json = {k: v.as_dict() for k, v in self._metrics.items()}
|
||||
with open(os.path.join(local_dir, "metrics.json"), "w") as f:
|
||||
json.dump(metrics_json, f)
|
||||
|
||||
def _report_checkpoint(self, metrics, checkpoint):
|
||||
logger.info(
|
||||
f"Uploading checkpoint @ epoch={self._train_epoch_idx}, "
|
||||
f"train_batch_idx={self._train_batch_idx}"
|
||||
)
|
||||
|
||||
checkpoint_dir_name = (
|
||||
f"checkpoint_epoch={self._train_epoch_idx}_batch={self._train_batch_idx}"
|
||||
)
|
||||
|
||||
ray.train.report(
|
||||
metrics,
|
||||
checkpoint=checkpoint,
|
||||
checkpoint_dir_name=checkpoint_dir_name,
|
||||
)
|
||||
|
||||
def run(self):
|
||||
starting_epoch = self._train_epoch_idx
|
||||
|
||||
for _ in range(starting_epoch, self.benchmark_config.num_epochs):
|
||||
with self._metrics["train/epoch"].timer():
|
||||
self._train_epoch()
|
||||
|
||||
if not self.benchmark_config.skip_validation_at_epoch_end:
|
||||
validation_metrics = self._validate()
|
||||
self._checkpoint(validation_metrics)
|
||||
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
logger.info(pprint.pformat(self.get_metrics(), indent=2))
|
||||
|
||||
self._cleanup()
|
||||
|
||||
def get_metrics(self, dataset_creation_time: float = 0.0) -> Dict[str, float]:
|
||||
# TODO: These metrics should be aggregated across training workers.
|
||||
metrics = {}
|
||||
for key, metric in self._metrics.items():
|
||||
metrics.update(
|
||||
{
|
||||
f"{key}-avg": metric.avg(),
|
||||
f"{key}-min": metric.min(),
|
||||
f"{key}-max": metric.max(),
|
||||
f"{key}-total": metric.get(),
|
||||
}
|
||||
)
|
||||
|
||||
metrics["train/dataset_creation_time"] = dataset_creation_time
|
||||
metrics["validation/dataset_creation_time"] = dataset_creation_time
|
||||
|
||||
# Throughput
|
||||
# TODO: Ray Data can provide these throughput metrics automatically.
|
||||
train_time = (
|
||||
metrics["train/dataset_creation_time"]
|
||||
+ self._metrics["train/step"].get()
|
||||
# Include the time it takes to get the first batch.
|
||||
+ self._metrics["train/iter_first_batch"].get()
|
||||
+ self._metrics["train/iter_batch"].get()
|
||||
)
|
||||
if train_time > 0:
|
||||
metrics["train/global_throughput"] = (
|
||||
self._metrics["train/rows_processed"].get() / train_time
|
||||
)
|
||||
|
||||
validation_time = (
|
||||
metrics["validation/dataset_creation_time"]
|
||||
+ self._metrics["validation/step"].get()
|
||||
# Include the time it takes to get the first batch.
|
||||
+ self._metrics["validation/iter_first_batch"].get()
|
||||
+ self._metrics["validation/iter_batch"].get()
|
||||
)
|
||||
if validation_time > 0:
|
||||
metrics["validation/global_throughput"] = (
|
||||
self._metrics["validation/rows_processed"].get() / validation_time
|
||||
)
|
||||
|
||||
# Extra time that each worker spends to restore from checkpoint,
|
||||
# which includes downloading the checkpoint, loading the checkpoint,
|
||||
# and skipping through batches that were already processed.
|
||||
restoration_time = (
|
||||
self._metrics["checkpoint/download"].get()
|
||||
+ self._metrics["checkpoint/load"].get()
|
||||
+ self._metrics["train/iter_skip_batch"].get()
|
||||
)
|
||||
if restoration_time > 0:
|
||||
metrics["checkpoint/restoration_time"] = restoration_time
|
||||
|
||||
# Dataloader metrics (ex: Ray Data stats)
|
||||
metrics.update(self.factory.get_dataloader_metrics())
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
class VanillaTorchRunner(TrainLoopRunner):
|
||||
"""A simple runner that uses a PyTorch model, optimizer, and loss function."""
|
||||
|
||||
def _setup(self):
|
||||
model = self.factory.get_model()
|
||||
self.model = ray.train.torch.prepare_model(model)
|
||||
self.loss_fn = self.factory.get_loss_fn()
|
||||
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
|
||||
|
||||
def _train_step(self, batch):
|
||||
self.model.train()
|
||||
|
||||
input_batch, labels = batch
|
||||
|
||||
self.model.train()
|
||||
self.optimizer.zero_grad()
|
||||
out = self.model(input_batch)
|
||||
loss = self.loss_fn(out, labels)
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
|
||||
def _validate_step(self, batch):
|
||||
self.model.eval()
|
||||
|
||||
input_batch, labels = batch
|
||||
|
||||
with torch.no_grad():
|
||||
out = self.model(input_batch)
|
||||
loss = self.loss_fn(out, labels)
|
||||
return loss
|
||||
|
||||
def _save_training_state(self, local_dir: str):
|
||||
# Standard DDP checkpointing.
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
torch.save(self.model.state_dict(), os.path.join(local_dir, "model.pt"))
|
||||
torch.save(
|
||||
self.optimizer.state_dict(), os.path.join(local_dir, "optimizer.pt")
|
||||
)
|
||||
|
||||
def _load_training_state(self, local_dir: str):
|
||||
self.model.load_state_dict(
|
||||
torch.load(os.path.join(local_dir, "model.pt"), map_location="cpu")
|
||||
)
|
||||
self.optimizer.load_state_dict(
|
||||
torch.load(os.path.join(local_dir, "optimizer.pt"), map_location="cpu")
|
||||
)
|
||||
Reference in New Issue
Block a user