wehub-resource-sync
154 lines
5.5 KiB
Python
154 lines
5.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import re
|
|
import torch
|
|
import pytest
|
|
import deepspeed
|
|
from types import SimpleNamespace
|
|
from deepspeed.profiling.flops_profiler import get_model_profile, FlopsProfiler
|
|
from unit.simple_model import SimpleModel, random_dataloader
|
|
from unit.common import DistributedTest
|
|
from deepspeed.utils.torch import required_torch_version
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
if torch.half not in get_accelerator().supported_dtypes():
|
|
pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True)
|
|
|
|
pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.3),
|
|
reason='requires Pytorch version 1.3 or above')
|
|
|
|
|
|
def within_range(val, target, tolerance):
|
|
return abs(val - target) / target < tolerance
|
|
|
|
|
|
TOLERANCE = 0.05
|
|
|
|
|
|
class LeNet5(torch.nn.Module):
|
|
|
|
def __init__(self, n_classes):
|
|
super(LeNet5, self).__init__()
|
|
|
|
self.feature_extractor = torch.nn.Sequential(
|
|
torch.nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1),
|
|
torch.nn.Tanh(),
|
|
torch.nn.AvgPool2d(kernel_size=2),
|
|
torch.nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1),
|
|
torch.nn.Tanh(),
|
|
torch.nn.AvgPool2d(kernel_size=2),
|
|
torch.nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5, stride=1),
|
|
torch.nn.Tanh(),
|
|
)
|
|
|
|
self.classifier = torch.nn.Sequential(
|
|
torch.nn.Linear(in_features=120, out_features=84),
|
|
torch.nn.Tanh(),
|
|
torch.nn.Linear(in_features=84, out_features=n_classes),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = self.feature_extractor(x)
|
|
x = torch.flatten(x, 1)
|
|
logits = self.classifier(x)
|
|
probs = torch.nn.functional.softmax(logits, dim=1)
|
|
return logits, probs
|
|
|
|
|
|
class TestFlopsProfiler(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self):
|
|
config_dict = {
|
|
"train_batch_size": 1,
|
|
"steps_per_print": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.001,
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": 0
|
|
},
|
|
"fp16": {
|
|
"enabled": True,
|
|
},
|
|
"flops_profiler": {
|
|
"enabled": True,
|
|
"step": 1,
|
|
"module_depth": -1,
|
|
"top_modules": 3,
|
|
},
|
|
}
|
|
hidden_dim = 10
|
|
model = SimpleModel(hidden_dim, empty_grad=False)
|
|
|
|
model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
|
|
|
|
data_loader = random_dataloader(model=model,
|
|
total_samples=50,
|
|
hidden_dim=hidden_dim,
|
|
device=model.device,
|
|
dtype=torch.half)
|
|
for n, batch in enumerate(data_loader):
|
|
loss = model(batch[0], batch[1])
|
|
model.backward(loss)
|
|
model.step()
|
|
if n == 3: break
|
|
assert within_range(model.flops_profiler.flops, 200, tolerance=TOLERANCE)
|
|
assert model.flops_profiler.params == 110
|
|
|
|
def test_flops_profiler_in_inference(self):
|
|
mod = LeNet5(10)
|
|
batch_size = 1024
|
|
input = torch.randn(batch_size, 1, 32, 32)
|
|
flops, macs, params = get_model_profile(
|
|
mod,
|
|
tuple(input.shape),
|
|
print_profile=True,
|
|
detailed=True,
|
|
module_depth=-1,
|
|
top_modules=3,
|
|
warm_up=1,
|
|
as_string=False,
|
|
ignore_modules=None,
|
|
)
|
|
print(flops, macs, params)
|
|
assert within_range(flops, 866076672, TOLERANCE)
|
|
assert within_range(macs, 426516480, TOLERANCE)
|
|
assert params == 61706
|
|
|
|
|
|
def test_print_model_profile_with_none_dp_world_size(capsys):
|
|
# Regression test for https://github.com/deepspeedai/DeepSpeed/issues/7483
|
|
# Under sequence parallelism (Ulysses) the engine reports dp_world_size as None, which used to
|
|
# crash print_model_profile with "unsupported format string passed to NoneType.__format__".
|
|
model = torch.nn.Sequential(torch.nn.Linear(128, 128))
|
|
prof = FlopsProfiler(model)
|
|
# Mimic a DeepSpeed engine configured with Ulysses sequence parallelism, where dp_world_size is
|
|
# None and the effective data-parallel replication is the sequence-data-parallel group.
|
|
prof.ds_engine = SimpleNamespace(world_size=8,
|
|
dp_world_size=None,
|
|
seq_dp_world_size=4,
|
|
mp_world_size=1,
|
|
has_moe_layers=False,
|
|
train_micro_batch_size_per_gpu=lambda: 1,
|
|
wall_clock_breakdown=lambda: False)
|
|
|
|
prof.start_profile()
|
|
# A few forward passes give the profiled modules a non-zero measured duration.
|
|
for _ in range(3):
|
|
model(torch.randn(64, 128))
|
|
prof.print_model_profile(profile_step=1, detailed=False)
|
|
prof.end_profile()
|
|
|
|
out = capsys.readouterr().out
|
|
match = re.search(r"data parallel size:\s+(\S+)", out)
|
|
assert match is not None
|
|
# The sequence-data-parallel world size is reported in place of the None dp_world_size.
|
|
assert match.group(1) == "4"
|