wehub-resource-sync
148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
import pytest
|
|
import deepspeed.comm as dist
|
|
import torch
|
|
|
|
from unit.common import DistributedTest
|
|
from unit.simple_model import random_dataloader, SimpleModel
|
|
|
|
import deepspeed
|
|
from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
|
|
from deepspeed.runtime.zero.partition_parameters import Init
|
|
from deepspeed.ops.aio import AsyncIOBuilder
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
|
|
@pytest.mark.sequential
|
|
class TestNVMeCheckpointing(DistributedTest):
|
|
world_size = 1
|
|
|
|
@pytest.mark.parametrize('param_offload_device, optim_offload_device',
|
|
[(OffloadDeviceEnum.none, OffloadDeviceEnum.nvme),
|
|
(OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme),
|
|
(OffloadDeviceEnum.nvme, OffloadDeviceEnum.none),
|
|
(OffloadDeviceEnum.nvme, OffloadDeviceEnum.cpu),
|
|
(OffloadDeviceEnum.nvme, OffloadDeviceEnum.nvme)])
|
|
def test_nvme_checkpointing(self, tmpdir, param_offload_device, optim_offload_device):
|
|
zero_dir, ckpt_dir = os.path.join(tmpdir, "zero"), os.path.join(tmpdir, "checkpoint")
|
|
|
|
first_stage_steps, second_stage_steps = 2, 2
|
|
|
|
if not get_accelerator().is_fp16_supported():
|
|
pytest.skip("fp16 is not supported")
|
|
|
|
if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
|
|
pytest.skip('Skip tests since async-io is not compatible')
|
|
|
|
torch.manual_seed(123)
|
|
|
|
config_dict = {
|
|
"train_micro_batch_size_per_gpu": 1,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00015,
|
|
}
|
|
},
|
|
"fp16": {
|
|
"enabled": True,
|
|
"initial_scale_power": 8
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 3,
|
|
"offload_param": {
|
|
"device": param_offload_device,
|
|
"nvme_path": str(zero_dir)
|
|
},
|
|
"offload_optimizer": {
|
|
"device": optim_offload_device,
|
|
"nvme_path": str(zero_dir)
|
|
},
|
|
"sub_group_size": 100,
|
|
"stage3_max_live_parameters": 100,
|
|
"stage3_param_persistence_threshold": 0,
|
|
},
|
|
"aio": {
|
|
"block_size": 1048576 # Minimum AIO bytes, anything smaller than this will not be offloaded
|
|
}
|
|
}
|
|
|
|
hidden_dim, nlayers = 2048, 2
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
model = SimpleModel(hidden_dim, nlayers=nlayers, empty_grad=False)
|
|
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
model.empty_partition_cache()
|
|
|
|
assert first_stage_steps > 0
|
|
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=first_stage_steps,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.float16)
|
|
dist.barrier()
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
|
|
dist.barrier()
|
|
model.save_checkpoint(ckpt_dir)
|
|
|
|
if second_stage_steps > 0:
|
|
second_stage_batches = list(
|
|
random_dataloader(model=model,
|
|
total_samples=second_stage_steps,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.float16))
|
|
dist.barrier()
|
|
for n, batch in enumerate(second_stage_batches):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
dist.barrier()
|
|
|
|
final_batch = next(
|
|
iter(
|
|
random_dataloader(model=model,
|
|
total_samples=1,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.float16)))
|
|
dist.barrier()
|
|
loss_before = float(model(final_batch[0], final_batch[1]))
|
|
|
|
# Needed in ZeRO 3. Not doing so can give memory leak
|
|
model.destroy()
|
|
|
|
# TODO: This should be on the engine? There needs to be a better way.
|
|
Init.param_id = 0
|
|
|
|
with deepspeed.zero.Init(config_dict_or_path=config_dict):
|
|
model = SimpleModel(hidden_dim, nlayers=nlayers, empty_grad=False)
|
|
|
|
model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
|
|
|
|
model.load_checkpoint(ckpt_dir)
|
|
|
|
if second_stage_steps > 0:
|
|
dist.barrier()
|
|
for n, batch in enumerate(second_stage_batches):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
dist.barrier()
|
|
|
|
dist.barrier()
|
|
loss_after = float(model(final_batch[0], final_batch[1]))
|
|
|
|
assert loss_before == loss_after
|