wehub-resource-sync
226 lines
7.0 KiB
Python
226 lines
7.0 KiB
Python
import os
|
|
from tempfile import TemporaryDirectory
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from accelerate import Accelerator
|
|
|
|
import ray
|
|
import ray.train as train
|
|
from ray.train import Checkpoint, ScalingConfig
|
|
from ray.train.examples.pytorch.torch_linear_example import LinearDataset
|
|
from ray.train.torch import TorchTrainer
|
|
|
|
DEEPSPEED_CONFIG = {
|
|
"fp16": {
|
|
"enabled": "auto",
|
|
"loss_scale": 0,
|
|
"loss_scale_window": 1000,
|
|
"initial_scale_power": 16,
|
|
"hysteresis": 2,
|
|
"min_loss_scale": 1,
|
|
},
|
|
"bf16": {"enabled": "auto"},
|
|
"optimizer": {
|
|
"type": "AdamW",
|
|
"params": {
|
|
"lr": "auto",
|
|
"weight_decay": "auto",
|
|
"torch_adam": True,
|
|
"adam_w_mode": True,
|
|
},
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 2,
|
|
"offload_optimizer": {"device": "cpu", "pin_memory": True},
|
|
"allgather_partitions": True,
|
|
"allgather_bucket_size": 2e8,
|
|
"overlap_comm": True,
|
|
"reduce_scatter": True,
|
|
"contiguous_gradients": True,
|
|
},
|
|
"gradient_accumulation_steps": 1,
|
|
"gradient_clipping": "auto",
|
|
"steps_per_print": 2000,
|
|
"train_batch_size": "auto",
|
|
"train_micro_batch_size_per_gpu": "auto",
|
|
"wall_clock_breakdown": False,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
def linear_train_func(accelerator: Accelerator, config):
|
|
from accelerate.utils import DummyOptim
|
|
from deepspeed.ops.adam import DeepSpeedCPUAdam
|
|
|
|
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)
|
|
|
|
model = nn.Linear(1, hidden_size)
|
|
|
|
loss_fn = nn.MSELoss()
|
|
if (
|
|
accelerator.state.deepspeed_plugin
|
|
and "optimizer" in accelerator.state.deepspeed_plugin.deepspeed_config
|
|
):
|
|
optimizer_cls = DummyOptim
|
|
elif accelerator.state.deepspeed_plugin:
|
|
optimizer_cls = DeepSpeedCPUAdam
|
|
else:
|
|
optimizer_cls = torch.optim.SGD
|
|
|
|
# Accelerate boilerplate
|
|
no_decay = ["bias", "LayerNorm.weight"]
|
|
optimizer_grouped_parameters = [
|
|
{
|
|
"params": [
|
|
p
|
|
for n, p in model.named_parameters()
|
|
if not any(nd in n for nd in no_decay)
|
|
],
|
|
"weight_decay": 0.0,
|
|
},
|
|
{
|
|
"params": [
|
|
p
|
|
for n, p in model.named_parameters()
|
|
if any(nd in n for nd in no_decay)
|
|
],
|
|
"weight_decay": 0.0,
|
|
},
|
|
]
|
|
optimizer = optimizer_cls(optimizer_grouped_parameters, lr=lr)
|
|
train_loader, validation_loader, model, optimizer = accelerator.prepare(
|
|
train_loader, validation_loader, model, optimizer
|
|
)
|
|
|
|
results = []
|
|
for _ in range(epochs):
|
|
for X, y in train_loader:
|
|
# Compute prediction error
|
|
pred = model(X)
|
|
loss = loss_fn(pred, y)
|
|
|
|
# Backpropagation
|
|
accelerator.backward(loss)
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
|
|
num_batches = len(validation_loader)
|
|
model.eval()
|
|
loss = 0
|
|
with torch.no_grad():
|
|
for X, y in validation_loader:
|
|
pred = model(X)
|
|
loss += loss_fn(pred, y).item()
|
|
loss /= num_batches
|
|
import copy
|
|
|
|
model_copy = copy.deepcopy(accelerator.unwrap_model(model))
|
|
state_dict, loss = model_copy.cpu().state_dict(), loss
|
|
|
|
result = dict(loss=loss)
|
|
results.append(result)
|
|
|
|
with TemporaryDirectory() as tmpdir:
|
|
torch.save(state_dict, os.path.join(tmpdir, "checkpoint.pt"))
|
|
train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
|
|
|
|
return results
|
|
|
|
|
|
@pytest.mark.parametrize("use_gpu", [True, False])
|
|
def test_accelerate_base(ray_2_node_2_gpu, use_gpu):
|
|
def train_func(config):
|
|
accelerator = Accelerator(cpu=not use_gpu)
|
|
assert accelerator.device == train.torch.get_device()
|
|
assert accelerator.process_index == train.get_context().get_world_rank()
|
|
if accelerator.device.type != "cpu":
|
|
assert (
|
|
accelerator.local_process_index == train.get_context().get_local_rank()
|
|
)
|
|
result = linear_train_func(accelerator, config)
|
|
assert len(result) == epochs
|
|
assert result[-1]["loss"] < result[0]["loss"]
|
|
|
|
epochs = 3
|
|
scaling_config = ScalingConfig(num_workers=2, use_gpu=use_gpu)
|
|
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=scaling_config,
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_accelerate_deepspeed(ray_2_node_2_gpu):
|
|
from accelerate import DeepSpeedPlugin
|
|
|
|
def train_func(config):
|
|
deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=DEEPSPEED_CONFIG)
|
|
accelerator = Accelerator(deepspeed_plugin=deepspeed_plugin)
|
|
|
|
assert accelerator.device == train.torch.get_device()
|
|
assert accelerator.process_index == train.get_context().get_world_rank()
|
|
assert accelerator.local_process_index == train.get_context().get_local_rank()
|
|
result = linear_train_func(accelerator, config)
|
|
assert len(result) == epochs
|
|
assert result[-1]["loss"] < result[0]["loss"]
|
|
|
|
epochs = 3
|
|
scaling_config = ScalingConfig(num_workers=2, use_gpu=True)
|
|
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=scaling_config,
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
# Using CPU on purpose
|
|
@pytest.mark.parametrize("num_workers", [1, 2])
|
|
def test_accelerate_e2e(ray_start_4_cpus, num_workers):
|
|
def train_func():
|
|
accelerator = Accelerator(cpu=True)
|
|
assert accelerator.device == train.torch.get_device()
|
|
assert accelerator.process_index == train.get_context().get_world_rank()
|
|
model = torch.nn.Linear(3, 1)
|
|
model = accelerator.prepare(model)
|
|
with TemporaryDirectory() as tmpdir:
|
|
torch.save(model, os.path.join(tmpdir, "checkpoint.pt"))
|
|
train.report({}, checkpoint=Checkpoint.from_directory(tmpdir))
|
|
|
|
scaling_config = ScalingConfig(num_workers=num_workers)
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=scaling_config,
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|