chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,153 @@
+# 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"