wehub-resource-sync
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import deepspeed
|
|
import pytest
|
|
from unit.common import DistributedTest, is_rocm_pytorch
|
|
from unit.util import skip_on_arch
|
|
|
|
try:
|
|
import transformer_engine.pytorch as transformer_engine
|
|
from transformer_engine.common import recipe
|
|
except ImportError:
|
|
pytest.skip("Transformer Engine package is missing, skipping tests", allow_module_level=True)
|
|
|
|
|
|
@pytest.mark.parametrize("base_datatype", ["fp16", "bf16", "fp32"])
|
|
class TestFp8ComposabilityAcrossZero(DistributedTest):
|
|
world_size = 1
|
|
|
|
def test(self, base_datatype):
|
|
skip_on_arch(min_arch=9)
|
|
|
|
def run_zero(stage, model_dtype):
|
|
num_batches = 128
|
|
batch_size = 16
|
|
hidden_dim = 768
|
|
# Have to set seed before model
|
|
torch.random.manual_seed(42)
|
|
enable_fp16 = model_dtype == torch.float16
|
|
enable_bf16 = model_dtype == torch.bfloat16
|
|
# TransformerEngine Model
|
|
model = transformer_engine.Linear(hidden_dim, hidden_dim, bias=True, params_dtype=model_dtype)
|
|
|
|
# Create FP8 recipe. Note: All input args are optional.
|
|
fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID,
|
|
amax_history_len=16,
|
|
amax_compute_algo="max")
|
|
config = {
|
|
"train_batch_size": batch_size,
|
|
"gradient_accumulation_steps": 1,
|
|
"optimizer": {
|
|
"type": "Adam",
|
|
"params": {
|
|
"lr": 0.00001
|
|
}
|
|
},
|
|
"zero_optimization": {
|
|
"stage": stage
|
|
},
|
|
"fp16": {
|
|
"enabled": enable_fp16,
|
|
"loss_scale": 0.1
|
|
},
|
|
"bf16": {
|
|
"enabled": enable_bf16
|
|
}
|
|
}
|
|
# Init DeepSpeed
|
|
model, optimizer, _, _ = deepspeed.initialize(args=None,
|
|
model=model,
|
|
model_parameters=model.parameters(),
|
|
config=config)
|
|
|
|
batches = torch.randn(num_batches, batch_size, hidden_dim, device=model.device, dtype=model_dtype)
|
|
for batch in batches:
|
|
# Enables autocasting for the forward pass
|
|
with transformer_engine.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
|
|
out = model(batch)
|
|
loss = out.mean()
|
|
model.backward(loss)
|
|
model.step()
|
|
return loss
|
|
|
|
if base_datatype == "fp16":
|
|
model_dtype = torch.float16
|
|
elif base_datatype == "bf16":
|
|
model_dtype = torch.bfloat16
|
|
else:
|
|
model_dtype = torch.float32
|
|
|
|
# Set default tolerances
|
|
rtol, atol = 1e-07, 1e-05
|
|
|
|
# Relax tolerance only for ROCm + FP16
|
|
if is_rocm_pytorch() and base_datatype in ["fp16", "bf16"]:
|
|
rtol, atol = 1e-07, 1e-04
|
|
|
|
# config
|
|
zero_stage = [0, 1, 2, 3]
|
|
losses = []
|
|
for stage in zero_stage:
|
|
loss = run_zero(stage, model_dtype)
|
|
losses.append(loss)
|
|
all_equal = all(torch.allclose(loss, losses[0], rtol, atol) for loss in losses)
|
|
assert (all_equal)
|