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,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import deepspeed
+import deepspeed.comm as dist
+import torch
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel, random_dataloader
+
+
+def create_model(config_dict):
+    hidden_dim = 64
+    model = SimpleModel(hidden_dim)
+    model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict)
+    return model
+
+
+def train_shared_loss(num_models, config_dict, dtype):
+    hidden_dim = 64
+
+    models = [create_model(config_dict) for _ in range(num_models)]
+    data_loader = random_dataloader(model=models[0],
+                                    total_samples=4,
+                                    hidden_dim=hidden_dim,
+                                    device=models[0].device,
+                                    dtype=dtype)
+    dist.barrier()
+    for _, batch in enumerate(data_loader):
+        losses = [m.module(batch[0], batch[1]) for m in models]
+        loss = sum(l / (i + 1) for i, l in enumerate(losses))
+        loss.backward()
+
+        for m in models:
+            m._backward_epilogue()
+
+        for m in models:
+            m.step()
+
+        for m in models:
+            m.optimizer.zero_grad()
+
+    for m in models:
+        m.destroy()
+
+
+def train_independent_loss(num_models, config_dict, dtype):
+    hidden_dim = 64
+
+    models = [create_model(config_dict) for _ in range(num_models)]
+    data_loader = random_dataloader(model=models[0],
+                                    total_samples=4,
+                                    hidden_dim=hidden_dim,
+                                    device=models[0].device,
+                                    dtype=dtype)
+    dist.barrier()
+    for _, batch in enumerate(data_loader):
+        losses = [m.module(batch[0], batch[1]) for m in models]
+        for m, loss in zip(models, losses):
+            m.backward(loss)
+            m.step()
+
+    for m in models:
+        m.destroy()
+
+
+@pytest.mark.parametrize('num_models', [1, 2, 3])
+class TestMultipleModels(DistributedTest):
+    world_size = 2
+    reuse_dist_env = True
+
+    @pytest.mark.parametrize('shared_loss', [False, True])
+    @pytest.mark.parametrize('zero_stage', [1, 2, 3])
+    @pytest.mark.parametrize('fp32_grad_accum', [False, True])
+    @pytest.mark.parametrize('contiguous_gradients', [False, True])
+    @pytest.mark.parametrize('overlap_comm', [False, True])
+    def test_zero_optimizer(self, num_models, shared_loss, zero_stage, fp32_grad_accum, contiguous_gradients,
+                            overlap_comm):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": zero_stage,
+                "contiguous_gradients": contiguous_gradients,
+                "overlap_comm": overlap_comm,
+            },
+            "fp16": {
+                "initial_scale_power": 8,
+                "enabled": True
+            },
+        }
+        if fp32_grad_accum:
+            config_dict["data_types"] = {"grad_accum_dtype": "fp32"}
+
+        if shared_loss:
+            train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
+        else:
+            train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.float16)
+
+    # TODO: Combination of shared_loss==True and bf16.immediate_grad_update==False is currently broken
+    @pytest.mark.parametrize('shared_loss', [False, True])
+    def test_bf16_optimizer(self, num_models, shared_loss):
+        config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-4
+                }
+            },
+            "zero_optimization": {
+                "stage": 1,
+            },
+            "bf16": {
+                "enabled": True,
+                "immediate_grad_update": True,
+            },
+            "data_types": {
+                "grad_accum_dtype": "fp32"
+            }
+        }
+
+        if shared_loss:
+            train_shared_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)
+        else:
+            train_independent_loss(num_models=num_models, config_dict=config_dict, dtype=torch.bfloat16)