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,129 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from pytest import approx
+from unit.util import torch_assert_close
+from unit.common import DistributedTest, preferred_dtype
+from unit.multi_output_model import MultiOutputModel, multi_output_dataloader
+
+
+class TestTwoOutputModel(DistributedTest):
+    world_size = 1
+
+    def test(self, tmpdir):
+        grad_accumulation_steps = 2
+        micro_batch_size = 1
+        world_size = self.world_size
+        config_dict = {
+            "train_micro_batch_size_per_gpu": micro_batch_size,
+            "gradient_accumulation_steps": grad_accumulation_steps,
+            "train_batch_size": micro_batch_size * grad_accumulation_steps * world_size,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+
+        hidden_dim = 10
+        weight_value = 0.1
+
+        model = MultiOutputModel(hidden_dim, weight_value)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+        total_samples = 4
+        data_loader = multi_output_dataloader(model=model,
+                                              total_samples=total_samples,
+                                              hidden_dim=hidden_dim,
+                                              device=model.device,
+                                              inputs=[1.0, 2.0],
+                                              targets=[1, 2])
+        for n, batch in enumerate(data_loader):
+            assert len(batch) % 2 == 0, \
+                 "multi_output_dataloader failed to return even number of data samples (input+target)"
+
+            midpoint = len(batch) // 2
+            inputs, targets = batch[:midpoint], batch[midpoint:]
+            loss_tuple = model(inputs, targets)
+
+            expected_loss = torch.tensor(2.302734375, dtype=preferred_dtype(), device=model.device)
+            for loss in loss_tuple:
+                assert loss.shape == torch.Size([])
+                assert loss.item() == approx(expected_loss.item())
+
+            summed_loss = sum(loss_tuple)
+            scaled_loss = model.backward(summed_loss)
+            expected_scaled_loss = summed_loss / grad_accumulation_steps
+            torch_assert_close(scaled_loss, expected_scaled_loss)
+
+            model.step()
+
+
+class TestThreeOutputModel(DistributedTest):
+    world_size = 1
+
+    def test(self, tmpdir):
+        grad_accumulation_steps = 3
+        micro_batch_size = 1
+        world_size = 1
+        config_dict = {
+            "train_micro_batch_size_per_gpu": micro_batch_size,
+            "gradient_accumulation_steps": grad_accumulation_steps,
+            "train_batch_size": micro_batch_size * grad_accumulation_steps * world_size,
+            "steps_per_print": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 0.00015
+                }
+            },
+        }
+        if get_accelerator().is_bf16_supported():
+            config_dict["bf16"] = {"enabled": True}
+        elif get_accelerator().is_fp16_supported():
+            config_dict["fp16"] = {"enabled": True}
+
+        hidden_dim = 10
+        weight_value = 0.1
+
+        model = MultiOutputModel(hidden_dim, weight_value)
+        model, _, _, _ = deepspeed.initialize(config=config_dict, model=model, model_parameters=model.parameters())
+
+        total_samples = grad_accumulation_steps * micro_batch_size * 2
+        data_loader = multi_output_dataloader(model=model,
+                                              total_samples=total_samples,
+                                              hidden_dim=hidden_dim,
+                                              device=model.device,
+                                              inputs=[1.0, 2.0, 3.0],
+                                              targets=[1, 2, 3])
+        for n, batch in enumerate(data_loader):
+            assert len(batch) % 2 == 0, \
+                 "multi_output_dataloader failed to return even number of data samples (input+target)"
+
+            midpoint = len(batch) // 2
+            inputs, targets = batch[:midpoint], batch[midpoint:]
+            loss_tuple = model(inputs, targets)
+            assert len(loss_tuple) == 3
+
+            expected_loss = torch.tensor(2.302734375, dtype=preferred_dtype(), device=model.device)
+
+            for loss in loss_tuple:
+                assert loss.shape == torch.Size([])
+                assert loss.item() == approx(expected_loss.item())
+
+            summed_loss = sum(loss_tuple)
+            scaled_loss = model.backward(summed_loss)
+            expected_scaled_loss = summed_loss / grad_accumulation_steps
+            torch_assert_close(scaled_loss, expected_scaled_loss)
+
+            model.step()