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

Code before:


def step_amp(enabled, baseline_model, baseline_optimizer, target_engine, dtype, enable_autocast_outside,
             baseline_scaler, step, x, y, rtol, atol, expect_match):
    device_type = get_accelerator().device_name()

    # Runs the forward pass with autocasting.
    with torch.autocast(device_type=device_type, dtype=dtype, enabled=enabled):
        baseline_optimizer.zero_grad()
        baseline_loss = baseline_model(x, y)

    baseline_scaler.scale(baseline_loss).backward()
    baseline_scaler.step(baseline_optimizer)
    baseline_scaler.update()

    # We don't need torch.autocast here in real applications, but want to test the behavior of nested autocast.
    with torch.autocast(device_type=device_type, dtype=dtype, enabled=enable_autocast_outside):
        target_loss = target_engine(x, y)

    # reduce-scatter in `dtype` makes a difference in the loss.
    if step <= 1 and expect_match:
        allclose_local = torch.allclose(baseline_loss.float(), target_loss.float(), rtol=rtol, atol=atol)
        if not allclose_local:
            print(f"Losses do not match: baseline_loss={baseline_loss}, target_loss={target_loss}")
        # Ensure all ranks either pass or fail together.
        # If some rank fail while others pass, subsequent tests or iterations may hang.
        if not reduce_boolean_flags(allclose_local, all):
            assert False, f"Losses do not match on one or more ranks."

    target_engine.backward(target_loss)
    target_engine.step()


@enable_determinism(123)
def compare_loss(model_cls,
                 enable,
                 zero_stage,
                 model_dtype,
                 dtype,
                 autocast_conf,
                 enable_autocast_outside,
                 lower_precision_safe_modules,
                 expect_match=True):
    iteration = 5
    hidden_dim = 10
    lr = 0.001

    if dtype == torch.bfloat16 and not bf16_required_version_check():
        raise ValueError(
            "DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
        )

Review comment (bug):
```suggestion
        # If some ranks fail while others pass, subsequent tests or iterations may hang.
```

Code after:


def step_amp(enabled, baseline_model, baseline_optimizer, target_engine, dtype, enable_autocast_outside,
             baseline_scaler, step, x, y, rtol, atol, expect_match):
    device_type = get_accelerator().device_name()

    # Runs the forward pass with autocasting.
    with torch.autocast(device_type=device_type, dtype=dtype, enabled=enabled):
        baseline_optimizer.zero_grad()
        baseline_loss = baseline_model(x, y)

    baseline_scaler.scale(baseline_loss).backward()
    baseline_scaler.step(baseline_optimizer)
    baseline_scaler.update()

    # We don't need torch.autocast here in real applications, but want to test the behavior of nested autocast.
    with torch.autocast(device_type=device_type, dtype=dtype, enabled=enable_autocast_outside):
        target_loss = target_engine(x, y)

    # reduce-scatter in `dtype` makes a difference in the loss.
    if step <= 1 and expect_match:
        allclose_local = torch.allclose(baseline_loss.float(), target_loss.float(), rtol=rtol, atol=atol)
        if not allclose_local:
            print(f"Losses do not match: baseline_loss={baseline_loss}, target_loss={target_loss}")
        # Ensure all ranks either pass or fail together.
        # If some ranks fail while others pass, subsequent tests or iterations may hang.
        if not reduce_boolean_flags(allclose_local, all):
            assert False, f"Losses do not match on one or more ranks."

    target_engine.backward(target_loss)
    target_engine.step()


@enable_determinism(123)
def compare_loss(model_cls,
                 enable,
                 zero_stage,
                 model_dtype,
                 dtype,
                 autocast_conf,
                 enable_autocast_outside,
                 lower_precision_safe_modules,
                 expect_match=True):
    iteration = 5
    hidden_dim = 10
    lr = 0.001

    if dtype == torch.bfloat16 and not bf16_required_version_check():
        raise ValueError(
            "DeepSpeed BFloat16 tests need torch >= 1.10, NCCL >= 2.10.3, CUDA > =11.0 and HW support for BFloat16 to run correctly"
        )