Repository: microsoft/DeepSpeed
File: tests/unit/runtime/zero/test_zero_leaf_module.py (Python)

Code before:
class TestZ3LeafOptimization(DistributedTest):
    world_size = 2
    reuse_dist_env = True

    def test_finegrained_optimization(self):
        hidden_dim = 128
        num_block = 16
        stage3_coalesced_fetch_threshold = 12000
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
                "stage3_prefetch_bucket_size": hidden_dim**2,
                "stage3_param_persistence_threshold": 0,
                "stage3_max_reuse_distance": 0,
                "stage3_coalesced_fetch_threshold": stage3_coalesced_fetch_threshold
            }
        }
        if get_accelerator().is_fp16_supported():
            config_dict["fp16"] = {"enabled": True}
        elif get_accelerator().is_bf16_supported():
            config_dict["bf16"] = {"enabled": True}

        def bench_loss_and_time(config):
            warm_up_step = 10
            model = modelWithFineGrainedBlock(hidden_dim, num_block)
            model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
            data_loader = random_dataloader(model=model,
                                            total_samples=20,
                                            hidden_dim=hidden_dim,
                                            device=model.device,
                                            dtype=preferred_dtype())
            dist.barrier()
            loss_list = []

            for i, batch in enumerate(data_loader):
                if i == warm_up_step:
                    dist.barrier()
                    get_accelerator().synchronize()
                    start_time = time.time()
                batch[0].requires_grad = True
                loss = model(batch[0], batch[1])
                loss = loss[1]
                loss_list.append(loss)

Review comment (suggestion):
Please use  `preferred_dtype()` instead of `is_*_supported()` to ensure consistency with line 231. 

Code after:

@pytest.mark.parametrize("module_granularity_threshold", [0, 100, 12100, 10000000])
class TestZ3LeafOptimization(DistributedTest):
    world_size = 2
    reuse_dist_env = True

    def test_finegrained_optimization(self, module_granularity_threshold: int):
        hidden_dim = 128
        num_block = 16
        config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-6
                }
            },
            "zero_optimization": {
                "stage": 3,
                "stage3_prefetch_bucket_size": hidden_dim**2,
                "stage3_param_persistence_threshold": 0,
                "stage3_max_reuse_distance": 0,
            }
        }
        if preferred_dtype() is torch.float16:
            config_dict["fp16"] = {"enabled": True}
        elif preferred_dtype() is torch.bfloat16:
            config_dict["bf16"] = {"enabled": True}

        def bench_loss_and_time(config):
            warm_up_step = 10
            model = modelWithFineGrainedBlock(hidden_dim, num_block)
            model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
            data_loader = random_dataloader(model=model,
                                            total_samples=20,
                                            hidden_dim=hidden_dim,
                                            device=model.device,
                                            dtype=preferred_dtype())
            dist.barrier()
            loss_list = []

            for i, batch in enumerate(data_loader):
                if i == warm_up_step:
                    dist.barrier()
                    get_accelerator().synchronize()
                    start_time = time.time()
                batch[0].requires_grad = True
                loss = model(batch[0], batch[1])
                loss = loss[1]
                loss_list.append(loss)