chore: import upstream snapshot with attribution

python/ray/train/examples/accelerate/accelerate_torch_trainer.py

@@ -0,0 +1,164 @@
+# __accelerate_torch_basic_example_start__
+"""
+Minimal Ray Train and Accelerate example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with Hugging Face Accelerate and Ray Train and Ray Data
+"""
+
+from tempfile import TemporaryDirectory
+
+import evaluate
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset
+from torch.optim import AdamW
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+import ray
+import ray.train
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that launches on each worker."""
+
+    # Unpack training configs
+    lr = config["lr"]
+    seed = config["seed"]
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+    train_ds_size = config["train_dataset_size"]
+
+    set_seed(seed)
+
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Load datasets and metrics
+    metric = evaluate.load("glue", "mrpc")
+
+    # Prepare Ray Data loaders
+    # ====================================================
+    train_ds = ray.train.get_dataset_shard("train")
+    eval_ds = ray.train.get_dataset_shard("validation")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            list(batch["sentence1"]),
+            list(batch["sentence2"]),
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor(batch["label"])
+        outputs = {k: v.to(accelerator.device) for k, v in outputs.items()}
+        return outputs
+
+    train_dataloader = train_ds.iter_torch_batches(
+        batch_size=train_batch_size, collate_fn=collate_fn
+    )
+    eval_dataloader = eval_ds.iter_torch_batches(
+        batch_size=eval_batch_size, collate_fn=collate_fn
+    )
+    # ====================================================
+
+    # Instantiate the model, optimizer, lr_scheduler
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    steps_per_epoch = train_ds_size // (accelerator.num_processes * train_batch_size)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(steps_per_epoch * num_epochs),
+    )
+
+    # Prepare everything with accelerator
+    model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            predictions, references = accelerator.gather_for_metrics(
+                (predictions, batch["labels"])
+            )
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==========================================
+        with TemporaryDirectory() as tmpdir:
+            if accelerator.is_main_process:
+                unwrapped_model = accelerator.unwrap_model(model)
+                accelerator.save(unwrapped_model, f"{tmpdir}/ckpt_{epoch}.bin")
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+            ray.train.report(metrics=eval_metric, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    config = {
+        "lr": 2e-5,
+        "num_epochs": 3,
+        "seed": 42,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+    }
+
+    # Prepare Ray Datasets
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    ray_datasets = {
+        "train": ray.data.from_huggingface(hf_datasets["train"]),
+        "validation": ray.data.from_huggingface(hf_datasets["validation"]),
+    }
+    config["train_dataset_size"] = ray_datasets["train"].count()
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        datasets=ray_datasets,
+        dataset_config=DataConfig(datasets_to_split=["train", "validation"]),
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+# __accelerate_torch_basic_example_end__

python/ray/train/examples/accelerate/accelerate_torch_trainer_no_raydata.py

@@ -0,0 +1,169 @@
+# __accelerate_torch_basic_example_no_raydata_start__
+"""
+Minimal Ray Train + Accelerate example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with Hugging Face Accelerate and Ray Train
+"""
+
+from tempfile import TemporaryDirectory
+
+import evaluate
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+import ray.train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    lr = config["lr"]
+    seed = config["seed"]
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+
+    set_seed(seed)
+
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Load datasets and metrics
+    metric = evaluate.load("glue", "mrpc")
+
+    # Prepare PyTorch DataLoaders
+    # ====================================================
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            [sample["sentence1"] for sample in batch],
+            [sample["sentence2"] for sample in batch],
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor([sample["label"] for sample in batch])
+        outputs = {k: v.to(accelerator.device) for k, v in outputs.items()}
+        return outputs
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        hf_datasets["train"],
+        shuffle=True,
+        collate_fn=collate_fn,
+        batch_size=train_batch_size,
+        drop_last=True,
+    )
+    eval_dataloader = DataLoader(
+        hf_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=eval_batch_size,
+        drop_last=True,
+    )
+    # ====================================================
+
+    # Instantiate the model, optimizer, lr_scheduler
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    optimizer = AdamW(params=model.parameters(), lr=lr)
+
+    steps_per_epoch = len(train_dataloader)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=100,
+        num_training_steps=(steps_per_epoch * num_epochs),
+    )
+
+    # Prepare everything with accelerator
+    (
+        model,
+        optimizer,
+        train_dataloader,
+        eval_dataloader,
+        lr_scheduler,
+    ) = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            predictions, references = accelerator.gather_for_metrics(
+                (predictions, batch["labels"])
+            )
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+
+        # Report Checkpoint and metrics to Ray Train
+        # ==========================================
+        with TemporaryDirectory() as tmpdir:
+            if accelerator.is_main_process:
+                unwrapped_model = accelerator.unwrap_model(model)
+                accelerator.save(unwrapped_model, f"{tmpdir}/ckpt_{epoch}.bin")
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+            ray.train.report(metrics=eval_metric, checkpoint=checkpoint)
+
+
+if __name__ == "__main__":
+    config = {
+        "lr": 2e-5,
+        "num_epochs": 3,
+        "seed": 42,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+# __accelerate_torch_basic_example_no_raydata_end__

python/ray/train/examples/deepspeed/deepspeed_torch_trainer.py

@@ -0,0 +1,185 @@
+# __deepspeed_torch_basic_example_start__
+"""
+Minimal Ray Train + DeepSpeed example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with DeepSpeed ZeRO-3 and Ray Train and Ray Data
+"""
+
+from tempfile import TemporaryDirectory
+
+import deepspeed
+import torch
+from datasets import load_dataset
+from deepspeed.accelerator import get_accelerator
+from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed
+
+import ray
+import ray.train
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    set_seed(config["seed"])
+    num_epochs = config["num_epochs"]
+    train_batch_size = config["train_batch_size"]
+    eval_batch_size = config["eval_batch_size"]
+
+    # Instantiate the Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    # Prepare Ray Data Loaders
+    # ====================================================
+    train_ds = ray.train.get_dataset_shard("train")
+    eval_ds = ray.train.get_dataset_shard("validation")
+
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            list(batch["sentence1"]),
+            list(batch["sentence2"]),
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor(batch["label"])
+        return outputs
+
+    train_dataloader = train_ds.iter_torch_batches(
+        batch_size=train_batch_size, collate_fn=collate_fn
+    )
+    eval_dataloader = eval_ds.iter_torch_batches(
+        batch_size=eval_batch_size, collate_fn=collate_fn
+    )
+    # ====================================================
+
+    # Initialize DeepSpeed Engine
+    model, optimizer, _, lr_scheduler = deepspeed.initialize(
+        model=model,
+        model_parameters=model.parameters(),
+        config=deepspeed_config,
+    )
+    device = get_accelerator().device_name(model.local_rank)
+
+    # Initialize Evaluation Metrics
+    f1 = BinaryF1Score().to(device)
+    accuracy = BinaryAccuracy().to(device)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            f1.update(predictions, batch["labels"])
+            accuracy.update(predictions, batch["labels"])
+
+        # torchmetrics will aggregate the metrics across all workers
+        eval_metric = {
+            "f1": f1.compute().item(),
+            "accuracy": accuracy.compute().item(),
+        }
+        f1.reset()
+        accuracy.reset()
+
+        if model.global_rank == 0:
+            print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==============================================================
+        with TemporaryDirectory() as tmpdir:
+            # Each worker saves its own checkpoint shard
+            model.save_checkpoint(tmpdir)
+
+            # Ensure all workers finished saving their checkpoint shard
+            torch.distributed.barrier()
+
+            # Report checkpoint shards from each worker in parallel
+            ray.train.report(
+                metrics=eval_metric, checkpoint=Checkpoint.from_directory(tmpdir)
+            )
+        # ==============================================================
+
+
+if __name__ == "__main__":
+    deepspeed_config = {
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": 2e-5,
+            },
+        },
+        "scheduler": {"type": "WarmupLR", "params": {"warmup_num_steps": 100}},
+        "fp16": {"enabled": True},
+        "bf16": {"enabled": False},  # Turn this on if using AMPERE GPUs.
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": "none",
+            },
+            "offload_param": {
+                "device": "none",
+            },
+        },
+        "gradient_accumulation_steps": 1,
+        "gradient_clipping": True,
+        "steps_per_print": 10,
+        "train_micro_batch_size_per_gpu": 16,
+        "wall_clock_breakdown": False,
+    }
+
+    training_config = {
+        "seed": 42,
+        "num_epochs": 3,
+        "train_batch_size": 16,
+        "eval_batch_size": 32,
+        "deepspeed_config": deepspeed_config,
+    }
+
+    # Prepare Ray Datasets
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    ray_datasets = {
+        "train": ray.data.from_huggingface(hf_datasets["train"]),
+        "validation": ray.data.from_huggingface(hf_datasets["validation"]),
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=training_config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        datasets=ray_datasets,
+        dataset_config=DataConfig(datasets_to_split=["train", "validation"]),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+    # Retrieve the best checkponints from results
+    _ = result.best_checkpoints
+
+# __deepspeed_torch_basic_example_end__

python/ray/train/examples/deepspeed/deepspeed_torch_trainer_no_raydata.py

@@ -0,0 +1,178 @@
+# __deepspeed_torch_basic_example_no_raydata_start__
+"""
+Minimal Ray Train + DeepSpeed example adapted from
+https://github.com/huggingface/accelerate/blob/main/examples/nlp_example.py
+
+Fine-tune a BERT model with DeepSpeed ZeRO-3 and Ray Train
+"""
+
+from tempfile import TemporaryDirectory
+
+import deepspeed
+import torch
+from datasets import load_dataset
+from deepspeed.accelerator import get_accelerator
+from torch.utils.data import DataLoader
+from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, set_seed
+
+import ray
+import ray.train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func(config):
+    """Your training function that will be launched on each worker."""
+
+    # Unpack training configs
+    set_seed(config["seed"])
+    num_epochs = config["num_epochs"]
+    eval_batch_size = config["eval_batch_size"]
+
+    # Instantiate the Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", return_dict=True
+    )
+
+    # Prepare PyTorch Data Loaders
+    # ====================================================
+    hf_datasets = load_dataset("nyu-mll/glue", "mrpc")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def collate_fn(batch):
+        outputs = tokenizer(
+            [sample["sentence1"] for sample in batch],
+            [sample["sentence2"] for sample in batch],
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+        )
+        outputs["labels"] = torch.LongTensor([sample["label"] for sample in batch])
+        return outputs
+
+    # Instantiate dataloaders.
+    # The train_dataloader already created by `deepspeed.initialize`
+    eval_dataloader = DataLoader(
+        hf_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=eval_batch_size,
+        drop_last=True,
+    )
+    # ====================================================
+
+    # Initialize DeepSpeed Engine
+    model, optimizer, train_dataloader, lr_scheduler = deepspeed.initialize(
+        model=model,
+        model_parameters=model.parameters(),
+        training_data=hf_datasets["train"],
+        collate_fn=collate_fn,
+        config=deepspeed_config,
+    )
+    device = get_accelerator().device_name(model.local_rank)
+
+    # Initialize Evaluation Metrics
+    f1 = BinaryF1Score().to(device)
+    accuracy = BinaryAccuracy().to(device)
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        for batch in train_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            model.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+
+            f1.update(predictions, batch["labels"])
+            accuracy.update(predictions, batch["labels"])
+
+        # torchmetrics will aggregate the metrics across all workers
+        eval_metric = {
+            "f1": f1.compute().item(),
+            "accuracy": accuracy.compute().item(),
+        }
+        f1.reset()
+        accuracy.reset()
+
+        if model.global_rank == 0:
+            print(f"epoch {epoch}:", eval_metric)
+
+        # Report checkpoint and metrics to Ray Train
+        # ==============================================================
+        with TemporaryDirectory() as tmpdir:
+            # Each worker saves its own checkpoint shard
+            model.save_checkpoint(tmpdir)
+
+            # Ensure all workers finished saving their checkpoint shard
+            torch.distributed.barrier()
+
+            # Report checkpoint shards from each worker in parallel
+            ray.train.report(
+                metrics=eval_metric, checkpoint=Checkpoint.from_directory(tmpdir)
+            )
+        # ==============================================================
+
+
+if __name__ == "__main__":
+    deepspeed_config = {
+        "optimizer": {
+            "type": "AdamW",
+            "params": {
+                "lr": 2e-5,
+            },
+        },
+        "scheduler": {"type": "WarmupLR", "params": {"warmup_num_steps": 100}},
+        "fp16": {"enabled": True},
+        "bf16": {"enabled": False},  # Turn this on if using AMPERE GPUs.
+        "zero_optimization": {
+            "stage": 3,
+            "offload_optimizer": {
+                "device": "none",
+            },
+            "offload_param": {
+                "device": "none",
+            },
+        },
+        "gradient_accumulation_steps": 1,
+        "gradient_clipping": True,
+        "steps_per_print": 10,
+        "train_micro_batch_size_per_gpu": 16,
+        "wall_clock_breakdown": False,
+    }
+
+    training_config = {
+        "seed": 42,
+        "num_epochs": 3,
+        "eval_batch_size": 32,
+        "deepspeed_config": deepspeed_config,
+    }
+
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=training_config,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        # run_config=ray.train.RunConfig(storage_path="s3://..."),
+    )
+
+    result = trainer.fit()
+
+    # Retrieve the best checkponints from results
+    _ = result.best_checkpoints
+
+# __deepspeed_torch_basic_example_no_raydata_end__

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_comet.py

@@ -0,0 +1,41 @@
+# isort: skip_file
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_comet_start__
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import CometLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = CometLogger(api_key=os.environ["COMET_API_KEY"])
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "COMET_API_KEY" in os.environ
+), 'Please do COMET_API_KEY="abcde" when running this script.'
+# This makes sure that all workers have this env var set.
+ray.init(runtime_env={"env_vars": {"COMET_API_KEY": os.environ["COMET_API_KEY"]}})
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_mlflow.py

@@ -0,0 +1,63 @@
+# ruff: noqa
+# isort: skip_file
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+
+# __lightning_experiment_tracking_mlflow_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import MLFlowLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = MLFlowLogger(
+            experiment_name="demo-project",
+            tracking_uri=f"file:{save_dir}",
+        )
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_mlflow_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py

@@ -0,0 +1,46 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+
+# __model_dl_start__
+import lightning.pytorch as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+# Create dummy data
+X = torch.randn(128, 3)  # 128 samples, 3 features
+y = torch.randint(0, 2, (128,))  # 128 binary labels
+
+# Create a TensorDataset to wrap the data
+dataset = TensorDataset(X, y)
+
+# Create a DataLoader to iterate over the dataset
+batch_size = 8
+dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+
+# Define a dummy model
+class DummyModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.layer = torch.nn.Linear(3, 1)
+
+    def forward(self, x):
+        return self.layer(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float())
+
+        # The metrics below will be reported to Loggers
+        self.log("train_loss", loss)
+        self.log_dict({
+            "metric_1": 1 / (batch_idx + 1), "metric_2": batch_idx * 100
+        })
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-3)

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_tensorboard.py

@@ -0,0 +1,60 @@
+# ruff: noqa
+# isort: skip_file
+
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+# __lightning_experiment_tracking_tensorboard_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import TensorBoardLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = TensorBoardLogger(name="demo-run", save_dir=f"file:{save_dir}")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "tensorboard")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_tensorboard_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_wandb.py

@@ -0,0 +1,50 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_wandb_start__
+import os
+import wandb
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import WandbLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = WandbLogger(name="demo-run", project="demo-project")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "WANDB_API_KEY" in os.environ
+), 'Please set WANDB_API_KEY="abcde" when running this script.'
+
+# This ensures that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_mlflow.py

@@ -0,0 +1,85 @@
+# ruff: noqa
+# isort: skip_file
+from filelock import FileLock
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+# __start__
+# Run the following script with the SHARED_STORAGE_PATH env var set.
+# The MLflow offline logs are saved to SHARED_STORAGE_PATH/mlruns.
+
+import mlflow
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+import torch
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+from torch.utils.data import DataLoader
+
+assert os.environ.get(
+    "SHARED_STORAGE_PATH", None
+), "Please set SHARED_STORAGE_PATH env var."
+
+
+# Assumes you are passing a `save_dir` in `config`
+def train_func(config):
+    save_dir = config["save_dir"]
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.set_tracking_uri(f"file:{save_dir}")
+        mlflow.set_experiment("my_experiment")
+        mlflow.start_run()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                mlflow.log_metrics({"loss": loss.item(), "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.end_run()
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()
+# __end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_wandb.py

@@ -0,0 +1,75 @@
+# ruff: noqa
+# fmt: off
+# isort: off
+
+# __start__
+from filelock import FileLock
+import os
+
+import torch
+import wandb
+
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+# Run the following script with the WANDB_API_KEY env var set.
+assert os.environ.get("WANDB_API_KEY", None), "Please set WANDB_API_KEY env var."
+
+# This makes sure that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+
+
+def train_func(config):
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.init()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                wandb.log({"loss": loss, "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()

python/ray/train/examples/horovod/horovod-cluster.yaml

@@ -0,0 +1,55 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: horovod-cluster
+
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 3
+max_workers: 3
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-west-2
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+available_node_types:
+    ray.head.default:
+        min_workers: 0
+        max_workers: 0
+        resources: {}
+        node_config:
+            InstanceType: g3.8xlarge
+            ImageId: latest_dlami
+            InstanceMarketOptions:
+                MarketType: spot
+            BlockDeviceMappings:
+                - DeviceName: /dev/sda1
+                  Ebs:
+                      VolumeSize: 300
+
+
+    ray.worker.default:
+        min_workers: 3
+        max_workers: 3
+        resources: {}
+        node_config:
+            InstanceType: g3.8xlarge
+            ImageId: latest_dlami
+            InstanceMarketOptions:
+                MarketType: spot
+            BlockDeviceMappings:
+                - DeviceName: /dev/sda1
+                  Ebs:
+                      VolumeSize: 300
+
+
+setup_commands:
+    # This replaces the standard anaconda Ray installation
+    - pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
+    - pip install ray[tune]
+
+    # Install Horovod
+    - HOROVOD_WITH_GLOO=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_TENSORFLOW=1 HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITH_PYTORCH=1 pip install torch torchvision horovod

python/ray/train/examples/horovod/horovod_example.py

@@ -0,0 +1,286 @@
+import argparse
+import os
+
+import horovod.torch as hvd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data.distributed
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.horovod import HorovodTrainer
+
+
+def metric_average(val, name):
+    tensor = torch.tensor(val)
+    avg_tensor = hvd.allreduce(tensor, name=name)
+    return avg_tensor.item()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x)
+
+
+def setup(config):
+    data_dir = config.get("data_dir", None)
+    seed = config.get("seed", 42)
+    batch_size = config.get("batch_size", 64)
+    use_adasum = config.get("use_adasum", False)
+    lr = config.get("lr", 0.01)
+    momentum = config.get("momentum", 0.5)
+    use_cuda = config.get("use_cuda", False)
+
+    # Horovod: initialize library.
+    hvd.init()
+    torch.manual_seed(seed)
+
+    if use_cuda:
+        # Horovod: pin GPU to local rank.
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(seed)
+
+    # Horovod: limit # of CPU threads to be used per worker.
+    torch.set_num_threads(1)
+
+    kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
+    data_dir = data_dir or "~/data"
+    with FileLock(os.path.expanduser("~/.horovod_lock")):
+        train_dataset = datasets.MNIST(
+            data_dir,
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+            ),
+        )
+    # Horovod: use DistributedSampler to partition the training data.
+    train_sampler = torch.utils.data.distributed.DistributedSampler(
+        train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
+    )
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+    )
+
+    model = Net()
+
+    # By default, Adasum doesn't need scaling up learning rate.
+    lr_scaler = hvd.size() if not use_adasum else 1
+
+    if use_cuda:
+        # Move model to GPU.
+        model.cuda()
+        # If using GPU Adasum allreduce, scale learning rate by local_size.
+        if use_adasum and hvd.nccl_built():
+            lr_scaler = hvd.local_size()
+
+    # Horovod: scale learning rate by lr_scaler.
+    optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Adasum if use_adasum else hvd.Average,
+    )
+
+    return model, optimizer, train_loader, train_sampler
+
+
+def train_epoch(
+    model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+):
+    loss = None
+    model.train()
+    # Horovod: set epoch to sampler for shuffling.
+    train_sampler.set_epoch(epoch)
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if use_cuda:
+            data, target = data.cuda(), target.cuda()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            # Horovod: use train_sampler to determine the number of
+            # examples in this worker's partition.
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_sampler),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+    return loss.item() if loss else None
+
+
+# Horovod function API.
+
+
+def train_func(config):
+    num_epochs = config.get("num_epochs", 10)
+    log_interval = config.get("log_interval", 10)
+    use_cuda = config.get("use_cuda", False)
+
+    model, optimizer, train_loader, train_sampler = setup(config)
+
+    for epoch in range(num_epochs):
+        loss = train_epoch(
+            model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+        )
+        train.report(dict(loss=loss))
+
+
+def main(num_workers, use_gpu, kwargs):
+    trainer = HorovodTrainer(
+        train_func,
+        train_loop_config=kwargs,
+        scaling_config=ScalingConfig(use_gpu=use_gpu, num_workers=num_workers),
+    )
+    results = trainer.fit()
+    print(results.metrics)
+
+
+# Horovod Class API.
+
+
+class HorovodTrainClass:
+    def __init__(self, config):
+        self.log_interval = config.get("log_interval", 10)
+        self.use_cuda = config.get("use_cuda", False)
+
+        if self.use_cuda:
+            torch.cuda.set_device(hvd.local_rank())
+
+        self.model, self.optimizer, self.train_loader, self.train_sampler = setup(
+            config
+        )
+
+    def train(self, epoch):
+        loss = train_epoch(
+            self.model,
+            self.optimizer,
+            self.train_sampler,
+            self.train_loader,
+            epoch,
+            self.log_interval,
+            self.use_cuda,
+        )
+        return loss
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(
+        description="PyTorch MNIST Example",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=5,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--momentum",
+        type=float,
+        default=0.5,
+        metavar="M",
+        help="SGD momentum (default: 0.5)",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--use-adasum",
+        action="store_true",
+        default=False,
+        help="use adasum algorithm to do reduction",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=2,
+        help="Number of Ray workers to use for training.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="location of the training dataset in the local filesystem ("
+        "will be downloaded if needed)",
+    )
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        default=None,
+        help="Address of Ray cluster.",
+    )
+
+    args = parser.parse_args()
+
+    if args.address:
+        ray.init(args.address)
+    else:
+        ray.init()
+
+    use_cuda = args.use_gpu if args.use_gpu is not None else False
+
+    kwargs = {
+        "data_dir": args.data_dir,
+        "seed": args.seed,
+        "use_cuda": use_cuda,
+        "batch_size": args.batch_size,
+        "use_adasum": args.use_adasum if args.use_adasum else False,
+        "lr": args.lr,
+        "momentum": args.momentum,
+        "num_epochs": args.num_epochs,
+        "log_interval": args.log_interval,
+    }
+
+    main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs)

python/ray/train/examples/horovod/horovod_pytorch_example.py

@@ -0,0 +1,270 @@
+import argparse
+import os
+import tempfile
+
+import horovod.torch as hvd
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.utils.data.distributed
+from filelock import FileLock
+from torchvision import datasets, transforms
+
+import ray.train.torch
+from ray import train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.horovod import HorovodTrainer
+
+
+def metric_average(val, name):
+    tensor = torch.tensor(val)
+    avg_tensor = hvd.allreduce(tensor, name=name)
+    return avg_tensor.item()
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
+        self.conv2_drop = nn.Dropout2d()
+        self.fc1 = nn.Linear(320, 50)
+        self.fc2 = nn.Linear(50, 10)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x)
+
+
+def setup(config):
+    data_dir = config.get("data_dir", None)
+    seed = config.get("seed", 42)
+    batch_size = config.get("batch_size", 64)
+    use_adasum = config.get("use_adasum", False)
+    lr = config.get("lr", 0.01)
+    momentum = config.get("momentum", 0.5)
+    use_cuda = config.get("use_cuda", False)
+
+    # Horovod: initialize library.
+    hvd.init()
+    torch.manual_seed(seed)
+
+    if use_cuda:
+        # Horovod: pin GPU to local rank.
+        torch.cuda.set_device(hvd.local_rank())
+        torch.cuda.manual_seed(seed)
+
+    # Horovod: limit # of CPU threads to be used per worker.
+    torch.set_num_threads(1)
+
+    kwargs = {"pin_memory": True} if use_cuda else {}
+    data_dir = data_dir or "~/data"
+    with FileLock(os.path.expanduser("~/.horovod_lock")):
+        train_dataset = datasets.MNIST(
+            data_dir,
+            train=True,
+            download=True,
+            transform=transforms.Compose(
+                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+            ),
+        )
+    # Horovod: use DistributedSampler to partition the training data.
+    train_sampler = torch.utils.data.distributed.DistributedSampler(
+        train_dataset, num_replicas=hvd.size(), rank=hvd.rank()
+    )
+    # Note, don't set `num_workers` in DataLoader (not even 1),
+    # as that will separately start multiple processes (each corresponding to 1 worker)
+    # to load the data. This is known to cause issues with Ray.
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+    )
+
+    model = Net()
+
+    # By default, Adasum doesn't need scaling up learning rate.
+    lr_scaler = hvd.size() if not use_adasum else 1
+
+    if use_cuda:
+        # Move model to GPU.
+        model.cuda()
+        # If using GPU Adasum allreduce, scale learning rate by local_size.
+        if use_adasum and hvd.nccl_built():
+            lr_scaler = hvd.local_size()
+
+    # Horovod: scale learning rate by lr_scaler.
+    optimizer = optim.SGD(model.parameters(), lr=lr * lr_scaler, momentum=momentum)
+
+    # Horovod: wrap optimizer with DistributedOptimizer.
+    optimizer = hvd.DistributedOptimizer(
+        optimizer,
+        named_parameters=model.named_parameters(),
+        op=hvd.Adasum if use_adasum else hvd.Average,
+    )
+
+    return model, optimizer, train_loader, train_sampler
+
+
+def train_epoch(
+    model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+):
+    loss = None
+    model.train()
+    # Horovod: set epoch to sampler for shuffling.
+    train_sampler.set_epoch(epoch)
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if use_cuda:
+            data, target = data.cuda(), target.cuda()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            # Horovod: use train_sampler to determine the number of
+            # examples in this worker's partition.
+            print(
+                "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
+                    epoch,
+                    batch_idx * len(data),
+                    len(train_sampler),
+                    100.0 * batch_idx / len(train_loader),
+                    loss.item(),
+                )
+            )
+    return loss.item() if loss else None
+
+
+def train_func(config):
+    num_epochs = config.get("num_epochs", 10)
+    log_interval = config.get("log_interval", 10)
+    use_cuda = config.get("use_cuda", False)
+
+    model, optimizer, train_loader, train_sampler = setup(config)
+
+    results = []
+    for epoch in range(num_epochs):
+        loss = train_epoch(
+            model, optimizer, train_sampler, train_loader, epoch, log_interval, use_cuda
+        )
+        results.append(loss)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report({"loss": loss}, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    # Only used for testing.
+    return results
+
+
+def main(num_workers, use_gpu, kwargs):
+    trainer = HorovodTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config={
+            "num_epochs": kwargs["num_epochs"],
+            "log_interval": kwargs["log_interval"],
+            "use_cuda": kwargs["use_cuda"],
+        },
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    result = trainer.fit()
+    print(result)
+
+
+if __name__ == "__main__":
+    # Training settings
+    parser = argparse.ArgumentParser(
+        description="PyTorch MNIST Example",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=64,
+        metavar="N",
+        help="input batch size for training (default: 64)",
+    )
+    parser.add_argument(
+        "--num-epochs",
+        type=int,
+        default=5,
+        metavar="N",
+        help="number of epochs to train (default: 10)",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.01,
+        metavar="LR",
+        help="learning rate (default: 0.01)",
+    )
+    parser.add_argument(
+        "--momentum",
+        type=float,
+        default=0.5,
+        metavar="M",
+        help="SGD momentum (default: 0.5)",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="enables CUDA training"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
+    )
+    parser.add_argument(
+        "--log-interval",
+        type=int,
+        default=10,
+        metavar="N",
+        help="how many batches to wait before logging training status",
+    )
+    parser.add_argument(
+        "--use-adasum",
+        action="store_true",
+        default=False,
+        help="use adasum algorithm to do reduction",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=2,
+        help="Number of Ray workers to use for training.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="location of the training dataset in the local filesystem ("
+        "will be downloaded if needed)",
+    )
+    parser.add_argument(
+        "--address",
+        required=False,
+        type=str,
+        default=None,
+        help="Address of Ray cluster.",
+    )
+
+    args = parser.parse_args()
+
+    if args.address:
+        ray.init(args.address)
+    else:
+        ray.init()
+
+    use_cuda = args.use_gpu if args.use_gpu is not None else False
+
+    kwargs = {
+        "data_dir": args.data_dir,
+        "seed": args.seed,
+        "use_cuda": use_cuda,
+        "batch_size": args.batch_size,
+        "use_adasum": args.use_adasum if args.use_adasum else False,
+        "lr": args.lr,
+        "momentum": args.momentum,
+        "num_epochs": args.num_epochs,
+        "log_interval": args.log_interval,
+    }
+
+    main(num_workers=args.num_workers, use_gpu=use_cuda, kwargs=kwargs)

python/ray/train/examples/horovod/horovod_tune_example.py

@@ -0,0 +1,139 @@
+import time
+
+import numpy as np
+import torch
+
+import ray
+import ray.train.torch
+from ray import train, tune
+from ray.train import ScalingConfig
+from ray.train.horovod import HorovodTrainer
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+
+
+def sq(x):
+    m2 = 1.0
+    m1 = -20.0
+    m0 = 50.0
+    return m2 * x * x + m1 * x + m0
+
+
+def qu(x):
+    m3 = 10.0
+    m2 = 5.0
+    m1 = -20.0
+    m0 = -5.0
+    return m3 * x * x * x + m2 * x * x + m1 * x + m0
+
+
+class Net(torch.nn.Module):
+    def __init__(self, mode="sq"):
+        super(Net, self).__init__()
+
+        if mode == "square":
+            self.mode = 0
+            self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0]))
+        else:
+            self.mode = 1
+            self.param = torch.nn.Parameter(torch.FloatTensor([1.0, -1.0, 1.0]))
+
+    def forward(self, x):
+        if ~self.mode:
+            return x * x + self.param[0] * x + self.param[1]
+        else:
+            return_val = 10 * x * x * x
+            return_val += self.param[0] * x * x
+            return_val += self.param[1] * x + self.param[2]
+            return return_val
+
+
+def train_loop_per_worker(config):
+    import horovod.torch as hvd
+    import torch
+
+    hvd.init()
+    device = ray.train.torch.get_device()
+    mode = config["mode"]
+    net = Net(mode).to(device)
+    optimizer = torch.optim.SGD(
+        net.parameters(),
+        lr=config["lr"],
+    )
+    optimizer = hvd.DistributedOptimizer(optimizer)
+
+    num_steps = 5
+    print(hvd.size())
+    np.random.seed(1 + hvd.rank())
+    torch.manual_seed(1234)
+    # To ensure consistent initialization across workers,
+    hvd.broadcast_parameters(net.state_dict(), root_rank=0)
+    hvd.broadcast_optimizer_state(optimizer, root_rank=0)
+
+    start = time.time()
+    x_max = config["x_max"]
+    for step in range(1, num_steps + 1):
+        features = torch.Tensor(np.random.rand(1) * 2 * x_max - x_max).to(device)
+        if mode == "square":
+            labels = sq(features)
+        else:
+            labels = qu(features)
+        optimizer.zero_grad()
+        outputs = net(features)
+        loss = torch.nn.MSELoss()(outputs, labels)
+        loss.backward()
+
+        optimizer.step()
+        time.sleep(0.1)
+        train.report(dict(loss=loss.item()))
+    total = time.time() - start
+    print(f"Took {total:0.3f} s. Avg: {total / num_steps:0.3f} s.")
+
+
+def tune_horovod(num_workers, num_samples, use_gpu, mode="square", x_max=1.0):
+    horovod_trainer = HorovodTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        train_loop_config={"mode": mode, "x_max": x_max},
+    )
+
+    tuner = Tuner(
+        horovod_trainer,
+        param_space={"train_loop_config": {"lr": tune.uniform(0.1, 1)}},
+        tune_config=TuneConfig(mode="min", metric="loss", num_samples=num_samples),
+        _tuner_kwargs={"fail_fast": True},
+    )
+
+    result_grid = tuner.fit()
+
+    print("Best hyperparameters found were: ", result_grid.get_best_result().config)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--mode", type=str, default="square", choices=["square", "cubic"]
+    )
+    parser.add_argument(
+        "--learning_rate", type=float, default=0.1, dest="learning_rate"
+    )
+    parser.add_argument("--x_max", type=float, default=1.0, dest="x_max")
+    parser.add_argument("--gpu", action="store_true")
+    parser.add_argument(
+        "--smoke-test", action="store_true", help=("Finish quickly for testing.")
+    )
+    parser.add_argument("--num-workers", type=int, default=2)
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        ray.init(num_cpus=3)
+
+    tune_horovod(
+        num_workers=args.num_workers,
+        num_samples=2 if args.smoke_test else 10,
+        use_gpu=args.gpu,
+        mode=args.mode,
+        x_max=args.x_max,
+    )

python/ray/train/examples/pytorch/torch_data_prefetch_benchmark/auto_pipeline_for_host_to_device_data_transfer.py

@@ -0,0 +1,161 @@
+# The PyTorch data transfer benchmark script.
+import argparse
+import warnings
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ray.train as train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class Net(nn.Module):
+    def __init__(self, in_d, hidden):
+        # output dim = 1
+        super(Net, self).__init__()
+        dims = [in_d] + hidden + [1]
+        self.layers = nn.ModuleList(
+            [nn.Linear(dims[i - 1], dims[i]) for i in range(len(dims))]
+        )
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class BenchmarkDataset(torch.utils.data.Dataset):
+    """Create a naive dataset for the benchmark"""
+
+    def __init__(self, dim, size=1000):
+        self.x = torch.from_numpy(np.random.normal(size=(size, dim))).float()
+        self.y = torch.from_numpy(np.random.normal(size=(size, 1))).float()
+        self.size = size
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return self.size
+
+
+def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
+    if train.get_context().get_world_size() > 1:
+        dataloader.sampler.set_epoch(epoch)
+
+    for X, y in dataloader:
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def train_func(config):
+    data_size = config.get("data_size", 4096 * 50)
+    batch_size = config.get("batch_size", 4096)
+    hidden_size = config.get("hidden_size", 1)
+    use_auto_transfer = config.get("use_auto_transfer", False)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 10)
+
+    train_dataset = BenchmarkDataset(4096, size=data_size)
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True
+    )
+
+    train_loader = train.torch.prepare_data_loader(
+        data_loader=train_loader, move_to_device=True, auto_transfer=use_auto_transfer
+    )
+
+    model = Net(in_d=4096, hidden=[4096] * hidden_size)
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+
+    choice = "with" if use_auto_transfer else "without"
+    print(f"Starting the torch data prefetch benchmark {choice} auto pipeline...")
+
+    torch.cuda.synchronize()
+    start.record()
+    for epoch in range(epochs):
+        train_epoch(epoch, train_loader, model, loss_fn, optimizer)
+    end.record()
+    torch.cuda.synchronize()
+
+    print(
+        f"Finished the torch data prefetch benchmark {choice} "
+        f"auto pipeline: {start.elapsed_time(end)} ms."
+    )
+
+    return "Experiment done."
+
+
+def train_linear(num_workers=1, num_hidden_layers=1, use_auto_transfer=True, epochs=3):
+    config = {
+        "lr": 1e-2,
+        "hidden_size": num_hidden_layers,
+        "batch_size": 4096,
+        "epochs": epochs,
+        "use_auto_transfer": use_auto_transfer,
+    }
+    trainer = TorchTrainer(
+        train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(use_gpu=True, num_workers=num_workers),
+    )
+    results = trainer.fit()
+
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=1, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--num_hidden_layers",
+        type=int,
+        default=1,
+        help="Number of epochs to train for.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    ray.init(address=args.address)
+
+    if not torch.cuda.is_available():
+        warnings.warn("GPU is not available. Skip the test using auto pipeline.")
+    else:
+        train_linear(
+            num_workers=1,
+            num_hidden_layers=args.num_hidden_layers,
+            use_auto_transfer=True,
+            epochs=args.epochs,
+        )
+
+    torch.cuda.empty_cache()
+    train_linear(
+        num_workers=1,
+        num_hidden_layers=args.num_hidden_layers,
+        use_auto_transfer=False,
+        epochs=args.epochs,
+    )
+
+    ray.shutdown()

python/ray/train/examples/pytorch/torch_fashion_mnist_example.py

@@ -0,0 +1,154 @@
+import os
+from typing import Dict
+
+import torch
+from filelock import FileLock
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.transforms import Normalize, ToTensor
+from tqdm import tqdm
+
+import ray.train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def get_dataloaders(batch_size):
+    # Transform to normalize the input images
+    transform = transforms.Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
+
+    with FileLock(os.path.expanduser("~/data.lock")):
+        # Download training data from open datasets
+        training_data = datasets.FashionMNIST(
+            root="~/data",
+            train=True,
+            download=True,
+            transform=transform,
+        )
+
+        # Download test data from open datasets
+        test_data = datasets.FashionMNIST(
+            root="~/data",
+            train=False,
+            download=True,
+            transform=transform,
+        )
+
+    # Create data loaders
+    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
+    test_dataloader = DataLoader(test_data, batch_size=batch_size)
+
+    return train_dataloader, test_dataloader
+
+
+# Model Definition
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Dropout(0.25),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Dropout(0.25),
+            nn.Linear(512, 10),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+def train_func_per_worker(config: Dict):
+    ray.train.torch.enable_reproducibility()
+
+    lr = config["lr"]
+    epochs = config["epochs"]
+    batch_size = config["batch_size_per_worker"]
+
+    # Get dataloaders inside the worker training function
+    train_dataloader, test_dataloader = get_dataloaders(batch_size=batch_size)
+
+    # [1] Prepare Dataloader for distributed training
+    # Shard the datasets among workers and move batches to the correct device
+    # =======================================================================
+    train_dataloader = ray.train.torch.prepare_data_loader(train_dataloader)
+    test_dataloader = ray.train.torch.prepare_data_loader(test_dataloader)
+
+    model = NeuralNetwork()
+
+    # [2] Prepare and wrap your model with DistributedDataParallel
+    # Move the model to the correct GPU/CPU device
+    # ============================================================
+    model = ray.train.torch.prepare_model(model)
+
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
+
+    # Model training loop
+    for epoch in range(epochs):
+        if ray.train.get_context().get_world_size() > 1:
+            # Required for the distributed sampler to shuffle properly across epochs.
+            train_dataloader.sampler.set_epoch(epoch)
+
+        model.train()
+        for X, y in tqdm(train_dataloader, desc=f"Train Epoch {epoch}"):
+            pred = model(X)
+            loss = loss_fn(pred, y)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        model.eval()
+        test_loss, num_correct, num_total = 0, 0, 0
+        with torch.no_grad():
+            for X, y in tqdm(test_dataloader, desc=f"Test Epoch {epoch}"):
+                pred = model(X)
+                loss = loss_fn(pred, y)
+
+                test_loss += loss.item()
+                num_total += y.shape[0]
+                num_correct += (pred.argmax(1) == y).sum().item()
+
+        test_loss /= len(test_dataloader)
+        accuracy = num_correct / num_total
+
+        # [3] Report metrics to Ray Train
+        # ===============================
+        ray.train.report(metrics={"loss": test_loss, "accuracy": accuracy})
+
+
+def train_fashion_mnist(num_workers=2, use_gpu=False):
+    global_batch_size = 32
+
+    train_config = {
+        "lr": 1e-3,
+        "epochs": 10,
+        "batch_size_per_worker": global_batch_size // num_workers,
+    }
+
+    # Configure computation resources
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+
+    # Initialize a Ray TorchTrainer
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func_per_worker,
+        train_loop_config=train_config,
+        scaling_config=scaling_config,
+    )
+
+    # [4] Start distributed training
+    # Run `train_func_per_worker` on all workers
+    # =============================================
+    result = trainer.fit()
+    print(f"Training result: {result}")
+
+
+if __name__ == "__main__":
+    train_fashion_mnist(num_workers=4, use_gpu=True)

python/ray/train/examples/pytorch/torch_linear_example.py

@@ -0,0 +1,147 @@
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ray.train as train
+from ray.train import Checkpoint, RunConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class LinearDataset(torch.utils.data.Dataset):
+    """y = a * x + b"""
+
+    def __init__(self, a, b, size=1000):
+        x = np.arange(0, 10, 10 / size, dtype=np.float32)
+        self.x = torch.from_numpy(x)
+        self.y = torch.from_numpy(a * x + b)
+
+    def __getitem__(self, index):
+        return self.x[index, None], self.y[index, None]
+
+    def __len__(self):
+        return len(self.x)
+
+
+def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
+    if train.get_context().get_world_size() > 1:
+        dataloader.sampler.set_epoch(epoch)
+
+    for X, y in dataloader:
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def validate_epoch(dataloader, model, loss_fn):
+    num_batches = len(dataloader)
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for X, y in dataloader:
+            pred = model(X)
+            loss += loss_fn(pred, y).item()
+    loss /= num_batches
+    import copy
+
+    model_copy = copy.deepcopy(model)
+    return model_copy.cpu().state_dict(), loss
+
+
+def train_func(config):
+    data_size = config.get("data_size", 1000)
+    val_size = config.get("val_size", 400)
+    batch_size = config.get("batch_size", 32)
+    hidden_size = config.get("hidden_size", 1)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 3)
+
+    train_dataset = LinearDataset(2, 5, size=data_size)
+    val_dataset = LinearDataset(2, 5, size=val_size)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
+    validation_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
+
+    train_loader = train.torch.prepare_data_loader(train_loader)
+    validation_loader = train.torch.prepare_data_loader(validation_loader)
+
+    model = nn.Linear(1, hidden_size)
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.MSELoss()
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    results = []
+    for epoch in range(epochs):
+        train_epoch(epoch, train_loader, model, loss_fn, optimizer)
+        state_dict, loss = validate_epoch(validation_loader, model, loss_fn)
+        result = dict(loss=loss)
+        results.append(result)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(state_dict, os.path.join(tmpdir, "model.pt"))
+            train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    return results
+
+
+def train_linear(num_workers=2, use_gpu=False, epochs=3, storage_path=None):
+    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        run_config=RunConfig(storage_path=storage_path),
+    )
+    result = trainer.fit()
+
+    print(result.metrics)
+    return result.metrics
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", help="Whether to use GPU for training."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    if args.smoke_test:
+        # 2 workers + 1 for trainer.
+        ray.init(num_cpus=3)
+        train_linear()
+    else:
+        ray.init(address=args.address)
+        train_linear(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )

python/ray/train/examples/pytorch/torch_quick_start.py

@@ -0,0 +1,110 @@
+# ruff: noqa
+# fmt: off
+# isort: skip_file
+
+# __torch_setup_begin__
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
+def get_dataset():
+    return datasets.FashionMNIST(
+        root="/tmp/data",
+        train=True,
+        download=True,
+        transform=ToTensor(),
+    )
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+        )
+
+    def forward(self, inputs):
+        inputs = self.flatten(inputs)
+        logits = self.linear_relu_stack(inputs)
+        return logits
+# __torch_setup_end__
+
+# __torch_single_begin__
+def train_func():
+    num_epochs = 3
+    batch_size = 64
+
+    dataset = get_dataset()
+    dataloader = DataLoader(dataset, batch_size=batch_size)
+
+    model = NeuralNetwork()
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(num_epochs):
+        for inputs, labels in dataloader:
+            optimizer.zero_grad()
+            pred = model(inputs)
+            loss = criterion(pred, labels)
+            loss.backward()
+            optimizer.step()
+        print(f"epoch: {epoch}, loss: {loss.item()}")
+# __torch_single_end__
+
+# __torch_distributed_begin__
+import ray.train.torch
+
+def train_func_distributed():
+    num_epochs = 3
+    batch_size = 64
+
+    dataset = get_dataset()
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    dataloader = ray.train.torch.prepare_data_loader(dataloader)
+
+    model = NeuralNetwork()
+    model = ray.train.torch.prepare_model(model)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(num_epochs):
+        if ray.train.get_context().get_world_size() > 1:
+            dataloader.sampler.set_epoch(epoch)
+
+        for inputs, labels in dataloader:
+            optimizer.zero_grad()
+            pred = model(inputs)
+            loss = criterion(pred, labels)
+            loss.backward()
+            optimizer.step()
+        print(f"epoch: {epoch}, loss: {loss.item()}")
+# __torch_distributed_end__
+
+
+if __name__ == "__main__":
+    # __torch_single_run_begin__
+    train_func()
+    # __torch_single_run_end__
+
+    # __torch_trainer_begin__
+    from ray.train.torch import TorchTrainer
+    from ray.train import ScalingConfig
+
+    # For GPU Training, set `use_gpu` to True.
+    use_gpu = False
+
+    trainer = TorchTrainer(
+        train_func_distributed,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu)
+    )
+
+    results = trainer.fit()
+    # __torch_trainer_end__

python/ray/train/examples/pytorch/torch_regression_example.py

@@ -0,0 +1,159 @@
+import argparse
+import os
+import tempfile
+from typing import Tuple
+
+import pandas as pd
+import torch
+import torch.nn as nn
+
+import ray
+import ray.train as train
+from ray.data import Dataset
+from ray.train import Checkpoint, DataConfig, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def get_datasets(split: float = 0.7) -> Tuple[Dataset]:
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+
+    def combine_x(batch):
+        return pd.DataFrame(
+            {
+                "x": batch[[f"x{i:03d}" for i in range(100)]].values.tolist(),
+                "y": batch["y"],
+            }
+        )
+
+    dataset = dataset.map_batches(combine_x, batch_format="pandas")
+    train_dataset, validation_dataset = dataset.repartition(
+        num_blocks=4
+    ).train_test_split(split, shuffle=True)
+    return train_dataset, validation_dataset
+
+
+def train_epoch(iterable_dataset, model, loss_fn, optimizer, device):
+    model.train()
+    for X, y in iterable_dataset:
+        X = X.to(device)
+        y = y.to(device)
+
+        # Compute prediction error
+        pred = model(X)
+        loss = loss_fn(pred, y)
+
+        # Backpropagation
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+
+def validate_epoch(iterable_dataset, model, loss_fn, device):
+    num_batches = 0
+    model.eval()
+    loss = 0
+    with torch.no_grad():
+        for X, y in iterable_dataset:
+            X = X.to(device)
+            y = y.to(device)
+            num_batches += 1
+            pred = model(X)
+            loss += loss_fn(pred, y).item()
+    loss /= num_batches
+    result = {"loss": loss}
+    return result
+
+
+def train_func(config):
+    batch_size = config.get("batch_size", 32)
+    hidden_size = config.get("hidden_size", 10)
+    lr = config.get("lr", 1e-2)
+    epochs = config.get("epochs", 3)
+
+    train_dataset_shard = train.get_dataset_shard("train")
+    validation_dataset = train.get_dataset_shard("validation")
+
+    model = nn.Sequential(
+        nn.Linear(100, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1)
+    )
+    model = train.torch.prepare_model(model)
+
+    loss_fn = nn.L1Loss()
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    results = []
+
+    def create_torch_iterator(shard):
+        iterator = shard.iter_torch_batches(batch_size=batch_size)
+        for batch in iterator:
+            yield batch["x"].float(), batch["y"].float()
+
+    for _ in range(epochs):
+        train_torch_dataset = create_torch_iterator(train_dataset_shard)
+        validation_torch_dataset = create_torch_iterator(validation_dataset)
+
+        device = train.torch.get_device()
+
+        train_epoch(train_torch_dataset, model, loss_fn, optimizer, device)
+        if train.get_context().get_world_rank() == 0:
+            result = validate_epoch(validation_torch_dataset, model, loss_fn, device)
+        else:
+            result = {}
+        results.append(result)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
+            train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
+
+    return results
+
+
+def train_regression(num_workers=2, use_gpu=False):
+    train_dataset, val_dataset = get_datasets()
+    config = {"lr": 1e-2, "hidden_size": 20, "batch_size": 4, "epochs": 3}
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        datasets={"train": train_dataset, "validation": val_dataset},
+        dataset_config=DataConfig(datasets_to_split=["train"]),
+    )
+
+    result = trainer.fit()
+    return result
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Use GPU for training."
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        ray.init(num_cpus=4)
+        result = train_regression()
+    else:
+        ray.init(address=args.address)
+        result = train_regression(num_workers=args.num_workers, use_gpu=args.use_gpu)
+    print(result)

python/ray/train/examples/pytorch_geometric/distributed_sage_example.py

@@ -0,0 +1,228 @@
+# Adapted from https://github.com/pyg-team/pytorch_geometric/blob/2.1.0
+# /examples/multi_gpu/distributed_sampling.py
+
+import argparse
+import os
+
+import torch
+import torch.nn.functional as F
+from filelock import FileLock
+from torch_geometric.datasets import FakeDataset, Reddit
+from torch_geometric.loader import NeighborSampler
+from torch_geometric.nn import SAGEConv
+from torch_geometric.transforms import RandomNodeSplit
+
+from ray import train
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+class SAGE(torch.nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=2):
+        super().__init__()
+        self.num_layers = num_layers
+
+        self.convs = torch.nn.ModuleList()
+        self.convs.append(SAGEConv(in_channels, hidden_channels))
+        for _ in range(self.num_layers - 2):
+            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
+        self.convs.append(SAGEConv(hidden_channels, out_channels))
+
+    def forward(self, x, adjs):
+        for i, (edge_index, _, size) in enumerate(adjs):
+            x_target = x[: size[1]]  # Target nodes are always placed first.
+            x = self.convs[i]((x, x_target), edge_index)
+            if i != self.num_layers - 1:
+                x = F.relu(x)
+                x = F.dropout(x, p=0.5, training=self.training)
+        return x.log_softmax(dim=-1)
+
+    @torch.no_grad()
+    def test(self, x_all, subgraph_loader):
+        for i in range(self.num_layers):
+            xs = []
+            for batch_size, n_id, adj in subgraph_loader:
+                edge_index, _, size = adj
+                x = x_all[n_id.to(x_all.device)].to(train.torch.get_device())
+                x_target = x[: size[1]]
+                x = self.convs[i]((x, x_target), edge_index)
+                if i != self.num_layers - 1:
+                    x = F.relu(x)
+                xs.append(x.cpu())
+
+            x_all = torch.cat(xs, dim=0)
+
+        return x_all
+
+
+def train_loop_per_worker(train_loop_config):
+    dataset = train_loop_config["dataset_fn"]()
+    batch_size = train_loop_config["batch_size"]
+    num_epochs = train_loop_config["num_epochs"]
+
+    data = dataset[0]
+    train_idx = data.train_mask.nonzero(as_tuple=False).view(-1)
+    train_idx = train_idx.split(
+        train_idx.size(0) // train.get_context().get_world_size()
+    )[train.get_context().get_world_rank()]
+
+    train_loader = NeighborSampler(
+        data.edge_index,
+        node_idx=train_idx,
+        sizes=[25, 10],
+        batch_size=batch_size,
+        shuffle=True,
+    )
+
+    # Disable distributed sampler since the train_loader has already been split above.
+    train_loader = train.torch.prepare_data_loader(train_loader, add_dist_sampler=False)
+
+    # Do validation on rank 0 worker only.
+    if train.get_context().get_world_rank() == 0:
+        subgraph_loader = NeighborSampler(
+            data.edge_index, node_idx=None, sizes=[-1], batch_size=2048, shuffle=False
+        )
+        subgraph_loader = train.torch.prepare_data_loader(
+            subgraph_loader, add_dist_sampler=False
+        )
+
+    model = SAGE(dataset.num_features, 256, dataset.num_classes)
+    model = train.torch.prepare_model(model)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+    x, y = data.x.to(train.torch.get_device()), data.y.to(train.torch.get_device())
+
+    for epoch in range(num_epochs):
+        model.train()
+
+        # ``batch_size`` is the number of samples in the current batch.
+        # ``n_id`` are the ids of all the nodes used in the computation. This is
+        # needed to pull in the necessary features just for the current batch that is
+        # being trained on.
+        # ``adjs`` is a list of 3 element tuple consisting of ``(edge_index, e_id,
+        # size)`` for each sample in the batch, where ``edge_index``represent the
+        # edges of the sampled subgraph, ``e_id`` are the ids of the edges in the
+        # sample, and ``size`` holds the shape of the subgraph.
+        # See ``torch_geometric.loader.neighbor_sampler.NeighborSampler`` for more info.
+        for batch_size, n_id, adjs in train_loader:
+            optimizer.zero_grad()
+            out = model(x[n_id], adjs)
+            loss = F.nll_loss(out, y[n_id[:batch_size]])
+            loss.backward()
+            optimizer.step()
+
+        if train.get_context().get_world_rank() == 0:
+            print(f"Epoch: {epoch:03d}, Loss: {loss:.4f}")
+
+        train_accuracy = validation_accuracy = test_accuracy = None
+
+        # Do validation on rank 0 worker only.
+        if train.get_context().get_world_rank() == 0:
+            model.eval()
+            with torch.no_grad():
+                out = model.module.test(x, subgraph_loader)
+            res = out.argmax(dim=-1) == data.y
+            train_accuracy = int(res[data.train_mask].sum()) / int(
+                data.train_mask.sum()
+            )
+            validation_accuracy = int(res[data.val_mask].sum()) / int(
+                data.val_mask.sum()
+            )
+            test_accuracy = int(res[data.test_mask].sum()) / int(data.test_mask.sum())
+
+        train.report(
+            dict(
+                train_accuracy=train_accuracy,
+                validation_accuracy=validation_accuracy,
+                test_accuracy=test_accuracy,
+            )
+        )
+
+
+def gen_fake_dataset():
+    """Returns a function to be called on each worker that returns a Fake Dataset."""
+
+    # For fake dataset, since the dataset is randomized, we create it once on the
+    # driver, and then send the same dataset to all the training workers.
+    # Use 10% of nodes for validation and 10% for testing.
+    fake_dataset = FakeDataset(transform=RandomNodeSplit(num_val=0.1, num_test=0.1))
+
+    def gen_dataset():
+        return fake_dataset
+
+    return gen_dataset
+
+
+def gen_reddit_dataset():
+    """Returns a function to be called on each worker that returns Reddit Dataset."""
+
+    # For Reddit dataset, we have to download the data on each node, so we create the
+    # dataset on each training worker.
+    with FileLock(os.path.expanduser("~/.reddit_dataset_lock")):
+        dataset = Reddit("./data/Reddit")
+    return dataset
+
+
+def train_gnn(
+    num_workers=2, use_gpu=False, epochs=3, global_batch_size=32, dataset="reddit"
+):
+    per_worker_batch_size = global_batch_size // num_workers
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_loop_per_worker,
+        train_loop_config={
+            "num_epochs": epochs,
+            "batch_size": per_worker_batch_size,
+            "dataset_fn": gen_reddit_dataset
+            if dataset == "reddit"
+            else gen_fake_dataset(),
+        },
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+    )
+    result = trainer.fit()
+    print(result.metrics)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", help="Whether to use GPU for training."
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--global-batch-size",
+        "-b",
+        type=int,
+        default=32,
+        help="Global batch size to use for training.",
+    )
+    parser.add_argument(
+        "--dataset",
+        "-d",
+        type=str,
+        choices=["reddit", "fake"],
+        default="reddit",
+        help="The dataset to use. Either 'reddit' or 'fake' Defaults to 'reddit'.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    train_gnn(
+        num_workers=args.num_workers,
+        use_gpu=args.use_gpu,
+        epochs=args.epochs,
+        global_batch_size=args.global_batch_size,
+        dataset=args.dataset,
+    )

python/ray/train/examples/tf/tensorflow_autoencoder_example.py

@@ -0,0 +1,177 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+import ray
+from ray import train
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.data.datasource import SimpleTensorFlowDatasource
+from ray.data.extensions import TensorArray
+from ray.train import Result, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard
+
+
+def get_dataset(split_type="train"):
+    def dataset_factory():
+        return tfds.load("mnist", split=[split_type], as_supervised=True)[0].take(128)
+
+    dataset = ray.data.read_datasource(
+        SimpleTensorFlowDatasource(), dataset_factory=dataset_factory
+    )
+
+    def normalize_images(x):
+        x = np.float32(x.numpy()) / 255.0
+        x = np.reshape(x, (-1,))
+        return x
+
+    def preprocess_dataset(batch):
+        return [
+            (normalize_images(image), normalize_images(image)) for image, _ in batch
+        ]
+
+    dataset = dataset.map_batches(preprocess_dataset)
+
+    def convert_batch_to_pandas(batch):
+
+        images = [TensorArray(image) for image, _ in batch]
+        # because we did autoencoder here
+        df = pd.DataFrame({"image": images, "label": images})
+        return df
+
+    dataset = dataset.map_batches(convert_batch_to_pandas)
+    return dataset
+
+
+def build_autoencoder_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(784,)),
+            # encoder
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(32, activation="relu"),
+            # decoder
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(784, activation="sigmoid"),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    dataset_shard = train.get_dataset_shard("train")
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_autoencoder_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.BinaryCrossentropy(),
+            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
+            metrics=[
+                "binary_crossentropy",
+            ],
+        )
+
+    def to_tf_dataset(dataset, batch_size):
+        def to_tensor_iterator():
+            for batch in dataset.iter_tf_batches(
+                batch_size=batch_size, dtypes=tf.float32
+            ):
+                yield batch["image"], batch["label"]
+
+        output_signature = (
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+        )
+        tf_dataset = tf.data.Dataset.from_generator(
+            to_tensor_iterator, output_signature=output_signature
+        )
+        return prepare_dataset_shard(tf_dataset)
+
+    results = []
+    for epoch in range(epochs):
+        tf_dataset = to_tf_dataset(
+            dataset=dataset_shard,
+            batch_size=per_worker_batch_size,
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2, use_gpu: bool = False, epochs: int = 4
+) -> Result:
+    train_dataset = get_dataset(split_type="train")
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        datasets={"train": train_dataset},
+        scaling_config=scaling_config,
+    )
+
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )
+    print(result)

python/ray/train/examples/tf/tensorflow_mnist_example.py

@@ -0,0 +1,138 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import json
+
+import numpy as np
+import tensorflow as tf
+from filelock import FileLock
+
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.train import Result, RunConfig, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer
+
+
+def mnist_dataset(batch_size: int) -> tf.data.Dataset:
+    with FileLock(os.path.expanduser("~/.mnist_lock")):
+        (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = (
+        tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        .shuffle(60000)
+        .repeat()
+        .batch(batch_size)
+    )
+    return train_dataset
+
+
+def build_cnn_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(28, 28)),
+            tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+            tf.keras.layers.Conv2D(32, 3, activation="relu"),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(10),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+    steps_per_epoch = config.get("steps_per_epoch", 70)
+
+    tf_config = json.loads(os.environ["TF_CONFIG"])
+    num_workers = len(tf_config["cluster"]["worker"])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_cnn_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
+            metrics=["accuracy"],
+        )
+
+    history = multi_worker_model.fit(
+        multi_worker_dataset,
+        epochs=epochs,
+        steps_per_epoch=steps_per_epoch,
+        callbacks=[ReportCheckpointCallback()],
+    )
+    results = history.history
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2,
+    use_gpu: bool = False,
+    epochs: int = 4,
+    storage_path: str = None,
+) -> Result:
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        run_config=RunConfig(storage_path=storage_path),
+    )
+    results = trainer.fit()
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )

python/ray/train/examples/tf/tensorflow_quick_start.py

@@ -0,0 +1,90 @@
+# ruff: noqa
+# fmt: off
+# isort: skip_file
+
+# __tf_setup_begin__
+import os
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import sys
+import numpy as np
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+def mnist_dataset(batch_size):
+    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = tf.data.Dataset.from_tensor_slices(
+        (x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
+    return train_dataset
+
+
+def build_and_compile_cnn_model():
+    model = tf.keras.Sequential([
+        tf.keras.layers.InputLayer(input_shape=(28, 28)),
+        tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+        tf.keras.layers.Conv2D(32, 3, activation='relu'),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(128, activation='relu'),
+        tf.keras.layers.Dense(10)
+    ])
+    model.compile(
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
+        metrics=['accuracy'])
+    return model
+# __tf_setup_end__
+
+# __tf_single_begin__
+def train_func():
+    batch_size = 64
+    single_worker_dataset = mnist_dataset(batch_size)
+    single_worker_model = build_and_compile_cnn_model()
+    single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_single_end__
+
+# __tf_distributed_begin__
+import json
+import os
+
+def train_func_distributed():
+    per_worker_batch_size = 64
+    # This environment variable will be set by Ray Train.
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    num_workers = len(tf_config['cluster']['worker'])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_and_compile_cnn_model()
+
+    multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_distributed_end__
+
+if __name__ == "__main__":
+    # __tf_single_run_begin__
+    train_func()
+    # __tf_single_run_end__
+
+    # __tf_trainer_begin__
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train import ScalingConfig
+
+    # For GPU Training, set `use_gpu` to True.
+    use_gpu = False
+
+    trainer = TensorflowTrainer(train_func_distributed, scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu))
+
+    trainer.fit()
+    # __tf_trainer_end__

python/ray/train/examples/tf/tensorflow_regression_example.py

@@ -0,0 +1,115 @@
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import sys
+
+import ray
+from ray import train
+from ray.data.preprocessors import Concatenator
+from ray.train import Result, ScalingConfig
+
+if sys.version_info >= (3, 12):
+    # Skip this test in Python 3.12+ because TensorFlow is not supported.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train.tensorflow.keras import ReportCheckpointCallback
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=(100,)),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_model()
+        multi_worker_model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
+            loss=tf.keras.losses.mean_absolute_error,
+            metrics=[tf.keras.metrics.mean_squared_error],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    results = []
+    for _ in range(epochs):
+        tf_dataset = dataset.to_tf(
+            feature_columns="x", label_columns="y", batch_size=batch_size
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_regression(num_workers: int = 2, use_gpu: bool = False) -> Result:
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+    columns_to_concatenate = [f"x{i:03}" for i in range(100)]
+    preprocessor = Concatenator(columns=columns_to_concatenate, output_column_name="x")
+    dataset = preprocessor.fit_transform(dataset)
+
+    config = {"lr": 1e-3, "batch_size": 32, "epochs": 4}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+        datasets={"train": dataset},
+    )
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_regression(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_regression(
+            num_workers=args.num_workers, use_gpu=args.use_gpu
+        )
+    print(result)

python/ray/train/examples/transformers/transformers_torch_trainer_basic.py

@@ -0,0 +1,78 @@
+import evaluate
+import numpy as np
+
+# Minimal Example adapted from https://huggingface.co/docs/transformers/training
+from datasets import load_dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+
+from ray.train import ScalingConfig
+from ray.train.huggingface.transformers import RayTrainReportCallback, prepare_trainer
+from ray.train.torch import TorchTrainer
+
+
+# [1] Define a training function that includes all your training logic
+# ====================================================================
+def train_func(config):
+    # Datasets
+    dataset = load_dataset("Yelp/yelp_review_full")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    tokenized_ds = dataset.map(tokenize_function, batched=True)
+
+    small_train_ds = tokenized_ds["train"].shuffle(seed=42).select(range(1000))
+    small_eval_ds = tokenized_ds["test"].shuffle(seed=42).select(range(1000))
+
+    # Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", num_labels=5
+    )
+
+    # Evaluation metrics
+    metric = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        logits, labels = eval_pred
+        predictions = np.argmax(logits, axis=-1)
+        return metric.compute(predictions=predictions, references=labels)
+
+    # Hugging Face Trainer
+    training_args = TrainingArguments(
+        output_dir="test_trainer", eval_strategy="epoch", report_to="none"
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=small_train_ds,
+        eval_dataset=small_eval_ds,
+        compute_metrics=compute_metrics,
+    )
+
+    # [2] Report metrics and checkpoints to Ray Train
+    # ===============================================
+    trainer.add_callback(RayTrainReportCallback())
+
+    # [3] Prepare your trainer for Ray Data integration
+    # =================================================
+    trainer = prepare_trainer(trainer)
+
+    # Start Training
+    trainer.train()
+
+
+if __name__ == "__main__":
+    # [4] Build a Ray TorchTrainer to launch `train_func` on all workers
+    # ==================================================================
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    trainer.fit()